@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 81bc196913

[Model Runner V2] Cache draft logits in model's LM head dtype (#50910)

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

Giancarlo Delfin committed Aug 5, 2026 at 19:12 UTC 81bc196913653c9e06b957e39a606002a09db171

6 files changed +170 -76

tests/v1/spec_decode/test_rejection_sampler_utils.py

+21 -7

@@ -25,6 +25,12 @@ def _build_rejection_sample_inputs(

25 temperature: float,

26 num_trials: int,

27 ) -> dict:

28 + """Build rejection_sample kwargs from a fixed target and draft distribution.

29 +

30 + target_logits_1d must already have temperature applied (the sampler applies

31 + sampling params before verification), whereas draft_logits_1d must not:

32 + rejection_sample divides the draft logits by the temperature on load.

33 + """

34 device = target_logits_1d.device

35 vocab_size = target_logits_1d.shape[0]

36 K = num_speculative_steps

@@ -35,7 +41,10 @@ def _build_rejection_sample_inputs(

41 draft_logits_1d.view(1, 1, vocab_size).expand(num_trials, K, -1).contiguous()

42 )

38 - draft_probs = torch.softmax(draft_logits_1d, dim=0)

44 + scaled_draft_logits_1d = draft_logits_1d.float()

45 + if temperature > 0:

46 + scaled_draft_logits_1d = scaled_draft_logits_1d / temperature

47 + draft_probs = torch.softmax(scaled_draft_logits_1d, dim=0)

48 draft_tokens = torch.multinomial(

49 draft_probs.expand(num_trials, -1), K, replacement=True

50 )

@@ -138,7 +147,10 @@ def _assert_distribution_match(

147 (3, 1.0),

148 ],

149 )

141 -def test_stochastic_rejection_sample(num_speculative_steps: int, temperature: float):

150 +@pytest.mark.parametrize("draft_logits_dtype", [torch.float32, torch.bfloat16])

151 +def test_stochastic_rejection_sample(

152 + num_speculative_steps: int, temperature: float, draft_logits_dtype: torch.dtype

153 +):

154 """

155 Verify that rejection sampling produces the target distribution.

156 This is done by simulating many independent trials of speculative

@@ -146,6 +158,9 @@ def test_stochastic_rejection_sample(num_speculative_steps: int, temperature: fl

158 run rejection sample on all of the trials (requests), and verify

159 that the sampled tokens at every position follow the target

160 distribution p(x).

161 +

162 + Parametrized over the draft-logits dtype: storing them in the draft head's

163 + dtype must not bias the output distribution.

164 """

165

166 torch.manual_seed(42)

@@ -153,11 +168,12 @@ def test_stochastic_rejection_sample(num_speculative_steps: int, temperature: fl

168 num_trials = 10 * VOCAB_SIZE

169

170 target_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)

156 - draft_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)

171 + draft_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32).to(

172 + draft_logits_dtype

173 + )

174

175 if temperature > 0:

176 target_logits_1d /= temperature

160 - draft_logits_1d /= temperature

177

178 inputs = _build_rejection_sample_inputs(

179 target_logits_1d,

@@ -251,7 +267,6 @@ def test_synthetic_rejection_sample(

267

268 if temperature > 0:

269 target_logits_1d /= temperature

254 - draft_logits_1d /= temperature

270

271 inputs = _build_rejection_sample_inputs(

272 target_logits_1d,

@@ -334,7 +349,7 @@ def test_placeholder_draft_token_rejected():

349 temperature = 0.6

350

351 target_logits_1d = torch.randn(VOCAB_SIZE, device=device) / temperature

337 - draft_logits_1d = torch.randn(VOCAB_SIZE, device=device) / temperature

352 + draft_logits_1d = torch.randn(VOCAB_SIZE, device=device)

353

354 inputs = _build_rejection_sample_inputs(

355 target_logits_1d,

@@ -383,7 +398,6 @@ def test_block_verification_rejection_sample(

398

399 if temperature > 0:

400 target_logits_1d /= temperature

386 - draft_logits_1d /= temperature

401

402 inputs = _build_rejection_sample_inputs(

403 target_logits_1d,

tests/v1/worker/test_gpu_gumbel_sample.py

+66

@@ -225,3 +225,69 @@ def test_vocab_size_not_multiple_of_block(vocab_size: int):

225 df = vocab_size - 1

226 if df >= 1:

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

228 +

229 +

230 +# ----------------------------- Logits cache --------------------------------

231 +

232 +

233 +def _float_bits(t: torch.Tensor) -> torch.Tensor:

234 + """Reinterpret floats as same-width ints, so equality is truly bitwise."""

235 + int_dtype = torch.int32 if t.dtype == torch.float32 else torch.int16

236 + return t.view(int_dtype)

237 +

238 +

239 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])

240 +@pytest.mark.parametrize("per_token_col", [False, True])

241 +def test_logits_cache_stores_input_logits_bitwise(

242 + dtype: torch.dtype, per_token_col: bool

243 +):

244 + """`logits_cache` must receive the input logits, pre-temperature and bit-exact.

245 +

246 + The cache is kept in the head's dtype, which is only lossless because the

247 + stored value is the input logit itself. Storing `logit / temp` instead would

248 + generally not be representable there, and the rejection sampler -- which

249 + divides by the same temperature on load -- would then verify against a `q`

250 + the draft never sampled from. Temperatures 0.0 and 1.0 are included because

251 + both make the divide a no-op and would mask such a bug.

252 +

253 + Both column-addressing modes are covered: a 0-d column (one draft step per

254 + call, as MTP/EAGLE do) and a [num_tokens] column (as DFlash does, sampling

255 + every step in one call).

256 + """

257 + torch.manual_seed(0)

258 + num_reqs, vocab_size, num_steps = 8, 4099, 3

259 + logits = torch.randn(num_reqs, vocab_size, device=DEVICE, dtype=dtype)

260 + idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE)

261 + temp = torch.tensor(

262 + [0.0, 0.1, 0.5, 1.0, 1.5, 2.0, 0.7, 1.0], dtype=torch.float32, device=DEVICE

263 + )

264 + seed = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE)

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

266 +

267 + cache = torch.zeros(num_reqs, num_steps, vocab_size, device=DEVICE, dtype=dtype)

268 + if per_token_col:

269 + # Each token lands in its own column, cycling through the steps.

270 + cols = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE) % num_steps

271 + else:

272 + cols = torch.tensor(1, dtype=torch.int32, device=DEVICE)

273 +

274 + gumbel_sample(

275 + logits,

276 + idx_mapping,

277 + temp,

278 + seed,

279 + pos,

280 + apply_temperature=True,

281 + logits_cache=cache,

282 + logits_cache_col=cols,

283 + )

284 +

285 + if per_token_col:

286 + stored = cache[torch.arange(num_reqs, device=DEVICE), cols.long()]

287 + else:

288 + stored = cache[:, 1]

289 + # Untouched columns must stay untouched.

290 + assert not cache[:, 0].any() and not cache[:, 2].any()

291 + assert torch.equal(_float_bits(stored), _float_bits(logits)), (

292 + "cached logits differ from the input logits"

293 + )

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

+33 -36

@@ -91,9 +91,9 @@ def gumbel_block_argmax(

91 temp_ptr,

92 seeds_ptr,

93 pos_ptr,

94 - processed_logits_ptr,

95 - processed_logits_stride,

96 - processed_logits_col_ptr,

94 + logits_cache_ptr,

95 + logits_cache_stride,

96 + logits_cache_col_ptr,

97 vocab_size,

98 APPLY_TEMPERATURE: tl.constexpr,

99 USE_FP64: tl.constexpr,

@@ -104,30 +104,30 @@ def gumbel_block_argmax(

104 temp = tl.load(temp_ptr + req_state_idx, mask=is_valid_req, other=0.0).to(

105 tl.float32

106 )

107 - if temp != 0.0 and APPLY_TEMPERATURE:

108 - # Apply temperature.

109 - # NOTE(woosuk): Match the behavior of _temperature_kernel.

110 - # E.g., if the kernel uses tl.div_rn, we should use tl.div_rn here too.

111 - logits = logits / temp

112 -

113 - if processed_logits_ptr is not None:

114 - # Store the temperature-applied logits.

115 - if processed_logits_col_ptr is not None:

116 - if PER_TOKEN_COL:

117 - col = tl.load(processed_logits_col_ptr + token_idx)

118 - else:

119 - col = tl.load(processed_logits_col_ptr)

107 + if logits_cache_ptr is not None:

108 + # Store the logits *before* temperature. Dividing first would produce a

109 + # value that is generally not representable in the cache's dtype, forcing

110 + # it to be fp32. Consumers (the rejection sampler) divide by the same

111 + # temperature on load, which reproduces the value used below bitwise.

112 + if PER_TOKEN_COL:

113 + col = tl.load(logits_cache_col_ptr + token_idx)

114 else:

121 - col = 0

115 + col = tl.load(logits_cache_col_ptr)

116 tl.store(

123 - processed_logits_ptr

124 - + req_state_idx * processed_logits_stride

117 + logits_cache_ptr

118 + + req_state_idx * logits_cache_stride

119 + col * vocab_size

120 + block,

121 logits,

122 mask=mask & is_valid_req,

123 )

124

125 + if temp != 0.0 and APPLY_TEMPERATURE:

126 + # Apply temperature.

127 + # NOTE(woosuk): Match the behavior of _temperature_kernel.

128 + # E.g., if the kernel uses tl.div_rn, we should use tl.div_rn here too.

129 + logits = logits / temp

130 +

131 # fp32 is the default reduction dtype; fp64 is ~1/32–1/64x the throughput

132 # on H100/Ada/Blackwell and empirically indistinguishable for Gumbel-max.

133 if USE_FP64:

@@ -164,9 +164,9 @@ def _gumbel_sample_kernel(

164 local_argmax_stride,

165 local_max_ptr,

166 local_max_stride,

167 - processed_logits_ptr,

168 - processed_logits_stride,

169 - processed_logits_col_ptr,

167 + logits_cache_ptr,

168 + logits_cache_stride,

169 + logits_cache_col_ptr,

170 logits_ptr,

171 logits_stride,

172 expanded_idx_mapping_ptr,

@@ -199,9 +199,9 @@ def _gumbel_sample_kernel(

199 temp_ptr,

200 seeds_ptr,

201 pos_ptr,

202 - processed_logits_ptr,

203 - processed_logits_stride,

204 - processed_logits_col_ptr,

202 + logits_cache_ptr,

203 + logits_cache_stride,

204 + logits_cache_col_ptr,

205 vocab_size,

206 APPLY_TEMPERATURE=APPLY_TEMPERATURE,

207 USE_FP64=USE_FP64,

@@ -219,33 +219,30 @@ def gumbel_sample(

219 seed: torch.Tensor, # [max_num_reqs]

220 pos: torch.Tensor, # [num_tokens]

221 apply_temperature: bool,

222 - output_processed_logits: torch.Tensor | None = None,

223 - output_processed_logits_col: torch.Tensor | None = None,

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

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

224 use_fp64: bool = False,

225 ) -> torch.Tensor:

226 # Enforce contiguity on non-strided input tensors

227 expanded_idx_mapping = expanded_idx_mapping.contiguous()

228 pos = pos.contiguous()

229 - if output_processed_logits_col is not None:

230 - output_processed_logits_col = output_processed_logits_col.contiguous()

229 + if logits_cache_col is not None:

230 + logits_cache_col = logits_cache_col.contiguous()

231 num_tokens, vocab_size = logits.shape

232 BLOCK_SIZE = 1024

233 num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)

234 local_argmax = logits.new_empty(num_tokens, num_blocks, dtype=torch.int64)

235 local_max_dtype = torch.float64 if use_fp64 else torch.float32

236 local_max = logits.new_empty(num_tokens, num_blocks, dtype=local_max_dtype)

237 - per_token_col = (

238 - output_processed_logits_col is not None

239 - and output_processed_logits_col.dim() > 0

240 - )

237 + per_token_col = logits_cache_col is not None and logits_cache_col.dim() > 0

238 _gumbel_sample_kernel[(num_tokens, num_blocks)](

239 local_argmax,

240 local_argmax.stride(0),

241 local_max,

242 local_max.stride(0),

246 - output_processed_logits,

247 - output_processed_logits.stride(0) if output_processed_logits is not None else 0,

248 - output_processed_logits_col,

243 + logits_cache,

244 + logits_cache.stride(0) if logits_cache is not None else 0,

245 + logits_cache_col,

246 logits,

247 logits.stride(0),

248 expanded_idx_mapping,

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

+2 -2

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

127 self.seeds,

128 sample_pos - 1,

129 apply_temperature=True,

130 - output_processed_logits=self.draft_logits,

131 - output_processed_logits_col=self._step_cols[step],

130 + logits_cache=self.draft_logits,

131 + logits_cache_col=self._step_cols[step],

132 use_fp64=self.use_fp64_gumbel,

133 )

134

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

+44 -28

@@ -123,6 +123,7 @@ def _compute_global_logprobs_and_logsumexp(

123 logit_idx,

124 req_state_idx,

125 draft_step,

126 + temp,

127 # [num_logits, V]

128 target_logits_ptr,

129 target_logits_stride,

@@ -160,14 +161,18 @@ def _compute_global_logprobs_and_logsumexp(

161 )

162 target_log_prob = target_logit - target_lse

163 if HAS_DRAFT_LOGITS:

163 - draft_logit = tl.load(

164 - draft_logits_ptr

165 - + req_state_idx * draft_logits_stride_0

166 - + draft_step * draft_logits_stride_1

167 - + token,

168 - mask=mask,

169 - other=float("-inf"),

170 - ).to(tl.float32)

164 + # draft_logits is stored pre-temperature, so apply scale first.

165 + draft_logit = (

166 + tl.load(

167 + draft_logits_ptr

168 + + req_state_idx * draft_logits_stride_0

169 + + draft_step * draft_logits_stride_1

170 + + token,

171 + mask=mask,

172 + other=float("-inf"),

173 + ).to(tl.float32)

174 + / temp

175 + )

176 draft_lse = _compute_global_logsumexp(

177 draft_local_max_ptr,

178 draft_local_max_stride,

@@ -272,15 +277,19 @@ def _compute_local_logits_stats_kernel(

277 target_sumexp,

278 )

279 if HAS_DRAFT_LOGITS:

275 - # Get local draft max and summed exponentials.

276 - draft_logits = tl.load(

277 - draft_logits_ptr

278 - + req_state_idx * draft_logits_stride_0

279 - + draft_step_idx * draft_logits_stride_1

280 - + block_offsets,

281 - mask=mask,

282 - other=float("-inf"),

283 - ).to(tl.float32)

280 + # Get local draft max and summed exponentials. draft_logits is

281 + # stored pre-temperature, so apply scale first.

282 + draft_logits = (

283 + tl.load(

284 + draft_logits_ptr

285 + + req_state_idx * draft_logits_stride_0

286 + + draft_step_idx * draft_logits_stride_1

287 + + block_offsets,

288 + mask=mask,

289 + other=float("-inf"),

290 + ).to(tl.float32)

291 + / temp

292 + )

293 draft_max, draft_sumexp = _compute_max_and_sumexp(draft_logits)

294 tl.store(

295 draft_local_max_ptr + logit_idx * draft_local_max_stride + block_idx,

@@ -348,6 +357,7 @@ def _compute_cumulative_log_p_kernel(

357 logit_idx,

358 req_state_idx,

359 step,

360 + temp,

361 target_logits_ptr,

362 target_logits_stride,

363 target_local_max_ptr,

@@ -429,6 +439,7 @@ def _compute_local_residual_mass_kernel(

439 logit_idx,

440 req_state_idx,

441 draft_step_idx,

442 + temp,

443 target_logits_ptr,

444 target_logits_stride,

445 target_local_max_ptr,

@@ -599,6 +610,7 @@ def _rejection_kernel(

610 logit_idx,

611 req_state_idx,

612 i,

613 + temp,

614 target_logits_ptr,

615 target_logits_stride,

616 target_local_max_ptr,

@@ -724,14 +736,18 @@ def _resample_kernel(

736 # Bonus token (no rejections). Directly use the target logits.

737 residual_logits = target_logits

738 elif HAS_DRAFT_LOGITS:

727 - draft_logits = tl.load(

728 - draft_logits_ptr

729 - + req_state_idx * draft_logits_stride_0

730 - + resample_idx * draft_logits_stride_1

731 - + block,

732 - mask=mask,

733 - other=float("-inf"),

734 - ).to(tl.float32)

739 + # draft_logits is stored pre-temperature, so apply scale first.

740 + draft_logits = (

741 + tl.load(

742 + draft_logits_ptr

743 + + req_state_idx * draft_logits_stride_0

744 + + resample_idx * draft_logits_stride_1

745 + + block,

746 + mask=mask,

747 + other=float("-inf"),

748 + ).to(tl.float32)

749 + / temp

750 + )

751 target_lse = tl.load(target_rejected_logsumexp_ptr + req_idx)

752 draft_lse = tl.load(draft_rejected_logsumexp_ptr + req_idx)

753 target_log_probs = target_logits - target_lse

@@ -784,9 +800,9 @@ def _resample_kernel(

800 temp_ptr,

801 seed_ptr,

802 pos_ptr,

787 - None, # processed_logits_ptr

788 - 0, # processed_logits_stride

789 - None, # processed_logits_col_ptr

803 + None, # logits_cache_ptr

804 + 0, # logits_cache_stride

805 + None, # logits_cache_col_ptr

806 vocab_size,

807 APPLY_TEMPERATURE=False,

808 USE_FP64=USE_FP64,

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

+4 -3

@@ -131,11 +131,12 @@ class DraftModelSpeculator(BaseSpeculator):

131

132 self.draft_logits: torch.Tensor | None = None

133 if self.speculative_config.draft_sample_method == "probabilistic":

134 + # Pre-temperature logits, cached from the previous decode step.

135 self.draft_logits = torch.zeros(

136 self.max_num_reqs,

137 self.num_speculative_steps,

138 self.vocab_size,

138 - dtype=torch.float32,

139 + dtype=vllm_config.model_config.head_dtype,

140 device=device,

141 )

142

@@ -332,8 +333,8 @@ class DraftModelSpeculator(BaseSpeculator):

333 seeds,

334 positions + 1,

335 apply_temperature=True,

335 - output_processed_logits=draft_logits,

336 - output_processed_logits_col=draft_step,

336 + logits_cache=draft_logits,

337 + logits_cache_col=draft_step,

338 use_fp64=self.use_fp64_gumbel,

339 )

340 return self._greedy_sample_draft(hidden_states)