[MRV2][Spec Decode] Close DFlash/DSpark and adaptive-verification JIT warmup gaps (#63819)

Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: vllm-agent <claw@inferact.ai>

tests/v1/spec_decode/test_acceptance_estimator.py

@@ -15,6 +15,8 @@ import torch
 from vllm.platforms import current_platform
 from vllm.v1.worker.gpu.spec_decode.acceptance_estimator import (
     OnlineAcceptanceEstimator,
+    _accumulate_kernel,
+    _refit_kernel,
 )
 
 pytest.importorskip("triton")
@@ -295,3 +297,34 @@ def test_feature_ignores_which_token_was_drawn():
         atol=2e-3,
         rtol=2e-3,
     )
+
+
+def test_warmup_covers_every_batch_size_and_the_refit():
+    """After warmup, no batch size and no refit may JIT-compile during serving."""
+    device = current_platform.device_type
+
+    def compiled_variants() -> set:
+        return {
+            (kernel.fn.__name__, key)
+            for kernel in (_accumulate_kernel, _refit_kernel)
+            for caches in kernel.device_caches.values()
+            for key in caches[0]
+        }
+
+    estimator = OnlineAcceptanceEstimator(REFIT_NUM_REQS, NUM_STEPS, device)
+    estimator.REFIT_INTERVAL = 2
+    slots = torch.arange(REFIT_NUM_REQS, device=device, dtype=torch.int32)
+    num_accepted = torch.ones(REFIT_NUM_REQS, device=device)
+    features = torch.rand(REFIT_NUM_REQS, NUM_STEPS, device=device)
+    _drive_round(estimator, features[:1], num_accepted[:1], slots[:1])
+    estimator.warmup()
+    warmed = compiled_variants()
+
+    for num_reqs in (8, 300):
+        _drive_round(
+            estimator, features[:num_reqs], num_accepted[:num_reqs], slots[:num_reqs]
+        )
+    torch.accelerator.synchronize()
+
+    assert estimator._refits == 1
+    assert compiled_variants() == warmed

tests/v1/spec_decode/test_dflash_prepare_inputs.py

@@ -12,6 +12,7 @@ from vllm.platforms import current_platform
 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
 from vllm.v1.worker.gpu.spec_decode.dflash.speculator import (
     DFlashSpeculator,
+    _prepare_dflash_inputs_kernel,
 )
 
 pytestmark = pytest.mark.skipif(
@@ -28,6 +29,7 @@ def _run_prepare(
     cp_size: int = 1,
     cp_interleave: int = 1,
     draft_dcp_size: int | None = None,
+    num_scheduled_tokens: int = 4,
 ):
     device = torch.device(current_platform.device_type)
     max_num_reqs = 4
@@ -46,7 +48,7 @@ def _run_prepare(
     input_batch = SimpleNamespace(
         num_reqs=1,
         num_tokens=4,
-        num_scheduled_tokens=np.array([4], dtype=np.int32),
+        num_scheduled_tokens=np.array([num_scheduled_tokens], dtype=np.int32),
         seq_lens_cpu_upper_bound=torch.tensor([target_positions[-1] + 1]),
         positions=torch.tensor(target_positions, dtype=torch.int64, device=device),
         query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device=device),
@@ -220,3 +222,18 @@ def test_prepare_dflash_inputs_never_writes_the_null_block():
         PAD_SLOT_ID,
         PAD_SLOT_ID,
     ]
+
+
+def test_prepare_dflash_inputs_compiles_once_for_all_query_lengths():
+    def compiled_variants() -> set:
+        return {
+            key
+            for caches in _prepare_dflash_inputs_kernel.device_caches.values()
+            for key in caches[0]
+        }
+
+    kwargs = dict(target_positions=[10, 11, 12, 13], block_table_values=[0, 7, 8])
+    _run_prepare(**kwargs)
+    warmed = compiled_variants()
+    _run_prepare(**kwargs, num_scheduled_tokens=300)
+    assert compiled_variants() == warmed

tests/v1/worker/test_gpu_warmup_blocks.py

@@ -90,6 +90,7 @@ def _make_runner(
         num_speculative_steps=num_spec_steps,
         speculative_config=SimpleNamespace() if num_spec_steps else None,
         adaptive_verification=None,
+        speculator=None,
         rejection_sampler=None,
         decode_query_len=num_spec_steps + 1,
         is_pooling_model=False,

tests/v1/worker/test_mixed_warmup_gate.py

@@ -76,6 +76,7 @@ def test_kernel_warmup_restores_uncalibrated_adaptive_manager(monkeypatch, fail_
     rejection_sampler = SimpleNamespace(enable_adaptive_verification=True)
     runner = SimpleNamespace(
         adaptive_verification=manager,
+        speculator=None,
         rejection_sampler=rejection_sampler,
         speculative_config=SimpleNamespace(),
     )
@@ -116,6 +117,7 @@ def test_kernel_warmup_runs_same_rounds_on_every_pp_stage(monkeypatch, speculati
         rounds.clear()
         runner = SimpleNamespace(
             adaptive_verification=None,
+            speculator=None,
             speculative_config=SimpleNamespace() if speculative else None,
             rejection_sampler=(
                 SimpleNamespace(enable_adaptive_verification=False)

vllm/v1/worker/gpu/spec_decode/acceptance_estimator.py

@@ -347,6 +347,7 @@ class OnlineAcceptanceEstimator:
         self.device = device
         self._steps_since_refit = 0
         self._refits = 0
+        self._accumulate_block_r = max(256, triton.next_power_of_2(max_num_reqs))
 
         # The learned coefficients used for predicting acceptance.
         self.slope = torch.full((1,), _INIT_SLOPE, dtype=torch.float32, device=device)
@@ -422,8 +423,8 @@ class OnlineAcceptanceEstimator:
             self.predictions.stride(0),
             self.counts,
             num_reqs,
-            BLOCK_R=max(256, triton.next_power_of_2(num_reqs)),
-            num_warps=1 if num_reqs <= 256 else 4,
+            BLOCK_R=self._accumulate_block_r,
+            num_warps=1 if self._accumulate_block_r <= 256 else 4,
         )
 
         self._steps_since_refit += 1
@@ -433,6 +434,19 @@ class OnlineAcceptanceEstimator:
 
         # Fit the coefficients to the accumulated statistics gathered over the
         # course of the last REFIT_INTERVAL steps.
+        self._refit()
+        self._refits += 1
+
+    def warmup(self) -> None:
+        """Drop warmup statistics and compile the refit kernel. A refit on
+        zeroed statistics leaves the coefficients unchanged."""
+        self.info.zero_()
+        self.grad.zero_()
+        self.counts.zero_()
+        self._steps_since_refit = 0
+        self._refit()
+
+    def _refit(self) -> None:
         _refit_kernel[(1,)](
             self.slope,
             self.intercepts,
@@ -446,7 +460,6 @@ class OnlineAcceptanceEstimator:
             DAMPING=self.DAMPING_OBSERVATIONS,
             BLOCK=triton.next_power_of_2(self.num_speculative_steps),
         )
-        self._refits += 1
 
     def predict(
         self,

vllm/v1/worker/gpu/spec_decode/dflash/speculator.py

@@ -888,7 +888,7 @@ def prepare_dflash_inputs(
     # per-request query length, not the total token count across the batch.
     max_target_query_len = int(input_batch.num_scheduled_tokens[:num_reqs].max())
     max_tokens_per_req = max_target_query_len + num_query_per_req
-    BLOCK_SIZE = min(256, triton.next_power_of_2(max(1, max_tokens_per_req)))
+    BLOCK_SIZE = 128
     num_blocks = triton.cdiv(max_tokens_per_req, BLOCK_SIZE)
     _prepare_dflash_inputs_kernel[(num_reqs, num_blocks)](
         input_buffers.input_ids,

vllm/v1/worker/gpu/warmup.py

@@ -252,6 +252,11 @@ def warmup_kernels(
             rejection_sampler.enable_adaptive_verification = True
     if adaptive_verification is not None:
         adaptive_verification.warmup_budget_kernel()
+    acceptance_estimator = getattr(
+        model_runner.speculator, "acceptance_estimator", None
+    )
+    if acceptance_estimator is not None:
+        acceptance_estimator.warmup()
 
 
 def _warmup_kernels(