diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index 4dd43503daed..86787b53a105 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/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,101 @@ 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]) +def test_gumbel_drafted_rejection_sample_is_unbiased(num_speculative_steps: int): + """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 + + inputs = _build_rejection_sample_inputs( + target_logits_1d, + draft_logits_1d, + num_speculative_steps, + temperature=1.0, + 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): """ diff --git a/tests/v1/worker/test_gpu_gumbel_sample.py b/tests/v1/worker/test_gpu_gumbel_sample.py index a87c12416eb5..10a09b9de4a6 100644 --- a/tests/v1/worker/test_gpu_gumbel_sample.py +++ b/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), ) diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index a600ab9bc6f3..4bf1f47280cf 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -13,6 +13,12 @@ # attribute is `None`, and `tl.constexpr(...)` would crash at import time. _TL_RAND_MIN = tl.constexpr(4.6566127342e-10) if HAS_TRITON else 4.6566127342e-10 +# Offset salt keeping the draft's Gumbel noise disjoint from the target's. +# Verification is a probability-ratio test, not a Gumbel coupling, so a proposal +# and the residual it is resampled from must not share a noise vector. +# Positions are int64 and never approach 2**30, so the streams cannot collide. +_DRAFT_NOISE_SALT = tl.constexpr(1 << 30) if HAS_TRITON else (1 << 30) + @triton.jit def _temperature_kernel( @@ -89,6 +95,7 @@ def gumbel_noised_argmax( seed, pos, temp, + IS_DRAFTING: tl.constexpr, USE_FP64: tl.constexpr, APPLY_TEMPERATURE: tl.constexpr = True, ): @@ -108,6 +115,8 @@ def gumbel_noised_argmax( if USE_FP64: logits = logits.to(tl.float64) if temp != 0.0: + if IS_DRAFTING: + pos = pos + _DRAFT_NOISE_SALT gumbel_seed = tl.randint(seed, pos) if USE_FP64: u = tl_rand64(gumbel_seed, keys, includes_zero=False) @@ -137,6 +146,7 @@ def gumbel_block_argmax( logits_cache_stride_1, logits_cache_col_ptr, vocab_size, + IS_DRAFTING: tl.constexpr, APPLY_TEMPERATURE: tl.constexpr, USE_FP64: tl.constexpr, PER_TOKEN_COL: tl.constexpr = False, @@ -173,6 +183,7 @@ def gumbel_block_argmax( seed, pos, temp, + IS_DRAFTING=IS_DRAFTING, USE_FP64=USE_FP64, APPLY_TEMPERATURE=APPLY_TEMPERATURE, ) @@ -197,6 +208,7 @@ def _gumbel_sample_kernel( temp_ptr, vocab_size, BLOCK_SIZE: tl.constexpr, + IS_DRAFTING: tl.constexpr, APPLY_TEMPERATURE: tl.constexpr, USE_FP64: tl.constexpr, PER_TOKEN_COL: tl.constexpr, @@ -226,6 +238,7 @@ def _gumbel_sample_kernel( logits_cache_stride_1, logits_cache_col_ptr, vocab_size, + IS_DRAFTING=IS_DRAFTING, APPLY_TEMPERATURE=APPLY_TEMPERATURE, USE_FP64=USE_FP64, PER_TOKEN_COL=PER_TOKEN_COL, @@ -242,6 +255,7 @@ def gumbel_sample( seed: torch.Tensor, # [max_num_reqs] pos: torch.Tensor, # [num_tokens] apply_temperature: bool, + is_drafting: bool, logits_cache: torch.Tensor | None = None, # [max_num_reqs, num_cols, vocab_size] logits_cache_col: torch.Tensor | None = None, # scalar or [num_tokens] use_fp64: bool = False, @@ -281,6 +295,7 @@ def gumbel_sample( temperature, vocab_size, BLOCK_SIZE=BLOCK_SIZE, + IS_DRAFTING=is_drafting, APPLY_TEMPERATURE=apply_temperature, USE_FP64=use_fp64, PER_TOKEN_COL=per_token_col, diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index d3205a3dae83..3ee2c132940e 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -317,6 +317,7 @@ def sample( self.sampling_states.seeds.gpu, pos, apply_temperature=False, + is_drafting=False, use_fp64=self.use_fp64_gumbel, ) return sampled, processed_logits diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index dc5f115c70b1..8cdbc7322cb5 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -37,6 +37,9 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.last_token_indices = torch.zeros( self.max_num_reqs, dtype=torch.int64, device=device ) + self.sample_src_positions = torch.zeros( + self.max_num_reqs, dtype=torch.int64, device=device + ) self.inputs_embeds: torch.Tensor | None = None @@ -327,6 +330,7 @@ def propose( input_batch.seq_lens, num_rejected, self.input_buffers, + self.sample_src_positions, self.max_model_len, self.max_num_reqs, advance_draft_positions=self.advance_draft_positions, @@ -439,6 +443,10 @@ def _prefill( ) -> None: last_token_indices = self.last_token_indices[:num_reqs] positions = self.input_buffers.positions[last_token_indices] + # The output hidden state at position P (= positions) and the token id + # at P+1 are used to draft the token at P+2. Sampling keys a draw by the + # position before the sampled token, so the net adjustment is +1. + sample_src_positions = positions + 1 idx_mapping = self.idx_mapping[:num_reqs] last_hidden_states, hidden_states = self._run_model( @@ -449,11 +457,11 @@ def _prefill( cudagraph_runtime_mode=cudagraph_runtime_mode, mm_inputs=mm_inputs, ) - sample_hidden_states = last_hidden_states[last_token_indices] + sample_hidden_states = last_hidden_states[last_token_indices] self.draft_tokens[:num_reqs, 0] = self.sample_draft( sample_hidden_states, - positions, + sample_src_positions, idx_mapping, self.temperature, self.seeds, @@ -465,6 +473,7 @@ def _prefill( else: self.hidden_states[:num_reqs] = hidden_states[last_token_indices] self.input_buffers.positions[:num_reqs] = positions + self.sample_src_positions[:num_reqs] = sample_src_positions def _multi_step_decode( self, @@ -617,7 +626,6 @@ def _generate_draft( self._prepare_eplb_forward(num_reqs) idx_mapping = self.idx_mapping[:num_reqs] - positions = self.input_buffers.positions[:num_reqs] # Run the draft model forward pass. last_hidden_states, hidden_states = self._run_model( num_tokens_padded, @@ -626,18 +634,13 @@ def _generate_draft( num_tokens_across_dp, cudagraph_runtime_mode, ) - last_hidden_states = last_hidden_states[:num_reqs] - - sample_positions = positions - if not self.advance_draft_positions: - # The forward pass holds positions fixed (Q-only, shared target KV), - # but Gumbel sampling still needs the absolute draft position. - sample_positions = positions + self.current_draft_step # Sample the draft tokens. + sample_hidden_states = last_hidden_states[:num_reqs] + sample_src_positions = self.sample_src_positions[:num_reqs] draft_tokens = self.sample_draft( - last_hidden_states, - sample_positions, + sample_hidden_states, + sample_src_positions, idx_mapping, self.temperature, self.seeds, @@ -653,6 +656,7 @@ def _generate_draft( self.draft_tokens, self.hidden_states, self.input_buffers, + self.sample_src_positions, num_reqs, self.max_model_len, self.num_speculative_steps, @@ -790,6 +794,7 @@ def _prepare_decode_inputs_kernel( num_rejected_ptr, input_ids_ptr, positions_ptr, + sample_src_positions_ptr, query_start_loc_ptr, seq_lens_ptr, max_model_len, @@ -818,6 +823,10 @@ def _prepare_decode_inputs_kernel( draft_token = tl.load(draft_tokens_ptr + req_idx * draft_tokens_stride) tl.store(input_ids_ptr + req_idx, draft_token) + # Advance the draft sampling key. + sample_position = tl.load(sample_src_positions_ptr + req_idx) + tl.store(sample_src_positions_ptr + req_idx, sample_position + 1) + target_seq_len = tl.load(target_seq_lens_ptr + req_idx) num_rejected = tl.load(num_rejected_ptr + req_idx) seq_len = target_seq_len - num_rejected @@ -837,6 +846,7 @@ def prepare_decode_inputs( target_seq_lens: torch.Tensor, num_rejected: torch.Tensor, input_buffers: InputBuffers, + sample_src_positions: torch.Tensor, max_model_len: int, max_num_reqs: int, advance_draft_positions: bool = True, @@ -849,6 +859,7 @@ def prepare_decode_inputs( num_rejected, input_buffers.input_ids, input_buffers.positions, + sample_src_positions, input_buffers.query_start_loc, input_buffers.seq_lens, max_model_len, @@ -866,6 +877,7 @@ def _update_draft_inputs_kernel( next_input_hidden_states_stride, input_ids_ptr, positions_ptr, + sample_src_positions_ptr, seq_lens_ptr, draft_tokens_ptr, current_draft_step_ptr, @@ -891,6 +903,10 @@ def _update_draft_inputs_kernel( # This is the final step. Skip updating draft forward inputs. return + # Advance the draft sampling key. + sample_position = tl.load(sample_src_positions_ptr + req_idx) + tl.store(sample_src_positions_ptr + req_idx, sample_position + 1) + # Write the sampled draft token into the input ids tensor for the next # forward pass. tl.store(input_ids_ptr + req_idx, draft_token) @@ -932,6 +948,7 @@ def update_draft_inputs( output_draft_tokens: torch.Tensor, next_input_hidden_states: torch.Tensor, input_buffers: InputBuffers, + sample_src_positions: torch.Tensor, num_reqs: int, max_model_len: int, num_speculative_steps: int, @@ -945,6 +962,7 @@ def update_draft_inputs( next_input_hidden_states.stride(0), input_buffers.input_ids, input_buffers.positions, + sample_src_positions, input_buffers.seq_lens, draft_tokens, current_draft_step, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 2ff7b22d0636..67378b424adf 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -261,11 +261,11 @@ def _generate_draft( ) num_sample = num_reqs * self.num_speculative_steps sample_hidden_states = last_hidden_states[self.sample_indices[:num_sample]] - # sample_pos is the predicted token's position Q; verification keys - # Gumbel by the predecessor (Q-1). sample_draft adds +1, so pass Q-2. + # sample_pos is the predicted token's position P. Sampling keys a draw + # by the position before the sampled token, P-1. draft_tokens = self.sample_draft( sample_hidden_states, - self.sample_pos[:num_sample] - 2, + self.sample_pos[:num_sample] - 1, self.sample_idx_mapping[:num_sample], self.temperature, self.seeds, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py index 621afc38c716..1cd19acdab23 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash2/speculator.py @@ -51,15 +51,17 @@ def _selector_walk_kernel( other=0, ) - # Candidate ids key the noise, matching the target's own sampling. - position = tl.load(sample_pos_ptr + flat) - 1 + # sample_pos is the predicted token's position P. Sampling keys a draw + # by the position before the sampled token, P-1. + sample_pos = tl.load(sample_pos_ptr + flat) - 1 _, index = gumbel_noised_argmax( scores, candidates, mask & valid, seed, - position, + sample_pos, temperature if SAMPLE_PROBABILISTIC else 0.0, + IS_DRAFTING=True, USE_FP64=USE_FP64, ) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index 19fc45201846..e9eca550ca9a 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -135,8 +135,8 @@ def _sample_logits( buf.index_copy_(1, self._d2t_scatter_index, logits.to(buf.dtype)) logits = buf - # sample_pos is the predicted token's position Q; the target verifies - # it with the predecessor's Gumbel key (Q-1). Pass Q-1. + # sample_pos is the predicted token's position P. Sampling keys a draw + # by the position before the sampled token, P-1. return gumbel_sample( logits, idx_map, @@ -144,6 +144,7 @@ def _sample_logits( self.seeds, sample_pos - 1, apply_temperature=True, + is_drafting=True, logits_cache=self.draft_logits, logits_cache_col=self._step_cols[step], use_fp64=self.use_fp64_gumbel, diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py index 82835c7b3586..90572c03ccd5 100644 --- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py @@ -379,7 +379,11 @@ def _generate_drafts( cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, ) -> None: last_token_indices = self.last_token_indices[:num_reqs] - sample_positions = self.input_buffers.positions[last_token_indices] + positions = self.input_buffers.positions[last_token_indices] + # The output hidden state at position P (= positions) and the token id + # at P+1 are used to draft the token at P+2. Sampling keys a draw by the + # position before the sampled token, so the net adjustment is +1. + sample_src_positions = positions + 1 idx_mapping = self.idx_mapping[:num_reqs] # Cache the trailing token's ids, hidden states (and embeddings for @@ -417,7 +421,7 @@ def _generate_drafts( sample_hidden_states = last_hidden_states[last_token_indices] draft_tokens = self.sample_draft( sample_hidden_states, - sample_positions, + sample_src_positions, idx_mapping, self.temperature, self.seeds, @@ -448,7 +452,8 @@ def _generate_drafts( idx_mapping, num_reqs, ) - sample_positions += 1 + # Advance the draft sampling key. + sample_src_positions += 1 @triton.jit 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 abbf0ee96ef4..29b20afea69a 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -837,6 +837,7 @@ def _resample_kernel( 0, # logits_cache_stride_1 None, # logits_cache_col_ptr vocab_size, + IS_DRAFTING=False, APPLY_TEMPERATURE=False, USE_FP64=USE_FP64, ) diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 303f16b79f3c..bb833ab25e52 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -333,7 +333,7 @@ def _greedy_sample_draft(self, hidden_states: torch.Tensor) -> torch.Tensor: def sample_draft( self, hidden_states: torch.Tensor, - positions: torch.Tensor, + sample_src_positions: torch.Tensor, idx_mapping: torch.Tensor, temperature: torch.Tensor, seeds: torch.Tensor, @@ -342,15 +342,14 @@ def sample_draft( ) -> torch.Tensor: if draft_logits is not None: logits = self.model.compute_logits(hidden_states) - # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise - # used for draft and target sampling. return gumbel_sample( logits, idx_mapping, temperature, seeds, - positions + 1, + sample_src_positions, apply_temperature=True, + is_drafting=True, logits_cache=draft_logits, logits_cache_col=draft_step, use_fp64=self.use_fp64_gumbel,