@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 1,325 lines 44.7 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import torch
4
5 from vllm.triton_utils import tl, tldevice, triton
6 from vllm.v1.worker.gpu.sample.gumbel import gumbel_block_argmax, tl_rand32
7 from vllm.v1.worker.gpu.sample.watermark import philox_gumbel_block_argmax
8
9
10 @triton.jit
11 def _compute_max_and_sumexp(logits):
12 max = tl.max(logits, axis=0)
13 sumexp = tl.where(
14 max > float("-inf"),
15 tl.sum(tl.exp(logits - max)),
16 0.0,
17 )
18 return max, sumexp
19
20
21 @triton.jit
22 def _compute_global_logsumexp(
23 local_max_ptr,
24 local_max_stride,
25 local_sumexp_ptr,
26 local_sumexp_stride,
27 logit_idx,
28 vocab_num_blocks,
29 PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
30 ):
31 blocks = tl.arange(0, PADDED_VOCAB_NUM_BLOCKS)
32 blocks_mask = blocks < vocab_num_blocks
33 maxes = tl.load(
34 local_max_ptr + logit_idx * local_max_stride + blocks,
35 mask=blocks_mask,
36 other=float("-inf"),
37 )
38 sumexps = tl.load(
39 local_sumexp_ptr + logit_idx * local_sumexp_stride + blocks,
40 mask=blocks_mask,
41 other=0.0,
42 )
43 global_max = tl.max(maxes, axis=0)
44 global_lse = global_max + tl.log(tl.sum(sumexps * tl.exp(maxes - global_max)))
45 return global_lse
46
47
48 @triton.jit
49 def _compute_global_residual_mass(
50 local_residual_mass_ptr,
51 local_residual_mass_stride,
52 prefix_joint_ratio,
53 target_logits_ptr,
54 target_logits_stride,
55 target_local_max_ptr,
56 target_local_max_stride,
57 target_local_sumexp_ptr,
58 target_local_sumexp_stride,
59 draft_token,
60 logit_idx,
61 vocab_num_blocks,
62 PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
63 HAS_DRAFT_LOGITS: tl.constexpr,
64 ):
65 if HAS_DRAFT_LOGITS:
66 blocks = tl.arange(0, PADDED_VOCAB_NUM_BLOCKS)
67 mask = blocks < vocab_num_blocks
68 partials = tl.load(
69 local_residual_mass_ptr + logit_idx * local_residual_mass_stride + blocks,
70 mask=mask,
71 other=0.0,
72 )
73 return tl.sum(partials, axis=0)
74 else:
75 # One-hot draft. M_s is a point mass at this draft token
76 # so the residual mass reduces to the closed form:
77 # p * (1 - M_b(draft_token)).
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 # See _insert_resampled_kernel: NaN breaks tl.argmax index bounds.
112 local_max = tl.where(local_max != local_max, float("-inf"), local_max)
113 max_block_idx = tl.argmax(local_max, axis=0)
114 return tl.load(
115 target_local_argmax_ptr + logit_idx * target_local_argmax_stride + max_block_idx
116 ).to(tl.int64)
117
118
119 @triton.jit
120 def _compute_global_logprobs_and_logsumexp(
121 token,
122 mask,
123 logit_idx,
124 req_state_idx,
125 draft_step,
126 temp,
127 # [num_logits, V]
128 target_logits_ptr,
129 target_logits_stride,
130 # [num_logits, num_blocks]
131 target_local_max_ptr,
132 target_local_max_stride,
133 target_local_sumexp_ptr,
134 target_local_sumexp_stride,
135 # [max_num_reqs, num_speculative_steps, V]
136 draft_logits_ptr,
137 draft_logits_stride_0,
138 draft_logits_stride_1,
139 # [num_logits, num_blocks]
140 draft_local_max_ptr,
141 draft_local_max_stride,
142 draft_local_sumexp_ptr,
143 draft_local_sumexp_stride,
144 vocab_num_blocks,
145 PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
146 HAS_DRAFT_LOGITS: tl.constexpr,
147 ):
148 target_logit = tl.load(
149 target_logits_ptr + logit_idx * target_logits_stride + token,
150 mask=mask,
151 other=float("-inf"),
152 ).to(tl.float32)
153 target_lse = _compute_global_logsumexp(
154 target_local_max_ptr,
155 target_local_max_stride,
156 target_local_sumexp_ptr,
157 target_local_sumexp_stride,
158 logit_idx,
159 vocab_num_blocks,
160 PADDED_VOCAB_NUM_BLOCKS,
161 )
162 target_log_prob = target_logit - target_lse
163 if HAS_DRAFT_LOGITS:
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,
179 draft_local_sumexp_ptr,
180 draft_local_sumexp_stride,
181 logit_idx,
182 vocab_num_blocks,
183 PADDED_VOCAB_NUM_BLOCKS,
184 )
185 draft_log_prob = draft_logit - draft_lse
186 else:
187 # One-hot draft: q(token) = 1, log_q = 0.
188 draft_log_prob = 0.0
189 draft_lse = 0.0
190 return target_log_prob, draft_log_prob, target_lse, draft_lse
191
192
193 @triton.jit
194 def _compute_local_logits_stats_kernel(
195 # [num_logits, num_blocks]
196 target_local_argmax_ptr,
197 target_local_argmax_stride,
198 # [num_logits, num_blocks]
199 target_local_max_ptr,
200 target_local_max_stride,
201 # [num_logits, num_blocks]
202 target_local_sumexp_ptr,
203 target_local_sumexp_stride,
204 # [num_logits, num_blocks]
205 draft_local_max_ptr,
206 draft_local_max_stride,
207 # [num_logits, num_blocks]
208 draft_local_sumexp_ptr,
209 draft_local_sumexp_stride,
210 # [num_logits, V]
211 target_logits_ptr,
212 target_logits_stride,
213 # [max_num_reqs, num_speculative_steps, V]
214 draft_logits_ptr,
215 draft_logits_stride_0,
216 draft_logits_stride_1,
217 # [num_logits]
218 expanded_idx_mapping_ptr,
219 # [num_logits]
220 expanded_local_pos_ptr,
221 # [max_num_reqs]
222 temp_ptr,
223 vocab_size,
224 num_speculative_steps,
225 BLOCK_SIZE: tl.constexpr,
226 HAS_DRAFT_LOGITS: tl.constexpr,
227 ):
228 logit_idx = tl.program_id(0).to(tl.int64)
229 draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx)
230
231 if draft_step_idx >= num_speculative_steps:
232 # Bonus token. Max/argmax and summed exponentials are not needed.
233 return
234
235 req_state_idx = tl.load(expanded_idx_mapping_ptr + logit_idx).to(tl.int64)
236 temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
237
238 block_idx = tl.program_id(1)
239 block_offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
240 mask = block_offsets < vocab_size
241
242 if temp == 0.0:
243 # Greedy sampling. Only the target max/argmax are needed.
244 target_logits = tl.load(
245 target_logits_ptr + logit_idx * target_logits_stride + block_offsets,
246 mask=mask,
247 other=float("-inf"),
248 ).to(tl.float32)
249 value, idx = tl.max(target_logits, axis=0, return_indices=True)
250 token_id = block_idx * BLOCK_SIZE + idx
251 tl.store(
252 target_local_argmax_ptr
253 + logit_idx * target_local_argmax_stride
254 + block_idx,
255 token_id,
256 )
257 tl.store(
258 target_local_max_ptr + logit_idx * target_local_max_stride + block_idx,
259 value,
260 )
261 else:
262 # Get local target max and summed exponentials.
263 target_logits = tl.load(
264 target_logits_ptr + logit_idx * target_logits_stride + block_offsets,
265 mask=mask,
266 other=float("-inf"),
267 ).to(tl.float32)
268 target_max, target_sumexp = _compute_max_and_sumexp(target_logits)
269 tl.store(
270 target_local_max_ptr + logit_idx * target_local_max_stride + block_idx,
271 target_max,
272 )
273 tl.store(
274 target_local_sumexp_ptr
275 + logit_idx * target_local_sumexp_stride
276 + block_idx,
277 target_sumexp,
278 )
279 if HAS_DRAFT_LOGITS:
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,
296 draft_max,
297 )
298 tl.store(
299 draft_local_sumexp_ptr
300 + logit_idx * draft_local_sumexp_stride
301 + block_idx,
302 draft_sumexp,
303 )
304
305
306 @triton.jit
307 def _compute_cumulative_log_p_kernel(
308 # [num_logits]
309 cumulative_log_p_ptr,
310 # [num_logits, V]
311 target_logits_ptr,
312 target_logits_stride,
313 # [num_logits, num_blocks]
314 target_local_max_ptr,
315 target_local_max_stride,
316 # [num_logits, num_blocks]
317 target_local_sumexp_ptr,
318 target_local_sumexp_stride,
319 # [num_logits]
320 draft_sampled_ptr,
321 # [max_num_reqs, num_speculative_steps, V]
322 draft_logits_ptr,
323 draft_logits_stride_0,
324 draft_logits_stride_1,
325 # [num_logits, num_blocks]
326 draft_local_max_ptr,
327 draft_local_max_stride,
328 # [num_logits, num_blocks]
329 draft_local_sumexp_ptr,
330 draft_local_sumexp_stride,
331 # [num_reqs + 1]
332 cu_num_logits_ptr,
333 # [num_reqs]
334 idx_mapping_ptr,
335 # [max_num_reqs]
336 temp_ptr,
337 vocab_num_blocks,
338 PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
339 HAS_DRAFT_LOGITS: tl.constexpr,
340 ):
341 req_idx = tl.program_id(0)
342 req_state_idx = tl.load(idx_mapping_ptr + req_idx).to(tl.int64)
343 start_idx = tl.load(cu_num_logits_ptr + req_idx).to(tl.int64)
344 end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
345 num_draft_tokens = end_idx - start_idx - 1
346 temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
347 if temp == 0.0:
348 return
349
350 log_p = tl.zeros((), tl.float32)
351 for step in range(num_draft_tokens):
352 logit_idx = start_idx + step
353 draft_token = tl.load(draft_sampled_ptr + logit_idx + 1).to(tl.int64)
354 # -1 placeholder tokens can never be accepted. Skip their reductions
355 # and carry the last valid cumulative value.
356 if draft_token >= 0:
357 target_logprob, draft_logprob, _, _ = (
358 _compute_global_logprobs_and_logsumexp(
359 draft_token,
360 True, # mask
361 logit_idx,
362 req_state_idx,
363 step,
364 temp,
365 target_logits_ptr,
366 target_logits_stride,
367 target_local_max_ptr,
368 target_local_max_stride,
369 target_local_sumexp_ptr,
370 target_local_sumexp_stride,
371 draft_logits_ptr,
372 draft_logits_stride_0,
373 draft_logits_stride_1,
374 draft_local_max_ptr,
375 draft_local_max_stride,
376 draft_local_sumexp_ptr,
377 draft_local_sumexp_stride,
378 vocab_num_blocks,
379 PADDED_VOCAB_NUM_BLOCKS,
380 HAS_DRAFT_LOGITS,
381 )
382 )
383 log_p = tl.minimum(log_p + (target_logprob - draft_logprob), 0.0)
384 tl.store(cumulative_log_p_ptr + logit_idx, log_p)
385
386
387 @triton.jit
388 def _compute_local_residual_mass_kernel(
389 # [num_logits, num_blocks]
390 local_residual_mass_ptr,
391 local_residual_mass_stride,
392 # [num_logits]
393 cumulative_log_p_ptr,
394 # [num_logits, V]
395 target_logits_ptr,
396 target_logits_stride,
397 # [num_logits, num_blocks]
398 target_local_max_ptr,
399 target_local_max_stride,
400 # [num_logits, num_blocks]
401 target_local_sumexp_ptr,
402 target_local_sumexp_stride,
403 # [max_num_reqs, num_speculative_steps, V]
404 draft_logits_ptr,
405 draft_logits_stride_0,
406 draft_logits_stride_1,
407 # [num_logits, num_blocks]
408 draft_local_max_ptr,
409 draft_local_max_stride,
410 # [num_logits, num_blocks]
411 draft_local_sumexp_ptr,
412 draft_local_sumexp_stride,
413 # [num_logits]
414 draft_sampled_ptr,
415 # [num_logits]
416 expanded_idx_mapping_ptr,
417 # [num_logits]
418 expanded_local_pos_ptr,
419 # [max_num_reqs]
420 temp_ptr,
421 vocab_size,
422 num_speculative_steps,
423 vocab_num_blocks,
424 BLOCK_SIZE: tl.constexpr,
425 PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
426 ):
427 logit_idx = tl.program_id(0).to(tl.int64)
428 draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx)
429 if draft_step_idx == 0 or draft_step_idx >= num_speculative_steps:
430 # The acceptance threshold, h, looks one position ahead and sums
431 # over: max(p_i * M_b(x|x_{<i}) - M_s(x|x_{<i}), 0). Tokens at the
432 # first and last (bonus) positions aren't needed for this computation.
433 return
434
435 if tl.load(draft_sampled_ptr + logit_idx + 1) < 0:
436 # -1 placeholder token. The rejection kernel treats the preceding token
437 # as the end of the block, so this position's residual mass is unused.
438 return
439
440 req_state_idx = tl.load(expanded_idx_mapping_ptr + logit_idx).to(tl.int64)
441 temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
442 if temp == 0.0:
443 return
444
445 block_idx = tl.program_id(1)
446 block_offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
447 mask = block_offsets < vocab_size
448 target_log_probs, draft_log_probs, _, _ = _compute_global_logprobs_and_logsumexp(
449 block_offsets,
450 mask,
451 logit_idx,
452 req_state_idx,
453 draft_step_idx,
454 temp,
455 target_logits_ptr,
456 target_logits_stride,
457 target_local_max_ptr,
458 target_local_max_stride,
459 target_local_sumexp_ptr,
460 target_local_sumexp_stride,
461 draft_logits_ptr,
462 draft_logits_stride_0,
463 draft_logits_stride_1,
464 draft_local_max_ptr,
465 draft_local_max_stride,
466 draft_local_sumexp_ptr,
467 draft_local_sumexp_stride,
468 vocab_num_blocks,
469 PADDED_VOCAB_NUM_BLOCKS,
470 True, # HAS_DRAFT_LOGITS
471 )
472
473 # Compute the residual mass: max(p_i * M_b(x|x_{<i}) - M_s(x|x_{<i}), 0)
474 p = tl.exp(tl.load(cumulative_log_p_ptr + logit_idx - 1).to(tl.float32))
475 m_b = tl.exp(target_log_probs)
476 m_s = tl.exp(draft_log_probs)
477 partial = tl.sum(tl.maximum(p * m_b - m_s, 0.0), axis=0)
478 tl.store(
479 local_residual_mass_ptr + logit_idx * local_residual_mass_stride + block_idx,
480 partial,
481 )
482
483
484 @triton.jit
485 def _rejection_kernel(
486 # [num_reqs, num_speculative_steps + 1]
487 sampled_ptr,
488 sampled_stride,
489 # [num_reqs]
490 rejected_steps_ptr,
491 # [num_reqs]
492 target_rejected_logsumexp_ptr,
493 # [num_reqs]
494 draft_rejected_logsumexp_ptr,
495 # [num_logits, V]
496 target_logits_ptr,
497 target_logits_stride,
498 # [num_logits, num_blocks]
499 target_local_argmax_ptr,
500 target_local_argmax_stride,
501 # [num_logits, num_blocks]
502 target_local_max_ptr,
503 target_local_max_stride,
504 # [num_logits, num_blocks]
505 target_local_sumexp_ptr,
506 target_local_sumexp_stride,
507 # [num_logits]
508 draft_sampled_ptr,
509 # [max_num_reqs, num_speculative_steps, V]
510 draft_logits_ptr,
511 draft_logits_stride_0,
512 draft_logits_stride_1,
513 # [num_logits, num_blocks]
514 draft_local_max_ptr,
515 draft_local_max_stride,
516 # [num_logits, num_blocks]
517 draft_local_sumexp_ptr,
518 draft_local_sumexp_stride,
519 # [num_reqs + 1]
520 cu_num_logits_ptr,
521 # [num_reqs]
522 idx_mapping_ptr,
523 # [max_num_reqs]
524 temp_ptr,
525 # [max_num_reqs]
526 seed_ptr,
527 # [num_logits]
528 pos_ptr,
529 # [num_speculative_steps]
530 synthetic_conditional_rates_ptr,
531 # [num_logits]
532 cumulative_log_p_ptr,
533 # [num_logits, num_blocks]
534 local_residual_mass_ptr,
535 local_residual_mass_stride,
536 vocab_num_blocks,
537 PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
538 HAS_DRAFT_LOGITS: tl.constexpr,
539 SYNTHETIC_MODE: tl.constexpr,
540 USE_BLOCK_VERIFICATION: tl.constexpr,
541 ):
542 req_idx = tl.program_id(0)
543 req_state_idx = tl.load(idx_mapping_ptr + req_idx).to(tl.int64)
544 start_idx = tl.load(cu_num_logits_ptr + req_idx).to(tl.int64)
545 end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
546 num_draft_tokens = end_idx - start_idx - 1
547 seed = tl.load(seed_ptr + req_state_idx)
548 temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
549 is_greedy = temp == 0.0
550
551 accepted_length = tl.zeros((), tl.int64)
552 target_lse = 0.0
553 draft_lse = 0.0
554 verifying = True
555 for i in range(num_draft_tokens):
556 logit_idx = start_idx + i
557 draft_sampled = tl.load(draft_sampled_ptr + logit_idx + 1).to(tl.int64)
558 # -1 is used for placeholder draft token ids that should be rejected.
559 is_valid_draft = draft_sampled >= 0
560 # Avoid possible OOB ptr access.
561 draft_sampled = tl.maximum(0, draft_sampled)
562 if not is_greedy:
563 # A -1 placeholder ends verification. Greedy is excluded because it
564 # stores the target argmax upon first rejection, so it rejects the
565 # placeholder via `accepted` instead.
566 verifying &= is_valid_draft
567 if verifying:
568 pos = tl.load(pos_ptr + logit_idx)
569 u = tl_rand32(seed, pos, includes_zero=False)
570 if is_greedy:
571 # Greedy sampling. Accept IFF draft matches target argmax.
572 # NOTE: Target argmax is stored directly so that resampling
573 # can be skipped upon rejection.
574 target_argmax = _compute_global_target_argmax(
575 target_local_max_ptr,
576 target_local_max_stride,
577 target_local_argmax_ptr,
578 target_local_argmax_stride,
579 logit_idx,
580 vocab_num_blocks,
581 PADDED_VOCAB_NUM_BLOCKS,
582 )
583 if SYNTHETIC_MODE:
584 rate = tl.load(synthetic_conditional_rates_ptr + i)
585 accepted = u < rate
586 else:
587 accepted = target_argmax == draft_sampled
588 accepted &= is_valid_draft
589 verifying = accepted
590 accepted_length += accepted
591 tl.store(
592 sampled_ptr + req_idx * sampled_stride + i,
593 draft_sampled if accepted else target_argmax,
594 )
595 elif USE_BLOCK_VERIFICATION:
596 # Block verification (Sun et al., 2024): https://arxiv.org/abs/2403.10444
597 prefix_joint_ratio = tl.exp(
598 tl.load(cumulative_log_p_ptr + logit_idx).to(tl.float32)
599 )
600 next_draft_token = tl.load(
601 draft_sampled_ptr + logit_idx + 2,
602 mask=i < num_draft_tokens - 1,
603 other=-1,
604 ).to(tl.int64)
605 if next_draft_token >= 0:
606 residual_mass = _compute_global_residual_mass(
607 local_residual_mass_ptr,
608 local_residual_mass_stride,
609 prefix_joint_ratio,
610 target_logits_ptr,
611 target_logits_stride,
612 target_local_max_ptr,
613 target_local_max_stride,
614 target_local_sumexp_ptr,
615 target_local_sumexp_stride,
616 next_draft_token,
617 logit_idx + 1,
618 vocab_num_blocks,
619 PADDED_VOCAB_NUM_BLOCKS,
620 HAS_DRAFT_LOGITS,
621 )
622 denom = residual_mass + 1.0 - prefix_joint_ratio
623 h = tl.where(denom > 0.0, residual_mass / denom, 1.0)
624 else:
625 h = prefix_joint_ratio
626 accepted_length = tl.where(u <= h, i + 1, accepted_length)
627 tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled)
628 else:
629 # Speculative decoding (Leviathan et al., 2023): https://arxiv.org/abs/2211.17192
630 target_logprob, draft_logprob, target_lse, draft_lse = (
631 _compute_global_logprobs_and_logsumexp(
632 draft_sampled,
633 True, # mask
634 logit_idx,
635 req_state_idx,
636 i,
637 temp,
638 target_logits_ptr,
639 target_logits_stride,
640 target_local_max_ptr,
641 target_local_max_stride,
642 target_local_sumexp_ptr,
643 target_local_sumexp_stride,
644 draft_logits_ptr,
645 draft_logits_stride_0,
646 draft_logits_stride_1,
647 draft_local_max_ptr,
648 draft_local_max_stride,
649 draft_local_sumexp_ptr,
650 draft_local_sumexp_stride,
651 vocab_num_blocks,
652 PADDED_VOCAB_NUM_BLOCKS,
653 HAS_DRAFT_LOGITS,
654 )
655 )
656 if SYNTHETIC_MODE:
657 rate = tl.load(synthetic_conditional_rates_ptr + i)
658 accepted = u < rate
659 else:
660 # Probability ratio test: p(x) > u * q(x)
661 # Equivalent log form: log_p(x) > log(u) + log_q(x)
662 accepted = target_logprob > tl.log(u) + draft_logprob
663 verifying = accepted
664 accepted_length += accepted
665 tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled)
666
667 tl.store(rejected_steps_ptr + req_idx, accepted_length)
668 if USE_BLOCK_VERIFICATION and not is_greedy and accepted_length < num_draft_tokens:
669 # Compute the target and draft log exponential sums for the
670 # rejected token.
671 rejected_idx = start_idx + accepted_length
672 target_lse = _compute_global_logsumexp(
673 target_local_max_ptr,
674 target_local_max_stride,
675 target_local_sumexp_ptr,
676 target_local_sumexp_stride,
677 rejected_idx,
678 vocab_num_blocks,
679 PADDED_VOCAB_NUM_BLOCKS,
680 )
681 if HAS_DRAFT_LOGITS:
682 draft_lse = _compute_global_logsumexp(
683 draft_local_max_ptr,
684 draft_local_max_stride,
685 draft_local_sumexp_ptr,
686 draft_local_sumexp_stride,
687 rejected_idx,
688 vocab_num_blocks,
689 PADDED_VOCAB_NUM_BLOCKS,
690 )
691 tl.store(target_rejected_logsumexp_ptr + req_idx, target_lse)
692 tl.store(draft_rejected_logsumexp_ptr + req_idx, draft_lse)
693
694
695 @triton.jit
696 def _seeded_resample_argmax(
697 residual_logits,
698 block,
699 mask,
700 resample_token_idx,
701 expanded_idx_mapping_ptr,
702 temp_ptr,
703 seed_ptr,
704 pos_ptr,
705 vocab_size,
706 USE_FP64: tl.constexpr,
707 ):
708 return gumbel_block_argmax(
709 residual_logits,
710 block,
711 mask,
712 resample_token_idx,
713 expanded_idx_mapping_ptr,
714 temp_ptr,
715 seed_ptr,
716 pos_ptr,
717 None, # logits_cache_ptr
718 0, # logits_cache_stride_0
719 0, # logits_cache_stride_1
720 None, # logits_cache_col_ptr
721 None, # logits_cache_source_ptr
722 0, # logits_cache_source_stride
723 vocab_size,
724 IS_DRAFTING=False,
725 APPLY_TEMPERATURE=False,
726 USE_FP64=USE_FP64,
727 )
728
729
730 @triton.jit(do_not_specialize=["watermark_key_0", "watermark_key_1"])
731 def _resample_kernel(
732 # [num_reqs, num_blocks]
733 resampled_local_argmax_ptr,
734 resampled_local_argmax_stride,
735 # [num_reqs, num_blocks]
736 resampled_local_max_ptr,
737 resampled_local_max_stride,
738 # [num_logits, V]
739 target_logits_ptr,
740 target_logits_stride,
741 # [num_reqs]
742 target_rejected_logsumexp_ptr,
743 # [max_num_reqs, num_speculative_steps, V]
744 draft_logits_ptr,
745 draft_logits_stride_0,
746 draft_logits_stride_1,
747 # [num_reqs]
748 draft_rejected_logsumexp_ptr,
749 # [num_reqs]
750 rejected_step_ptr,
751 # [num_reqs + 1]
752 cu_num_logits_ptr,
753 # [num_logits]
754 expanded_idx_mapping_ptr,
755 # [num_logits]
756 draft_sampled_ptr,
757 # [max_num_reqs]
758 temp_ptr,
759 # [max_num_reqs]
760 seed_ptr,
761 # [num_logits]
762 pos_ptr,
763 # [num_logits]
764 cumulative_log_p_ptr,
765 # [num_logits, CONTEXT_WIDTH]
766 contexts_ptr,
767 contexts_stride,
768 # [max_num_reqs], uint8 view of a bool tensor
769 watermarking_ptr,
770 watermarking_skip_mask_ptr,
771 watermark_key_0,
772 watermark_key_1,
773 vocab_size,
774 BLOCK_SIZE: tl.constexpr,
775 HAS_DRAFT_LOGITS: tl.constexpr,
776 USE_FP64: tl.constexpr,
777 USE_BLOCK_VERIFICATION: tl.constexpr,
778 CONTEXT_WIDTH: tl.constexpr,
779 WATERMARK: tl.constexpr,
780 DEDUPLICATE_CONTEXTS: tl.constexpr,
781 ):
782 req_idx = tl.program_id(0)
783 resample_idx = tl.load(rejected_step_ptr + req_idx)
784 start_idx = tl.load(cu_num_logits_ptr + req_idx).to(tl.int64)
785 end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
786 resample_token_idx = start_idx + resample_idx
787 req_state_idx = tl.load(expanded_idx_mapping_ptr + resample_token_idx).to(tl.int64)
788
789 temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
790 is_bonus = resample_token_idx == end_idx - 1
791 if temp == 0.0 and not is_bonus:
792 # Greedy + non-bonus token. No resampling needed because
793 # the target argmax is already in the sampled tensor.
794 return
795
796 rejected_draft_token = tl.load(
797 draft_sampled_ptr + resample_token_idx + 1,
798 mask=not is_bonus,
799 other=0,
800 )
801 is_valid_rejected_draft = rejected_draft_token >= 0
802
803 block_idx = tl.program_id(1)
804 block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
805 mask = block < vocab_size
806 target_logits = tl.load(
807 target_logits_ptr + resample_token_idx * target_logits_stride + block,
808 mask=mask,
809 other=float("-inf"),
810 ).to(tl.float32)
811
812 # Compute the residual logits to resample the rejected token from.
813 is_watermarked = False
814 if WATERMARK:
815 is_watermarked = (
816 tl.load(
817 watermarking_ptr + req_state_idx,
818 mask=req_state_idx >= 0,
819 other=0,
820 )
821 != 0
822 )
823
824 skip_watermarking = False
825 if DEDUPLICATE_CONTEXTS:
826 skip_watermarking = (
827 (tl.load(watermarking_skip_mask_ptr + resample_token_idx) != 0)
828 & is_watermarked
829 & (temp != 0.0)
830 )
831
832 if is_bonus or not is_valid_rejected_draft:
833 # Bonus token (no rejections) or -1 placeholder token. In either case,
834 # directly use the target logits.
835 residual_logits = target_logits
836 elif HAS_DRAFT_LOGITS:
837 # draft_logits is stored pre-temperature, so apply scale first.
838 draft_logits = (
839 tl.load(
840 draft_logits_ptr
841 + req_state_idx * draft_logits_stride_0
842 + resample_idx * draft_logits_stride_1
843 + block,
844 mask=mask,
845 other=float("-inf"),
846 ).to(tl.float32)
847 / temp
848 )
849 target_lse = tl.load(target_rejected_logsumexp_ptr + req_idx)
850 draft_lse = tl.load(draft_rejected_logsumexp_ptr + req_idx)
851 target_log_probs = target_logits - target_lse
852 if USE_BLOCK_VERIFICATION:
853 # Block residual is:
854 # max(p_tau * M_b(x) - M_s(x), 0) / Z.
855 # Scale the target logprobs by log(p_tau). p_0 = 1, so skip
856 # shifting when nothing was accepted (tau == 0).
857 log_p_tau = 0.0
858 if resample_idx > 0:
859 log_p_tau = tl.load(cumulative_log_p_ptr + resample_token_idx - 1).to(
860 tl.float32
861 )
862 target_log_probs += log_p_tau
863 draft_log_probs = draft_logits - draft_lse
864 # Compute the residual:
865 # r(x) = max(p(x) - q(x), 0)
866 # Gumbel sampling needs logits, so we compute it in log space:
867 # log(r(x)) = log(max(exp(log_p(x)) - exp(log_q(x)), 0))
868 # The more numerically stable form is:
869 # log(max(exp(a) - exp(b), 0)) = a + log(max(1 - exp(b - a), 0))
870 ratio = tl.exp(draft_log_probs - target_log_probs)
871 residual_logits = tl.where(
872 ratio < 1.0,
873 target_log_probs + tldevice.log1p(-ratio),
874 float("-inf"),
875 ).to(tl.float32)
876 else:
877 # One-hot draft. The residual is just the target distribution with
878 # the rejected draft token probability zeroed out.
879 # NOTE: During block verification, the residual becomes:
880 # 0 if x == rejected_draft_token
881 # p_tau * M_b(x) / Z otherwise
882 # Therefore p_tau is a constant that cancels under normalization,
883 # and does not need to be applied.
884 residual_logits = tl.where(
885 block != rejected_draft_token,
886 target_logits,
887 float("-inf"),
888 ).to(tl.float32)
889
890 # Resample the rejected/bonus token.
891 if WATERMARK:
892 # Padded and greedy rows retain the stock draw.
893 if is_watermarked & (temp != 0.0) & ~skip_watermarking:
894 watermark_value, idx = philox_gumbel_block_argmax(
895 residual_logits,
896 mask,
897 block_idx,
898 contexts_ptr + resample_token_idx * contexts_stride,
899 watermark_key_0.to(tl.uint32),
900 watermark_key_1.to(tl.uint32),
901 CONTEXT_WIDTH,
902 BLOCK_SIZE,
903 )
904 # Detector compatibility fixes the keyed draw at fp32.
905 value = watermark_value.to(tl.float64) if USE_FP64 else watermark_value
906 else:
907 value, idx = _seeded_resample_argmax(
908 residual_logits,
909 block,
910 mask,
911 resample_token_idx,
912 expanded_idx_mapping_ptr,
913 temp_ptr,
914 seed_ptr,
915 pos_ptr,
916 vocab_size,
917 USE_FP64=USE_FP64,
918 )
919 else:
920 value, idx = _seeded_resample_argmax(
921 residual_logits,
922 block,
923 mask,
924 resample_token_idx,
925 expanded_idx_mapping_ptr,
926 temp_ptr,
927 seed_ptr,
928 pos_ptr,
929 vocab_size,
930 USE_FP64=USE_FP64,
931 )
932 token_id = block_idx * BLOCK_SIZE + idx
933 tl.store(
934 resampled_local_argmax_ptr
935 + req_idx * resampled_local_argmax_stride
936 + block_idx,
937 token_id,
938 )
939 tl.store(
940 resampled_local_max_ptr + req_idx * resampled_local_max_stride + block_idx,
941 value,
942 )
943
944
945 @triton.jit
946 def _insert_resampled_kernel(
947 # [num_reqs, num_speculative_steps + 1]
948 sampled_ptr,
949 sampled_stride,
950 # [num_reqs]
951 num_sampled_ptr,
952 # [num_reqs, num_blocks]
953 resampled_local_argmax_ptr,
954 resampled_local_argmax_stride,
955 # [num_reqs, num_blocks]
956 resampled_local_max_ptr,
957 resampled_local_max_stride,
958 resample_num_blocks,
959 # [num_reqs + 1]
960 cu_num_logits_ptr,
961 # [num_reqs]
962 expanded_idx_mapping_ptr,
963 # [max_num_reqs]
964 temp_ptr,
965 PADDED_RESAMPLE_NUM_BLOCKS: tl.constexpr,
966 ):
967 req_idx = tl.program_id(0)
968 num_sampled = tl.load(num_sampled_ptr + req_idx)
969 start_idx = tl.load(cu_num_logits_ptr + req_idx)
970 end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
971 resample_token_idx = start_idx + num_sampled
972 req_state_idx = tl.load(expanded_idx_mapping_ptr + resample_token_idx)
973
974 # Increment the number of sampled tokens.
975 tl.store(num_sampled_ptr + req_idx, num_sampled + 1)
976
977 temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
978 is_bonus = resample_token_idx == end_idx - 1
979 if temp == 0.0 and not is_bonus:
980 # Greedy + non-bonus token. The target argmax is already
981 # in the sampled tensor.
982 return
983
984 # Insert the resampled token.
985 block = tl.arange(0, PADDED_RESAMPLE_NUM_BLOCKS)
986 mask = block < resample_num_blocks
987 resampled_local_max = tl.load(
988 resampled_local_max_ptr + req_idx * resampled_local_max_stride + block,
989 mask=mask,
990 other=float("-inf"),
991 )
992 # NaN max values (from NaN target logits) make tl.argmax return an
993 # out-of-range block index (into the padded region), causing an OOB read
994 # of resampled_local_argmax. Map NaN to -inf so argmax stays in range.
995 resampled_local_max = tl.where(
996 resampled_local_max != resampled_local_max,
997 float("-inf"),
998 resampled_local_max,
999 )
1000 resampled_max_block_idx = tl.argmax(resampled_local_max, axis=0)
1001 resampled = tl.load(
1002 resampled_local_argmax_ptr
1003 + req_idx * resampled_local_argmax_stride
1004 + resampled_max_block_idx,
1005 )
1006 tl.store(
1007 sampled_ptr + req_idx * sampled_stride + num_sampled,
1008 resampled,
1009 )
1010
1011
1012 def rejection_sample(
1013 # [num_logits, V]
1014 target_logits: torch.Tensor,
1015 # [max_num_reqs, num_speculative_steps, V]
1016 draft_logits: torch.Tensor | None,
1017 # [num_logits]
1018 draft_sampled: torch.Tensor,
1019 # [num_reqs + 1]
1020 cu_num_logits: torch.Tensor,
1021 # [num_logits]
1022 pos: torch.Tensor,
1023 # [num_reqs]
1024 idx_mapping: torch.Tensor,
1025 # [num_logits]
1026 expanded_idx_mapping: torch.Tensor,
1027 # [num_logits]
1028 expanded_local_pos: torch.Tensor,
1029 # [max_num_reqs]
1030 temperature: torch.Tensor,
1031 # [max_num_reqs]
1032 seed: torch.Tensor,
1033 num_speculative_steps: int,
1034 # [num_speculative_steps]
1035 synthetic_conditional_rates: torch.Tensor | None = None,
1036 use_fp64: bool = False,
1037 use_block_verification: bool = False,
1038 contexts: torch.Tensor | None = None,
1039 watermarking: torch.Tensor | None = None,
1040 watermarking_skip_mask: torch.Tensor | None = None,
1041 watermark_key: int | None = None,
1042 ) -> tuple[torch.Tensor, torch.Tensor]:
1043 assert target_logits.ndim == 2 and target_logits.stride(-1) == 1
1044 assert draft_logits is None or (
1045 draft_logits.ndim == 3 and draft_logits.stride(-1) == 1
1046 )
1047 num_reqs = cu_num_logits.shape[0] - 1
1048 num_logits, vocab_size = target_logits.shape
1049
1050 watermark = contexts is not None
1051 assert watermark == (watermarking is not None) == (watermark_key is not None), (
1052 "contexts, watermarking and watermark_key must be set together."
1053 )
1054 assert watermarking_skip_mask is None or watermark, (
1055 "watermarking_skip_mask requires watermarking."
1056 )
1057 contexts_stride = 0
1058 context_width = 1
1059 watermarking_bytes: torch.Tensor | None = None
1060 watermarking_skip_mask_bytes: torch.Tensor | None = None
1061 watermark_key_0 = 0
1062 watermark_key_1 = 0
1063 if contexts is not None:
1064 assert watermarking is not None and watermark_key is not None
1065 assert contexts.ndim == 2 and contexts.shape[0] == num_logits
1066 # Context words are hashed as uint32, so int32 (the request-state token
1067 # dtype) and int64 both land on the same PRF stream, -1 included.
1068 assert contexts.dtype in (torch.int32, torch.int64)
1069 if contexts.stride(-1) != 1:
1070 contexts = contexts.contiguous()
1071 assert watermarking.ndim == 1 and watermarking.dtype == torch.bool
1072 contexts_stride = contexts.stride(0)
1073 context_width = contexts.shape[-1]
1074 watermarking_bytes = watermarking.view(torch.uint8)
1075 if watermarking_skip_mask is not None:
1076 assert watermarking_skip_mask.shape == (num_logits,)
1077 assert watermarking_skip_mask.dtype == torch.bool
1078 if not watermarking_skip_mask.is_contiguous():
1079 watermarking_skip_mask = watermarking_skip_mask.contiguous()
1080 watermarking_skip_mask_bytes = watermarking_skip_mask.view(torch.uint8)
1081 watermark_key_0 = watermark_key & 0xFFFFFFFF
1082 watermark_key_1 = watermark_key >> 32
1083 draft_logits_stride_0 = 0
1084 draft_logits_stride_1 = 0
1085 if has_draft_logits := draft_logits is not None:
1086 draft_logits_stride_0 = draft_logits.stride(0)
1087 draft_logits_stride_1 = draft_logits.stride(1)
1088 # In some cases (e.g. MiMo v2.5 Pro + DFlash) the target model's
1089 # vocab size is larger than the draft's due to padding.
1090 vocab_size = min(vocab_size, draft_logits.size(-1))
1091
1092 # Compute the per-vocab-block logits stats, such as target argmax
1093 # (for greedy requests), and target max + softmax exponential
1094 # (for non-greedy requests).
1095 VOCAB_BLOCK_SIZE = 8192
1096 vocab_num_blocks = triton.cdiv(vocab_size, VOCAB_BLOCK_SIZE)
1097 padded_vocab_num_blocks = triton.next_power_of_2(vocab_num_blocks)
1098 target_local_argmax = target_logits.new_empty(
1099 num_logits, vocab_num_blocks, dtype=torch.int64
1100 )
1101 target_local_max = target_logits.new_empty(
1102 num_logits, vocab_num_blocks, dtype=torch.float32
1103 )
1104 target_local_sumexp = target_logits.new_empty(
1105 num_logits, vocab_num_blocks, dtype=torch.float32
1106 )
1107 draft_local_max = target_logits.new_empty(
1108 num_logits, vocab_num_blocks, dtype=torch.float32
1109 )
1110 draft_local_sumexp = target_logits.new_empty(
1111 num_logits, vocab_num_blocks, dtype=torch.float32
1112 )
1113 _compute_local_logits_stats_kernel[(num_logits, vocab_num_blocks)](
1114 target_local_argmax,
1115 target_local_argmax.stride(0),
1116 target_local_max,
1117 target_local_max.stride(0),
1118 target_local_sumexp,
1119 target_local_sumexp.stride(0),
1120 draft_local_max,
1121 draft_local_max.stride(0),
1122 draft_local_sumexp,
1123 draft_local_sumexp.stride(0),
1124 target_logits,
1125 target_logits.stride(0),
1126 draft_logits,
1127 draft_logits_stride_0,
1128 draft_logits_stride_1,
1129 expanded_idx_mapping,
1130 expanded_local_pos,
1131 temperature,
1132 vocab_size,
1133 num_speculative_steps,
1134 BLOCK_SIZE=VOCAB_BLOCK_SIZE,
1135 HAS_DRAFT_LOGITS=has_draft_logits,
1136 )
1137
1138 # Precompute the running joint ratio and residual mass for block
1139 # verification.
1140 if use_block_verification:
1141 assert synthetic_conditional_rates is None, (
1142 "Block verification is incompatible with synthetic acceptance rates."
1143 )
1144
1145 # Compute the log of the running joint ratio, p_i.
1146 # cumulative_log_p[start + i] = log(p_{i+1}), the cumulative ratio after
1147 # the (i+1)-th draft token.
1148 cumulative_log_p = target_logits.new_empty(num_logits, dtype=torch.float32)
1149 _compute_cumulative_log_p_kernel[(num_reqs,)](
1150 cumulative_log_p,
1151 target_logits,
1152 target_logits.stride(0),
1153 target_local_max,
1154 target_local_max.stride(0),
1155 target_local_sumexp,
1156 target_local_sumexp.stride(0),
1157 draft_sampled,
1158 draft_logits,
1159 draft_logits_stride_0,
1160 draft_logits_stride_1,
1161 draft_local_max,
1162 draft_local_max.stride(0),
1163 draft_local_sumexp,
1164 draft_local_sumexp.stride(0),
1165 cu_num_logits,
1166 idx_mapping,
1167 temperature,
1168 vocab_num_blocks,
1169 PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks,
1170 HAS_DRAFT_LOGITS=has_draft_logits,
1171 num_warps=1,
1172 )
1173
1174 # Compute the per-vocab-block partials of the residual mass, later reduced
1175 # to the total by _compute_global_residual_mass. Only launched for full
1176 # draft logits distributions. One-hot drafts used a closed-form residual
1177 # mass instead.
1178 if has_draft_logits:
1179 local_residual_mass = target_logits.new_empty(
1180 num_logits, vocab_num_blocks, dtype=torch.float32
1181 )
1182 _compute_local_residual_mass_kernel[(num_logits, vocab_num_blocks)](
1183 local_residual_mass,
1184 local_residual_mass.stride(0),
1185 cumulative_log_p,
1186 target_logits,
1187 target_logits.stride(0),
1188 target_local_max,
1189 target_local_max.stride(0),
1190 target_local_sumexp,
1191 target_local_sumexp.stride(0),
1192 draft_logits,
1193 draft_logits_stride_0,
1194 draft_logits_stride_1,
1195 draft_local_max,
1196 draft_local_max.stride(0),
1197 draft_local_sumexp,
1198 draft_local_sumexp.stride(0),
1199 draft_sampled,
1200 expanded_idx_mapping,
1201 expanded_local_pos,
1202 temperature,
1203 vocab_size,
1204 num_speculative_steps,
1205 vocab_num_blocks,
1206 BLOCK_SIZE=VOCAB_BLOCK_SIZE,
1207 PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks,
1208 )
1209 else:
1210 local_residual_mass = None
1211 else:
1212 cumulative_log_p = None
1213 local_residual_mass = None
1214
1215 # Sample up until the first rejected/bonus token, and store
1216 # the step.
1217 sampled = draft_sampled.new_empty(
1218 num_reqs, num_speculative_steps + 1, dtype=torch.int64
1219 )
1220 num_sampled = sampled.new_empty(num_reqs, dtype=torch.int32)
1221 target_rejected_logsumexp = target_logits.new_empty(num_reqs, dtype=torch.float32)
1222 draft_rejected_logsumexp = target_logits.new_empty(num_reqs, dtype=torch.float32)
1223 _rejection_kernel[(num_reqs,)](
1224 sampled,
1225 sampled.stride(0),
1226 num_sampled,
1227 target_rejected_logsumexp,
1228 draft_rejected_logsumexp,
1229 target_logits,
1230 target_logits.stride(0),
1231 target_local_argmax,
1232 target_local_argmax.stride(0),
1233 target_local_max,
1234 target_local_max.stride(0),
1235 target_local_sumexp,
1236 target_local_sumexp.stride(0),
1237 draft_sampled,
1238 draft_logits,
1239 draft_logits_stride_0,
1240 draft_logits_stride_1,
1241 draft_local_max,
1242 draft_local_max.stride(0),
1243 draft_local_sumexp,
1244 draft_local_sumexp.stride(0),
1245 cu_num_logits,
1246 idx_mapping,
1247 temperature,
1248 seed,
1249 pos,
1250 synthetic_conditional_rates,
1251 cumulative_log_p,
1252 local_residual_mass,
1253 local_residual_mass.stride(0) if local_residual_mass is not None else 0,
1254 vocab_num_blocks,
1255 PADDED_VOCAB_NUM_BLOCKS=padded_vocab_num_blocks,
1256 HAS_DRAFT_LOGITS=has_draft_logits,
1257 SYNTHETIC_MODE=synthetic_conditional_rates is not None,
1258 USE_BLOCK_VERIFICATION=use_block_verification,
1259 num_warps=1,
1260 )
1261
1262 # Resample the rejected/bonus tokens.
1263 RESAMPLE_BLOCK_SIZE = 1024
1264 resample_num_blocks = triton.cdiv(vocab_size, RESAMPLE_BLOCK_SIZE)
1265 padded_resample_num_blocks = triton.next_power_of_2(resample_num_blocks)
1266 resampled_local_argmax = target_logits.new_empty(
1267 num_reqs, resample_num_blocks, dtype=torch.int64
1268 )
1269 resampled_local_max = target_logits.new_empty(
1270 num_reqs,
1271 resample_num_blocks,
1272 dtype=torch.float64 if use_fp64 else torch.float32,
1273 )
1274 _resample_kernel[(num_reqs, resample_num_blocks)](
1275 resampled_local_argmax,
1276 resampled_local_argmax.stride(0),
1277 resampled_local_max,
1278 resampled_local_max.stride(0),
1279 target_logits,
1280 target_logits.stride(0),
1281 target_rejected_logsumexp,
1282 draft_logits,
1283 draft_logits_stride_0,
1284 draft_logits_stride_1,
1285 draft_rejected_logsumexp,
1286 num_sampled,
1287 cu_num_logits,
1288 expanded_idx_mapping,
1289 draft_sampled,
1290 temperature,
1291 seed,
1292 pos,
1293 cumulative_log_p,
1294 contexts,
1295 contexts_stride,
1296 watermarking_bytes,
1297 watermarking_skip_mask_bytes,
1298 watermark_key_0,
1299 watermark_key_1,
1300 vocab_size,
1301 BLOCK_SIZE=RESAMPLE_BLOCK_SIZE,
1302 HAS_DRAFT_LOGITS=has_draft_logits,
1303 USE_FP64=use_fp64,
1304 USE_BLOCK_VERIFICATION=use_block_verification,
1305 CONTEXT_WIDTH=context_width,
1306 WATERMARK=watermark,
1307 DEDUPLICATE_CONTEXTS=watermarking_skip_mask is not None,
1308 )
1309
1310 # Insert the resampled tokens into the output sampled.
1311 _insert_resampled_kernel[(num_reqs,)](
1312 sampled,
1313 sampled.stride(0),
1314 num_sampled,
1315 resampled_local_argmax,
1316 resampled_local_argmax.stride(0),
1317 resampled_local_max,
1318 resampled_local_max.stride(0),
1319 resample_num_blocks,
1320 cu_num_logits,
1321 expanded_idx_mapping,
1322 temperature,
1323 PADDED_RESAMPLE_NUM_BLOCKS=padded_resample_num_blocks,
1324 )
1325 return sampled, num_sampled