@tduncan / Nvidia-Optimization-Stack-VLLM / commits / fe755c8899

[Bugfix][Model Runner V2][Spec Decode] Decouple the draft's gumbel noise stream from the target's (#54282)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>

Giancarlo Delfin committed Aug 29, 2026 at 09:58 UTC fe755c88995ad468882517b6c4bdd60138d46a3a

11 files changed +232 -30

tests/v1/spec_decode/test_rejection_sampler_utils.py

+96

@@ -6,6 +6,7 @@ import math

6 import pytest

7 import torch

9 +from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample

10 from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (

11 rejection_sample,

12 )

@@ -195,6 +196,101 @@ def test_stochastic_rejection_sample(

196 )

197

198

199 +# The test above spreads its samples too thin to resolve a small distributional

200 +# bias: VOCAB_SIZE bins over 10 * VOCAB_SIZE trials is ~10 samples per bin.

201 +# Sixteen bins over 200K trials is ~12K per bin, which resolves a few percent.

202 +NARROW_VOCAB_SIZE = 16

203 +NARROW_NUM_TRIALS = 200_000

204 +

205 +

206 +def _gumbel_drafted_tokens(

207 + inputs: dict,

208 + draft_logits_1d: torch.Tensor,

209 + num_trials: int,

210 + num_speculative_steps: int,

211 +) -> torch.Tensor:

212 + """Proposals drawn with gumbel_sample, shaped like inputs["draft_sampled"].

213 +

214 + _build_rejection_sample_inputs draws them with torch.multinomial, which is

215 + independent of the resample noise by construction. Production drafts come

216 + from gumbel_sample keyed by pos[t * (K + 1) + i] for step i of trial t --

217 + the same entry _rejection_kernel and _resample_kernel read for that token --

218 + so the draft and the residual compete for one noise stream.

219 + """

220 + k = num_speculative_steps

221 + vocab_size = draft_logits_1d.shape[0]

222 + device = draft_logits_1d.device

223 + draft_tokens = gumbel_sample(

224 + draft_logits_1d.unsqueeze(0).expand(num_trials * k, vocab_size).float(),

225 + inputs["expanded_idx_mapping"]

226 + .view(num_trials, k + 1)[:, :k]

227 + .reshape(-1)

228 + .contiguous(),

229 + inputs["temperature"],

230 + inputs["seed"],

231 + inputs["pos"].view(num_trials, k + 1)[:, :k].reshape(-1).contiguous(),

232 + apply_temperature=True,

233 + is_drafting=True,

234 + )

235 + draft_sampled = torch.zeros(num_trials * (k + 1), dtype=torch.int64, device=device)

236 + draft_sampled.view(num_trials, k + 1)[:, 1:] = draft_tokens.view(num_trials, k)

237 + return draft_sampled

238 +

239 +

240 +@pytest.mark.parametrize("num_speculative_steps", [1, 3])

241 +def test_gumbel_drafted_rejection_sample_is_unbiased(num_speculative_steps: int):

242 + """The proposal and the residual resample must not share a noise vector.

243 +

244 + Draws proposals on the same (seed, pos) stream the sampler verifies and

245 + resamples with, then checks the output still follows the target. Conditioned

246 + on a proposal winning the argmax, every other token's Gumbel is truncated

247 + below that max -- most tightly for the tokens the draft ranked highest -- so

248 + a shared stream makes the residual under-weight exactly those tokens.

249 +

250 + Runs narrow because the wide-vocab test above cannot resolve this: dropping

251 + `is_drafting=True` in _gumbel_drafted_tokens takes position 0 from chi2 ~12 to

252 + ~1500 against a threshold of ~70 here, while leaving that test passing.

253 + """

254 + torch.manual_seed(42)

255 + device = "cuda"

256 +

257 + # A draft that ranks tokens in exactly the opposite order rejects ~73% of

258 + # proposals, so most trials reach the residual resample where the bias

259 + # lives. The disagreement has to be constructed rather than sampled: two

260 + # independent randn draws land close together often enough that the signal

261 + # swings between chi2 ~14 and ~1900 depending on the seed. At temperature

262 + # 1.0 the target needs no scaling before being passed in.

263 + target_logits_1d = torch.randn(

264 + NARROW_VOCAB_SIZE, device=device, dtype=torch.float32

265 + )

266 + draft_logits_1d = -target_logits_1d

267 +

268 + inputs = _build_rejection_sample_inputs(

269 + target_logits_1d,

270 + draft_logits_1d,

271 + num_speculative_steps,

272 + temperature=1.0,

273 + num_trials=NARROW_NUM_TRIALS,

274 + )

275 + inputs["draft_sampled"] = _gumbel_drafted_tokens(

276 + inputs, draft_logits_1d, NARROW_NUM_TRIALS, num_speculative_steps

277 + )

278 +

279 + sampled, num_sampled = rejection_sample(

280 + **inputs, num_speculative_steps=num_speculative_steps

281 + )

282 +

283 + # Position 0 carries the power: every trial reaches it, while later

284 + # positions are only reached on acceptance, which is rare by construction.

285 + assert (num_sampled >= 1).all()

286 + target_probs = torch.softmax(target_logits_1d, dim=0)

287 + for pos in range(num_speculative_steps + 1):

288 + accepted_mask = num_sampled >= pos + 1

289 + _assert_distribution_match(

290 + sampled[accepted_mask, pos], target_probs, device, label=f"position {pos}"

291 + )

292 +

293 +

294 @pytest.mark.parametrize("num_speculative_steps", [1, 3])

295 def test_greedy_rejection_sample(num_speculative_steps: int):

296 """

tests/v1/worker/test_gpu_gumbel_sample.py

+67 -3

@@ -54,6 +54,7 @@ def _sample(

54 *,

55 use_fp64: bool = False,

56 temperature: float = 1.0,

57 + is_drafting: bool = False,

58 ) -> torch.Tensor:

59 """Sample `num_samples` tokens from one logit vector.

@@ -74,6 +75,7 @@ def _sample(

75 seed,

76 pos,

77 apply_temperature=True,

78 + is_drafting=is_drafting,

79 use_fp64=use_fp64,

80 )

@@ -84,7 +86,11 @@ def _z_score(observed: int, expected: float, num_trials: int) -> float:

88 def _sample_histogram(

87 - logits_1d: torch.Tensor, num_samples: int, *, chunk: int = 1_000_000

89 + logits_1d: torch.Tensor,

90 + num_samples: int,

91 + *,

92 + chunk: int = 1_000_000,

93 + is_drafting: bool = False,

94 ) -> torch.Tensor:

95 """Histogram of `num_samples` draws, accumulated in chunks.

@@ -101,7 +107,13 @@ def _sample_histogram(

107 seed = torch.tensor([0xABCD], dtype=torch.int64, device=DEVICE)

108 pos = torch.arange(start, start + size, dtype=torch.int64, device=DEVICE)

109 out = gumbel_sample(

104 - logits, idx_mapping, temp, seed, pos, apply_temperature=True

110 + logits,

111 + idx_mapping,

112 + temp,

113 + seed,

114 + pos,

115 + apply_temperature=True,

116 + is_drafting=is_drafting,

117 )

118 hist += torch.bincount(out, minlength=vocab_size).double()

119 return hist

@@ -164,6 +176,55 @@ def test_full_vocab_distribution_fidelity():

176 assert chi2 < df + 10 * math.sqrt(2 * df), f"chi2={chi2:.0f}, df={df}"

177

178

179 +# --------------------------- Noise streams ---------------------------------

180 +

181 +

182 +def test_drafting_uses_a_separate_noise_stream():

183 + """is_drafting salts the Philox offset: same inputs, different draws.

184 +

185 + The draft proposal and the residual resample after a rejection must be

186 + independent. They key noise by the same (seed, pos), so only the salt keeps

187 + them apart -- without it the resample inherits the very noise vector that

188 + picked the rejected proposal. See test_stochastic_rejection_sample in

189 + tests/v1/spec_decode/test_rejection_sampler_utils.py for the distributional

190 + consequence.

191 +

192 + Relocating the offset must not distort the draw either, so both streams are

193 + checked against the target's far-tail mass, which sits HEAD_LOG_GAP logits

194 + below the head -- the regime where fp32 Gumbel precision matters.

195 + """

196 + counts = _make_heavy_tailed_counts()

197 + total = counts.sum().item()

198 + logits = _counts_to_logits(counts)

199 + tail_prob = (total - counts[0].item()) / total

200 +

201 + target = _sample(logits, NUM_SAMPLES, is_drafting=False)

202 + draft = _sample(logits, NUM_SAMPLES, is_drafting=True)

203 +

204 + # The head dominates, so the streams agree on most draws by construction.

205 + # Compare instead which draws leave the head: shared noise makes that

206 + # identical, independent noise makes them differ on ~2p(1-p) of draws.

207 + target_tail = target != 0

208 + draft_tail = draft != 0

209 + disagree = (target_tail != draft_tail).double().mean().item()

210 + assert disagree > tail_prob, (

211 + f"streams leave the head on the same draws ({disagree:.3e} disagreement "

212 + f"vs tail mass {tail_prob:.3e}); the draft salt is not taking effect"

213 + )

214 +

215 + # Both streams must still reproduce the target's tail mass.

216 + for name, tail in (("target", target_tail), ("draft", draft_tail)):

217 + tail_count = tail.sum().item()

218 + z = _z_score(tail_count, NUM_SAMPLES * tail_prob, NUM_SAMPLES)

219 + assert abs(z) < Z_TOLERANCE, (

220 + f"{name} tail mass {tail_count / NUM_SAMPLES:.3e} != "

221 + f"{tail_prob:.3e} (z={z:.2f})"

222 + )

223 +

224 + # The draft stream is reproducible.

225 + assert torch.equal(draft, _sample(logits, NUM_SAMPLES, is_drafting=True))

226 +

227 +

228 # ----------------------------- Edge cases ----------------------------------

229

230

@@ -178,7 +239,7 @@ def test_greedy_temperature_zero_returns_argmax():

239 pos = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE)

240

241 sampled = gumbel_sample(

181 - logits, idx_mapping, temp, seed, pos, apply_temperature=True

242 + logits, idx_mapping, temp, seed, pos, apply_temperature=True, is_drafting=False

243 )

244 assert torch.equal(sampled, logits.argmax(dim=-1))

245

@@ -278,6 +339,7 @@ def test_logits_cache_stores_input_logits_bitwise(

339 seed,

340 pos,

341 apply_temperature=True,

342 + is_drafting=True,

343 logits_cache=cache,

344 logits_cache_col=cols,

345 )

@@ -326,6 +388,7 @@ def test_logits_cache_columns_stay_separate_across_steps(extra_cache_cols: int):

388 seed,

389 pos,

390 apply_temperature=True,

391 + is_drafting=True,

392 logits_cache=cache,

393 logits_cache_col=cols[step],

394 )

@@ -357,6 +420,7 @@ def test_logits_cache_narrower_than_logits_is_rejected():

420 seed,

421 pos,

422 apply_temperature=True,

423 + is_drafting=True,

424 logits_cache=cache,

425 logits_cache_col=torch.tensor(0, dtype=torch.int32, device=DEVICE),

426 )

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

+15

@@ -13,6 +13,12 @@ from vllm.triton_utils import HAS_TRITON, tl, tldevice, triton

13 # attribute is `None`, and `tl.constexpr(...)` would crash at import time.

14 _TL_RAND_MIN = tl.constexpr(4.6566127342e-10) if HAS_TRITON else 4.6566127342e-10

16 +# Offset salt keeping the draft's Gumbel noise disjoint from the target's.

17 +# Verification is a probability-ratio test, not a Gumbel coupling, so a proposal

18 +# and the residual it is resampled from must not share a noise vector.

19 +# Positions are int64 and never approach 2**30, so the streams cannot collide.

20 +_DRAFT_NOISE_SALT = tl.constexpr(1 << 30) if HAS_TRITON else (1 << 30)

21 +

23 @triton.jit

24 def _temperature_kernel(

@@ -89,6 +95,7 @@ def gumbel_noised_argmax(

95 seed,

96 pos,

97 temp,

98 + IS_DRAFTING: tl.constexpr,

99 USE_FP64: tl.constexpr,

100 APPLY_TEMPERATURE: tl.constexpr = True,

101 ):

@@ -108,6 +115,8 @@ def gumbel_noised_argmax(

115 if USE_FP64:

116 logits = logits.to(tl.float64)

117 if temp != 0.0:

118 + if IS_DRAFTING:

119 + pos = pos + _DRAFT_NOISE_SALT

120 gumbel_seed = tl.randint(seed, pos)

121 if USE_FP64:

122 u = tl_rand64(gumbel_seed, keys, includes_zero=False)

@@ -137,6 +146,7 @@ def gumbel_block_argmax(

146 logits_cache_stride_1,

147 logits_cache_col_ptr,

148 vocab_size,

149 + IS_DRAFTING: tl.constexpr,

150 APPLY_TEMPERATURE: tl.constexpr,

151 USE_FP64: tl.constexpr,

152 PER_TOKEN_COL: tl.constexpr = False,

@@ -173,6 +183,7 @@ def gumbel_block_argmax(

183 seed,

184 pos,

185 temp,

186 + IS_DRAFTING=IS_DRAFTING,

187 USE_FP64=USE_FP64,

188 APPLY_TEMPERATURE=APPLY_TEMPERATURE,

189 )

@@ -197,6 +208,7 @@ def _gumbel_sample_kernel(

208 temp_ptr,

209 vocab_size,

210 BLOCK_SIZE: tl.constexpr,

211 + IS_DRAFTING: tl.constexpr,

212 APPLY_TEMPERATURE: tl.constexpr,

213 USE_FP64: tl.constexpr,

214 PER_TOKEN_COL: tl.constexpr,

@@ -226,6 +238,7 @@ def _gumbel_sample_kernel(

238 logits_cache_stride_1,

239 logits_cache_col_ptr,

240 vocab_size,

241 + IS_DRAFTING=IS_DRAFTING,

242 APPLY_TEMPERATURE=APPLY_TEMPERATURE,

243 USE_FP64=USE_FP64,

244 PER_TOKEN_COL=PER_TOKEN_COL,

@@ -242,6 +255,7 @@ def gumbel_sample(

255 seed: torch.Tensor, # [max_num_reqs]

256 pos: torch.Tensor, # [num_tokens]

257 apply_temperature: bool,

258 + is_drafting: bool,

259 logits_cache: torch.Tensor | None = None, # [max_num_reqs, num_cols, vocab_size]

260 logits_cache_col: torch.Tensor | None = None, # scalar or [num_tokens]

261 use_fp64: bool = False,

@@ -281,6 +295,7 @@ def gumbel_sample(

295 temperature,

296 vocab_size,

297 BLOCK_SIZE=BLOCK_SIZE,

298 + IS_DRAFTING=is_drafting,

299 APPLY_TEMPERATURE=apply_temperature,

300 USE_FP64=use_fp64,

301 PER_TOKEN_COL=per_token_col,

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

@@ -317,6 +317,7 @@ class Sampler:

317 self.sampling_states.seeds.gpu,

318 pos,

319 apply_temperature=False,

320 + is_drafting=False,

321 use_fp64=self.use_fp64_gumbel,

322 )

323 return sampled, processed_logits

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

+30 -12

@@ -37,6 +37,9 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

37 self.last_token_indices = torch.zeros(

38 self.max_num_reqs, dtype=torch.int64, device=device

39 )

40 + self.sample_src_positions = torch.zeros(

41 + self.max_num_reqs, dtype=torch.int64, device=device

42 + )

44 self.inputs_embeds: torch.Tensor | None = None

@@ -327,6 +330,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

330 input_batch.seq_lens,

331 num_rejected,

332 self.input_buffers,

333 + self.sample_src_positions,

334 self.max_model_len,

335 self.max_num_reqs,

336 advance_draft_positions=self.advance_draft_positions,

@@ -439,6 +443,10 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

443 ) -> None:

444 last_token_indices = self.last_token_indices[:num_reqs]

445 positions = self.input_buffers.positions[last_token_indices]

446 + # The output hidden state at position P (= positions) and the token id

447 + # at P+1 are used to draft the token at P+2. Sampling keys a draw by the

448 + # position before the sampled token, so the net adjustment is +1.

449 + sample_src_positions = positions + 1

450 idx_mapping = self.idx_mapping[:num_reqs]

451

452 last_hidden_states, hidden_states = self._run_model(

@@ -449,11 +457,11 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

457 cudagraph_runtime_mode=cudagraph_runtime_mode,

458 mm_inputs=mm_inputs,

459 )

452 - sample_hidden_states = last_hidden_states[last_token_indices]

460

461 + sample_hidden_states = last_hidden_states[last_token_indices]

462 self.draft_tokens[:num_reqs, 0] = self.sample_draft(

463 sample_hidden_states,

456 - positions,

464 + sample_src_positions,

465 idx_mapping,

466 self.temperature,

467 self.seeds,

@@ -465,6 +473,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

473 else:

474 self.hidden_states[:num_reqs] = hidden_states[last_token_indices]

475 self.input_buffers.positions[:num_reqs] = positions

476 + self.sample_src_positions[:num_reqs] = sample_src_positions

477

478 def _multi_step_decode(

479 self,

@@ -617,7 +626,6 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

626 self._prepare_eplb_forward(num_reqs)

627

628 idx_mapping = self.idx_mapping[:num_reqs]

620 - positions = self.input_buffers.positions[:num_reqs]

629 # Run the draft model forward pass.

630 last_hidden_states, hidden_states = self._run_model(

631 num_tokens_padded,

@@ -626,18 +634,13 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

634 num_tokens_across_dp,

635 cudagraph_runtime_mode,

636 )

629 - last_hidden_states = last_hidden_states[:num_reqs]

630 -

631 - sample_positions = positions

632 - if not self.advance_draft_positions:

633 - # The forward pass holds positions fixed (Q-only, shared target KV),

634 - # but Gumbel sampling still needs the absolute draft position.

635 - sample_positions = positions + self.current_draft_step

637

638 # Sample the draft tokens.

639 + sample_hidden_states = last_hidden_states[:num_reqs]

640 + sample_src_positions = self.sample_src_positions[:num_reqs]

641 draft_tokens = self.sample_draft(

639 - last_hidden_states,

640 - sample_positions,

642 + sample_hidden_states,

643 + sample_src_positions,

644 idx_mapping,

645 self.temperature,

646 self.seeds,

@@ -653,6 +656,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

656 self.draft_tokens,

657 self.hidden_states,

658 self.input_buffers,

659 + self.sample_src_positions,

660 num_reqs,

661 self.max_model_len,

662 self.num_speculative_steps,

@@ -790,6 +794,7 @@ def _prepare_decode_inputs_kernel(

794 num_rejected_ptr,

795 input_ids_ptr,

796 positions_ptr,

797 + sample_src_positions_ptr,

798 query_start_loc_ptr,

799 seq_lens_ptr,

800 max_model_len,

@@ -818,6 +823,10 @@ def _prepare_decode_inputs_kernel(

823 draft_token = tl.load(draft_tokens_ptr + req_idx * draft_tokens_stride)

824 tl.store(input_ids_ptr + req_idx, draft_token)

825

826 + # Advance the draft sampling key.

827 + sample_position = tl.load(sample_src_positions_ptr + req_idx)

828 + tl.store(sample_src_positions_ptr + req_idx, sample_position + 1)

829 +

830 target_seq_len = tl.load(target_seq_lens_ptr + req_idx)

831 num_rejected = tl.load(num_rejected_ptr + req_idx)

832 seq_len = target_seq_len - num_rejected

@@ -837,6 +846,7 @@ def prepare_decode_inputs(

846 target_seq_lens: torch.Tensor,

847 num_rejected: torch.Tensor,

848 input_buffers: InputBuffers,

849 + sample_src_positions: torch.Tensor,

850 max_model_len: int,

851 max_num_reqs: int,

852 advance_draft_positions: bool = True,

@@ -849,6 +859,7 @@ def prepare_decode_inputs(

859 num_rejected,

860 input_buffers.input_ids,

861 input_buffers.positions,

862 + sample_src_positions,

863 input_buffers.query_start_loc,

864 input_buffers.seq_lens,

865 max_model_len,

@@ -866,6 +877,7 @@ def _update_draft_inputs_kernel(

877 next_input_hidden_states_stride,

878 input_ids_ptr,

879 positions_ptr,

880 + sample_src_positions_ptr,

881 seq_lens_ptr,

882 draft_tokens_ptr,

883 current_draft_step_ptr,

@@ -891,6 +903,10 @@ def _update_draft_inputs_kernel(

903 # This is the final step. Skip updating draft forward inputs.

904 return

905

906 + # Advance the draft sampling key.

907 + sample_position = tl.load(sample_src_positions_ptr + req_idx)

908 + tl.store(sample_src_positions_ptr + req_idx, sample_position + 1)

909 +

910 # Write the sampled draft token into the input ids tensor for the next

911 # forward pass.

912 tl.store(input_ids_ptr + req_idx, draft_token)

@@ -932,6 +948,7 @@ def update_draft_inputs(

948 output_draft_tokens: torch.Tensor,

949 next_input_hidden_states: torch.Tensor,

950 input_buffers: InputBuffers,

951 + sample_src_positions: torch.Tensor,

952 num_reqs: int,

953 max_model_len: int,

954 num_speculative_steps: int,

@@ -945,6 +962,7 @@ def update_draft_inputs(

962 next_input_hidden_states.stride(0),

963 input_buffers.input_ids,

964 input_buffers.positions,

965 + sample_src_positions,

966 input_buffers.seq_lens,

967 draft_tokens,

968 current_draft_step,

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

+3 -3

@@ -261,11 +261,11 @@ class DFlashSpeculator(DraftModelSpeculator):

261 )

262 num_sample = num_reqs * self.num_speculative_steps

263 sample_hidden_states = last_hidden_states[self.sample_indices[:num_sample]]

264 - # sample_pos is the predicted token's position Q; verification keys

265 - # Gumbel by the predecessor (Q-1). sample_draft adds +1, so pass Q-2.

264 + # sample_pos is the predicted token's position P. Sampling keys a draw

265 + # by the position before the sampled token, P-1.

266 draft_tokens = self.sample_draft(

267 sample_hidden_states,

268 - self.sample_pos[:num_sample] - 2,

268 + self.sample_pos[:num_sample] - 1,

269 self.sample_idx_mapping[:num_sample],

270 self.temperature,

271 self.seeds,

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

+5 -3

@@ -51,15 +51,17 @@ def _selector_walk_kernel(

51 other=0,

52 )

54 - # Candidate ids key the noise, matching the target's own sampling.

55 - position = tl.load(sample_pos_ptr + flat) - 1

54 + # sample_pos is the predicted token's position P. Sampling keys a draw

55 + # by the position before the sampled token, P-1.

56 + sample_pos = tl.load(sample_pos_ptr + flat) - 1

57 _, index = gumbel_noised_argmax(

58 scores,

59 candidates,

60 mask & valid,

61 seed,

61 - position,

62 + sample_pos,

63 temperature if SAMPLE_PROBABILISTIC else 0.0,

64 + IS_DRAFTING=True,

65 USE_FP64=USE_FP64,

66 )

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

+3 -2

@@ -135,8 +135,8 @@ class DSparkSpeculator(DFlashSpeculator):

135 buf.index_copy_(1, self._d2t_scatter_index, logits.to(buf.dtype))

136 logits = buf

137

138 - # sample_pos is the predicted token's position Q; the target verifies

139 - # it with the predecessor's Gumbel key (Q-1). Pass Q-1.

138 + # sample_pos is the predicted token's position P. Sampling keys a draw

139 + # by the position before the sampled token, P-1.

140 return gumbel_sample(

141 logits,

142 idx_map,

@@ -144,6 +144,7 @@ class DSparkSpeculator(DFlashSpeculator):

144 self.seeds,

145 sample_pos - 1,

146 apply_temperature=True,

147 + is_drafting=True,

148 logits_cache=self.draft_logits,

149 logits_cache_col=self._step_cols[step],

150 use_fp64=self.use_fp64_gumbel,

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

+8 -3

@@ -379,7 +379,11 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):

379 cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,

380 ) -> None:

381 last_token_indices = self.last_token_indices[:num_reqs]

382 - sample_positions = self.input_buffers.positions[last_token_indices]

382 + positions = self.input_buffers.positions[last_token_indices]

383 + # The output hidden state at position P (= positions) and the token id

384 + # at P+1 are used to draft the token at P+2. Sampling keys a draw by the

385 + # position before the sampled token, so the net adjustment is +1.

386 + sample_src_positions = positions + 1

387 idx_mapping = self.idx_mapping[:num_reqs]

388

389 # Cache the trailing token's ids, hidden states (and embeddings for

@@ -417,7 +421,7 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):

421 sample_hidden_states = last_hidden_states[last_token_indices]

422 draft_tokens = self.sample_draft(

423 sample_hidden_states,

420 - sample_positions,

424 + sample_src_positions,

425 idx_mapping,

426 self.temperature,

427 self.seeds,

@@ -448,7 +452,8 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):

452 idx_mapping,

453 num_reqs,

454 )

451 - sample_positions += 1

455 + # Advance the draft sampling key.

456 + sample_src_positions += 1

457

458

459 @triton.jit

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

@@ -837,6 +837,7 @@ def _resample_kernel(

837 0, # logits_cache_stride_1

838 None, # logits_cache_col_ptr

839 vocab_size,

840 + IS_DRAFTING=False,

841 APPLY_TEMPERATURE=False,

842 USE_FP64=USE_FP64,

843 )

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

+3 -4

@@ -364,7 +364,7 @@ class DraftModelSpeculator(BaseSpeculator):

364 def sample_draft(

365 self,

366 hidden_states: torch.Tensor,

367 - positions: torch.Tensor,

367 + sample_src_positions: torch.Tensor,

368 idx_mapping: torch.Tensor,

369 temperature: torch.Tensor,

370 seeds: torch.Tensor,

@@ -373,15 +373,14 @@ class DraftModelSpeculator(BaseSpeculator):

373 ) -> torch.Tensor:

374 if draft_logits is not None:

375 logits = self.model.compute_logits(hidden_states)

376 - # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise

377 - # used for draft and target sampling.

376 return gumbel_sample(

377 logits,

378 idx_mapping,

379 temperature,

380 seeds,

383 - positions + 1,

381 + sample_src_positions,

382 apply_temperature=True,

383 + is_drafting=True,

384 logits_cache=draft_logits,

385 logits_cache_col=draft_step,

386 use_fp64=self.use_fp64_gumbel,