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 12, 01:00 UTC
[Model Runner V2] Implement get_num_sampled_and_rejected kernel (#30029)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Dec 4, 2025 at 07:19 UTC
cc050558f424714f9548774cc2c661b3916d96ca
3 files changed
+65
-29
vllm/v1/worker/gpu/input_batch.py
+49
@@ -354,6 +354,55 @@ def combine_sampled_and_draft_tokens(
354
return logits_indices
355
356
357
+@triton.jit
358
+def _get_num_sampled_and_rejected_kernel(
359
+ num_sampled_ptr,
360
+ num_rejected_ptr,
361
+ seq_lens_ptr,
362
+ cu_num_logits_ptr,
363
+ idx_mapping_ptr,
364
+ prefill_len_ptr,
365
+):
366
+ batch_idx = tl.program_id(0)
367
+ req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
368
+
369
+ seq_len = tl.load(seq_lens_ptr + batch_idx)
370
+ prefill_len = tl.load(prefill_len_ptr + req_state_idx)
371
+ is_chunked_prefilling = seq_len < prefill_len
372
+
373
+ num_sampled = tl.load(num_sampled_ptr + batch_idx)
374
+ num_sampled = tl.where(is_chunked_prefilling, 0, num_sampled)
375
+ tl.store(num_sampled_ptr + batch_idx, num_sampled)
376
+
377
+ logits_start = tl.load(cu_num_logits_ptr + batch_idx)
378
+ logits_end = tl.load(cu_num_logits_ptr + batch_idx + 1)
379
+ num_logits = logits_end - logits_start
380
+
381
+ num_rejected = num_logits - num_sampled
382
+ num_rejected = tl.where(is_chunked_prefilling, 0, num_rejected)
383
+ tl.store(num_rejected_ptr + batch_idx, num_rejected)
384
+
385
+
386
+def get_num_sampled_and_rejected(
387
+ num_sampled: torch.Tensor,
388
+ seq_lens: torch.Tensor,
389
+ cu_num_logits: torch.Tensor,
390
+ idx_mapping: torch.Tensor,
391
+ prefill_len: torch.Tensor,
392
+) -> tuple[torch.Tensor, torch.Tensor]:
393
+ num_reqs = idx_mapping.shape[0]
394
+ num_rejected = torch.empty_like(num_sampled)
395
+ _get_num_sampled_and_rejected_kernel[(num_reqs,)](
396
+ num_sampled,
397
+ num_rejected,
398
+ seq_lens,
399
+ cu_num_logits,
400
+ idx_mapping,
401
+ prefill_len,
402
+ )
403
+ return num_sampled, num_rejected
404
+
405
+
406
@triton.jit
407
def _post_update_kernel(
408
idx_mapping_ptr,
vllm/v1/worker/gpu/model_runner.py
+16
-17
@@ -43,6 +43,7 @@ from vllm.v1.worker.gpu.input_batch import (
43
InputBatch,
44
InputBuffers,
45
combine_sampled_and_draft_tokens,
46
+ get_num_sampled_and_rejected,
47
post_update,
48
prepare_pos_seq_lens,
49
prepare_prefill_inputs,
@@ -54,10 +55,7 @@ from vllm.v1.worker.gpu.sample.metadata import (
55
)
56
from vllm.v1.worker.gpu.sample.sampler import Sampler
57
from vllm.v1.worker.gpu.spec_decode import init_speculator
57
-from vllm.v1.worker.gpu.spec_decode.rejection_sample import (
58
- get_num_rejected,
59
- rejection_sample,
60
-)
58
+from vllm.v1.worker.gpu.spec_decode.rejection_sample import rejection_sample
59
from vllm.v1.worker.gpu.states import RequestState
60
from vllm.v1.worker.gpu.structured_outputs import apply_grammar_bitmask
61
from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
@@ -621,16 +619,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
619
# Sample tokens and compute logprobs (if needed).
620
sampler_output = self.sampler(logits, sampling_metadata)
621
624
- # Get the number of sampled tokens.
625
- prefill_len = self.req_states.prefill_len.gpu[input_batch.idx_mapping]
626
- is_chunked_prefilling = input_batch.seq_lens < prefill_len
622
if input_batch.num_draft_tokens == 0:
623
# No draft tokens (common case).
629
- # 0 if chunked-prefilling, 1 if not.
630
- num_sampled = (~is_chunked_prefilling).int()
631
- num_rejected = torch.zeros_like(num_sampled)
624
+ num_sampled = torch.ones(
625
+ input_batch.num_reqs, dtype=torch.int32, device=self.device
626
+ )
627
else:
633
- # Draft tokens for spec decoding.
628
+ # Rejection sampling for spec decoding.
629
input_ids = input_batch.input_ids[input_batch.logits_indices]
630
sampled_tokens, num_sampled = rejection_sample(
631
sampler_output.sampled_token_ids,
@@ -638,13 +633,17 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
633
input_batch.cu_num_logits,
634
self.num_speculative_steps,
635
)
641
- num_sampled *= ~is_chunked_prefilling
642
- num_rejected = get_num_rejected(
643
- input_batch.cu_num_logits,
644
- num_sampled,
645
- )
636
sampler_output.sampled_token_ids = sampled_tokens
647
- # TODO(woosuk): Support logprobs with spec decoding.
637
+
638
+ # Get the number of sampled and rejected tokens.
639
+ # For chunked prefills, num_sampled and num_rejected are both 0.
640
+ num_sampled, num_rejected = get_num_sampled_and_rejected(
641
+ num_sampled,
642
+ input_batch.seq_lens,
643
+ input_batch.cu_num_logits,
644
+ input_batch.idx_mapping,
645
+ self.req_states.prefill_len.gpu,
646
+ )
647
return sampler_output, num_sampled, num_rejected
648
649
def compute_prompt_logprobs(
vllm/v1/worker/gpu/spec_decode/rejection_sample.py
-12
@@ -69,15 +69,3 @@ def rejection_sample(
69
num_warps=1,
70
)
71
return sampled, num_sampled
72
-
73
-
74
-@torch.compile(dynamic=True)
75
-def get_num_rejected(
76
- cu_num_logits: torch.Tensor,
77
- num_sampled: torch.Tensor,
78
-) -> torch.Tensor:
79
- num_logits = cu_num_logits[1:] - cu_num_logits[:-1]
80
- num_rejected = num_logits - num_sampled
81
- # No token is rejected for chunked prefills.
82
- num_rejected *= num_sampled > 0
83
- return num_rejected