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, 14:45 UTC
[Model Runner V2] Simplify Eagle bookkeeping with num_rejected (#29347)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Nov 24, 2025 at 13:54 UTC
f32c7d6f5455de2684686c7238f9c7ecca6b58b7
4 files changed
+50
-30
vllm/v1/worker/gpu/input_batch.py
+6
-13
@@ -344,8 +344,8 @@ def _post_update_kernel(
344
sampled_tokens_ptr,
345
sampled_tokens_stride,
346
num_sampled_ptr,
347
+ num_rejected_ptr,
348
query_start_loc_ptr,
348
- cu_num_logits_ptr,
349
):
350
req_id = tl.program_id(0)
351
req_state_idx = tl.load(idx_mapping_ptr + req_id)
@@ -360,17 +360,10 @@ def _post_update_kernel(
360
query_start = tl.load(query_start_loc_ptr + req_id)
361
query_end = tl.load(query_start_loc_ptr + req_id + 1)
362
query_len = query_end - query_start
363
+ num_rejected = tl.load(num_rejected_ptr + req_id)
364
365
num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
365
- num_computed += query_len
366
- # Consider the rejected tokens in spec decoding.
367
- if num_sampled > 0:
368
- # NOTE(woosuk): We must skip num_sampled == 0 to account for chunked prefills.
369
- logits_start = tl.load(cu_num_logits_ptr + req_id)
370
- logits_end = tl.load(cu_num_logits_ptr + req_id + 1)
371
- num_logits = logits_end - logits_start
372
- num_rejected = num_logits - num_sampled
373
- num_computed -= num_rejected
366
+ num_computed += query_len - num_rejected
367
tl.store(num_computed_tokens_ptr + req_state_idx, num_computed)
368
369
@@ -385,10 +378,10 @@ def post_update(
378
sampled_tokens: torch.Tensor,
379
# [num_reqs]
380
num_sampled: torch.Tensor,
381
+ # [num_reqs]
382
+ num_rejected: torch.Tensor,
383
# [num_reqs + 1]
384
query_start_loc: torch.Tensor,
390
- # [num_reqs + 1]
391
- cu_num_logits: torch.Tensor,
385
) -> None:
386
num_reqs = idx_mapping.shape[0]
387
_post_update_kernel[(num_reqs,)](
@@ -398,7 +391,7 @@ def post_update(
391
sampled_tokens,
392
sampled_tokens.stride(0),
393
num_sampled,
394
+ num_rejected,
395
query_start_loc,
402
- cu_num_logits,
396
num_warps=1,
397
)
vllm/v1/worker/gpu/model_runner.py
+22
-8
@@ -46,7 +46,10 @@ from vllm.v1.worker.gpu.input_batch import (
46
)
47
from vllm.v1.worker.gpu.sampler import Sampler, compute_prompt_logprobs
48
from vllm.v1.worker.gpu.spec_decode import init_speculator
49
-from vllm.v1.worker.gpu.spec_decode.rejection_sample import rejection_sample
49
+from vllm.v1.worker.gpu.spec_decode.rejection_sample import (
50
+ get_num_rejected,
51
+ rejection_sample,
52
+)
53
from vllm.v1.worker.gpu.states import RequestState, SamplingMetadata
54
from vllm.v1.worker.gpu.structured_outputs import apply_grammar_bitmask
55
from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
@@ -311,12 +314,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
314
device=self.device,
315
)
316
num_sampled = torch.ones(num_reqs, dtype=torch.int32, device=self.device)
317
+ num_rejected = torch.zeros(num_reqs, dtype=torch.int32, device=self.device)
318
self.propose_draft(
319
input_batch=input_batch,
320
sampling_metadata=sampling_metadata,
321
last_hidden_states=hidden_states,
322
aux_hidden_states=aux_hidden_states,
323
num_sampled=num_sampled,
324
+ num_rejected=num_rejected,
325
)
326
327
@torch.inference_mode()
@@ -606,7 +611,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
611
input_batch: InputBatch,
612
sampling_metadata: SamplingMetadata,
613
grammar_output: GrammarOutput | None,
609
- ) -> tuple[SamplerOutput, torch.Tensor]:
614
+ ) -> tuple[SamplerOutput, torch.Tensor, torch.Tensor]:
615
sample_hidden_states = hidden_states[input_batch.logits_indices]
616
logits = self.model.compute_logits(sample_hidden_states)
617
if grammar_output is not None:
@@ -632,6 +637,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
637
# No draft tokens (common case).
638
# 0 if chunked-prefilling, 1 if not.
639
num_sampled = (~is_chunked_prefilling).int()
640
+ num_rejected = torch.zeros_like(num_sampled)
641
else:
642
# Draft tokens for spec decoding.
643
input_ids = input_batch.input_ids[input_batch.logits_indices]
@@ -642,9 +648,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
648
self.num_speculative_steps,
649
)
650
num_sampled *= ~is_chunked_prefilling
651
+ num_rejected = get_num_rejected(
652
+ input_batch.cu_num_logits,
653
+ num_sampled,
654
+ )
655
sampler_output.sampled_token_ids = sampled_tokens
656
# TODO(woosuk): Support logprobs with spec decoding.
647
- return sampler_output, num_sampled
657
+ return sampler_output, num_sampled, num_rejected
658
659
def compute_prompt_logprobs(
660
self,
@@ -750,6 +760,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
760
input_batch: InputBatch,
761
sampled_tokens: torch.Tensor,
762
num_sampled: torch.Tensor,
763
+ num_rejected: torch.Tensor,
764
) -> None:
765
# Update the number of computed tokens.
766
post_update(
@@ -758,8 +769,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
769
self.req_states.last_sampled_tokens,
770
sampled_tokens,
771
num_sampled,
772
+ num_rejected,
773
input_batch.query_start_loc,
762
- input_batch.cu_num_logits,
774
)
775
776
# Update the number of computed prefill tokens.
@@ -779,6 +790,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
790
last_hidden_states: torch.Tensor,
791
aux_hidden_states: list[torch.Tensor] | None,
792
num_sampled: torch.Tensor,
793
+ num_rejected: torch.Tensor,
794
) -> torch.Tensor:
795
num_reqs = input_batch.num_reqs
796
idx_mapping_np = input_batch.idx_mapping_np
@@ -800,6 +812,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
812
last_hidden_states,
813
aux_hidden_states,
814
num_sampled,
815
+ num_rejected,
816
self.req_states.last_sampled_tokens,
817
next_prefill_tokens,
818
)
@@ -958,7 +971,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
971
self.execute_model_state = None # type: ignore
972
assert sampling_metadata is not None
973
961
- sampler_output, num_sampled_tokens = self.sample(
974
+ sampler_output, num_sampled, num_rejected = self.sample(
975
hidden_states, input_batch, sampling_metadata, grammar_output
976
)
977
prompt_logprobs_dict = self.compute_prompt_logprobs(hidden_states, input_batch)
@@ -979,7 +992,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
992
async_output = AsyncOutput(
993
model_runner_output=model_runner_output,
994
sampler_output=sampler_output,
982
- num_sampled_tokens=num_sampled_tokens,
995
+ num_sampled_tokens=num_sampled,
996
copy_stream=self.output_copy_stream,
997
copy_event=self.output_copy_event,
998
)
@@ -990,7 +1003,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1003
# This sequencing may slightly reduce latency as async D2H copy does not
1004
# need to wait for the postprocess to finish.
1005
self.postprocess(
993
- input_batch, sampler_output.sampled_token_ids, num_sampled_tokens
1006
+ input_batch, sampler_output.sampled_token_ids, num_sampled, num_rejected
1007
)
1008
if self.do_spec_decode:
1009
_ = self.propose_draft(
@@ -998,7 +1011,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1011
sampling_metadata,
1012
hidden_states,
1013
None, # aux_hidden_states
1001
- num_sampled_tokens,
1014
+ num_sampled,
1015
+ num_rejected,
1016
)
1017
1018
if self.use_async_scheduling:
vllm/v1/worker/gpu/spec_decode/eagle.py
+10
-9
@@ -60,6 +60,8 @@ class EagleSpeculator:
60
aux_hidden_states: list[torch.Tensor] | None,
61
# [num_reqs]
62
num_sampled: torch.Tensor,
63
+ # [num_reqs]
64
+ num_rejected: torch.Tensor,
65
# [max_num_reqs, 1]
66
last_sampled: torch.Tensor,
67
# [num_reqs]
@@ -84,6 +86,7 @@ class EagleSpeculator:
86
self.input_ids,
87
input_batch,
88
num_sampled,
89
+ num_rejected,
90
last_sampled,
91
next_prefill_tokens,
92
)
@@ -139,8 +142,8 @@ def _prepare_eagle_inputs_kernel(
142
last_sampled_ptr,
143
next_prefill_tokens_ptr,
144
num_sampled_ptr,
145
+ num_rejected_ptr,
146
query_start_loc_ptr,
143
- cu_num_logits_ptr,
147
BLOCK_SIZE: tl.constexpr,
148
):
149
batch_idx = tl.program_id(0)
@@ -149,17 +152,13 @@ def _prepare_eagle_inputs_kernel(
152
query_len = query_end - query_start
153
154
# Get the true query length and next token after accounting for rejected tokens.
155
+ num_rejected = tl.load(num_rejected_ptr + batch_idx)
156
+ query_len -= num_rejected
157
+
158
num_sampled = tl.load(num_sampled_ptr + batch_idx)
159
if num_sampled > 0:
160
req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
161
next_token = tl.load(last_sampled_ptr + req_state_idx).to(tl.int32)
156
-
157
- logits_start = tl.load(cu_num_logits_ptr + batch_idx)
158
- logits_end = tl.load(cu_num_logits_ptr + batch_idx + 1)
159
- num_logits = logits_end - logits_start
160
-
161
- num_rejected = num_logits - num_sampled
162
- query_len -= num_rejected
162
else:
163
# Chunked prefilling.
164
# Get the next prefill token.
@@ -182,6 +181,8 @@ def prepare_eagle_inputs(
181
input_batch: InputBatch,
182
# [num_reqs]
183
num_sampled: torch.Tensor,
184
+ # [num_reqs]
185
+ num_rejected: torch.Tensor,
186
# [max_num_reqs, 1]
187
last_sampled: torch.Tensor,
188
# [max_num_reqs]
@@ -201,8 +202,8 @@ def prepare_eagle_inputs(
202
last_sampled,
203
next_prefill_tokens,
204
num_sampled,
205
+ num_rejected,
206
input_batch.query_start_loc,
205
- input_batch.cu_num_logits,
207
BLOCK_SIZE=1024,
208
)
209
return last_token_indices
vllm/v1/worker/gpu/spec_decode/rejection_sample.py
+12
@@ -69,3 +69,15 @@ 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