[Spec Decode] Enable adaptive DSpark on SM100 sparse MLA (#52783)
Signed-off-by: mgoin <mgoin64@gmail.com>
Misha Goin committed
Aug 25, 2026 at 13:31 UTC
7de96050c530009eb8e81065aea76570fc6ed15b
7 files changed
+293
-18
tests/v1/attention/test_dspark_noncausal_sparse_mla.py
+60
-2
@@ -55,6 +55,7 @@ if not current_platform.is_cuda():
55
)
56
57
from vllm.utils.math_utils import cdiv
58
+from vllm.v1.attention.backend import AttentionCGSupport
59
from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
60
FlashInferMLASparseTRTLLMBackend,
61
)
@@ -118,6 +119,9 @@ def _run_sparse_backend_vs_sdpa(
119
num_heads: int,
120
device: torch.device,
121
force_future_dominance: bool = False,
122
+ qk_nope_head_dim: int = 128,
123
+ v_head_dim: int = 128,
124
+ stale_cpu_query_lens: list[int] | None = None,
125
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
126
"""Run a sparse-MLA backend with the given per-token indices and compute a
127
dense per-token SDPA reference over the SAME indices.
@@ -145,9 +149,7 @@ def _run_sparse_backend_vs_sdpa(
149
use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla"
150
151
kv_lora_rank = 512
148
- qk_nope_head_dim = 128
152
qk_rope_head_dim = 64
150
- v_head_dim = 128
153
head_size = kv_lora_rank + qk_rope_head_dim
154
155
max_seqlen = max(seq_lens)
@@ -290,6 +292,16 @@ def _run_sparse_backend_vs_sdpa(
292
common_attn_metadata = create_common_attn_metadata(
293
batch_spec, block_size, device, arange_block_indices=True
294
)
295
+ if stale_cpu_query_lens is not None:
296
+ # Adaptive verification updates only device boundaries; the stale CPU
297
+ # copy proves token_to_req_indices consumes the device layout.
298
+ assert len(stale_cpu_query_lens) == len(query_lens)
299
+ cpu_query_start_loc = [0]
300
+ for query_len in stale_cpu_query_lens:
301
+ cpu_query_start_loc.append(cpu_query_start_loc[-1] + query_len)
302
+ common_attn_metadata.query_start_loc_cpu = torch.tensor(
303
+ cpu_query_start_loc, dtype=torch.int32
304
+ )
305
kv_cache = create_and_prepopulate_kv_cache(
306
kv_c_contexts=kv_c_contexts,
307
k_pe_contexts=k_pe_contexts,
@@ -409,6 +421,52 @@ def _skip_if_backend_unavailable(backend_cls, kv_cache_dtype: str, block_size: i
421
pytest.skip("FlashInferMLASparseTRTLLMBackend requires SM 10.x capability")
422
423
424
+def test_flashinfer_sparse_mla_adaptive_varlen_matches_sdpa(
425
+ default_vllm_config,
426
+ dist_init,
427
+ workspace_init,
428
+):
429
+ """Adaptive request boundaries must drive SM100 sparse index conversion."""
430
+ backend_cls = FlashInferMLASparseTRTLLMBackend
431
+ _skip_if_backend_unavailable(backend_cls, "fp8", 64)
432
+ assert (
433
+ backend_cls.get_builder_cls().get_cudagraph_support(None, None)
434
+ == AttentionCGSupport.ALWAYS
435
+ )
436
+
437
+ device = torch.device(DEVICE_TYPE)
438
+ seq_lens = [257, 270, 265, 276]
439
+ query_lens = [1, 7, 3, 5]
440
+ sparse_indices = _build_dspark_noncausal_indices(
441
+ seq_lens,
442
+ query_lens,
443
+ window=128,
444
+ topk_width=256,
445
+ device=device,
446
+ )
447
+
448
+ backend_output, sdpa_reference, _ = _run_sparse_backend_vs_sdpa(
449
+ backend_cls,
450
+ seq_lens,
451
+ query_lens,
452
+ sparse_indices,
453
+ "fp8",
454
+ 64,
455
+ 16,
456
+ device,
457
+ qk_nope_head_dim=192,
458
+ v_head_dim=256,
459
+ stale_cpu_query_lens=[4, 4, 4, 4],
460
+ )
461
+
462
+ torch.testing.assert_close(
463
+ backend_output,
464
+ sdpa_reference,
465
+ rtol=0.065,
466
+ atol=0.05,
467
+ )
468
+
469
+
470
@pytest.mark.parametrize(
471
"backend_cls",
472
[FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend],
tests/v1/spec_decode/test_adaptive_verification.py
+58
@@ -5,9 +5,13 @@ from types import SimpleNamespace
5
6
import numpy as np
7
8
+from vllm.v1.attention.backend import AttentionCGSupport
9
from vllm.v1.worker.gpu.async_utils import StepTimingSample
10
+from vllm.v1.worker.gpu.attn_utils import AttentionCGSupportInfo
11
+from vllm.v1.worker.gpu.spec_decode import adaptive_verification as adaptive_module
12
from vllm.v1.worker.gpu.spec_decode.adaptive_verification import (
13
AdaptiveVerificationManager,
14
+ maybe_create_adaptive_verification_manager,
15
)
16
from vllm.v1.worker.gpu.structured_outputs import _build_grammar_mapping
17
@@ -31,6 +35,60 @@ def make_manager(
35
return manager
36
37
38
+def test_manager_scopes_varlen_check_without_weakening_runner_cg_mode(monkeypatch):
39
+ class Backend:
40
+ @classmethod
41
+ def supports_device_cpu_query_lens_mismatch(cls):
42
+ return True
43
+
44
+ class Builder:
45
+ def __init__(self, support):
46
+ self.support = support
47
+
48
+ def get_cudagraph_support(self, *_args):
49
+ return self.support
50
+
51
+ def group(layer_name, support):
52
+ builder = Builder(support)
53
+ return SimpleNamespace(
54
+ layer_names=[layer_name],
55
+ backend=Backend,
56
+ kv_cache_spec=None,
57
+ get_metadata_builder=lambda _index: builder,
58
+ )
59
+
60
+ groups = [
61
+ [
62
+ group("target", AttentionCGSupport.ALWAYS),
63
+ group("draft", AttentionCGSupport.UNIFORM_BATCH),
64
+ ]
65
+ ]
66
+ runner_support = AttentionCGSupportInfo(
67
+ AttentionCGSupport.UNIFORM_BATCH, "DraftBackend"
68
+ )
69
+ created = object()
70
+ monkeypatch.setattr(
71
+ adaptive_module,
72
+ "AdaptiveVerificationManager",
73
+ lambda *_args, **_kwargs: created,
74
+ )
75
+
76
+ manager = maybe_create_adaptive_verification_manager(
77
+ enable_adaptive_verification=True,
78
+ attn_groups=groups,
79
+ attn_cg_support=runner_support,
80
+ req_states=object(),
81
+ query_start_loc=object(),
82
+ num_bonus_tokens=1,
83
+ max_total_logits=1,
84
+ vllm_config=None,
85
+ target_layer_names={"target"},
86
+ )
87
+
88
+ assert manager is created
89
+ assert runner_support.min_cg_support == AttentionCGSupport.UNIFORM_BATCH
90
+
91
+
92
def test_budget_stops_where_marginal_drafts_stop_paying_for_themselves():
93
# Verification is cheap up to two extra tokens, then jumps 100x; only the
94
# highest-confidence draft is worth the cheap slot.
tests/v1/worker/test_attn_utils.py
+91
@@ -11,6 +11,7 @@ import pytest
11
import torch
12
13
from tests.v1.attention.utils import dense_kv_cache_views
14
+from vllm.v1.attention.backend import AttentionCGSupport
15
from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
16
from vllm.v1.kv_cache_interface import (
17
FullAttentionSpec,
@@ -21,12 +22,102 @@ from vllm.v1.kv_cache_interface import (
22
MLAAttentionSpec,
23
compute_layout_strides,
24
)
25
+from vllm.v1.worker.gpu.attn_utils import (
26
+ get_attn_cg_support,
27
+ get_query_lens_mismatch_unsupported_backend,
28
+)
29
from vllm.v1.worker.utils import (
30
+ AttentionGroup,
31
allocate_kv_cache,
32
copy_kv_cache_blocks_inplace,
33
)
34
35
36
+class _FakeMetadataBuilder:
37
+ def __init__(self, support: AttentionCGSupport):
38
+ self.support = support
39
+
40
+ def get_cudagraph_support(self, *_args):
41
+ return self.support
42
+
43
+
44
+class _TargetBackend:
45
+ @classmethod
46
+ def supports_device_cpu_query_lens_mismatch(cls) -> bool:
47
+ return True
48
+
49
+
50
+class _DraftBackend:
51
+ @classmethod
52
+ def supports_device_cpu_query_lens_mismatch(cls) -> bool:
53
+ return False
54
+
55
+
56
+def test_attention_checks_preserve_global_and_target_scoped_support():
57
+ spec = FullAttentionSpec(
58
+ block_size=16,
59
+ num_kv_heads=1,
60
+ head_size=128,
61
+ dtype=torch.bfloat16,
62
+ )
63
+ target_group = AttentionGroup(
64
+ _TargetBackend,
65
+ ["target"],
66
+ spec,
67
+ 0, # type: ignore[arg-type]
68
+ )
69
+ target_group.metadata_builders = [
70
+ _FakeMetadataBuilder(AttentionCGSupport.ALWAYS) # type: ignore[list-item]
71
+ ]
72
+ draft_group = AttentionGroup(
73
+ _DraftBackend,
74
+ ["draft"],
75
+ spec,
76
+ 0, # type: ignore[arg-type]
77
+ )
78
+ draft_group.metadata_builders = [
79
+ _FakeMetadataBuilder(AttentionCGSupport.UNIFORM_BATCH) # type: ignore[list-item]
80
+ ]
81
+ groups = [[target_group, draft_group]]
82
+
83
+ # The runner-wide execution mode must still honor the drafter's limit.
84
+ unfiltered = get_attn_cg_support(groups, None) # type: ignore[arg-type]
85
+ assert unfiltered.min_cg_support == AttentionCGSupport.UNIFORM_BATCH
86
+ assert unfiltered.min_cg_attn_backend == "_DraftBackend"
87
+
88
+ # Adaptive verification validates only the target's varlen graphs.
89
+ target_only = get_attn_cg_support(
90
+ groups,
91
+ None, # type: ignore[arg-type]
92
+ checked_layer_names={"target"},
93
+ )
94
+ assert target_only.min_cg_support == AttentionCGSupport.ALWAYS
95
+ assert target_only.min_cg_attn_backend is None
96
+ assert (
97
+ get_query_lens_mismatch_unsupported_backend(
98
+ groups,
99
+ checked_layer_names={"target"},
100
+ )
101
+ is None
102
+ )
103
+
104
+ # Shared target/draft groups still participate in target-scoped checks.
105
+ draft_group.layer_names.append("target")
106
+ target_with_shared_group = get_attn_cg_support(
107
+ groups,
108
+ None, # type: ignore[arg-type]
109
+ checked_layer_names={"target"},
110
+ )
111
+ assert target_with_shared_group.min_cg_support == AttentionCGSupport.UNIFORM_BATCH
112
+ assert (
113
+ get_query_lens_mismatch_unsupported_backend(
114
+ groups,
115
+ checked_layer_names={"target"},
116
+ )
117
+ == "_DraftBackend"
118
+ )
119
+
120
+
121
def test_reshape_padded_kv_cache_strides_by_padded_page():
122
num_blocks = 3
123
spec = FullAttentionSpec(
vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
+17
@@ -35,6 +35,7 @@ from vllm.v1.kv_cache_interface import AttentionSpec
35
36
if TYPE_CHECKING:
37
from vllm.model_executor.models.deepseek_v2 import Indexer
38
+ from vllm.v1.attention.backend import CommonAttentionMetadata
39
40
logger = init_logger(__name__)
41
@@ -83,6 +84,10 @@ class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase):
84
def get_impl_cls() -> type[MLAAttentionImpl]:
85
return FlashInferMLASparseImpl
86
87
+ @staticmethod
88
+ def get_builder_cls() -> type["FlashInferMLASparseTRTLLMMetadataBuilder"]:
89
+ return FlashInferMLASparseTRTLLMMetadataBuilder
90
+
91
@classmethod
92
def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
93
return capability.major == 10
@@ -261,6 +266,18 @@ class FlashInferMLASparseMetadataBuilder(
266
)
267
268
269
+class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilder):
270
+ """Metadata builder for the SM100 TRT-LLM sparse MLA kernel."""
271
+
272
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
273
+
274
+ def _build_req_id_per_token(
275
+ self,
276
+ common_attn_metadata: "CommonAttentionMetadata",
277
+ ) -> torch.Tensor:
278
+ return common_attn_metadata.token_to_req_indices(self.req_id_per_token_buffer)
279
+
280
+
281
# Global workspace buffer (lazily initialized)
282
_fi_sparse_workspace: torch.Tensor | None = None
283
vllm/v1/worker/gpu/attn_utils.py
+27
-7
@@ -130,8 +130,6 @@ def init_attn_backend(
130
131
# Phase 3: create metadata builders and determine cudagraph support.
132
attn_backend_workspace: torch.Tensor | None = None
133
- min_cg_support = AttentionCGSupport.ALWAYS
134
- min_cg_attn_backend = None
133
for kv_cache_group_id, groups in enumerate(attn_groups):
134
kernel_block_size = None
135
if kv_cache_group_id < len(kernel_block_sizes):
@@ -150,7 +148,25 @@ def init_attn_backend(
148
else:
149
if hasattr(builder, "set_workspace_buffer"):
150
builder.set_workspace_buffer(attn_backend_workspace)
153
- # Check cudagraph support for the attention backend
151
+ attn_cg_support_info = get_attn_cg_support(attn_groups, vllm_config)
152
+ return attn_groups, attn_cg_support_info, kernel_block_sizes
153
+
154
+
155
+def get_attn_cg_support(
156
+ attn_groups: list[list[AttentionGroup]],
157
+ vllm_config: VllmConfig,
158
+ checked_layer_names: set[str] | None = None,
159
+) -> AttentionCGSupportInfo:
160
+ """Return the weakest CUDA graph support among the checked layers."""
161
+ min_cg_support = AttentionCGSupport.ALWAYS
162
+ min_cg_attn_backend = None
163
+ for groups in attn_groups:
164
+ for group in groups:
165
+ if checked_layer_names is not None and checked_layer_names.isdisjoint(
166
+ group.layer_names
167
+ ):
168
+ continue
169
+ builder = group.get_metadata_builder(0)
170
cg_support = builder.get_cudagraph_support(
171
vllm_config,
172
group.kv_cache_spec,
@@ -158,15 +174,15 @@ def init_attn_backend(
174
if cg_support.value < min_cg_support.value:
175
min_cg_support = cg_support
176
min_cg_attn_backend = group.backend.__name__
161
-
162
- attn_cg_support_info = AttentionCGSupportInfo(
163
- min_cg_support=min_cg_support, min_cg_attn_backend=min_cg_attn_backend
177
+ return AttentionCGSupportInfo(
178
+ min_cg_support=min_cg_support,
179
+ min_cg_attn_backend=min_cg_attn_backend,
180
)
165
- return attn_groups, attn_cg_support_info, kernel_block_sizes
181
182
183
def get_query_lens_mismatch_unsupported_backend(
184
attn_groups: list[list[AttentionGroup]],
185
+ checked_layer_names: set[str] | None = None,
186
) -> str | None:
187
"""Name the first backend needing the CPU query lengths to be exact, if any.
188
@@ -176,6 +192,10 @@ def get_query_lens_mismatch_unsupported_backend(
192
"""
193
for groups in attn_groups:
194
for group in groups:
195
+ if checked_layer_names is not None and checked_layer_names.isdisjoint(
196
+ group.layer_names
197
+ ):
198
+ continue
199
if not group.backend.supports_device_cpu_query_lens_mismatch():
200
return group.backend.__name__
201
return None
vllm/v1/worker/gpu/model_runner.py
+16
-4
@@ -559,12 +559,21 @@ class GPUModelRunner(LoRAModelRunnerMixin):
559
max_num_blocks = get_block_table_width(max_num_blocks, spec.block_size)
560
max_num_blocks_per_group.append(max_num_blocks)
561
562
+ target_attn_layer_names = None
563
+ if isinstance(self.speculator, DraftModelSpeculator):
564
+ # Adaptive verification validates target attention separately.
565
+ target_attn_layer_names = {
566
+ layer_name
567
+ for group in self.kv_cache_config.kv_cache_groups
568
+ for layer_name in group.layer_names
569
+ } - self.speculator.draft_attn_layer_names
570
self.attn_groups, attn_cg_support, self.kernel_block_sizes = init_attn_backend(
563
- self.kv_cache_config, self.vllm_config, self.device
564
- )
565
- attn_cg_support = attn_cg_support.narrow(
566
- *self.model_state.get_additional_cg_support()
571
+ self.kv_cache_config,
572
+ self.vllm_config,
573
+ self.device,
574
)
575
+ additional_attn_cg_support = self.model_state.get_additional_cg_support()
576
+ attn_cg_support = attn_cg_support.narrow(*additional_attn_cg_support)
577
# The speculator clears the flag at load time when the checkpoint has
578
# no confidence head, so it holds the effective value.
579
self.adaptive_verification = maybe_create_adaptive_verification_manager(
@@ -577,6 +586,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
586
query_start_loc=self.input_buffers.query_start_loc,
587
num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step,
588
max_total_logits=get_max_chunk_logits(self.vocab_size),
589
+ vllm_config=self.vllm_config,
590
+ target_layer_names=target_attn_layer_names,
591
+ additional_attn_cg_support=additional_attn_cg_support,
592
)
593
594
self.block_tables = BlockTables(
vllm/v1/worker/gpu/spec_decode/adaptive_verification.py
+24
-5
@@ -17,6 +17,7 @@ from vllm.v1.attention.backend import AttentionCGSupport
17
from vllm.v1.utils import CpuGpuBuffer
18
from vllm.v1.worker.gpu.async_utils import StepTimingSample, stream
19
from vllm.v1.worker.gpu.attn_utils import (
20
+ get_attn_cg_support,
21
get_query_lens_mismatch_unsupported_backend,
22
)
23
from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
@@ -25,6 +26,7 @@ logger = init_logger(__name__)
26
_PROFILE_REPLAYS = 5
27
28
if TYPE_CHECKING:
29
+ from vllm.config import VllmConfig
30
from vllm.v1.worker.gpu.attn_utils import AttentionCGSupportInfo
31
from vllm.v1.worker.gpu.input_batch import InputBatch
32
from vllm.v1.worker.gpu.states import RequestState
@@ -444,13 +446,19 @@ def maybe_create_adaptive_verification_manager(
446
query_start_loc: torch.Tensor,
447
num_bonus_tokens: int,
448
max_total_logits: int,
449
+ vllm_config: "VllmConfig",
450
+ target_layer_names: set[str] | None = None,
451
+ additional_attn_cg_support: tuple[AttentionCGSupport, str | None] | None = None,
452
) -> AdaptiveVerificationManager | None:
453
if not enable_adaptive_verification:
454
return None
455
456
# The selector rejects unsupported backends, but models that
457
# hard-wire theirs (e.g. DeepSeek-V4) never go through it.
453
- backend = get_query_lens_mismatch_unsupported_backend(attn_groups)
458
+ backend = get_query_lens_mismatch_unsupported_backend(
459
+ attn_groups,
460
+ checked_layer_names=target_layer_names,
461
+ )
462
if backend is not None:
463
raise ValueError(
464
"Adaptive verification trims verification requests on device, which"
@@ -459,12 +467,23 @@ def maybe_create_adaptive_verification_manager(
467
"use a backend that does."
468
)
469
462
- if attn_cg_support.min_cg_support != AttentionCGSupport.ALWAYS:
470
+ target_attn_cg_support = attn_cg_support
471
+ if target_layer_names is not None:
472
+ target_attn_cg_support = get_attn_cg_support(
473
+ attn_groups,
474
+ vllm_config,
475
+ checked_layer_names=target_layer_names,
476
+ )
477
+ if additional_attn_cg_support is not None:
478
+ target_attn_cg_support = target_attn_cg_support.narrow(
479
+ *additional_attn_cg_support
480
+ )
481
+ if target_attn_cg_support.min_cg_support != AttentionCGSupport.ALWAYS:
482
raise ValueError(
483
"Adaptive verification captures varlen decode cudagraphs, so every"
465
- " attention builder must report AttentionCGSupport.ALWAYS, but "
466
- f"{attn_cg_support.min_cg_attn_backend} reports "
467
- f"{attn_cg_support.min_cg_support}. Pass "
484
+ " target attention builder must report AttentionCGSupport.ALWAYS, but "
485
+ f"{target_attn_cg_support.min_cg_attn_backend} reports "
486
+ f"{target_attn_cg_support.min_cg_support}. Pass "
487
"enable_adaptive_verification=false in the speculative config, or "
488
"use a backend that does."
489
)