[Spec Decode] Enable adaptive verification on DSv4 + sm90 (#52795)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Lucas Wilkinson committed
Aug 20, 2026 at 22:44 UTC
5df31ea52d7d175dd920fce61a19017941b18b4b
1 file changed
+28
-4
vllm/v1/attention/backends/mla/indexer.py
+28
-4
@@ -157,8 +157,11 @@ class DeepseekV32IndexerBackend(AttentionBackend):
157
@classmethod
158
def supports_device_cpu_query_lens_mismatch(cls) -> bool:
159
# Only the varlen paged MQA logits kernel takes per-request query
160
- # lengths from device tensors; otherwise the indexer needs uniform ones.
161
- return _supports_varlen_paged_mqa_logits()
160
+ # lengths from device tensors natively. Hopper can instead flatten each
161
+ # query into a single-token row using device-built metadata.
162
+ return _supports_varlen_paged_mqa_logits() or (
163
+ _supports_flattened_device_query_lens()
164
+ )
165
166
@staticmethod
167
def get_name() -> str:
@@ -493,6 +496,14 @@ def _supports_varlen_paged_mqa_logits() -> bool:
496
)
497
498
499
+def _supports_flattened_device_query_lens() -> bool:
500
+ return (
501
+ current_platform.is_cuda()
502
+ and current_platform.is_device_capability_family(90)
503
+ and has_deep_gemm()
504
+ )
505
+
506
+
507
def _supports_native_decode(next_n: int) -> bool:
508
"""Whether decode can pass `next_n` Q rows per request to the kernel
509
instead of flattening to one single-token row per query, which re-reads
@@ -507,6 +518,16 @@ def _supports_native_decode(next_n: int) -> bool:
518
return next_n in (1, 2)
519
520
521
+def _use_flattening(vllm_config: VllmConfig) -> bool:
522
+ speculative_config = vllm_config.speculative_config
523
+ next_n = 1 + vllm_config.num_speculative_tokens
524
+ return not _supports_native_decode(next_n) or (
525
+ speculative_config is not None
526
+ and speculative_config.enable_adaptive_verification
527
+ and _supports_flattened_device_query_lens()
528
+ )
529
+
530
+
531
class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
532
# The indexer opts out of the shared reorder-threshold vote (see __init__),
533
# so this is None; its own split uses self.decode_threshold.
@@ -519,7 +540,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
540
vllm_config: VllmConfig,
541
kv_cache_spec: KVCacheSpec,
542
) -> AttentionCGSupport:
522
- if _supports_varlen_paged_mqa_logits():
543
+ if _supports_varlen_paged_mqa_logits() or _use_flattening(vllm_config):
544
return AttentionCGSupport.ALWAYS
545
return AttentionCGSupport.UNIFORM_BATCH
546
@@ -553,7 +574,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
574
next_n = self.num_speculative_tokens + 1
575
self.decode_threshold = next_n
576
self.reorder_batch_threshold = None
556
- self.use_flattening = not _supports_native_decode(next_n)
577
+ self.use_flattening = _use_flattening(self.vllm_config)
578
self.supports_varlen = _supports_varlen_paged_mqa_logits()
579
logger.info_once(
580
"DSA indexer decode path: use_flattening=%s supports_varlen=%s "
@@ -674,11 +695,14 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
695
max_decode_len: int,
696
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, bool]:
697
"""Prepare native or per-token flattened decode tensors."""
698
+ spec_config = self.vllm_config.speculative_config
699
+ adaptive = bool(spec_config and spec_config.enable_adaptive_verification)
700
min_decode_len = int(decode_lens_cpu.min().item())
701
if not use_native:
702
assert self.decode_seq_lens_buffer.dim() == 1
703
if (
704
not self.supports_varlen
705
+ and (num_decodes == 1 or not adaptive)
706
and min_decode_len == max_decode_len
707
and num_decodes * max_decode_len == num_decode_tokens
708
):