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

[Model Runner V2][Spec Decode] Fix draft logits cache column stride in gumbel_sample (#53017)

Giancarlo Delfin committed Aug 19, 2026 at 17:39 UTC d4f4d3f40fc5350a71777fcb0e5eb8a57bda631f

3 files changed +88 -6

tests/v1/worker/test_gpu_gumbel_sample.py

+69

@@ -291,3 +291,72 @@ def test_logits_cache_stores_input_logits_bitwise(

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

292 "cached logits differ from the input logits"

293 )

294 +

295 +

296 +@pytest.mark.parametrize("extra_cache_cols", [0, 1])

297 +def test_logits_cache_columns_stay_separate_across_steps(extra_cache_cols: int):

298 + """Each drafting step must land in its own cache column.

299 +

300 + `extra_cache_cols=1` is the shape a draft produces when it adds an

301 + input-only mask/noise row to its embedding table but keeps a full-width

302 + output head: N-wide logits cached into N+1-wide rows. Striding cache

303 + columns by the logits width instead of the cache's own row width leaves

304 + step 0 correct and silently misaligns every step after it.

305 + """

306 + torch.manual_seed(0)

307 + num_reqs, vocab_size, num_steps = 4, 1031, 3

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

309 + temp = torch.ones(num_reqs, dtype=torch.float32, device=DEVICE)

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

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

312 +

313 + cache = torch.zeros(

314 + num_reqs, num_steps, vocab_size + extra_cache_cols, device=DEVICE

315 + )

316 + cols = torch.arange(num_steps, dtype=torch.int32, device=DEVICE)

317 + per_step = [

318 + torch.randn(num_reqs, vocab_size, device=DEVICE) for _ in range(num_steps)

319 + ]

320 +

321 + for step, logits in enumerate(per_step):

322 + gumbel_sample(

323 + logits,

324 + idx_mapping,

325 + temp,

326 + seed,

327 + pos,

328 + apply_temperature=True,

329 + logits_cache=cache,

330 + logits_cache_col=cols[step],

331 + )

332 +

333 + for step, logits in enumerate(per_step):

334 + stored = cache[:, step, :vocab_size]

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

336 + f"step {step} was overwritten by a later step"

337 + )

338 + # Columns past the sampled width belong to no step and stay untouched.

339 + assert not cache[:, :, vocab_size:].any()

340 +

341 +

342 +def test_logits_cache_narrower_than_logits_is_rejected():

343 + """A cache too narrow to hold a step would silently drop its tail."""

344 + num_reqs, vocab_size, num_steps = 2, 64, 3

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

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

347 + temp = torch.ones(num_reqs, dtype=torch.float32, device=DEVICE)

348 + seed = torch.zeros(num_reqs, dtype=torch.int64, device=DEVICE)

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

350 + cache = torch.zeros(num_reqs, num_steps, vocab_size - 1, device=DEVICE)

351 +

352 + with pytest.raises(AssertionError, match="narrower"):

353 + gumbel_sample(

354 + logits,

355 + idx_mapping,

356 + temp,

357 + seed,

358 + pos,

359 + apply_temperature=True,

360 + logits_cache=cache,

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

362 + )

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

+17 -5

@@ -91,8 +91,10 @@ def gumbel_block_argmax(

91 temp_ptr,

92 seeds_ptr,

93 pos_ptr,

94 + # [max_num_reqs, num_cols, vocab_size]

95 logits_cache_ptr,

95 - logits_cache_stride,

96 + logits_cache_stride_0,

97 + logits_cache_stride_1,

98 logits_cache_col_ptr,

99 vocab_size,

100 APPLY_TEMPERATURE: tl.constexpr,

@@ -115,8 +117,8 @@ def gumbel_block_argmax(

117 col = tl.load(logits_cache_col_ptr)

118 tl.store(

119 logits_cache_ptr

118 - + req_state_idx * logits_cache_stride

119 - + col * vocab_size

120 + + req_state_idx * logits_cache_stride_0

121 + + col * logits_cache_stride_1

122 + block,

123 logits,

124 mask=mask & is_valid_req,

@@ -164,8 +166,10 @@ def _gumbel_sample_kernel(

166 local_argmax_stride,

167 local_max_ptr,

168 local_max_stride,

169 + # [max_num_reqs, num_cols, vocab_size]

170 logits_cache_ptr,

168 - logits_cache_stride,

171 + logits_cache_stride_0,

172 + logits_cache_stride_1,

173 logits_cache_col_ptr,

174 logits_ptr,

175 logits_stride,

@@ -200,7 +204,8 @@ def _gumbel_sample_kernel(

204 seeds_ptr,

205 pos_ptr,

206 logits_cache_ptr,

203 - logits_cache_stride,

207 + logits_cache_stride_0,

208 + logits_cache_stride_1,

209 logits_cache_col_ptr,

210 vocab_size,

211 APPLY_TEMPERATURE=APPLY_TEMPERATURE,

@@ -229,6 +234,12 @@ def gumbel_sample(

234 if logits_cache_col is not None:

235 logits_cache_col = logits_cache_col.contiguous()

236 num_tokens, vocab_size = logits.shape

237 + if logits_cache is not None:

238 + assert logits_cache.size(-1) >= vocab_size, (

239 + f"draft logits cache vocab dim ({logits_cache.size(-1)}) is narrower "

240 + f"than the sampled logits ({vocab_size}). Cached logits would be "

241 + "truncated."

242 + )

243 BLOCK_SIZE = 1024

244 num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)

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

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

253 local_max.stride(0),

254 logits_cache,

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

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

257 logits_cache_col,

258 logits,

259 logits.stride(0),

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

+2 -1

@@ -833,7 +833,8 @@ def _resample_kernel(

833 seed_ptr,

834 pos_ptr,

835 None, # logits_cache_ptr

836 - 0, # logits_cache_stride

836 + 0, # logits_cache_stride_0

837 + 0, # logits_cache_stride_1

838 None, # logits_cache_col_ptr

839 vocab_size,

840 APPLY_TEMPERATURE=False,