A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 10:45 UTC
[Model Runner V2] Spec decode rejection sampler logprobs support (#37237)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Giancarlo Delfin committed
Mar 18, 2026 at 18:36 UTC
053f3b6309a82cca092c1dc99320e2325a9c8ca0
1 file changed
+65
-2
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+65
-2
@@ -3,11 +3,14 @@
3
import torch
4
5
from vllm.triton_utils import tl, triton
6
+from vllm.v1.outputs import LogprobsTensors
7
from vllm.v1.worker.gpu.input_batch import InputBatch
8
from vllm.v1.worker.gpu.metrics.logits import get_num_nans
9
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
10
+from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
11
from vllm.v1.worker.gpu.sample.output import SamplerOutput
12
from vllm.v1.worker.gpu.sample.sampler import Sampler
13
+from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS
14
15
16
@triton.jit
@@ -418,6 +421,26 @@ def probabilistic_rejection_sample(
421
return sampled, rejected_steps + 1
422
423
424
+@triton.jit
425
+def _flatten_sampled_kernel(
426
+ # [num_logits]
427
+ flat_sampled_ptr,
428
+ # [num_reqs, num_speculative_steps + 1]
429
+ sampled_ptr,
430
+ sampled_stride,
431
+ # [num_reqs]
432
+ num_sampled_ptr,
433
+ # [num_reqs + 1]
434
+ cu_num_logits_ptr,
435
+):
436
+ req_idx = tl.program_id(0)
437
+ start_idx = tl.load(cu_num_logits_ptr + req_idx)
438
+ num_sampled = tl.load(num_sampled_ptr + req_idx)
439
+ for i in range(num_sampled):
440
+ token_id = tl.load(sampled_ptr + req_idx * sampled_stride + i)
441
+ tl.store(flat_sampled_ptr + start_idx + i, token_id)
442
+
443
+
444
class RejectionSampler:
445
def __init__(
446
self,
@@ -429,6 +452,40 @@ class RejectionSampler:
452
self.num_speculative_steps = num_speculative_steps
453
self.use_strict_rejection_sampling = use_strict_rejection_sampling
454
455
+ def _get_logprobs_tensors(
456
+ self,
457
+ input_batch: InputBatch,
458
+ sampled: torch.Tensor,
459
+ num_sampled: torch.Tensor,
460
+ logits: torch.Tensor,
461
+ ) -> LogprobsTensors | None:
462
+ max_num_logprobs = self.sampler.sampling_states.max_num_logprobs(
463
+ input_batch.idx_mapping_np
464
+ )
465
+ if max_num_logprobs == NO_LOGPROBS:
466
+ return None
467
+
468
+ num_reqs = input_batch.cu_num_logits.shape[0] - 1
469
+ num_logits = logits.shape[0]
470
+ flat_sampled = torch.zeros(
471
+ num_logits, dtype=sampled.dtype, device=sampled.device
472
+ )
473
+ _flatten_sampled_kernel[(num_reqs,)](
474
+ flat_sampled,
475
+ sampled,
476
+ sampled.stride(0),
477
+ num_sampled,
478
+ input_batch.cu_num_logits,
479
+ num_warps=1,
480
+ )
481
+ expanded_logits = num_logits != input_batch.idx_mapping.shape[0]
482
+ return compute_topk_logprobs(
483
+ logits,
484
+ max_num_logprobs,
485
+ flat_sampled,
486
+ input_batch.cu_num_logits_np.tolist() if expanded_logits else None,
487
+ )
488
+
489
def __call__(
490
self,
491
logits: torch.Tensor,
@@ -460,8 +517,6 @@ class RejectionSampler:
517
draft_sampled,
518
input_batch.expanded_local_pos,
519
)
463
- # TODO (TheEpicDolphin): Return logprobs for sampled token ids.
464
- logprobs_tensors = None
520
sampled, num_sampled = probabilistic_rejection_sample(
521
processed_logits,
522
draft_logits,
@@ -475,6 +530,14 @@ class RejectionSampler:
530
self.sampler.sampling_states.seeds.gpu,
531
self.num_speculative_steps,
532
)
533
+ logprobs_tensors = self._get_logprobs_tensors(
534
+ input_batch,
535
+ sampled,
536
+ num_sampled,
537
+ processed_logits
538
+ if self.sampler.logprobs_mode == "processed_logprobs"
539
+ else logits,
540
+ )
541
542
return SamplerOutput(
543
sampled_token_ids=sampled,