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

[Attention] Enable masked MHA for GLM-5 head dimensions (#53785)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: OpenAI Codex <noreply@openai.com>

Matthew Bonanni committed Aug 27, 2026 at 17:17 UTC de9250ac9e9b249133fff14e15d9248a1ebbcdb8

10 files changed +527 -74

benchmarks/attention_benchmarks/benchmark.py

+32 -6

@@ -1143,12 +1143,26 @@ def main():

1143 ]

1144 else:

1145 variants = available_variants

1146 + model_sweep = getattr(args, "model_parameter_sweep", None)

1147 + if model_sweep is None:

1148 + num_q_heads_values = [args.num_q_heads]

1149 + elif model_sweep.param_name == "num_q_heads":

1150 + num_q_heads_values = model_sweep.values

1151 + else:

1152 + raise ValueError(

1153 + "mha_vs_mqa mode only supports a num_q_heads model parameter sweep"

1154 + )

1155 + benchmark_specs = [

1156 + (spec, num_q_heads)

1157 + for spec in args.batch_specs

1158 + for num_q_heads in num_q_heads_values

1159 + ]

1160 masked_mha_max_seq_len = getattr(

1161 args, "sparse_mla_masked_mha_max_seq_len", None

1162 )

1163 formatter = ResultsFormatter(console)

1164 total = 0

1151 - for spec in args.batch_specs:

1165 + for spec, _ in benchmark_specs:

1166 q_len = max(request.q_len for request in parse_batch_spec(spec))

1167 for variant_label, _, _ in variants:

1168 if (

@@ -1164,7 +1178,7 @@ def main():

1178 total += len(backends)

1179

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

1167 - for spec in args.batch_specs:

1181 + for spec, num_q_heads in benchmark_specs:

1182 q_len = max(request.q_len for request in parse_batch_spec(spec))

1183 for backend in backends:

1184 for variant_label, force_mqa, mha_mode in variants:

@@ -1178,12 +1192,17 @@ def main():

1192 and q_len > masked_mha_max_seq_len

1193 ):

1194 continue

1195 + result_backend = f"{backend}_{variant_label}"

1196 + if model_sweep is not None:

1197 + result_backend = model_sweep.get_label(

1198 + result_backend, num_q_heads

1199 + )

1200 config = BenchmarkConfig(

1182 - backend=f"{backend}_{variant_label}",

1201 + backend=result_backend,

1202 batch_spec=spec,

1203 num_layers=args.num_layers,

1204 head_dim=args.head_dim,

1186 - num_q_heads=args.num_q_heads,

1205 + num_q_heads=num_q_heads,

1206 num_kv_heads=args.num_kv_heads,

1207 block_size=args.block_size,

1208 device=args.device,

@@ -1211,7 +1230,7 @@ def main():

1230 # run_mla_benchmark needs the real backend name

1231 from mla_runner import run_mla_benchmark as run_mla

1232

1214 - run_label = f"{backend}_{variant_label} {spec}"

1233 + run_label = f"{backend}_{variant_label} {num_q_heads}h {spec}"

1234 pbar.set_postfix_str(run_label)

1235

1236 try:

@@ -1248,7 +1267,14 @@ def main():

1267

1268 # Display results with variant labels as separate "backends"

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

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

1270 + variant_backends = []

1271 + for num_q_heads in num_q_heads_values:

1272 + for backend in backends:

1273 + for variant, _, _ in variants:

1274 + label = f"{backend}_{variant}"

1275 + if model_sweep is not None:

1276 + label = model_sweep.get_label(label, num_q_heads)

1277 + variant_backends.append(label)

1278 formatter.print_table(all_results, variant_backends)

1279

1280 # Handle model parameter sweep mode

benchmarks/attention_benchmarks/configs/mla_sparse_masked_mha_vs_mqa_glm5.yaml new

+59

@@ -0,0 +1,59 @@

1 +mode: mha_vs_mqa

2 +model:

3 + name: "glm-5.x"

4 + num_layers: 78

5 + num_q_heads: 64

6 + num_kv_heads: 1

7 + head_dim: 576

8 + kv_lora_rank: 512

9 + qk_nope_head_dim: 192

10 + qk_rope_head_dim: 64

11 + v_head_dim: 256

12 + block_size: 128

13 + max_model_len: 131072

14 +model_parameter_sweep:

15 + param_name: "num_q_heads"

16 + values: [64, 32, 16, 8]

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

18 +batch_specs:

19 + - "q256s2048"

20 + - "q512s2048"

21 + - "q1024s2048"

22 + - "q1536s2048"

23 + - "q2048"

24 + - "q256s4096"

25 + - "q512s4096"

26 + - "q1024s4096"

27 + - "q1536s4096"

28 + - "q2048s4096"

29 + - "q4096"

30 + - "q256s8192"

31 + - "q512s8192"

32 + - "q1024s8192"

33 + - "q1536s8192"

34 + - "q2048s8192"

35 + - "q4096s8192"

36 + - "q8192"

37 + - "q256s16384"

38 + - "q512s16384"

39 + - "q1024s16384"

40 + - "q1536s16384"

41 + - "q2048s16384"

42 + - "q4096s16384"

43 + - "q8192s16384"

44 + - "q16384"

45 + - "q256s32768"

46 + - "q512s32768"

47 + - "q1024s32768"

48 + - "q1536s32768"

49 + - "q2048s32768"

50 + - "q4096s32768"

51 + - "q8192s32768"

52 + - "q16384s32768"

53 + - "q32768"

54 +backends: ["FLASHMLA_SPARSE", "FLASHINFER_MLA_SPARSE"]

55 +device: "cuda:0"

56 +profile_memory: false

57 +sparse_mla_masked_mha_max_seq_len: 131072

58 +sparse_mla_mha_variants: ["masked_mha", "mqa"]

59 +sparse_mla_topk_pattern: "random"

benchmarks/attention_benchmarks/mla_runner.py

+32

@@ -313,6 +313,9 @@ _BACKEND_PROPERTIES = {

313 "FLASHMLA_SPARSE": {

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

315 },

316 + "FLASHINFER_MLA_SPARSE": {

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

318 + },

319 }

320

321

@@ -1184,6 +1187,35 @@ def _run_mla_benchmark_batched(

1187 )

1188 # Run each benchmark with the shared impl

1189 for config, threshold, num_splits in configs_with_params:

1190 + force_masked_mha = (

1191 + is_sparse and getattr(config, "sparse_mla_mha_mode", "auto") == "masked"

1192 + )

1193 + if force_masked_mha:

1194 + impl.masked_mha_available = True

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

1196 + GLOBAL_TOPK_MASK_MAX_BYTES,

1197 + _topk_mask_shape,

1198 + )

1199 +

1200 + requests = parse_batch_spec(config.batch_spec)

1201 + max_query_len = max(request.q_len for request in requests)

1202 + mask_shape = _topk_mask_shape(

1203 + len(requests), max_query_len, max_query_len

1204 + )

1205 + mask_words = max(

1206 + GLOBAL_TOPK_MASK_MAX_BYTES // torch.int32.itemsize,

1207 + mask_shape[0] * mask_shape[1] * mask_shape[2],

1208 + )

1209 + if (

1210 + builder_instance.topk_mask_workspace is None

1211 + or builder_instance.topk_mask_workspace.numel() != mask_words

1212 + ):

1213 + builder_instance.topk_mask_workspace = torch.zeros(

1214 + mask_words,

1215 + dtype=torch.int32,

1216 + device=device,

1217 + )

1218 +

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

1220 original_threshold = None

1221 effective_threshold = threshold

cmake/external_projects/vllm_flash_attn.cmake

+1 -1

@@ -39,7 +39,7 @@ else()

39 FetchContent_Declare(

40 vllm-flash-attn

41 GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git

42 - GIT_TAG 4eb29359a51b4fe67c8c0f0c5e931236971aa492

42 + GIT_TAG 06bdd47c0d0383daf6a2ff0c418faff9c6da16e5

43 GIT_PROGRESS TRUE

44 # Don't share the vllm-flash-attn build between build types

45 BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

tests/model_executor/layers/test_mla_short_prefill_indexer.py

+75 -7

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

9 import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer

10 from vllm.config import CUDAGraphMode

11 +from vllm.models.deepseek_v32 import attention as deepseek_v32_attention

12 from vllm.models.deepseek_v32.attention import DeepseekV32Attention

13 from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata

@@ -49,7 +50,6 @@ def make_mla_metadata(*, use_dense_mha: bool = True, num_decode_tokens: int = 0)

50 [

51 "short",

52 "threshold_mismatch",

52 - "mqa_only_layer",

53 "force_mqa",

54 "mla_decode",

55 "capture",

@@ -141,11 +141,7 @@ def test_short_prefill_updates_k_cache_before_scoring_decision(

141 topk_indices = torch.full((7, 2048), 17, dtype=torch.int32)

142

143 def run_indexer():

144 - if batch_kind == "mqa_only_layer":

145 - assert not DeepseekV32Attention.supports_dense_mha_prefill

146 - dense_mha_layer = ""

147 - else:

148 - dense_mha_layer = MLA_LAYER

144 + assert DeepseekV32Attention.supports_dense_mha_prefill

145 return sparse_indexer.sparse_attn_indexer(

146 hidden_states,

147 INDEXER_LAYER,

@@ -163,7 +159,7 @@ def test_short_prefill_updates_k_cache_before_scoring_decision(

159 topk_indices,

160 False,

161 False,

166 - dense_mha_layer,

162 + MLA_LAYER,

163 )

164

165 if should_skip:

@@ -222,3 +218,75 @@ def test_skipped_k_cache_insert_accepts_no_k(

218

219 assert result is topk_indices

220 assert torch.all(topk_indices == -1)

221 +

222 +

223 +@pytest.mark.parametrize("fp8_query", [False, True])

224 +def test_deepseek_v32_dispatches_selected_mha(

225 + monkeypatch: pytest.MonkeyPatch,

226 + fp8_query: bool,

227 +) -> None:

228 + attn_metadata = SimpleNamespace(num_actual_tokens=2)

229 + kv_cache = torch.empty(1)

230 + monkeypatch.setattr(

231 + deepseek_v32_attention,

232 + "get_attention_context",

233 + lambda _: (attn_metadata, None, kv_cache, None),

234 + )

235 +

236 + observed = {}

237 +

238 + def record_forward_impl(*args):

239 + observed["args"] = args

240 +

241 + layer = SimpleNamespace(

242 + indexer=None,

243 + skip_topk=False,

244 + layer_name=MLA_LAYER,

245 + use_pcp=False,

246 + _fp8_query=fp8_query,

247 + _use_sparse_mha=lambda _: True,

248 + rotary_emb=lambda _positions, q: (q + 1, None),

249 + forward_impl=record_forward_impl,

250 + )

251 + q_nope = torch.randn(2, 1, 2)

252 + q_pe = torch.randn(2, 1, 2)

253 + mqa_q = torch.randn(2, 1, 2)

254 + mha_q = torch.cat((q_nope, q_pe + 1 if fp8_query else mqa_q), dim=-1)

255 + kv_c = torch.empty(2, 2)

256 + k_pe = torch.empty(2, 2)

257 + output = torch.empty(2, 2)

258 +

259 + DeepseekV32Attention._sparse_indexer_and_attn(

260 + layer,

261 + torch.arange(2),

262 + torch.empty(2, 2),

263 + q_nope,

264 + q_pe,

265 + None,

266 + None,

267 + None,

268 + kv_c,

269 + k_pe,

270 + torch.empty(2, 1, 2),

271 + mqa_q,

272 + output,

273 + )

274 +

275 + expected_args = (

276 + mha_q,

277 + kv_c,

278 + kv_cache,

279 + attn_metadata,

280 + output,

281 + )

282 + actual_args = observed["args"]

283 + torch.testing.assert_close(actual_args[0], mha_q)

284 + assert actual_args[2].shape == (2, 1, 2)

285 + assert actual_args[2].data_ptr() == k_pe.data_ptr()

286 + assert all(

287 + actual is expected

288 + for actual, expected in zip(

289 + actual_args[1:2] + actual_args[3:],

290 + expected_args[1:],

291 + )

292 + )

tests/v1/attention/test_mla_backends.py

+119

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

27 MLAAttention,

28 QueryLenSupport,

29 _DecodeConcatQuantFP8,

30 + _use_masked_mha,

31 build_mla_chunked_context_metadata,

32 )

33 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape

@@ -68,6 +69,124 @@ if current_platform.is_rocm():

69 DEVICE_TYPE = current_platform.device_type

72 +@pytest.mark.parametrize(

73 + ("tensor_parallel_size", "query_len", "expected"),

74 + [

75 + (1, 8192, True),

76 + (1, 9216, False),

77 + (2, 20480, True),

78 + (2, 24576, False),

79 + (4, 48 * 1024, True),

80 + (4, 56 * 1024, False),

81 + (8, 112 * 1024, True),

82 + (8, 128 * 1024, False),

83 + ],

84 +)

85 +def test_glm5_masked_mha_pure_prefill_routing(

86 + tensor_parallel_size, query_len, expected

87 +):

88 + assert (

89 + _use_masked_mha(

90 + backend_name="FLASHMLA_SPARSE",

91 + tensor_parallel_size=tensor_parallel_size,

92 + qk_head_dim=256,

93 + v_head_dim=256,

94 + query_len=query_len,

95 + seq_len=query_len,

96 + has_context=False,

97 + )

98 + is expected

99 + )

100 +

101 +

102 +@pytest.mark.parametrize(

103 + ("tensor_parallel_size", "query_len", "seq_len", "expected"),

104 + [

105 + (1, 8192, 8192, False),

106 + (2, 4096, 6144, True),

107 + (2, 4096, 8192, False),

108 + (2, 8192, 10240, True),

109 + (4, 16384, 24576, True),

110 + (4, 16384, 32768, False),

111 + (4, 32768, 34 * 1024, True),

112 + (4, 32768, 36 * 1024, False),

113 + (8, 32768, 49152, True),

114 + (8, 32768, 65536, False),

115 + ],

116 +)

117 +def test_glm5_masked_mha_context_routing(

118 + tensor_parallel_size, query_len, seq_len, expected

119 +):

120 + assert (

121 + _use_masked_mha(

122 + backend_name="FLASHMLA_SPARSE",

123 + tensor_parallel_size=tensor_parallel_size,

124 + qk_head_dim=256,

125 + v_head_dim=256,

126 + query_len=query_len,

127 + seq_len=seq_len,

128 + has_context=True,

129 + )

130 + is expected

131 + )

132 +

133 +

134 +@pytest.mark.parametrize(

135 + ("tensor_parallel_size", "query_len", "seq_len", "has_context", "expected"),

136 + [

137 + (4, 36 * 1024, 36 * 1024, False, True),

138 + (4, 40 * 1024, 40 * 1024, False, False),

139 + (4, 4 * 1024, 12 * 1024, True, True),

140 + (4, 4 * 1024, 16 * 1024, True, False),

141 + (4, 24 * 1024, 28 * 1024, True, True),

142 + (4, 32 * 1024, 36 * 1024, True, False),

143 + (8, 64 * 1024, 64 * 1024, False, True),

144 + (8, 68 * 1024, 68 * 1024, False, False),

145 + (8, 4 * 1024, 20 * 1024, True, True),

146 + (8, 16 * 1024, 32 * 1024, True, True),

147 + (8, 48 * 1024, 64 * 1024, True, True),

148 + (8, 56 * 1024, 72 * 1024, True, True),

149 + (8, 63 * 1024, 79 * 1024, True, False),

150 + ],

151 +)

152 +def test_glm5_flashinfer_masked_mha_routing(

153 + tensor_parallel_size, query_len, seq_len, has_context, expected

154 +):

155 + assert (

156 + _use_masked_mha(

157 + backend_name="FLASHINFER_MLA_SPARSE",

158 + tensor_parallel_size=tensor_parallel_size,

159 + qk_head_dim=256,

160 + v_head_dim=256,

161 + query_len=query_len,

162 + seq_len=seq_len,

163 + has_context=has_context,

164 + )

165 + is expected

166 + )

167 +

168 +

169 +def test_masked_mha_routing_is_dimension_specific():

170 + assert _use_masked_mha(

171 + backend_name="FLASHMLA_SPARSE",

172 + tensor_parallel_size=1,

173 + qk_head_dim=192,

174 + v_head_dim=128,

175 + query_len=1536,

176 + seq_len=2048,

177 + has_context=True,

178 + )

179 + assert not _use_masked_mha(

180 + backend_name="FLASHMLA_SPARSE",

181 + tensor_parallel_size=1,

182 + qk_head_dim=128,

183 + v_head_dim=128,

184 + query_len=1536,

185 + seq_len=2048,

186 + has_context=True,

187 + )

188 +

189 +

190 @pytest.mark.parametrize(

191 ("cache_dtype", "expected_quant_mode"),

192 [

tests/v1/attention/test_sparse_mla_backends.py

+42 -5

@@ -21,6 +21,7 @@ from tests.v1.attention.utils import (

21 )

22 from vllm import _custom_ops as ops

23 from vllm.config import set_current_vllm_config

24 +from vllm.model_executor.layers.attention.mla_attention import _use_masked_mha

25 from vllm.model_executor.layers.attention.sparse_mla_attention import (

26 GLOBAL_TOPK_MASK_MAX_BYTES,

27 _masked_mha_workspace_fits,

@@ -884,6 +885,35 @@ def test_masked_mha_workspace_fits_single_request_boundary(max_query_len, expect

885 )

886

887

888 +@pytest.mark.parametrize(

889 + ("backend_name", "tensor_parallel_size", "query_len"),

890 + [

891 + ("FLASHMLA_SPARSE", 4, 48 * 1024),

892 + ("FLASHMLA_SPARSE", 8, 112 * 1024),

893 + ("FLASHINFER_MLA_SPARSE", 4, 36 * 1024),

894 + ("FLASHINFER_MLA_SPARSE", 8, 64 * 1024),

895 + ],

896 +)

897 +def test_masked_mha_workspace_guards_long_routing_policy(

898 + backend_name, tensor_parallel_size, query_len

899 +):

900 + assert _use_masked_mha(

901 + backend_name=backend_name,

902 + tensor_parallel_size=tensor_parallel_size,

903 + qk_head_dim=256,

904 + v_head_dim=256,

905 + query_len=query_len,

906 + seq_len=query_len,

907 + has_context=False,

908 + )

909 + assert not _masked_mha_workspace_fits(

910 + batch_size=1,

911 + max_query_len=query_len,

912 + max_context_chunk_seq_len=0,

913 + workspace_numel=GLOBAL_TOPK_MASK_MAX_BYTES // torch.int32.itemsize,

914 + )

915 +

916 +

917 def test_masked_mha_workspace_fits_accounts_for_batch_and_context():

918 """Request count and context chunk length are independent multipliers."""

919 base = dict(batch_size=2, max_query_len=2048, max_context_chunk_seq_len=2048)

@@ -912,14 +942,25 @@ PREFILL_BATCH_SPECS = {

942 )

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

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

945 +@pytest.mark.parametrize(

946 + ("num_heads", "qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim"),

947 + [

948 + pytest.param(128, 128, 64, 128, id="deepseek_hd192_v128"),

949 + pytest.param(64, 192, 64, 256, id="glm5_hd256_v256"),

950 + ],

951 +)

952 def test_sparse_backend_prefill_correctness(

953 default_vllm_config,

954 dist_init,

955 batch_name,

956 kv_cache_dtype,

957 + num_heads,

958 + qk_nope_head_dim,

959 + qk_rope_head_dim,

960 + v_head_dim,

961 workspace_init,

962 ):

922 - """Test single-pass dense and masked MHA for sparse MLA prefill."""

963 + """Test dense and masked MHA across supported sparse MLA dimensions."""

964 backend_cls = FlashMLASparseBackend

965 batch_spec = PREFILL_BATCH_SPECS[batch_name]

966

@@ -927,11 +968,7 @@ def test_sparse_backend_prefill_correctness(

968 dtype = torch.bfloat16

969 block_size = 64

970

930 - num_heads = 128

971 kv_lora_rank = 512

932 - qk_nope_head_dim = 128

933 - qk_rope_head_dim = 64

934 - v_head_dim = 128

972 head_size = kv_lora_rank + qk_rope_head_dim

973 masked_mha = batch_name.startswith("masked_mha")

974 topk_tokens = 200 if masked_mha else 512

vllm/model_executor/layers/attention/mla_attention.py

+126 -34

@@ -827,25 +827,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

827 use_mha = True

828

829 if self.impl.is_sparse and num_mha_tokens > 0:

830 - prefill = getattr(attn_metadata, "prefill", None)

831 - use_dense_mha = getattr(prefill, "use_dense_mha", False)

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

833 - use_masked_mha = (

834 - self.prefill_backend is not None

835 - and self.impl.masked_mha_available # type: ignore[attr-defined]

836 - and self.impl.dcp_world_size <= 1

837 - and prefill is not None

838 - and _use_masked_mha(

839 - backend_name=self.attn_backend.get_name(),

840 - tensor_parallel_size=self._vllm_config.parallel_config.tensor_parallel_size,

841 - query_len=prefill.max_query_len,

842 - seq_len=prefill_max_seq_len,

843 - )

844 - and self.impl.masked_mha_workspace_fits(prefill) # type: ignore[attr-defined]

845 - )

846 - use_mha = (use_dense_mha or use_masked_mha) and not (

847 - self._vllm_config.attention_config.sparse_mla_force_mqa

848 - )

830 + use_mha = self._use_sparse_mha(attn_metadata)

831 if not use_mha:

832 num_mqa_tokens = q.size(0)

833 num_mha_tokens = 0

@@ -1063,6 +1045,31 @@ class MLAAttention(nn.Module, AttentionLayerBase):

1045 output_padded[num_actual_toks:].zero_()

1046 return output_padded

1047

1048 + def _use_sparse_mha(self, attn_metadata: "MLACommonMetadata") -> bool:

1049 + prefill = attn_metadata.prefill

1050 + if prefill is None:

1051 + return False

1052 + use_masked_mha = (

1053 + self.prefill_backend is not None

1054 + and self.impl.masked_mha_available # type: ignore[attr-defined]

1055 + and self.impl.dcp_world_size <= 1

1056 + and _use_masked_mha(

1057 + backend_name=self.attn_backend.get_name(),

1058 + tensor_parallel_size=(

1059 + self._vllm_config.parallel_config.tensor_parallel_size

1060 + ),

1061 + qk_head_dim=self.qk_nope_head_dim + self.qk_rope_head_dim,

1062 + v_head_dim=self.v_head_dim,

1063 + query_len=prefill.max_query_len,

1064 + seq_len=attn_metadata.prefill_max_seq_len, # type: ignore[attr-defined]

1065 + has_context=prefill.chunked_context is not None,

1066 + )

1067 + and self.impl.masked_mha_workspace_fits(prefill) # type: ignore[attr-defined]

1068 + )

1069 + return (prefill.use_dense_mha or use_masked_mha) and not (

1070 + self._vllm_config.attention_config.sparse_mla_force_mqa

1071 + )

1072 +

1073 def process_weights_after_loading(self, act_dtype: torch.dtype):

1074 # Let per-backend impls do their own weight packing first (no-op

1075 # unless overridden), mirroring Attention.process_weights_after_loading.

@@ -1619,36 +1626,121 @@ def get_mla_dims(model_config: ModelConfig) -> MLADims:

1626 )

1627

1628

1622 -_DSV32_MASKED_MHA_THRESHOLDS: dict[str, dict[int, tuple[int | None, ...]]] = {

1623 - "FLASHMLA_SPARSE": {

1624 - 1: (1536, 4096, None, None, None),

1625 - 2: (512, 1024, 4096, None, None),

1626 - 4: (512, 1024, 1536, 8192, None),

1627 - 8: (512, 512, 1024, 2048, 16384),

1629 +_MaskedMHARanges = tuple[tuple[int, int], ...]

1630 +

1631 +_MASKED_MHA_THRESHOLDS: dict[

1632 + tuple[int, int],

1633 + dict[str, dict[int, tuple[_MaskedMHARanges, _MaskedMHARanges | None]]],

1634 +] = {

1635 + (192, 128): {

1636 + "FLASHMLA_SPARSE": {

1637 + 1: (((2048, 1536), (4096, 4096)), None),

1638 + 2: (((2048, 512), (4096, 1024), (8192, 4096)), None),

1639 + 4: (

1640 + ((2048, 512), (4096, 1024), (8192, 1536), (16384, 8192)),

1641 + None,

1642 + ),

1643 + 8: (

1644 + (

1645 + (2048, 512),

1646 + (4096, 512),

1647 + (8192, 1024),

1648 + (16384, 2048),

1649 + (32768, 16384),

1650 + ),

1651 + None,

1652 + ),

1653 + },

1654 + "FLASHINFER_MLA_SPARSE": {

1655 + 8: (

1656 + (

1657 + (2048, 512),

1658 + (4096, 1024),

1659 + (8192, 1024),

1660 + (16384, 2048),

1661 + (32768, 32768),

1662 + ),

1663 + None,

1664 + ),

1665 + },

1666 },

1629 - "FLASHINFER_MLA_SPARSE": {

1630 - 8: (512, 1024, 1024, 2048, 32768),

1667 + (256, 256): {

1668 + "FLASHMLA_SPARSE": {

1669 + 1: (((8 * 1024, 0),), ()),

1670 + 2: (

1671 + ((20 * 1024, 0),),

1672 + ((6 * 1024, 4 * 1024), (10 * 1024, 8 * 1024)),

1673 + ),

1674 + 4: (

1675 + ((48 * 1024, 0),),

1676 + (

1677 + (8 * 1024, 4 * 1024),

1678 + (16 * 1024, 8 * 1024),

1679 + (24 * 1024, 16 * 1024),

1680 + (34 * 1024, 32 * 1024),

1681 + ),

1682 + ),

1683 + 8: (

1684 + ((112 * 1024, 0),),

1685 + (

1686 + (8 * 1024, 4 * 1024),

1687 + (16 * 1024, 8 * 1024),

1688 + (32 * 1024, 16 * 1024),

1689 + (48 * 1024, 32 * 1024),

1690 + ),

1691 + ),

1692 + },

1693 + "FLASHINFER_MLA_SPARSE": {

1694 + 4: (

1695 + ((36 * 1024, 0),),

1696 + (

1697 + (12 * 1024, 4 * 1024),

1698 + (16 * 1024, 8 * 1024),

1699 + (24 * 1024, 16 * 1024),

1700 + (28 * 1024, 24 * 1024),

1701 + ),

1702 + ),

1703 + 8: (

1704 + ((64 * 1024, 0),),

1705 + (

1706 + (20 * 1024, 4 * 1024),

1707 + (24 * 1024, 8 * 1024),

1708 + (32 * 1024, 16 * 1024),

1709 + (48 * 1024, 32 * 1024),

1710 + (56 * 1024, 40 * 1024),

1711 + (64 * 1024, 48 * 1024),

1712 + (72 * 1024, 56 * 1024),

1713 + ),

1714 + ),

1715 + },

1716 },

1717 }

1633 -_DSV32_SEQ_LEN_BUCKETS = (2048, 4096, 8192, 16384, 32768)

1718

1719

1720 def _use_masked_mha(

1721 *,

1722 backend_name: str,

1723 tensor_parallel_size: int,

1724 + qk_head_dim: int,

1725 + v_head_dim: int,

1726 query_len: int,

1727 seq_len: int,

1728 + has_context: bool,

1729 ) -> bool:

1643 - thresholds = _DSV32_MASKED_MHA_THRESHOLDS.get(backend_name, {}).get(

1644 - tensor_parallel_size

1730 + thresholds = (

1731 + _MASKED_MHA_THRESHOLDS.get((qk_head_dim, v_head_dim), {})

1732 + .get(backend_name, {})

1733 + .get(tensor_parallel_size)

1734 )

1735 if thresholds is None:

1736 return False

1648 - for bucket_idx, bucket_seq_len in enumerate(_DSV32_SEQ_LEN_BUCKETS):

1649 - if seq_len <= bucket_seq_len:

1650 - min_query_len = thresholds[bucket_idx]

1651 - return min_query_len is not None and query_len >= min_query_len

1737 + pure_prefill_ranges, context_ranges = thresholds

1738 + ranges = pure_prefill_ranges

1739 + if has_context and context_ranges is not None:

1740 + ranges = context_ranges

1741 + for max_seq_len, min_query_len in ranges:

1742 + if seq_len <= max_seq_len:

1743 + return query_len >= min_query_len

1744 return False

1745

1746

vllm/model_executor/layers/attention/sparse_mla_attention.py

+18 -7

@@ -88,16 +88,27 @@ def _is_masked_mha_available(

88 """Check if masked MHA can ever fire for this model configuration."""

89 if not current_platform.is_device_capability_family(100):

90 return False

91 - if (

92 - num_heads_total != 128

93 - or kv_lora_rank != 512

94 - or qk_nope_head_dim != 128

95 - or qk_rope_head_dim != 64

96 - or v_head_dim != 128

91 + model_dims = (

92 + num_heads_total,

93 + kv_lora_rank,

94 + qk_nope_head_dim,

95 + qk_rope_head_dim,

96 + v_head_dim,

97 + )

98 + if model_dims not in (

99 + (128, 512, 128, 64, 128),

100 + (64, 512, 192, 64, 256),

101 ):

102 return False

103 qk_head_dim = qk_nope_head_dim + qk_rope_head_dim

100 - fa_version = get_flash_attn_version(head_size=qk_head_dim, head_size_v=v_head_dim)

104 + if qk_head_dim == 256 and v_head_dim == 256:

105 + # This path uses contiguous K/V, so it does not need the paged-KV

106 + # features that keep head-dim 256 disabled in the general FA selector.

107 + fa_version = get_flash_attn_version()

108 + else:

109 + fa_version = get_flash_attn_version(

110 + head_size=qk_head_dim, head_size_v=v_head_dim

111 + )

112 return fa_version == 4 and not is_quantized_kv_cache(kv_cache_dtype)

113

114

vllm/models/deepseek_v32/attention.py

+23 -14

@@ -119,8 +119,6 @@ class DeepseekV32Attention(MLAAttention):

119 indexer: "DeepseekV32Indexer | None"

120 indexer_cls: "type[DeepseekV32Indexer]" = DeepseekV32Indexer

121

122 - supports_dense_mha_prefill = False

123 -

122 def __init__(

123 self,

124 vllm_config: VllmConfig,

@@ -303,15 +301,6 @@ class DeepseekV32Attention(MLAAttention):

301 device=hidden_states.device,

302 )

303 forward_context = get_forward_context()

306 - attn_metadata_raw = forward_context.attn_metadata

307 - if isinstance(attn_metadata_raw, dict):

308 - attn_metadata = attn_metadata_raw.get(self.layer_name)

309 - elif isinstance(attn_metadata_raw, list):

310 - attn_metadata = attn_metadata_raw[0].get(self.layer_name)

311 - else:

312 - attn_metadata = attn_metadata_raw

313 - attn_metadata = cast("MLACommonMetadata | None", attn_metadata)

314 -

304 slot_mapping = forward_context.slot_mapping

305 assert isinstance(slot_mapping, dict)

306 mla_slot = slot_mapping.get(self.layer_name)

@@ -337,7 +326,7 @@ class DeepseekV32Attention(MLAAttention):

326 indexer_softmax_scale = 0.0

327 indexer_n_head_scale = 0.0

328

340 - if attn_metadata is None or self.use_pcp:

329 + if forward_context.attn_metadata is None or self.use_pcp:

330 mla_kv_cache = None

331 mla_k_scale = None

332 indexer_k_cache = None

@@ -346,8 +335,8 @@ class DeepseekV32Attention(MLAAttention):

335 mla_kv_cache = self.kv_cache

336 mla_k_scale = self._k_scale

337

349 - kv_c_out = torch.empty_like(kv_c) if self.use_pcp else None

350 - k_pe_out = torch.empty_like(k_pe) if self.use_pcp else None

338 + kv_c_out = torch.empty_like(kv_c)

339 + k_pe_out = torch.empty_like(k_pe)

340 q_c = fused_norm_rope(

341 positions,

342 q_c,

@@ -403,7 +392,10 @@ class DeepseekV32Attention(MLAAttention):

392 )

393

394 self._sparse_indexer_and_attn(

395 + positions,

396 q_c,

397 + q_nope,

398 + q_pe,

399 index_q_fp8,

400 index_k_out,

401 index_weights_out,

@@ -418,7 +410,10 @@ class DeepseekV32Attention(MLAAttention):

410 @eager_break_during_capture

411 def _sparse_indexer_and_attn(

412 self,

413 + positions: torch.Tensor,

414 q_c: torch.Tensor,

415 + q_nope: torch.Tensor,

416 + q_pe: torch.Tensor,

417 index_q_fp8: torch.Tensor | None,

418 index_k: torch.Tensor | None,

419 index_weights_out: torch.Tensor | None,

@@ -498,6 +493,20 @@ class DeepseekV32Attention(MLAAttention):

493 output.zero_()

494 return

495

496 + if self._use_sparse_mha(attn_metadata):

497 + assert kv_c is not None and k_pe is not None

498 + mha_q_pe = self.rotary_emb(positions, q_pe)[0] if self._fp8_query else mqa_q

499 + mha_q = torch.cat((q_nope, mha_q_pe), dim=-1)

500 + self.forward_impl(

501 + mha_q,

502 + kv_c,

503 + k_pe.unsqueeze(1),

504 + kv_cache,

505 + attn_metadata,

506 + output,

507 + )

508 + return

509 +

510 if self._fp8_kv_needs_view:

511 kv_cache = kv_cache.view(torch.float8_e4m3fn)

512 if self._fp8_query: