[Attention][DSA] Take the native decode path for MTP=3 on SM90 (#52164)
Signed-off-by: zobinHuang <zobin1999@gmail.com> Co-authored-by: nodeeeeee <zhangkai.nodeee@gmail.com>
Zhuobin Huang committed
Aug 14, 2026 at 18:09 UTC
63a9a5010a6d1539c52957646ef9d6bbcf7a4deb
4 files changed
+232
-107
tests/kernels/attention/test_deepgemm_attention.py
+96
-91
@@ -13,6 +13,7 @@ from vllm.utils.deep_gemm import (
13
fp8_fp4_paged_mqa_logits,
14
get_num_sms,
15
get_paged_mqa_logits_metadata,
16
+ native_next_n_supported,
17
)
18
from vllm.utils.import_utils import has_deep_gemm
19
from vllm.utils.math_utils import cdiv
@@ -205,7 +206,12 @@ def _ref_fp8_fp4_paged_mqa_logits(
206
@pytest.mark.skipif(
207
not current_platform.has_device_capability(90), reason="SM90 and SM100 only"
208
)
208
-def test_deepgemm_fp8_fp4_paged_mqa_logits():
209
+# next_n = 1 + num_speculative_tokens, so next_n=4 is MTP=3 (issue #35878).
210
+@pytest.mark.parametrize("batch_size,next_n", [(4, 1), (2, 2), (2, 4)])
211
+def test_deepgemm_fp8_fp4_paged_mqa_logits(batch_size: int, next_n: int):
212
+ if not native_next_n_supported(next_n):
213
+ pytest.skip(f"next_n={next_n} has no native kernel on this architecture")
214
+
215
# NOTE: clean_logits=True is incompatible with the 2D context_lens
216
# required by csrc/apis/attention.hpp; only the False path is exercised.
217
clean_logits = False
@@ -213,98 +219,97 @@ def test_deepgemm_fp8_fp4_paged_mqa_logits():
219
random.seed(0)
220
221
max_model_len = 4096
216
- for batch_size, next_n in [(4, 1), (2, 2)]:
217
- for heads, index_dim in [(32, 128)]:
218
- for avg_kv in (2048,):
219
- num_blocks, blocksize = max_model_len * 2, 64
220
-
221
- q = torch.randn(
222
- (batch_size, next_n, heads, index_dim),
223
- device="cuda",
224
- dtype=torch.bfloat16,
225
- )
226
- kv_cache = torch.randn(
227
- (num_blocks, blocksize, 1, index_dim),
228
- device="cuda",
229
- dtype=torch.bfloat16,
230
- )
231
- weights = torch.randn(
232
- (batch_size * next_n, heads),
233
- device="cuda",
234
- dtype=torch.float32,
235
- )
236
-
237
- context_lens = (
238
- torch.randint(int(0.8 * avg_kv), int(1.2 * avg_kv), (batch_size,))
239
- .cuda()
240
- .to(torch.int32)
241
- )
242
- max_block_len = (
243
- (context_lens.max().item() + blocksize - 1) // blocksize * blocksize
244
- )
245
- block_tables = torch.zeros(
246
- (batch_size, max_block_len),
247
- device="cuda",
248
- dtype=torch.int32,
249
- )
222
+ for heads, index_dim in [(32, 128)]:
223
+ for avg_kv in (2048,):
224
+ num_blocks, blocksize = max_model_len * 2, 64
225
251
- counter = 0
252
- block_idx_pool = list(range(num_blocks))
253
- random.shuffle(block_idx_pool)
254
- for i in range(batch_size):
255
- ctx_len = int(context_lens[i].item())
256
- for j in range((ctx_len + blocksize - 1) // blocksize):
257
- block_tables[i][j] = block_idx_pool[counter]
258
- counter += 1
226
+ q = torch.randn(
227
+ (batch_size, next_n, heads, index_dim),
228
+ device="cuda",
229
+ dtype=torch.bfloat16,
230
+ )
231
+ kv_cache = torch.randn(
232
+ (num_blocks, blocksize, 1, index_dim),
233
+ device="cuda",
234
+ dtype=torch.bfloat16,
235
+ )
236
+ weights = torch.randn(
237
+ (batch_size * next_n, heads),
238
+ device="cuda",
239
+ dtype=torch.float32,
240
+ )
241
260
- q_fp8 = q.to(torch.float8_e4m3fn)
261
- kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache)
262
-
263
- # deep_gemm paged MQA logits requires 2D context_lens of
264
- # shape (B, next_n) (csrc/apis/attention.hpp:332-335);
265
- # see indexer.py:607-608. For each batch/next_n token, the
266
- # effective context length is context_lens[b] - next_n + j + 1.
267
- next_n_arange = torch.arange(next_n, device="cuda", dtype=torch.int32)
268
- context_lens_2d = (
269
- context_lens.unsqueeze(-1) - next_n + 1 + next_n_arange
270
- ).contiguous()
271
- schedule_metadata = get_paged_mqa_logits_metadata(
272
- context_lens_2d, blocksize, get_num_sms()
273
- )
274
- logits = fp8_fp4_paged_mqa_logits(
275
- (q_fp8, None),
276
- kv_cache_fp8,
277
- weights,
278
- context_lens_2d,
279
- block_tables,
280
- schedule_metadata,
281
- max_model_len,
282
- clean_logits=clean_logits,
283
- )
242
+ context_lens = (
243
+ torch.randint(int(0.8 * avg_kv), int(1.2 * avg_kv), (batch_size,))
244
+ .cuda()
245
+ .to(torch.int32)
246
+ )
247
+ max_block_len = (
248
+ (context_lens.max().item() + blocksize - 1) // blocksize * blocksize
249
+ )
250
+ block_tables = torch.zeros(
251
+ (batch_size, max_block_len),
252
+ device="cuda",
253
+ dtype=torch.int32,
254
+ )
255
285
- ref_logits = _ref_fp8_fp4_paged_mqa_logits(
286
- q,
287
- kv_cache,
288
- weights,
289
- context_lens,
290
- block_tables,
291
- max_model_len,
292
- )
256
+ counter = 0
257
+ block_idx_pool = list(range(num_blocks))
258
+ random.shuffle(block_idx_pool)
259
+ for i in range(batch_size):
260
+ ctx_len = int(context_lens[i].item())
261
+ for j in range((ctx_len + blocksize - 1) // blocksize):
262
+ block_tables[i][j] = block_idx_pool[counter]
263
+ counter += 1
264
+
265
+ q_fp8 = q.to(torch.float8_e4m3fn)
266
+ kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache)
267
+
268
+ # deep_gemm paged MQA logits requires 2D context_lens of
269
+ # shape (B, next_n) (csrc/apis/attention.hpp:332-335);
270
+ # see indexer.py:607-608. For each batch/next_n token, the
271
+ # effective context length is context_lens[b] - next_n + j + 1.
272
+ next_n_arange = torch.arange(next_n, device="cuda", dtype=torch.int32)
273
+ context_lens_2d = (
274
+ context_lens.unsqueeze(-1) - next_n + 1 + next_n_arange
275
+ ).contiguous()
276
+ schedule_metadata = get_paged_mqa_logits_metadata(
277
+ context_lens_2d,
278
+ blocksize,
279
+ get_num_sms(),
280
+ )
281
+ logits = fp8_fp4_paged_mqa_logits(
282
+ (q_fp8, None),
283
+ kv_cache_fp8,
284
+ weights,
285
+ context_lens_2d,
286
+ block_tables,
287
+ schedule_metadata,
288
+ max_model_len,
289
+ clean_logits=clean_logits,
290
+ )
291
294
- positions = (
295
- torch.arange(max_model_len, device="cuda")
296
- .unsqueeze(0)
297
- .expand(batch_size * next_n, -1)
298
- )
299
- row_indices = torch.arange(batch_size * next_n, device="cuda") // next_n
300
- next_n_offset = (
301
- torch.arange(batch_size * next_n, device="cuda") % next_n
302
- )
303
- mask = positions <= (
304
- context_lens[row_indices] - next_n + next_n_offset
305
- ).unsqueeze(1)
292
+ ref_logits = _ref_fp8_fp4_paged_mqa_logits(
293
+ q,
294
+ kv_cache,
295
+ weights,
296
+ context_lens,
297
+ block_tables,
298
+ max_model_len,
299
+ )
300
307
- logits = logits.masked_fill(~mask, 0)
308
- ref_logits = ref_logits.masked_fill(~mask, 0)
309
- diff = calc_diff(logits, ref_logits)
310
- assert diff < 1e-3, f"{diff=}"
301
+ positions = (
302
+ torch.arange(max_model_len, device="cuda")
303
+ .unsqueeze(0)
304
+ .expand(batch_size * next_n, -1)
305
+ )
306
+ row_indices = torch.arange(batch_size * next_n, device="cuda") // next_n
307
+ next_n_offset = torch.arange(batch_size * next_n, device="cuda") % next_n
308
+ mask = positions <= (
309
+ context_lens[row_indices] - next_n + next_n_offset
310
+ ).unsqueeze(1)
311
+
312
+ logits = logits.masked_fill(~mask, 0)
313
+ ref_logits = ref_logits.masked_fill(~mask, 0)
314
+ diff = calc_diff(logits, ref_logits)
315
+ assert diff < 1e-3, f"{diff=}"
tests/v1/attention/test_indexer_native_next_n.py
new
+75
@@ -0,0 +1,75 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Which next_n the DSA indexer decode path may hand to DeepGEMM unflattened.
4
+
5
+Getting this wrong is not a slow path but a crash: `fp8_fp4_paged_mqa_logits`
6
+asserts both that the architecture implements the requested `next_n` and that
7
+the schedule metadata was sized for the matching slot count.
8
+"""
9
+
10
+import pytest
11
+
12
+from vllm.platforms import current_platform
13
+from vllm.utils.deep_gemm import _paged_mqa_logits_schedule_slots
14
+from vllm.v1.attention.backends.mla import indexer
15
+
16
+NUM_SMS = 114 # H100 PCIe
17
+
18
+
19
+def _set_arch(monkeypatch, family: int, *, cuda: bool = True, deep_gemm: bool = True):
20
+ monkeypatch.setattr(current_platform, "is_cuda", lambda: cuda)
21
+ monkeypatch.setattr(
22
+ current_platform,
23
+ "is_device_capability_family",
24
+ lambda capability, device_id=0: capability // 10 == family,
25
+ )
26
+ monkeypatch.setattr(indexer, "has_deep_gemm", lambda: deep_gemm)
27
+
28
+
29
+@pytest.mark.parametrize(
30
+ "family,expected_native",
31
+ [
32
+ # SM90 gained next_n=4 (MTP=3) via 2-CTA multicast, but never 3.
33
+ (9, {1, 2, 4}),
34
+ # SM100 schedules any next_n with multi-atom tiles.
35
+ (10, {1, 2, 3, 4, 5, 8}),
36
+ # SM120 advertises multi-atom too but is unvalidated on hardware, so
37
+ # it stays on the conservative gate. Loosen it only with measurements.
38
+ (12, {1, 2}),
39
+ ],
40
+)
41
+def test_native_decode_gate_per_architecture(monkeypatch, family, expected_native):
42
+ _set_arch(monkeypatch, family)
43
+ for next_n in (1, 2, 3, 4, 5, 8):
44
+ assert indexer._supports_native_decode(next_n) == (next_n in expected_native), (
45
+ f"family={family} next_n={next_n}"
46
+ )
47
+
48
+
49
+@pytest.mark.parametrize(
50
+ "cuda,deep_gemm", [(False, True), (True, False), (False, False)]
51
+)
52
+def test_native_decode_gate_without_deepgemm(monkeypatch, cuda, deep_gemm):
53
+ """Without the DeepGEMM kernels only the shapes every backend handles."""
54
+ _set_arch(monkeypatch, 9, cuda=cuda, deep_gemm=deep_gemm)
55
+ assert [indexer._supports_native_decode(n) for n in (1, 2, 3, 4)] == [
56
+ True,
57
+ True,
58
+ False,
59
+ False,
60
+ ]
61
+
62
+
63
+def test_sm90_next_n_4_halves_the_schedule_slots(monkeypatch):
64
+ """SM90 next_n=4 runs one scheduler task per 2-CTA cluster, not per SM."""
65
+ _set_arch(monkeypatch, 9)
66
+ assert _paged_mqa_logits_schedule_slots(NUM_SMS, 4) == NUM_SMS // 2
67
+ for next_n in (1, 2, 3):
68
+ assert _paged_mqa_logits_schedule_slots(NUM_SMS, next_n) == NUM_SMS
69
+
70
+
71
+@pytest.mark.parametrize("family", [10, 12])
72
+def test_multicast_is_sm90_only(monkeypatch, family):
73
+ _set_arch(monkeypatch, family)
74
+ for next_n in (1, 2, 3, 4):
75
+ assert _paged_mqa_logits_schedule_slots(NUM_SMS, next_n) == NUM_SMS
vllm/utils/deep_gemm.py
+31
-5
@@ -552,6 +552,29 @@ def fp8_fp4_mqa_logits(
552
)
553
554
555
+def native_next_n_supported(next_n: int) -> bool:
556
+ """Whether the paged MQA logits kernel takes `next_n` Q rows per request.
557
+
558
+ SM90 implements only {1, 2, 4}; SM100 and SM120 schedule any `next_n` via
559
+ multi-atom tiles. Unsupported values must be flattened to one row per query.
560
+ """
561
+ if current_platform.is_device_capability_family(90):
562
+ return next_n in (1, 2, 4)
563
+ return True
564
+
565
+
566
+def _paged_mqa_logits_schedule_slots(num_sms: int, next_n: int) -> int:
567
+ """Scheduler tasks the paged MQA logits kernel launches.
568
+
569
+ SM90 `next_n=4` runs one task per 2-CTA multicast cluster rather than per
570
+ SM, and `fp8_fp4_paged_mqa_logits` asserts its metadata is sized to match.
571
+ """
572
+ num_kv_multicast = (
573
+ 2 if next_n == 4 and current_platform.is_device_capability_family(90) else 1
574
+ )
575
+ return num_sms // num_kv_multicast
576
+
577
+
578
def get_paged_mqa_logits_metadata(
579
context_lens: torch.Tensor,
580
block_size: int,
@@ -561,22 +584,24 @@ def get_paged_mqa_logits_metadata(
584
"""Build scheduling metadata for paged MQA logits.
585
586
Args:
564
- context_lens: Tensor of shape [B], dtype int32; effective context length
565
- per batch element.
587
+ context_lens: Tensor of shape [B, next_n], dtype int32; effective
588
+ context length per Q row.
589
block_size: KV-cache block size in tokens (e.g., 64).
590
num_sms: Number of SMs available. 132 for Hopper
591
indices: Optional request index for each varlen row.
592
593
Returns:
571
- Backend-specific tensor consumed by `fp8_fp4_paged_mqa_logits` to
572
- schedule work across SMs.
594
+ Tensor of shape [slots + 1, 2] consumed by `fp8_fp4_paged_mqa_logits`
595
+ to schedule work across SMs.
596
"""
597
_lazy_init()
598
if _get_paged_mqa_logits_metadata_impl is None:
599
return _missing()
600
+ next_n = context_lens.shape[1] if context_lens.dim() == 2 else 1
601
+ num_slots = _paged_mqa_logits_schedule_slots(num_sms, next_n)
602
kwargs = {} if indices is None else {"indices": indices}
603
return _get_paged_mqa_logits_metadata_impl(
579
- context_lens, block_size, num_sms, **kwargs
604
+ context_lens, block_size, num_slots, **kwargs
605
)
606
607
@@ -752,6 +777,7 @@ __all__ = [
777
"fp8_fp4_mqa_logits",
778
"fp8_fp4_paged_mqa_logits",
779
"get_paged_mqa_logits_metadata",
780
+ "native_next_n_supported",
781
"per_block_cast_to_fp8",
782
"is_deep_gemm_e8m0_used",
783
"is_deep_gemm_supported",
vllm/v1/attention/backends/mla/indexer.py
+30
-11
@@ -21,6 +21,7 @@ from vllm.triton_utils import tl, triton
21
from vllm.utils.deep_gemm import (
22
get_paged_mqa_logits_metadata,
23
has_deep_gemm,
24
+ native_next_n_supported,
25
)
26
from vllm.utils.platform_utils import num_compute_units
27
from vllm.v1.attention.backend import (
@@ -467,6 +468,20 @@ def _supports_varlen_paged_mqa_logits() -> bool:
468
)
469
470
471
+def _supports_native_decode(next_n: int) -> bool:
472
+ """Whether decode can pass `next_n` Q rows per request to the kernel
473
+ instead of flattening to one single-token row per query, which re-reads
474
+ the KV tile once per row.
475
+ """
476
+ if not (current_platform.is_cuda() and has_deep_gemm()):
477
+ return next_n in (1, 2)
478
+ if current_platform.is_device_capability_family(100):
479
+ return True
480
+ if current_platform.is_device_capability_family(90):
481
+ return native_next_n_supported(next_n)
482
+ return next_n in (1, 2)
483
+
484
+
485
class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
486
# The indexer opts out of the shared reorder-threshold vote (see __init__),
487
# so this is None; its own split uses self.decode_threshold.
@@ -524,14 +539,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
539
next_n = self.num_speculative_tokens + 1
540
self.decode_threshold = next_n
541
self.reorder_batch_threshold = None
527
- # NOTE: SM100 datacenter GPUs support any next_n natively via the
528
- # multi-atom paged MQA logits kernels (FP8 and FP4 indexer
529
- # caches). Outside the SM100 family the FP8
530
- # paged MQA logits kernel only supports next_n in (1, 2)
531
- # (deepgemm smxx_fp8_fp4_paged_mqa_logits.hpp:233), so flatten there.
532
- self.use_flattening = not current_platform.is_device_capability_family(
533
- 100
534
- ) and next_n not in (1, 2)
542
+ self.use_flattening = not _supports_native_decode(next_n)
543
self.supports_varlen = _supports_varlen_paged_mqa_logits()
544
logger.info_once(
545
"DSA indexer decode path: use_flattening=%s supports_varlen=%s "
@@ -585,7 +593,8 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
593
device=self.device,
594
)
595
588
- # See: DeepGMM/csrc/apis/attention.hpp
596
+ # See: DeepGMM/csrc/apis/attention.hpp. Sized for one slot per SM;
597
+ # build() narrows it to whatever the kernel actually schedules.
598
self.scheduler_metadata_buffer = torch.empty(
599
(self.num_sms + 1, 2), dtype=torch.int32, device=self.device
600
)
@@ -924,9 +933,16 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
933
934
max_decode_len = int(decode_lens_cpu.max().item())
935
next_n = 1 + self.num_speculative_tokens
936
+ # The kernel sees max_decode_len Q rows, not the configured next_n,
937
+ # so legality is per-step: on SM90 a uniformly 3-deep batch has no
938
+ # native kernel. max_decode_len <= 1 always has one.
939
+ step_next_n_ok = max_decode_len <= 1 or _supports_native_decode(
940
+ max_decode_len
941
+ )
942
use_native = (
943
not (self.use_flattening or self.supports_varlen)
944
and max_decode_len <= next_n
945
+ and step_next_n_ok
946
)
947
948
global_seq_lens_for_decode = self._prepare_global_decode_seq_lens(
@@ -992,20 +1008,23 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1008
seq_lens = seq_lens.unsqueeze(-1)
1009
1010
# DeepGEMM is required for the paged MQA logits on CUDA devices
1011
+ schedule_metadata = self.scheduler_metadata_buffer
1012
if current_platform.is_cuda() and has_deep_gemm():
996
- self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata(
1013
+ metadata = get_paged_mqa_logits_metadata(
1014
seq_lens,
1015
self.kv_cache_spec.storage_block_size,
1016
self.num_sms,
1017
indices=decode_indices,
1018
)
1019
+ schedule_metadata = self.scheduler_metadata_buffer[: metadata.shape[0]]
1020
+ schedule_metadata[:] = metadata
1021
1022
decode_metadata = DeepSeekV32IndexerDecodeMetadata(
1023
block_table=block_table,
1024
seq_lens=seq_lens,
1025
decode_lens=decode_lens,
1026
requires_padding=requires_padding,
1008
- schedule_metadata=self.scheduler_metadata_buffer,
1027
+ schedule_metadata=schedule_metadata,
1028
indices=decode_indices,
1029
global_seq_lens=global_seq_lens_for_decode,
1030
)