[Model Runner V2][Spec Decode] Implement block verification for rejection sampling (#46781)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Giancarlo Delfin committed
Jun 30, 2026 at 10:07 UTC
db808b39614384a0349378268a46a1a0feabcec3
4 files changed
+641
-85
tests/v1/spec_decode/test_rejection_sampler_utils.py
+103
@@ -309,3 +309,106 @@ def test_placeholder_draft_token_rejected():
309
assert torch.equal(num_sampled, torch.ones_like(num_sampled))
310
recovered = sampled[:, 0]
311
assert (recovered >= 0).all() and (recovered < VOCAB_SIZE).all()
312
+
313
+
314
+@pytest.mark.parametrize(
315
+ "num_speculative_steps,temperature",
316
+ [
317
+ (1, 0.6),
318
+ (3, 0.6),
319
+ (1, 1.0),
320
+ (3, 1.0),
321
+ (5, 1.0),
322
+ ],
323
+)
324
+@pytest.mark.parametrize("has_draft_logits", [True, False])
325
+def test_block_verification_rejection_sample(
326
+ num_speculative_steps: int, temperature: float, has_draft_logits: bool
327
+):
328
+ """
329
+ Verify that block verification (Sun et al.) preserves the target
330
+ distribution at every accepted position, for both the full draft-logits
331
+ case and the one-hot (no draft logits) case. Block verification changes
332
+ *which* prefix is accepted, but the marginal of every output position must
333
+ still match the target distribution p(x).
334
+ """
335
+
336
+ torch.manual_seed(42)
337
+ device = "cuda"
338
+ num_trials = 10 * VOCAB_SIZE
339
+
340
+ target_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
341
+ draft_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
342
+
343
+ if temperature > 0:
344
+ target_logits_1d /= temperature
345
+ draft_logits_1d /= temperature
346
+
347
+ inputs = _build_rejection_sample_inputs(
348
+ target_logits_1d,
349
+ draft_logits_1d,
350
+ num_speculative_steps,
351
+ temperature=temperature,
352
+ num_trials=num_trials,
353
+ )
354
+ if not has_draft_logits:
355
+ inputs["draft_logits"] = None
356
+
357
+ sampled, num_sampled = rejection_sample(
358
+ **inputs,
359
+ num_speculative_steps=num_speculative_steps,
360
+ use_block_verification=True,
361
+ )
362
+
363
+ target_probs = torch.softmax(target_logits_1d, dim=0)
364
+ for pos in range(num_speculative_steps + 1):
365
+ accepted_mask = num_sampled >= pos + 1
366
+ _assert_distribution_match(
367
+ sampled[accepted_mask, pos], target_probs, device, label=f"position {pos}"
368
+ )
369
+
370
+
371
+@pytest.mark.parametrize("num_speculative_steps", [3, 5])
372
+def test_block_verification_accepts_at_least_as_many(num_speculative_steps: int):
373
+ """
374
+ Block verification is designed to accept at least as long a prefix as
375
+ token verification in expectation. Verify the mean accepted length is no
376
+ worse than the standard method on the same inputs.
377
+ """
378
+
379
+ torch.manual_seed(0)
380
+ device = "cuda"
381
+ num_trials = 20 * VOCAB_SIZE
382
+ temperature = 1.0
383
+
384
+ target_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
385
+ # A draft close to the target gives block verification room to recover
386
+ # prefixes that token verification would have truncated.
387
+ draft_logits_1d = target_logits_1d + 0.5 * torch.randn(
388
+ VOCAB_SIZE, device=device, dtype=torch.float32
389
+ )
390
+
391
+ inputs = _build_rejection_sample_inputs(
392
+ target_logits_1d,
393
+ draft_logits_1d,
394
+ num_speculative_steps,
395
+ temperature=temperature,
396
+ num_trials=num_trials,
397
+ )
398
+
399
+ _, num_sampled_standard = rejection_sample(
400
+ **inputs, num_speculative_steps=num_speculative_steps
401
+ )
402
+ _, num_sampled_block = rejection_sample(
403
+ **inputs,
404
+ num_speculative_steps=num_speculative_steps,
405
+ use_block_verification=True,
406
+ )
407
+
408
+ mean_standard = (num_sampled_standard - 1).float().mean().item()
409
+ mean_block = (num_sampled_block - 1).float().mean().item()
410
+ # Allow a small slack for sampling noise.
411
+ assert mean_block >= mean_standard - 1e-2, (
412
+ f"Block verification mean accepted length {mean_block:.4f} is worse "
413
+ f"than standard {mean_standard:.4f}."
414
+ )
vllm/config/speculative.py
+4
-2
@@ -67,7 +67,7 @@ SpeculativeMethod = Literal[
67
EagleModelTypes,
68
NgramGPUTypes,
69
]
70
-RejectionSampleMethod = Literal["standard", "synthetic"]
70
+RejectionSampleMethod = Literal["standard", "synthetic", "block"]
71
DraftSampleMethod = Literal["greedy", "probabilistic"]
72
73
@@ -201,7 +201,9 @@ class SpeculativeConfig:
201
"""The rejection sampling method to use. 'standard' uses probabilistic
202
rejection sampling (with or without cached draft logits, controlled by
203
draft_sample_method). 'synthetic' accepts draft tokens with a decaying
204
- probability calibrated to synthetic_acceptance_rate."""
204
+ probability calibrated to synthetic_acceptance_rate. 'block' uses block
205
+ verification (Sun et al.), which jointly verifies the draft tokens as a
206
+ block instead of one at a time."""
207
208
synthetic_acceptance_rates: list[float] | None = None
209
"""Per-position *unconditional* acceptance rates for synthetic rejection
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+6
-2
@@ -49,9 +49,10 @@ class RejectionSampler:
49
):
50
self.sampler = sampler
51
self.num_speculative_steps = spec_config.num_speculative_tokens
52
- self.rejection_sample_method = spec_config.rejection_sample_method
52
+ rejection_sample_method = spec_config.rejection_sample_method
53
+ self.use_block_verification: bool = False
54
self.synthetic_conditional_rates: torch.Tensor | None = None
54
- if self.rejection_sample_method == "synthetic":
55
+ if rejection_sample_method == "synthetic":
56
assert spec_config.synthetic_acceptance_rates is not None
57
self.synthetic_conditional_rates = torch.tensor(
58
unconditional_to_conditional_rates(
@@ -60,6 +61,8 @@ class RejectionSampler:
61
dtype=torch.float32,
62
device=device,
63
)
64
+ elif rejection_sample_method == "block":
65
+ self.use_block_verification = True
66
67
def _get_logprobs_tensors(
68
self,
@@ -129,6 +132,7 @@ class RejectionSampler:
132
self.num_speculative_steps,
133
self.synthetic_conditional_rates,
134
use_fp64=self.sampler.use_fp64_gumbel,
135
+ use_block_verification=self.use_block_verification,
136
)
137
logprobs_tensors = self._get_logprobs_tensors(
138
input_batch,
vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
+528
-81
@@ -7,18 +7,18 @@ from vllm.v1.worker.gpu.sample.gumbel import gumbel_block_argmax, tl_rand32
7
8
9
@triton.jit
10
-def _compute_block_max_and_sumexp(logits):
11
- block_max = tl.max(logits, axis=0)
12
- block_sumexp = tl.where(
13
- block_max > float("-inf"),
14
- tl.sum(tl.exp(logits - block_max)),
10
+def _compute_max_and_sumexp(logits):
11
+ max = tl.max(logits, axis=0)
12
+ sumexp = tl.where(
13
+ max > float("-inf"),
14
+ tl.sum(tl.exp(logits - max)),
15
0.0,
16
)
17
- return block_max, block_sumexp
17
+ return max, sumexp
18
19
20
@triton.jit
21
-def _compute_global_lse(
21
+def _compute_global_logsumexp(
22
local_max_ptr,
23
local_max_stride,
24
local_sumexp_ptr,
@@ -45,7 +45,146 @@ def _compute_global_lse(
45
46
47
@triton.jit
48
-def _compute_block_stats_kernel(
48
+def _compute_global_residual_mass(
49
+ local_residual_mass_ptr,
50
+ local_residual_mass_stride,
51
+ prefix_joint_ratio,
52
+ target_logits_ptr,
53
+ target_logits_stride,
54
+ target_local_max_ptr,
55
+ target_local_max_stride,
56
+ target_local_sumexp_ptr,
57
+ target_local_sumexp_stride,
58
+ draft_sampled_ptr,
59
+ logit_idx,
60
+ vocab_num_blocks,
61
+ PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
62
+ HAS_DRAFT_LOGITS: tl.constexpr,
63
+):
64
+ if HAS_DRAFT_LOGITS:
65
+ blocks = tl.arange(0, PADDED_VOCAB_NUM_BLOCKS)
66
+ mask = blocks < vocab_num_blocks
67
+ partials = tl.load(
68
+ local_residual_mass_ptr + logit_idx * local_residual_mass_stride + blocks,
69
+ mask=mask,
70
+ other=0.0,
71
+ )
72
+ return tl.sum(partials, axis=0)
73
+ else:
74
+ # One-hot draft. M_s is a point mass at this draft token
75
+ # so the residual mass reduces to the closed form:
76
+ # p * (1 - M_b(draft_token)).
77
+ draft_token = tl.load(draft_sampled_ptr + logit_idx + 1).to(tl.int64)
78
+ target_lse = _compute_global_logsumexp(
79
+ target_local_max_ptr,
80
+ target_local_max_stride,
81
+ target_local_sumexp_ptr,
82
+ target_local_sumexp_stride,
83
+ logit_idx,
84
+ vocab_num_blocks,
85
+ PADDED_VOCAB_NUM_BLOCKS,
86
+ )
87
+ target_logit = tl.load(
88
+ target_logits_ptr + logit_idx * target_logits_stride + draft_token,
89
+ ).to(tl.float32)
90
+ m_b = tl.exp(target_logit - target_lse)
91
+ return prefix_joint_ratio * (1.0 - m_b)
92
+
93
+
94
+@triton.jit
95
+def _compute_global_target_argmax(
96
+ target_local_max_ptr,
97
+ target_local_max_stride,
98
+ target_local_argmax_ptr,
99
+ target_local_argmax_stride,
100
+ logit_idx,
101
+ vocab_num_blocks,
102
+ PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
103
+):
104
+ blocks = tl.arange(0, PADDED_VOCAB_NUM_BLOCKS)
105
+ blocks_mask = blocks < vocab_num_blocks
106
+ local_max = tl.load(
107
+ target_local_max_ptr + logit_idx * target_local_max_stride + blocks,
108
+ mask=blocks_mask,
109
+ other=float("-inf"),
110
+ )
111
+ max_block_idx = tl.argmax(local_max, axis=0)
112
+ return tl.load(
113
+ target_local_argmax_ptr + logit_idx * target_local_argmax_stride + max_block_idx
114
+ ).to(tl.int64)
115
+
116
+
117
+@triton.jit
118
+def _compute_global_logprobs_and_logsumexp(
119
+ token,
120
+ mask,
121
+ logit_idx,
122
+ req_state_idx,
123
+ draft_step,
124
+ # [num_logits, V]
125
+ target_logits_ptr,
126
+ target_logits_stride,
127
+ # [num_logits, num_blocks]
128
+ target_local_max_ptr,
129
+ target_local_max_stride,
130
+ target_local_sumexp_ptr,
131
+ target_local_sumexp_stride,
132
+ # [max_num_reqs, num_speculative_steps, V]
133
+ draft_logits_ptr,
134
+ draft_logits_stride_0,
135
+ draft_logits_stride_1,
136
+ # [num_logits, num_blocks]
137
+ draft_local_max_ptr,
138
+ draft_local_max_stride,
139
+ draft_local_sumexp_ptr,
140
+ draft_local_sumexp_stride,
141
+ vocab_num_blocks,
142
+ PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
143
+ HAS_DRAFT_LOGITS: tl.constexpr,
144
+):
145
+ target_logit = tl.load(
146
+ target_logits_ptr + logit_idx * target_logits_stride + token,
147
+ mask=mask,
148
+ other=float("-inf"),
149
+ ).to(tl.float32)
150
+ target_lse = _compute_global_logsumexp(
151
+ target_local_max_ptr,
152
+ target_local_max_stride,
153
+ target_local_sumexp_ptr,
154
+ target_local_sumexp_stride,
155
+ logit_idx,
156
+ vocab_num_blocks,
157
+ PADDED_VOCAB_NUM_BLOCKS,
158
+ )
159
+ target_log_prob = target_logit - target_lse
160
+ if HAS_DRAFT_LOGITS:
161
+ draft_logit = tl.load(
162
+ draft_logits_ptr
163
+ + req_state_idx * draft_logits_stride_0
164
+ + draft_step * draft_logits_stride_1
165
+ + token,
166
+ mask=mask,
167
+ other=float("-inf"),
168
+ ).to(tl.float32)
169
+ draft_lse = _compute_global_logsumexp(
170
+ draft_local_max_ptr,
171
+ draft_local_max_stride,
172
+ draft_local_sumexp_ptr,
173
+ draft_local_sumexp_stride,
174
+ logit_idx,
175
+ vocab_num_blocks,
176
+ PADDED_VOCAB_NUM_BLOCKS,
177
+ )
178
+ draft_log_prob = draft_logit - draft_lse
179
+ else:
180
+ # One-hot draft: q(token) = 1, log_q = 0.
181
+ draft_log_prob = 0.0
182
+ draft_lse = 0.0
183
+ return target_log_prob, draft_log_prob, target_lse, draft_lse
184
+
185
+
186
+@triton.jit
187
+def _compute_local_logits_stats_kernel(
188
# [num_logits, num_blocks]
189
target_local_argmax_ptr,
190
target_local_argmax_stride,
@@ -119,7 +258,7 @@ def _compute_block_stats_kernel(
258
mask=mask,
259
other=float("-inf"),
260
).to(tl.float32)
122
- target_max, target_sumexp = _compute_block_max_and_sumexp(target_logits)
261
+ target_max, target_sumexp = _compute_max_and_sumexp(target_logits)
262
tl.store(
263
target_local_max_ptr + logit_idx * target_local_max_stride + block_idx,
264
target_max,
@@ -140,7 +279,7 @@ def _compute_block_stats_kernel(
279
mask=mask,
280
other=float("-inf"),
281
).to(tl.float32)
143
- draft_max, draft_sumexp = _compute_block_max_and_sumexp(draft_logits)
282
+ draft_max, draft_sumexp = _compute_max_and_sumexp(draft_logits)
283
tl.store(
284
draft_local_max_ptr + logit_idx * draft_local_max_stride + block_idx,
285
draft_max,
@@ -153,6 +292,170 @@ def _compute_block_stats_kernel(
292
)
293
294
295
+@triton.jit
296
+def _compute_cumulative_log_p_kernel(
297
+ # [num_logits]
298
+ cumulative_log_p_ptr,
299
+ # [num_logits, V]
300
+ target_logits_ptr,
301
+ target_logits_stride,
302
+ # [num_logits, num_blocks]
303
+ target_local_max_ptr,
304
+ target_local_max_stride,
305
+ # [num_logits, num_blocks]
306
+ target_local_sumexp_ptr,
307
+ target_local_sumexp_stride,
308
+ # [num_logits]
309
+ draft_sampled_ptr,
310
+ # [max_num_reqs, num_speculative_steps, V]
311
+ draft_logits_ptr,
312
+ draft_logits_stride_0,
313
+ draft_logits_stride_1,
314
+ # [num_logits, num_blocks]
315
+ draft_local_max_ptr,
316
+ draft_local_max_stride,
317
+ # [num_logits, num_blocks]
318
+ draft_local_sumexp_ptr,
319
+ draft_local_sumexp_stride,
320
+ # [num_reqs + 1]
321
+ cu_num_logits_ptr,
322
+ # [num_reqs]
323
+ idx_mapping_ptr,
324
+ # [max_num_reqs]
325
+ temp_ptr,
326
+ vocab_num_blocks,
327
+ PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
328
+ HAS_DRAFT_LOGITS: tl.constexpr,
329
+):
330
+ req_idx = tl.program_id(0)
331
+ req_state_idx = tl.load(idx_mapping_ptr + req_idx)
332
+ start_idx = tl.load(cu_num_logits_ptr + req_idx)
333
+ end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
334
+ num_draft_tokens = end_idx - start_idx - 1
335
+ temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
336
+ if temp == 0.0:
337
+ return
338
+
339
+ log_p = 0.0
340
+ for step in range(num_draft_tokens):
341
+ logit_idx = start_idx + step
342
+ draft_token = tl.load(draft_sampled_ptr + logit_idx + 1).to(tl.int64)
343
+ target_logprob, draft_logprob, _, _ = _compute_global_logprobs_and_logsumexp(
344
+ draft_token,
345
+ True, # mask
346
+ logit_idx,
347
+ req_state_idx,
348
+ step,
349
+ target_logits_ptr,
350
+ target_logits_stride,
351
+ target_local_max_ptr,
352
+ target_local_max_stride,
353
+ target_local_sumexp_ptr,
354
+ target_local_sumexp_stride,
355
+ draft_logits_ptr,
356
+ draft_logits_stride_0,
357
+ draft_logits_stride_1,
358
+ draft_local_max_ptr,
359
+ draft_local_max_stride,
360
+ draft_local_sumexp_ptr,
361
+ draft_local_sumexp_stride,
362
+ vocab_num_blocks,
363
+ PADDED_VOCAB_NUM_BLOCKS,
364
+ HAS_DRAFT_LOGITS,
365
+ )
366
+ log_p = tl.minimum(log_p + (target_logprob - draft_logprob), 0.0)
367
+ tl.store(cumulative_log_p_ptr + logit_idx, log_p)
368
+
369
+
370
+@triton.jit
371
+def _compute_local_residual_mass_kernel(
372
+ # [num_logits, num_blocks]
373
+ local_residual_mass_ptr,
374
+ local_residual_mass_stride,
375
+ # [num_logits]
376
+ cumulative_log_p_ptr,
377
+ # [num_logits, V]
378
+ target_logits_ptr,
379
+ target_logits_stride,
380
+ # [num_logits, num_blocks]
381
+ target_local_max_ptr,
382
+ target_local_max_stride,
383
+ # [num_logits, num_blocks]
384
+ target_local_sumexp_ptr,
385
+ target_local_sumexp_stride,
386
+ # [max_num_reqs, num_speculative_steps, V]
387
+ draft_logits_ptr,
388
+ draft_logits_stride_0,
389
+ draft_logits_stride_1,
390
+ # [num_logits, num_blocks]
391
+ draft_local_max_ptr,
392
+ draft_local_max_stride,
393
+ # [num_logits, num_blocks]
394
+ draft_local_sumexp_ptr,
395
+ draft_local_sumexp_stride,
396
+ # [num_logits]
397
+ expanded_idx_mapping_ptr,
398
+ # [num_logits]
399
+ expanded_local_pos_ptr,
400
+ # [max_num_reqs]
401
+ temp_ptr,
402
+ vocab_size,
403
+ num_speculative_steps,
404
+ vocab_num_blocks,
405
+ BLOCK_SIZE: tl.constexpr,
406
+ PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
407
+):
408
+ logit_idx = tl.program_id(0)
409
+ draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx)
410
+ if draft_step_idx == 0 or draft_step_idx >= num_speculative_steps:
411
+ # The acceptance threshold, h, looks one position ahead and sums
412
+ # over: max(p_i * M_b(x|x_{<i}) - M_s(x|x_{<i}), 0). Tokens at the
413
+ # first and last (bonus) positions aren't needed for this computation.
414
+ return
415
+
416
+ req_state_idx = tl.load(expanded_idx_mapping_ptr + logit_idx)
417
+ temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
418
+ if temp == 0.0:
419
+ return
420
+
421
+ block_idx = tl.program_id(1)
422
+ block_offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
423
+ mask = block_offsets < vocab_size
424
+ target_log_probs, draft_log_probs, _, _ = _compute_global_logprobs_and_logsumexp(
425
+ block_offsets,
426
+ mask,
427
+ logit_idx,
428
+ req_state_idx,
429
+ draft_step_idx,
430
+ target_logits_ptr,
431
+ target_logits_stride,
432
+ target_local_max_ptr,
433
+ target_local_max_stride,
434
+ target_local_sumexp_ptr,
435
+ target_local_sumexp_stride,
436
+ draft_logits_ptr,
437
+ draft_logits_stride_0,
438
+ draft_logits_stride_1,
439
+ draft_local_max_ptr,
440
+ draft_local_max_stride,
441
+ draft_local_sumexp_ptr,
442
+ draft_local_sumexp_stride,
443
+ vocab_num_blocks,
444
+ PADDED_VOCAB_NUM_BLOCKS,
445
+ True, # HAS_DRAFT_LOGITS
446
+ )
447
+
448
+ # Compute the residual mass: max(p_i * M_b(x|x_{<i}) - M_s(x|x_{<i}), 0)
449
+ p = tl.exp(tl.load(cumulative_log_p_ptr + logit_idx - 1).to(tl.float32))
450
+ m_b = tl.exp(target_log_probs)
451
+ m_s = tl.exp(draft_log_probs)
452
+ partial = tl.sum(tl.maximum(p * m_b - m_s, 0.0), axis=0)
453
+ tl.store(
454
+ local_residual_mass_ptr + logit_idx * local_residual_mass_stride + block_idx,
455
+ partial,
456
+ )
457
+
458
+
459
@triton.jit
460
def _rejection_kernel(
461
# [num_reqs, num_speculative_steps + 1]
@@ -200,50 +503,78 @@ def _rejection_kernel(
503
pos_ptr,
504
# [num_speculative_steps]
505
synthetic_conditional_rates_ptr,
506
+ # [num_logits]
507
+ cumulative_log_p_ptr,
508
+ # [num_logits, num_blocks]
509
+ local_residual_mass_ptr,
510
+ local_residual_mass_stride,
511
vocab_num_blocks,
512
PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
513
HAS_DRAFT_LOGITS: tl.constexpr,
514
SYNTHETIC_MODE: tl.constexpr,
515
+ USE_BLOCK_VERIFICATION: tl.constexpr,
516
):
517
req_idx = tl.program_id(0)
518
req_state_idx = tl.load(idx_mapping_ptr + req_idx).to(tl.int64)
519
start_idx = tl.load(cu_num_logits_ptr + req_idx).to(tl.int64)
520
end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
212
- num_tokens = end_idx - start_idx
521
+ num_draft_tokens = end_idx - start_idx - 1
522
seed = tl.load(seed_ptr + req_state_idx)
523
temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
524
+ is_greedy = temp == 0.0
525
216
- rejected_step = 0
526
+ accepted_length = tl.zeros((), tl.int64)
527
target_lse = 0.0
528
draft_lse = 0.0
529
accepted = True
220
- for i in range(num_tokens - 1):
221
- if accepted:
222
- logit_idx = start_idx + i
223
- draft_sampled = tl.load(draft_sampled_ptr + logit_idx + 1).to(tl.int64)
224
- if temp == 0.0:
530
+ for i in range(num_draft_tokens):
531
+ logit_idx = start_idx + i
532
+ draft_sampled = tl.load(draft_sampled_ptr + logit_idx + 1).to(tl.int64)
533
+ pos = tl.load(pos_ptr + logit_idx)
534
+ u = tl_rand32(seed, pos, includes_zero=False)
535
+ if USE_BLOCK_VERIFICATION and not is_greedy:
536
+ # Block verification (Sun et al., 2024): https://arxiv.org/abs/2403.10444
537
+ prefix_joint_ratio = tl.exp(
538
+ tl.load(cumulative_log_p_ptr + logit_idx).to(tl.float32)
539
+ )
540
+ if i < num_draft_tokens - 1:
541
+ residual_mass = _compute_global_residual_mass(
542
+ local_residual_mass_ptr,
543
+ local_residual_mass_stride,
544
+ prefix_joint_ratio,
545
+ target_logits_ptr,
546
+ target_logits_stride,
547
+ target_local_max_ptr,
548
+ target_local_max_stride,
549
+ target_local_sumexp_ptr,
550
+ target_local_sumexp_stride,
551
+ draft_sampled_ptr,
552
+ logit_idx + 1,
553
+ vocab_num_blocks,
554
+ PADDED_VOCAB_NUM_BLOCKS,
555
+ HAS_DRAFT_LOGITS,
556
+ )
557
+ denom = residual_mass + 1.0 - prefix_joint_ratio
558
+ h = tl.where(denom > 0.0, residual_mass / denom, 1.0)
559
+ else:
560
+ h = prefix_joint_ratio
561
+ accepted_length = tl.where(u <= h, i + 1, accepted_length)
562
+ tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled)
563
+ elif accepted:
564
+ if is_greedy:
565
# Greedy sampling. Accept IFF draft matches target argmax.
566
# NOTE: Target argmax is stored directly so that resampling
567
# can be skipped upon rejection.
228
- target_blocks = tl.arange(0, PADDED_VOCAB_NUM_BLOCKS)
229
- target_blocks_mask = target_blocks < vocab_num_blocks
230
- target_local_max = tl.load(
231
- target_local_max_ptr
232
- + logit_idx * target_local_max_stride
233
- + target_blocks,
234
- mask=target_blocks_mask,
235
- other=float("-inf"),
568
+ target_argmax = _compute_global_target_argmax(
569
+ target_local_max_ptr,
570
+ target_local_max_stride,
571
+ target_local_argmax_ptr,
572
+ target_local_argmax_stride,
573
+ logit_idx,
574
+ vocab_num_blocks,
575
+ PADDED_VOCAB_NUM_BLOCKS,
576
)
237
- max_target_block_idx = tl.argmax(target_local_max, axis=0)
238
- target_argmax = tl.load(
239
- target_local_argmax_ptr
240
- + logit_idx * target_local_argmax_stride
241
- + max_target_block_idx
242
- ).to(tl.int64)
243
-
577
if SYNTHETIC_MODE:
245
- pos = tl.load(pos_ptr + logit_idx)
246
- u = tl_rand32(seed, pos, includes_zero=False)
578
rate = tl.load(synthetic_conditional_rates_ptr + i)
579
# -1 is used for padded draft token ids that should be rejected.
580
accepted &= (u < rate) & (draft_sampled >= 0)
@@ -254,57 +585,70 @@ def _rejection_kernel(
585
draft_sampled if accepted else target_argmax,
586
)
587
else:
588
+ # Speculative decoding (Leviathan et al., 2023): https://arxiv.org/abs/2211.17192
589
# -1 is used for padded draft token ids that should be rejected.
590
is_valid_draft = draft_sampled >= 0
591
# Avoid possible OOB ptr access.
592
draft_sampled = tl.maximum(0, draft_sampled)
261
- target_logit = tl.load(
262
- target_logits_ptr + logit_idx * target_logits_stride + draft_sampled
263
- ).to(tl.float32)
264
- target_lse = _compute_global_lse(
265
- target_local_max_ptr,
266
- target_local_max_stride,
267
- target_local_sumexp_ptr,
268
- target_local_sumexp_stride,
269
- logit_idx,
270
- vocab_num_blocks,
271
- PADDED_VOCAB_NUM_BLOCKS,
272
- )
273
- target_log_prob = target_logit - target_lse
274
- pos = tl.load(pos_ptr + logit_idx)
275
- u = tl_rand32(seed, pos, includes_zero=False)
276
- if HAS_DRAFT_LOGITS:
277
- draft_logit = tl.load(
278
- draft_logits_ptr
279
- + req_state_idx * draft_logits_stride_0
280
- + i * draft_logits_stride_1
281
- + draft_sampled
282
- ).to(tl.float32)
283
- draft_lse = _compute_global_lse(
593
+ target_logprob, draft_logprob, target_lse, draft_lse = (
594
+ _compute_global_logprobs_and_logsumexp(
595
+ draft_sampled,
596
+ True, # mask
597
+ logit_idx,
598
+ req_state_idx,
599
+ i,
600
+ target_logits_ptr,
601
+ target_logits_stride,
602
+ target_local_max_ptr,
603
+ target_local_max_stride,
604
+ target_local_sumexp_ptr,
605
+ target_local_sumexp_stride,
606
+ draft_logits_ptr,
607
+ draft_logits_stride_0,
608
+ draft_logits_stride_1,
609
draft_local_max_ptr,
610
draft_local_max_stride,
611
draft_local_sumexp_ptr,
612
draft_local_sumexp_stride,
288
- logit_idx,
613
vocab_num_blocks,
614
PADDED_VOCAB_NUM_BLOCKS,
615
+ HAS_DRAFT_LOGITS,
616
)
292
- draft_log_prob = draft_logit - draft_lse
293
- else:
294
- # One-hot draft: q(draft_token) = 1, log_q = 0.
295
- draft_log_prob = 0
296
-
617
+ )
618
if SYNTHETIC_MODE:
619
rate = tl.load(synthetic_conditional_rates_ptr + i)
620
accepted &= u < rate
621
else:
622
# Probability ratio test: p(x) > u * q(x)
623
# Equivalent log form: log_p(x) > log(u) + log_q(x)
303
- accepted &= target_log_prob > tl.log(u) + draft_log_prob
624
+ accepted &= target_logprob > tl.log(u) + draft_logprob
625
accepted &= is_valid_draft
626
tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled)
306
- rejected_step += accepted
307
- tl.store(rejected_steps_ptr + req_idx, rejected_step)
627
+ accepted_length += accepted
628
+ tl.store(rejected_steps_ptr + req_idx, accepted_length)
629
+ if USE_BLOCK_VERIFICATION and not is_greedy and accepted_length < num_draft_tokens:
630
+ # Compute the target and draft log exponential sums for the
631
+ # rejected token.
632
+ rejected_idx = start_idx + accepted_length
633
+ target_lse = _compute_global_logsumexp(
634
+ target_local_max_ptr,
635
+ target_local_max_stride,
636
+ target_local_sumexp_ptr,
637
+ target_local_sumexp_stride,
638
+ rejected_idx,
639
+ vocab_num_blocks,
640
+ PADDED_VOCAB_NUM_BLOCKS,
641
+ )
642
+ if HAS_DRAFT_LOGITS:
643
+ draft_lse = _compute_global_logsumexp(
644
+ draft_local_max_ptr,
645
+ draft_local_max_stride,
646
+ draft_local_sumexp_ptr,
647
+ draft_local_sumexp_stride,
648
+ rejected_idx,
649
+ vocab_num_blocks,
650
+ PADDED_VOCAB_NUM_BLOCKS,
651
+ )
652
tl.store(target_rejected_logsumexp_ptr + req_idx, target_lse)
653
tl.store(draft_rejected_logsumexp_ptr + req_idx, draft_lse)
654
@@ -342,10 +686,13 @@ def _resample_kernel(
686
seed_ptr,
687
# [num_logits]
688
pos_ptr,
689
+ # [num_logits]
690
+ cumulative_log_p_ptr,
691
vocab_size,
692
BLOCK_SIZE: tl.constexpr,
693
HAS_DRAFT_LOGITS: tl.constexpr,
694
USE_FP64: tl.constexpr,
695
+ USE_BLOCK_VERIFICATION: tl.constexpr,
696
):
697
req_idx = tl.program_id(0)
698
resample_idx = tl.load(rejected_step_ptr + req_idx)
@@ -386,6 +733,17 @@ def _resample_kernel(
733
target_lse = tl.load(target_rejected_logsumexp_ptr + req_idx)
734
draft_lse = tl.load(draft_rejected_logsumexp_ptr + req_idx)
735
target_log_probs = target_logits - target_lse
736
+ if USE_BLOCK_VERIFICATION:
737
+ # Block residual is:
738
+ # max(p_tau * M_b(x) - M_s(x), 0) / Z.
739
+ # Scale the target logprobs by log(p_tau). p_0 = 1, so skip
740
+ # shifting when nothing was accepted (tau == 0).
741
+ log_p_tau = 0.0
742
+ if resample_idx > 0:
743
+ log_p_tau = tl.load(cumulative_log_p_ptr + resample_token_idx - 1).to(
744
+ tl.float32
745
+ )
746
+ target_log_probs += log_p_tau
747
draft_log_probs = draft_logits - draft_lse
748
# Compute the residual:
749
# r(x) = max(p(x) - q(x), 0)
@@ -402,6 +760,11 @@ def _resample_kernel(
760
else:
761
# One-hot draft. The residual is just the target distribution with
762
# the rejected draft token probability zeroed out.
763
+ # NOTE: During block verification, the residual becomes:
764
+ # 0 if x == rejected_draft_token
765
+ # p_tau * M_b(x) / Z otherwise
766
+ # Therefore p_tau is a constant that cancels under normalization,
767
+ # and does not need to be applied.
768
rejected_draft_token = tl.load(draft_sampled_ptr + resample_token_idx + 1)
769
residual_logits = tl.where(
770
block != rejected_draft_token,
@@ -523,18 +886,20 @@ def rejection_sample(
886
# [num_speculative_steps]
887
synthetic_conditional_rates: torch.Tensor | None = None,
888
use_fp64: bool = False,
889
+ use_block_verification: bool = False,
890
) -> tuple[torch.Tensor, torch.Tensor]:
891
num_reqs = cu_num_logits.shape[0] - 1
892
num_logits, vocab_size = target_logits.shape
529
- has_draft_logits = draft_logits is not None
530
-
531
- if draft_logits is None:
532
- # When draft_logits is None, create a dummy tensor so that Triton
533
- # kernel signatures receive valid pointers/strides. The kernels
534
- # will never read from it when HAS_DRAFT_LOGITS=False.
535
- draft_logits = target_logits.new_empty(1, 1, 1)
893
+ draft_logits_stride_0 = 0
894
+ draft_logits_stride_1 = 0
895
+ if has_draft_logits := draft_logits is not None:
896
+ draft_logits_stride_0 = draft_logits.stride(0)
897
+ draft_logits_stride_1 = draft_logits.stride(1)
898
+ # In some cases (e.g. MiMo v2.5 Pro + DFlash) the target model's
899
+ # vocab size is larger than the draft's due to padding.
900
+ vocab_size = min(vocab_size, draft_logits.size(-1))
901
537
- # Compute the block-level logits stats, such as target argmax
902
+ # Compute the per-vocab-block logits stats, such as target argmax
903
# (for greedy requests), and target max + softmax exponential
904
# (for non-greedy requests).
905
VOCAB_BLOCK_SIZE = 8192
@@ -555,7 +920,7 @@ def rejection_sample(
920
draft_local_sumexp = target_logits.new_empty(
921
num_logits, vocab_num_blocks, dtype=torch.float32
922
)
558
- _compute_block_stats_kernel[(num_logits, vocab_num_blocks)](
923
+ _compute_local_logits_stats_kernel[(num_logits, vocab_num_blocks)](
924
target_local_argmax,
925
target_local_argmax.stride(0),
926
target_local_max,
@@ -569,8 +934,8 @@ def rejection_sample(
934
target_logits,
935
target_logits.stride(0),
936
draft_logits,
572
- draft_logits.stride(0),
573
- draft_logits.stride(1),
937
+ draft_logits_stride_0,
938
+ draft_logits_stride_1,
939
expanded_idx_mapping,
940
expanded_local_pos,
941
temperature,
@@ -580,6 +945,82 @@ def rejection_sample(
945
HAS_DRAFT_LOGITS=has_draft_logits,
946
)
947
948
+ # Precompute the running joint ratio and residual mass for block
949
+ # verification.
950
+ if use_block_verification:
951
+ assert synthetic_conditional_rates is None, (
952
+ "Block verification is incompatible with synthetic acceptance rates."
953
+ )
954
+
955
+ # Compute the log of the running joint ratio, p_i.
956
+ # cumulative_log_p[start + i] = log(p_{i+1}), the cumulative ratio after
957
+ # the (i+1)-th draft token.
958
+ cumulative_log_p = target_logits.new_empty(num_logits, dtype=torch.float32)
959
+ _compute_cumulative_log_p_kernel[(num_reqs,)](
960
+ cumulative_log_p,
961
+ target_logits,
962
+ target_logits.stride(0),
963
+ target_local_max,
964
+ target_local_max.stride(0),
965
+ target_local_sumexp,
966
+ target_local_sumexp.stride(0),
967
+ draft_sampled,
968
+ draft_logits,
969
+ draft_logits_stride_0,
970
+ draft_logits_stride_1,
971
+ draft_local_max,
972
+ draft_local_max.stride(0),
973
+ draft_local_sumexp,
974
+ draft_local_sumexp.stride(0),
975
+ cu_num_logits,
976
+ idx_mapping,
977
+ temperature,
978
+ vocab_num_blocks,
979
+ PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks,
980
+ HAS_DRAFT_LOGITS=has_draft_logits,
981
+ num_warps=1,
982
+ )
983
+
984
+ # Compute the per-vocab-block partials of the residual mass, later reduced
985
+ # to the total by _compute_global_residual_mass. Only launched for full
986
+ # draft logits distributions. One-hot drafts used a closed-form residual
987
+ # mass instead.
988
+ if has_draft_logits:
989
+ local_residual_mass = target_logits.new_empty(
990
+ num_logits, vocab_num_blocks, dtype=torch.float32
991
+ )
992
+ _compute_local_residual_mass_kernel[(num_logits, vocab_num_blocks)](
993
+ local_residual_mass,
994
+ local_residual_mass.stride(0),
995
+ cumulative_log_p,
996
+ target_logits,
997
+ target_logits.stride(0),
998
+ target_local_max,
999
+ target_local_max.stride(0),
1000
+ target_local_sumexp,
1001
+ target_local_sumexp.stride(0),
1002
+ draft_logits,
1003
+ draft_logits_stride_0,
1004
+ draft_logits_stride_1,
1005
+ draft_local_max,
1006
+ draft_local_max.stride(0),
1007
+ draft_local_sumexp,
1008
+ draft_local_sumexp.stride(0),
1009
+ expanded_idx_mapping,
1010
+ expanded_local_pos,
1011
+ temperature,
1012
+ vocab_size,
1013
+ num_speculative_steps,
1014
+ vocab_num_blocks,
1015
+ BLOCK_SIZE=VOCAB_BLOCK_SIZE,
1016
+ PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks,
1017
+ )
1018
+ else:
1019
+ local_residual_mass = None
1020
+ else:
1021
+ cumulative_log_p = None
1022
+ local_residual_mass = None
1023
+
1024
# Sample up until the first rejected/bonus token, and store
1025
# the step.
1026
sampled = draft_sampled.new_empty(
@@ -604,8 +1045,8 @@ def rejection_sample(
1045
target_local_sumexp.stride(0),
1046
draft_sampled,
1047
draft_logits,
607
- draft_logits.stride(0),
608
- draft_logits.stride(1),
1048
+ draft_logits_stride_0,
1049
+ draft_logits_stride_1,
1050
draft_local_max,
1051
draft_local_max.stride(0),
1052
draft_local_sumexp,
@@ -616,10 +1057,14 @@ def rejection_sample(
1057
seed,
1058
pos,
1059
synthetic_conditional_rates,
1060
+ cumulative_log_p,
1061
+ local_residual_mass,
1062
+ local_residual_mass.stride(0) if local_residual_mass is not None else 0,
1063
vocab_num_blocks,
1064
PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks,
1065
HAS_DRAFT_LOGITS=has_draft_logits,
1066
SYNTHETIC_MODE=synthetic_conditional_rates is not None,
1067
+ USE_BLOCK_VERIFICATION=use_block_verification,
1068
num_warps=1,
1069
)
1070
@@ -644,8 +1089,8 @@ def rejection_sample(
1089
target_logits.stride(0),
1090
target_rejected_logsumexp,
1091
draft_logits,
647
- draft_logits.stride(0),
648
- draft_logits.stride(1),
1092
+ draft_logits_stride_0,
1093
+ draft_logits_stride_1,
1094
draft_rejected_logsumexp,
1095
num_sampled,
1096
cu_num_logits,
@@ -654,10 +1099,12 @@ def rejection_sample(
1099
temperature,
1100
seed,
1101
pos,
1102
+ cumulative_log_p,
1103
vocab_size,
1104
BLOCK_SIZE=RESAMPLE_BLOCK_SIZE,
1105
HAS_DRAFT_LOGITS=has_draft_logits,
1106
USE_FP64=use_fp64,
1107
+ USE_BLOCK_VERIFICATION=use_block_verification,
1108
)
1109
1110
# Insert the resampled tokens into the output sampled.