From 9edcc7fa7ec858e13d4dc10ed968b741e2c95d6a Mon Sep 17 00:00:00 2001 From: DimensionSTP <65501090+DimensionSTP@users.noreply.github.com> Date: Fri, 16 Oct 2026 06:46:38 +0900 Subject: [PATCH] [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 --- tests/v1/spec_decode/test_dspark_topk.py | 42 +++++++++++ tests/v1/spec_decode/test_lilicorr.py | 55 ++++++++++++++ .../test_rejection_sampler_utils.py | 73 +++++++++++++++++++ tests/watermarking/test_gumbel.py | 2 + vllm/v1/worker/gpu/sample/gumbel.py | 8 ++ vllm/v1/worker/gpu/sample/watermark.py | 4 +- .../gpu/spec_decode/dflash2/speculator.py | 4 +- .../spec_decode/rejection_sampler_utils.py | 9 ++- 8 files changed, 192 insertions(+), 5 deletions(-) diff --git a/tests/v1/spec_decode/test_dspark_topk.py b/tests/v1/spec_decode/test_dspark_topk.py index 2de1ff052a..164ef14ff4 100644 --- a/tests/v1/spec_decode/test_dspark_topk.py +++ b/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 diff --git a/tests/v1/spec_decode/test_lilicorr.py b/tests/v1/spec_decode/test_lilicorr.py index 46dbaf41af..95fac9f834 100644 --- a/tests/v1/spec_decode/test_lilicorr.py +++ b/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"), [ diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index f635a66a46..ab94ec54d2 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/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 diff --git a/tests/watermarking/test_gumbel.py b/tests/watermarking/test_gumbel.py index aa0ffa2799..b2d7aa006e 100644 --- a/tests/watermarking/test_gumbel.py +++ b/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) diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index a34169b733..9ab0579a6f 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/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, diff --git a/vllm/v1/worker/gpu/sample/watermark.py b/vllm/v1/worker/gpu/sample/watermark.py index 216d189f8d..d366b7622a 100644 --- a/vllm/v1/worker/gpu/sample/watermark.py +++ b/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 diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py index bb14621129..110a2b34a7 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py +++ b/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, diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index 4387b1d38a..f35468d62b 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/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