[Spec Decode] Add graph-aware adaptive K for DFlash

docs/features/speculative_decoding/adaptive_verification.md

@@ -4,15 +4,25 @@ Speculative decoding buys fewer decode steps with more compute. At batch size 1
 
 That matters because per-position acceptance decays fast. While the GPU is memory-bound that slot is effectively free and worth the gamble; once it saturates the gamble has a real throughput cost. The crossover moves with load and with workload-dependent acceptance rates, so no static `num_speculative_tokens` is right across concurrencies.
 
-Adaptive verification decides per step how much of the draft to verify instead. Every (request, position) draft slot is scored by its *survival probability*, the running product of that request's per-position confidences, and the highest-scoring slots are admitted until a global budget is spent. Slots compete across requests: position 5 of a confident request can outrank position 1 of a doubtful one, so one request keeps its full block while another could be trimmed after a token or two.
+Adaptive verification decides per step how much of the draft to verify instead.
+DSpark scores every (request, position) draft slot by its *survival
+probability*, the running product of that request's per-position confidences,
+and admits the highest-scoring slots until a global budget is spent. DFlash
+uses its observed accepted-prefix survival to choose one K for the next batch.
 
 The budget itself comes from a cost model profiled at startup. vLLM measures what a step costs at each shape, then picks the token count that maximizes expected accepted tokens per second.
 
 The practical effect is that one configuration holds up across the whole load range, which removes most of the need to tune `num_speculative_tokens` per deployment.
 
 ## Support
 
-Adaptive verification needs per-position acceptance estimates, so today it is only supported for DSpark with a **confidence head**.
+Adaptive verification supports:
+
+- DSpark checkpoints with a **confidence head**. It trims the current
+  verification batch per request.
+- DFlash checkpoints. It chooses a uniform K from batch size, accepted-prefix
+  history, and the profiled draft/verify costs. K=0 skips the DFlash forward and
+  runs ordinary target decoding for the next step.
 
 ## Usage
 
@@ -32,11 +42,23 @@ vllm serve deepseek-ai/DeepSeek-V4-Flash-DSpark \
 
 Set `enable_adaptive_verification: false` to verify the full block for every request.
 
+For DFlash, use the same flag with `"method": "dflash"`. The configured
+`num_speculative_tokens` remains the maximum K. DFlash drafts all mask positions
+in one parallel forward, so lowering a nonzero K reduces target verification
+work but not draft-model work; K=0 is the only choice that skips the drafter.
+The runtime considers compact graph buckets (K=0, 1, 3, 7, ... and the
+configured maximum) rather than capturing a graph for every integer K.
+
 ## Requirements and limitations
 
-- The attention backend must tolerate device-decided query lengths, since the CPU lengths only bound them from above. Backends that plan off the CPU lengths are excluded by the attention selector, and rejected at startup for models that hard-wire their backend.
-- Full cudagraphs are required: step costs are profiled from captured graphs, so `--enforce-eager` is rejected at startup.
-- Not supported with LoRA (the per-token LoRA mapping is built from CPU-side boundaries), pipeline parallelism (cost curves and confidences exist only on the last rank), or output logprobs (to be fixed).
+- DSpark requires an attention backend that tolerates device-decided query
+  lengths. Backends that plan off CPU lengths are rejected at startup.
+- Full cudagraphs are required: step costs are profiled from captured graphs,
+  so `--enforce-eager` is rejected at startup.
+- Adaptive verification is not supported with LoRA or pipeline parallelism.
+  DSpark additionally does not support output logprobs. DFlash automatic K
+  does not compact the current verification logits and supports output
+  logprobs.
 
 ## Tuning the cost profile
 

tests/test_config.py

@@ -1767,6 +1767,28 @@ def test_draft_sample_method_gumbel_is_rejected():
         )
 
 
+def test_adaptive_verification_accepts_dflash():
+    speculative_config = SpeculativeConfig(
+        method="ngram",
+        num_speculative_tokens=3,
+    )
+    speculative_config.method = "dflash"
+    speculative_config.enable_adaptive_verification = True
+
+    speculative_config._validate_adaptive_verification()
+
+
+def test_adaptive_verification_rejects_unrelated_methods():
+    speculative_config = SpeculativeConfig(
+        method="ngram",
+        num_speculative_tokens=3,
+    )
+    speculative_config.enable_adaptive_verification = True
+
+    with pytest.raises(ValueError, match="only supported with DSpark and DFlash"):
+        speculative_config._validate_adaptive_verification()
+
+
 def test_ir_op_priority_default():
     """Test that IR op priority defaults are set correctly."""
     from vllm.config.kernel import IrOpPriorityConfig

tests/test_sampling_params.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from dataclasses import dataclass
+from types import SimpleNamespace
 
 import pytest
 
@@ -49,3 +50,22 @@ def test_diffusion_accepts_top_k_top_p():
 def test_non_diffusion_models_unaffected():
     params = SamplingParams(temperature=0.7, top_k=10, seed=42)
     params.verify(MockModelConfig(), None, None, None)
+
+
+def test_dflash_adaptive_k_allows_output_logprobs():
+    params = SamplingParams(logprobs=1)
+    speculative_config = SimpleNamespace(
+        method="dflash", enable_adaptive_verification=True
+    )
+
+    params._validate_spec_decode(speculative_config)
+
+
+def test_dspark_adaptive_verification_rejects_output_logprobs():
+    params = SamplingParams(logprobs=1)
+    speculative_config = SimpleNamespace(
+        method="dspark", enable_adaptive_verification=True
+    )
+
+    with pytest.raises(ValueError, match="DSpark confidence-based"):
+        params._validate_spec_decode(speculative_config)