fix(spec_decode): backport independent draft RNG from upstream #54282

tests/v1/spec_decode/test_rejection_sampler_utils.py

@@ -6,6 +6,7 @@
 import pytest
 import torch
 
+from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
 from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (
     rejection_sample,
 )
@@ -195,6 +196,106 @@ def test_stochastic_rejection_sample(
         )
 
 
+# The test above spreads its samples too thin to resolve a small distributional
+# bias: VOCAB_SIZE bins over 10 * VOCAB_SIZE trials is ~10 samples per bin.
+# Sixteen bins over 200K trials is ~12K per bin, which resolves a few percent.
+NARROW_VOCAB_SIZE = 16
+NARROW_NUM_TRIALS = 200_000
+
+
+def _gumbel_drafted_tokens(
+    inputs: dict,
+    draft_logits_1d: torch.Tensor,
+    num_trials: int,
+    num_speculative_steps: int,
+) -> torch.Tensor:
+    """Proposals drawn with gumbel_sample, shaped like inputs["draft_sampled"].
+
+    _build_rejection_sample_inputs draws them with torch.multinomial, which is
+    independent of the resample noise by construction. Production drafts come
+    from gumbel_sample keyed by pos[t * (K + 1) + i] for step i of trial t --
+    the same entry _rejection_kernel and _resample_kernel read for that token --
+    so the draft and the residual compete for one noise stream.
+    """
+    k = num_speculative_steps
+    vocab_size = draft_logits_1d.shape[0]
+    device = draft_logits_1d.device
+    draft_tokens = gumbel_sample(
+        draft_logits_1d.unsqueeze(0).expand(num_trials * k, vocab_size).float(),
+        inputs["expanded_idx_mapping"]
+        .view(num_trials, k + 1)[:, :k]
+        .reshape(-1)
+        .contiguous(),
+        inputs["temperature"],
+        inputs["seed"],
+        inputs["pos"].view(num_trials, k + 1)[:, :k].reshape(-1).contiguous(),
+        apply_temperature=True,
+        is_drafting=True,
+    )
+    draft_sampled = torch.zeros(num_trials * (k + 1), dtype=torch.int64, device=device)
+    draft_sampled.view(num_trials, k + 1)[:, 1:] = draft_tokens.view(num_trials, k)
+    return draft_sampled
+
+
+@pytest.mark.parametrize("num_speculative_steps", [1, 3])
+@pytest.mark.parametrize("draft_dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("temperature", [0.6, 1.0])
+def test_gumbel_drafted_rejection_sample_is_unbiased(
+    num_speculative_steps: int, draft_dtype: torch.dtype, temperature: float
+):
+    """The proposal and the residual resample must not share a noise vector.
+
+    Draws proposals on the same (seed, pos) stream the sampler verifies and
+    resamples with, then checks the output still follows the target. Conditioned
+    on a proposal winning the argmax, every other token's Gumbel is truncated
+    below that max -- most tightly for the tokens the draft ranked highest -- so
+    a shared stream makes the residual under-weight exactly those tokens.
+
+    Runs narrow because the wide-vocab test above cannot resolve this: dropping
+    `is_drafting=True` in _gumbel_drafted_tokens takes position 0 from chi2 ~12 to
+    ~1500 against a threshold of ~70 here, while leaving that test passing.
+    """
+    torch.manual_seed(42)
+    device = "cuda"
+
+    # A draft that ranks tokens in exactly the opposite order rejects ~73% of
+    # proposals, so most trials reach the residual resample where the bias
+    # lives. The disagreement has to be constructed rather than sampled: two
+    # independent randn draws land close together often enough that the signal
+    # swings between chi2 ~14 and ~1900 depending on the seed. At temperature
+    # 1.0 the target needs no scaling before being passed in.
+    target_logits_1d = torch.randn(
+        NARROW_VOCAB_SIZE, device=device, dtype=torch.float32
+    )
+    draft_logits_1d = (-target_logits_1d).to(draft_dtype)
+    target_logits_1d = target_logits_1d / temperature
+
+    inputs = _build_rejection_sample_inputs(
+        target_logits_1d,
+        draft_logits_1d,
+        num_speculative_steps,
+        temperature=temperature,
+        num_trials=NARROW_NUM_TRIALS,
+    )
+    inputs["draft_sampled"] = _gumbel_drafted_tokens(
+        inputs, draft_logits_1d, NARROW_NUM_TRIALS, num_speculative_steps
+    )
+
+    sampled, num_sampled = rejection_sample(
+        **inputs, num_speculative_steps=num_speculative_steps
+    )
+
+    # Position 0 carries the power: every trial reaches it, while later
+    # positions are only reached on acceptance, which is rare by construction.
+    assert (num_sampled >= 1).all()
+    target_probs = torch.softmax(target_logits_1d, dim=0)
+    for pos in range(num_speculative_steps + 1):
+        accepted_mask = num_sampled >= pos + 1
+        _assert_distribution_match(
+            sampled[accepted_mask, pos], target_probs, device, label=f"position {pos}"
+        )
+
+
 @pytest.mark.parametrize("num_speculative_steps", [1, 3])
 def test_greedy_rejection_sample(num_speculative_steps: int):
     """

tests/v1/worker/test_gpu_autoregressive_speculator.py

@@ -258,11 +258,13 @@ def test_propose_restores_mtp_state_when_draft_decode_raises(monkeypatch) -> Non
     speculator.model = SimpleNamespace(model=lifecycle)
     speculator.rollback_qsa_interval_starts = True
     speculator.share_mtp_topk_indices = True
+    speculator.prefill_outputs_are_compact = False
     speculator.num_speculative_steps = 2
     speculator.max_model_len = 32
     speculator.max_num_reqs = 1
     speculator.hidden_states = torch.zeros(3, 2)
     speculator.last_token_indices = torch.zeros(1, dtype=torch.int64)
+    speculator.sample_src_positions = torch.zeros(1, dtype=torch.int64)
     speculator.current_draft_step = torch.tensor(0, dtype=torch.int64)
     speculator.input_buffers = SimpleNamespace()
     speculator.draft_tokens = torch.zeros((1, 2), dtype=torch.int64)
@@ -591,7 +593,11 @@ def test_mrope_profile_uses_scalar_positions_before_target_state_is_bound(monkey
 
 
 @pytest.mark.skipif(not torch.accelerator.is_available(), reason="accelerator required")
-def test_mrope_prefill_compaction_and_continuation_kernels():
+@pytest.mark.parametrize("max_model_len", [108, 4096])
+@pytest.mark.parametrize("advance_draft_positions", [False, True])
+def test_mrope_prefill_compaction_and_continuation_kernels(
+    max_model_len: int, advance_draft_positions: bool
+):
     device = torch.device("cuda")
     sentinel = -777
     input_buffers = SimpleNamespace(
@@ -660,19 +666,27 @@ def test_mrope_prefill_compaction_and_continuation_kernels():
         num_reqs=2,
     )
     input_buffers.positions[:2] = torch.tensor([102, 106], device=device)
+    sample_src_positions = torch.tensor([103, 107], device=device)
     prepare_decode_inputs(
         draft_tokens=torch.tensor([70, 71], dtype=torch.int64, device=device),
         target_seq_lens=input_batch.seq_lens,
         num_rejected=num_rejected,
         input_buffers=input_buffers,
-        max_model_len=4096,
+        sample_src_positions=sample_src_positions,
+        max_model_len=max_model_len,
         max_num_reqs=3,
+        advance_draft_positions=advance_draft_positions,
         mrope_positions=mrope_positions,
     )
     torch.accelerator.synchronize()
 
-    assert input_buffers.positions[:2].tolist() == [103, 107]
-    assert mrope_positions[:, :2].tolist() == [[13, 17], [13, 17], [13, 17]]
+    shift = int(advance_draft_positions)
+    assert input_buffers.positions[:2].tolist() == [102 + shift, 106 + shift]
+    expected_mrope = (
+        [[13, 17]] * 3 if advance_draft_positions else [[12, 16], [22, 26], [32, 36]]
+    )
+    assert mrope_positions[:, :2].tolist() == expected_mrope
+    assert sample_src_positions.tolist() == [104, 108]
     assert mrope_positions.data_ptr() == backing_ptr
 
     overlapping_positions = torch.tensor(
@@ -722,15 +736,25 @@ def test_mrope_prefill_compaction_and_continuation_kernels():
         output_draft_tokens=torch.zeros((2, 3), dtype=torch.int64, device=device),
         next_input_hidden_states=torch.zeros((2, 4), device=device),
         input_buffers=input_buffers,
+        sample_src_positions=sample_src_positions,
         num_reqs=2,
-        max_model_len=4096,
+        max_model_len=max_model_len,
         num_speculative_steps=3,
+        advance_draft_positions=advance_draft_positions,
         mrope_positions=mrope_positions,
     )
     torch.accelerator.synchronize()
 
-    assert input_buffers.positions[:2].tolist() == [104, 108]
-    assert mrope_positions[:, :2].tolist() == [[14, 18], [14, 18], [14, 18]]
+    assert input_buffers.positions[:2].tolist() == [
+        102 + 2 * shift,
+        min(106 + 2 * shift, max_model_len - 1),
+    ]
+    expected_mrope = (
+        [[14, 18]] * 3 if advance_draft_positions else [[12, 16], [22, 26], [32, 36]]
+    )
+    assert mrope_positions[:, :2].tolist() == expected_mrope
+    # Sampling positions must advance even for clamped or Q-only forward positions.
+    assert sample_src_positions.tolist() == [105, 109]
     assert mrope_positions.data_ptr() == backing_ptr
 
 

tests/v1/worker/test_gpu_gumbel_sample.py

@@ -54,6 +54,7 @@ def _sample(
     *,
     use_fp64: bool = False,
     temperature: float = 1.0,
+    is_drafting: bool = False,
 ) -> torch.Tensor:
     """Sample `num_samples` tokens from one logit vector.
 
@@ -74,6 +75,7 @@ def _sample(
         seed,
         pos,
         apply_temperature=True,
+        is_drafting=is_drafting,
         use_fp64=use_fp64,
     )
 
@@ -84,7 +86,11 @@ def _z_score(observed: int, expected: float, num_trials: int) -> float:
 
 
 def _sample_histogram(
-    logits_1d: torch.Tensor, num_samples: int, *, chunk: int = 1_000_000
+    logits_1d: torch.Tensor,
+    num_samples: int,
+    *,
+    chunk: int = 1_000_000,
+    is_drafting: bool = False,
 ) -> torch.Tensor:
     """Histogram of `num_samples` draws, accumulated in chunks.
 
@@ -101,7 +107,13 @@ def _sample_histogram(
         seed = torch.tensor([0xABCD], dtype=torch.int64, device=DEVICE)
         pos = torch.arange(start, start + size, dtype=torch.int64, device=DEVICE)
         out = gumbel_sample(
-            logits, idx_mapping, temp, seed, pos, apply_temperature=True
+            logits,
+            idx_mapping,
+            temp,
+            seed,
+            pos,
+            apply_temperature=True,
+            is_drafting=is_drafting,
         )
         hist += torch.bincount(out, minlength=vocab_size).double()
     return hist
@@ -164,6 +176,55 @@ def test_full_vocab_distribution_fidelity():
     assert chi2 < df + 10 * math.sqrt(2 * df), f"chi2={chi2:.0f}, df={df}"
 
 
+# --------------------------- Noise streams ---------------------------------
+
+
+def test_drafting_uses_a_separate_noise_stream():
+    """is_drafting salts the Philox offset: same inputs, different draws.
+
+    The draft proposal and the residual resample after a rejection must be
+    independent. They key noise by the same (seed, pos), so only the salt keeps
+    them apart -- without it the resample inherits the very noise vector that
+    picked the rejected proposal. See test_stochastic_rejection_sample in
+    tests/v1/spec_decode/test_rejection_sampler_utils.py for the distributional
+    consequence.
+
+    Relocating the offset must not distort the draw either, so both streams are
+    checked against the target's far-tail mass, which sits HEAD_LOG_GAP logits
+    below the head -- the regime where fp32 Gumbel precision matters.
+    """
+    counts = _make_heavy_tailed_counts()
+    total = counts.sum().item()
+    logits = _counts_to_logits(counts)
+    tail_prob = (total - counts[0].item()) / total
+
+    target = _sample(logits, NUM_SAMPLES, is_drafting=False)
+    draft = _sample(logits, NUM_SAMPLES, is_drafting=True)
+
+    # The head dominates, so the streams agree on most draws by construction.
+    # Compare instead which draws leave the head: shared noise makes that
+    # identical, independent noise makes them differ on ~2p(1-p) of draws.
+    target_tail = target != 0
+    draft_tail = draft != 0
+    disagree = (target_tail != draft_tail).double().mean().item()
+    assert disagree > tail_prob, (
+        f"streams leave the head on the same draws ({disagree:.3e} disagreement "
+        f"vs tail mass {tail_prob:.3e}); the draft salt is not taking effect"
+    )
+
+    # Both streams must still reproduce the target's tail mass.
+    for name, tail in (("target", target_tail), ("draft", draft_tail)):
+        tail_count = tail.sum().item()
+        z = _z_score(tail_count, NUM_SAMPLES * tail_prob, NUM_SAMPLES)
+        assert abs(z) < Z_TOLERANCE, (
+            f"{name} tail mass {tail_count / NUM_SAMPLES:.3e} != "
+            f"{tail_prob:.3e} (z={z:.2f})"
+        )
+
+    # The draft stream is reproducible.
+    assert torch.equal(draft, _sample(logits, NUM_SAMPLES, is_drafting=True))
+
+
 # ----------------------------- Edge cases ----------------------------------
 
 
@@ -178,7 +239,7 @@ def test_greedy_temperature_zero_returns_argmax():
     pos = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE)
 
     sampled = gumbel_sample(
-        logits, idx_mapping, temp, seed, pos, apply_temperature=True
+        logits, idx_mapping, temp, seed, pos, apply_temperature=True, is_drafting=False
     )
     assert torch.equal(sampled, logits.argmax(dim=-1))
 
@@ -278,6 +339,7 @@ def test_logits_cache_stores_input_logits_bitwise(
         seed,
         pos,
         apply_temperature=True,
+        is_drafting=True,
         logits_cache=cache,
         logits_cache_col=cols,
     )
@@ -326,6 +388,7 @@ def test_logits_cache_columns_stay_separate_across_steps(extra_cache_cols: int):
             seed,
             pos,
             apply_temperature=True,
+            is_drafting=True,
             logits_cache=cache,
             logits_cache_col=cols[step],
         )
@@ -357,6 +420,7 @@ def test_logits_cache_narrower_than_logits_is_rejected():
             seed,
             pos,
             apply_temperature=True,
+            is_drafting=True,
             logits_cache=cache,
             logits_cache_col=torch.tensor(0, dtype=torch.int32, device=DEVICE),
         )