[Bugfix][Spec Decode] Avoid reusing random draws across block verification rounds (#61161)

Signed-off-by: DimensionSTP <65501090+DimensionSTP@users.noreply.github.com>
Co-authored-by: Kevin H. Luu <khluu000@gmail.com>

tests/v1/spec_decode/test_dspark_topk.py

@@ -15,6 +15,8 @@ from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import
 )
 from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead
 from vllm.platforms import current_platform
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
+from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator
 
 
 def _markov_head(weight: torch.Tensor) -> DSparkMarkovHead:
@@ -70,6 +72,46 @@ def test_gathered_markov_bias_matches_dense_at_full_vocab():
     torch.testing.assert_close(result, expected)
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+@pytest.mark.parametrize("use_fp64", [False, True])
+def test_dspark_revisited_position_uses_step_seed(use_fp64):
+    """A cached draft at a revisited position must use its current step."""
+    device = torch.device("cuda")
+    batch = 4096
+    spec = DSparkSpeculator.__new__(DSparkSpeculator)
+    spec.acceptance_estimator = None
+    spec._d2t_scatter_index = None
+    spec.draft_watermarker = None
+    spec.use_fp64_gumbel = use_fp64
+    spec.temperature = torch.ones(batch, device=device)
+    spec.seeds = torch.arange(batch, dtype=torch.int64, device=device)
+    spec._step_cols = torch.arange(2, dtype=torch.int32, device=device)
+    spec.draft_logits = torch.empty(batch, 2, 2, device=device)
+    logits = torch.tensor([0.7, 0.3], device=device).log().expand(batch, -1)
+    mapping = torch.arange(batch, device=device)
+    positions = torch.full((batch,), 22, dtype=torch.int64, device=device)
+
+    first = spec._sample_logits(logits, mapping, positions, 1)
+    repeated = spec._sample_logits(logits, mapping, positions, 1)
+    next_round = spec._sample_logits(logits, mapping, positions, 0)
+    torch.testing.assert_close(first, repeated)
+    expected = gumbel_sample(
+        logits,
+        mapping,
+        spec.temperature,
+        spec.seeds,
+        positions - 1,
+        apply_temperature=True,
+        is_drafting=True,
+        logits_cache=spec.draft_logits,
+        logits_cache_col=spec._step_cols[0],
+        use_fp64=use_fp64,
+    )
+    torch.testing.assert_close(next_round, expected)
+    mismatch_rate = (first != next_round).float().mean().item()
+    assert 0.35 < mismatch_rate < 0.5
+
+
 @pytest.mark.parametrize("cuda_alike_platform", [False, True])
 def test_gathered_markov_bias_dequantizes_selected_w4a16_rows(
     monkeypatch, cuda_alike_platform

tests/v1/spec_decode/test_lilicorr.py

@@ -22,6 +22,7 @@ from vllm.model_executor.layers import logits_processor
 from vllm.model_executor.layers.linear import LinearBase
 from vllm.model_executor.models.lilicorr import LiLiCorrHead
 from vllm.transformers_utils.configs.eagle import EAGLEConfig
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
 from vllm.v1.worker.gpu.spec_decode.dflash2.speculator import CandidateSampler
 from vllm.v1.worker.gpu.spec_decode.lilicorr.speculator import LiLiCorrSpeculator
 
@@ -500,6 +501,60 @@ def test_candidate_walk_preserves_conditional_scores_and_padding(probabilistic):
         assert cache[1].isneginf().all()
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+@pytest.mark.parametrize("use_fp64", [False, True])
+def test_candidate_walk_revisited_position_matches_step_seed(use_fp64):
+    """The candidate walk must use the same step key as ordinary draft sampling."""
+    device = torch.device("cuda")
+    batch, steps, top_k = 4096, 2, 2
+    sampler = CandidateSampler(batch, steps, top_k, device)
+    candidates = torch.arange(top_k, device=device).expand(batch, steps, top_k)
+    logits = torch.tensor([0.7, 0.3], device=device).log()
+    scores = logits.expand(batch, steps, top_k, top_k)
+    mapping = torch.arange(batch, device=device)
+    temperature = torch.ones(batch, device=device)
+    seeds = mapping.to(torch.int64)
+    positions = torch.arange(21, 21 + steps, device=device).repeat(batch)
+    tokens = torch.empty(batch, steps, dtype=torch.int64, device=device)
+    cache = torch.empty(batch, steps, top_k, device=device)
+    first_round = None
+    for start in (21, 22):
+        positions.view(batch, steps).copy_(
+            torch.arange(start, start + steps, device=device)
+        )
+        sampler.sample(
+            candidates,
+            scores,
+            batch,
+            positions,
+            mapping.repeat_interleave(steps),
+            temperature,
+            seeds,
+            tokens,
+            cache,
+            use_fp64,
+        )
+        for step in range(steps):
+            expected = gumbel_sample(
+                logits.expand(batch, -1),
+                mapping,
+                temperature,
+                seeds,
+                positions.view(batch, steps)[:, step] - 1,
+                apply_temperature=True,
+                is_drafting=True,
+                logits_cache=cache,
+                logits_cache_col=torch.tensor(step, device=device),
+                use_fp64=use_fp64,
+            )
+            torch.testing.assert_close(tokens[:, step], expected)
+        if first_round is None:
+            first_round = tokens.clone()
+        else:
+            mismatch_rate = (first_round[:, 1] != tokens[:, 0]).float().mean().item()
+            assert 0.35 < mismatch_rate < 0.5
+
+
 @pytest.mark.parametrize(
     ("convolution", "mismatch", "quantized"),
     [

tests/v1/spec_decode/test_rejection_sampler_utils.py

@@ -557,6 +557,79 @@ def test_gumbel_drafted_rejection_sample_is_unbiased(num_speculative_steps: int)
         )
 
 
+@pytest.mark.parametrize("use_block_verification", [False, True])
+@pytest.mark.parametrize("probabilistic_draft", [False, True])
+def test_multiround_rejection_sample_is_unbiased(
+    use_block_verification: bool, probabilistic_draft: bool
+):
+    """Uncommitted positions must receive fresh noise on later verification rounds."""
+    torch.manual_seed(42)
+    num_trials, num_tokens, k = 32768, 6, 2
+    device = "cuda"
+    target_probs = torch.tensor(
+        [0.3, 0.7] if probabilistic_draft else [0.5, 0.5], device=device
+    )
+    draft_logits = torch.tensor(
+        [0.7, 0.3] if probabilistic_draft else [1.0, 0.0], device=device
+    ).log()
+    inputs = _build_rejection_sample_inputs(
+        target_probs.log(), draft_logits, k, 1.0, num_trials
+    )
+    seeds = inputs["seed"]
+    output = torch.full((num_trials, num_tokens + k), -1, device=device)
+    output[:, 0] = torch.multinomial(target_probs, num_trials, replacement=True)
+    lengths = torch.ones(num_trials, dtype=torch.int64, device=device)
+    accepted = 0
+    while (lengths < num_tokens).any():
+        active = (lengths < num_tokens).nonzero().flatten()
+        n = active.numel()
+        starts = lengths[active]
+        local_positions = torch.arange(k + 1, device=device).repeat(n)
+        positions = (starts[:, None] + torch.arange(k + 1, device=device)).flatten()
+        draft_sampled = torch.zeros((n, k + 1), dtype=torch.int64, device=device)
+        draft_sampled[:, 0] = output[active, starts - 1]
+        logits_cache = inputs["draft_logits"] if probabilistic_draft else None
+        if probabilistic_draft:
+            for step in range(k):
+                draft_sampled[:, step + 1] = gumbel_sample(
+                    draft_logits.expand(n, -1),
+                    active,
+                    inputs["temperature"],
+                    seeds,
+                    starts + step,
+                    apply_temperature=True,
+                    is_drafting=True,
+                    logits_cache=logits_cache,
+                    logits_cache_col=torch.tensor(step, device=device),
+                )
+        sampled, counts = rejection_sample(
+            target_logits=target_probs.log().expand(n * (k + 1), -1).contiguous(),
+            draft_logits=logits_cache,
+            draft_sampled=draft_sampled.flatten(),
+            cu_num_logits=torch.arange(n + 1, dtype=torch.int32, device=device)
+            * (k + 1),
+            pos=positions,
+            idx_mapping=active,
+            expanded_idx_mapping=active.repeat_interleave(k + 1),
+            expanded_local_pos=local_positions,
+            temperature=inputs["temperature"],
+            seed=seeds,
+            num_speculative_steps=k,
+            use_block_verification=use_block_verification,
+        )
+        for column in range(k + 1):
+            valid = counts > column
+            output[active[valid], starts[valid] + column] = sampled[valid, column]
+        lengths[active] += counts
+        accepted += (counts - 1).sum().item()
+
+    assert accepted > num_trials, "The test must exercise accepted draft tokens."
+    for position in range(num_tokens):
+        _assert_distribution_match(
+            output[:, position], target_probs, device, label=f"position {position}"
+        )
+
+
 @pytest.mark.parametrize("num_speculative_steps", [1, 3])
 def test_greedy_rejection_sample(num_speculative_steps: int):
     """Verify that greedy (temperature=0) always outputs the target argmax

tests/watermarking/test_gumbel.py

@@ -253,6 +253,8 @@ def test_philox_mixed_draft_sampling_caches_unprocessed_logits():
         positions,
         apply_temperature=True,
         is_drafting=True,
+        logits_cache=cache,
+        logits_cache_col=cols,
     )
     watermarked = philox_gumbel_sample(processed_logits, contexts, 42)
     expected = torch.where(skip_mask, ordinary, watermarked)

vllm/v1/worker/gpu/sample/gumbel.py

@@ -70,6 +70,12 @@ def tl_rand32(seed, offset, includes_zero: tl.constexpr):
     return u
 
 
+@triton.jit
+def spec_step_seed(seed, step):
+    """Mix a speculative step into the high bits of a request seed."""
+    return seed ^ (tl.cast(step, tl.int64) << 32)
+
+
 @triton.jit
 def _murmur3_rotl32(value, shift: tl.constexpr):
     return (value << shift) | (value >> (32 - shift))
@@ -257,6 +263,8 @@ def gumbel_block_argmax(
 
     seed = tl.load(seeds_ptr + req_state_idx, mask=is_valid_req, other=0)
     pos = tl.load(pos_ptr + token_idx)
+    if logits_cache_ptr is not None:
+        seed = spec_step_seed(seed, col)
     return gumbel_noised_argmax(
         logits,
         block,

vllm/v1/worker/gpu/sample/watermark.py

@@ -25,7 +25,7 @@ from vllm.v1.watermarking.prfs.philox import (
 from vllm.v1.watermarking.prfs.philox import (
     _UINT32_MASK as _UINT32_MASK_VALUE,
 )
-from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax, spec_step_seed
 
 if HAS_TRITON:
     from triton.language import math as tl_math
@@ -644,6 +644,8 @@ def _philox_gumbel_kernel(
             )
             seed = tl.load(seeds_ptr + req_state_idx, mask=valid_req, other=0)
             pos = tl.load(pos_ptr + row)
+            if logits_cache_ptr is not None:
+                seed = spec_step_seed(seed, col)
             logits_row = logits_ptr + row * logits_stride
             logits = tl.load(logits_row + candidate, mask=mask, other=float("-inf")).to(
                 tl.float32

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

@@ -8,7 +8,7 @@ import torch
 from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
 from vllm.triton_utils import tl, triton
-from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax, spec_step_seed
 from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator
 
 
@@ -58,7 +58,7 @@ def _selector_walk_kernel(
             scores,
             candidates,
             mask & valid,
-            seed,
+            spec_step_seed(seed, step),
             sample_pos,
             temperature if SAMPLE_PROBABILISTIC else 0.0,
             IS_DRAFTING=True,

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

@@ -3,7 +3,11 @@
 import torch
 
 from vllm.triton_utils import tl, tldevice, triton
-from vllm.v1.worker.gpu.sample.gumbel import gumbel_block_argmax, tl_rand32
+from vllm.v1.worker.gpu.sample.gumbel import (
+    gumbel_block_argmax,
+    spec_step_seed,
+    tl_rand32,
+)
 from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_block_argmax
 
 
@@ -569,7 +573,8 @@ def _rejection_kernel(
             verifying &= is_valid_draft
         if verifying:
             pos = tl.load(pos_ptr + logit_idx)
-            u = tl_rand32(seed, pos, includes_zero=False)
+            step_seed = spec_step_seed(seed, i) if USE_BLOCK_VERIFICATION else seed
+            u = tl_rand32(step_seed, pos, includes_zero=False)
             if is_greedy:
                 # Greedy sampling. Accept IFF draft matches target argmax.
                 # NOTE: Target argmax is stored directly so that resampling