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 10, 01:30 UTC
[Model Runner V2] Minor optimization for eagle input processing (#32535)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jan 17, 2026 at 21:55 UTC
963dc0b865a3b6011fde7e0d938f86245dccbfac
2 files changed
+12
-14
vllm/v1/worker/gpu/model_runner.py
+2
-8
@@ -827,20 +827,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
827
num_rejected: torch.Tensor,
828
) -> torch.Tensor:
829
assert self.speculator is not None
830
- last_sampled_tokens = self.req_states.last_sampled_tokens[
831
- input_batch.idx_mapping
832
- ]
833
- next_prefill_tokens = self.req_states.next_prefill_tokens[
834
- input_batch.idx_mapping
835
- ]
830
draft_tokens = self.speculator.propose(
831
input_batch,
832
last_hidden_states,
833
aux_hidden_states,
834
num_sampled,
835
num_rejected,
842
- last_sampled_tokens,
843
- next_prefill_tokens,
836
+ self.req_states.last_sampled_tokens,
837
+ self.req_states.next_prefill_tokens,
838
self.sampler.sampling_states.temperature.gpu,
839
self.sampler.sampling_states.seeds.gpu,
840
)
vllm/v1/worker/gpu/spec_decode/eagle.py
+10
-6
@@ -195,9 +195,9 @@ class EagleSpeculator:
195
num_sampled: torch.Tensor,
196
# [num_reqs]
197
num_rejected: torch.Tensor,
198
- # [num_reqs]
198
+ # [max_num_reqs]
199
last_sampled: torch.Tensor,
200
- # [num_reqs]
200
+ # [max_num_reqs]
201
next_prefill_tokens: torch.Tensor,
202
# [max_num_reqs]
203
temperature: torch.Tensor,
@@ -320,6 +320,7 @@ def _prepare_eagle_inputs_kernel(
320
eagle_positions_ptr,
321
target_input_ids_ptr,
322
target_positions_ptr,
323
+ idx_mapping_ptr,
324
last_sampled_ptr,
325
next_prefill_tokens_ptr,
326
num_sampled_ptr,
@@ -328,6 +329,8 @@ def _prepare_eagle_inputs_kernel(
329
BLOCK_SIZE: tl.constexpr,
330
):
331
batch_idx = tl.program_id(0)
332
+ req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
333
+
334
query_start = tl.load(query_start_loc_ptr + batch_idx)
335
query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
336
query_len = query_end - query_start
@@ -338,11 +341,11 @@ def _prepare_eagle_inputs_kernel(
341
342
num_sampled = tl.load(num_sampled_ptr + batch_idx)
343
if num_sampled > 0:
341
- next_token = tl.load(last_sampled_ptr + batch_idx).to(tl.int32)
344
+ next_token = tl.load(last_sampled_ptr + req_state_idx).to(tl.int32)
345
else:
346
# Chunked prefilling.
347
# Get the next prefill token.
345
- next_token = tl.load(next_prefill_tokens_ptr + batch_idx)
348
+ next_token = tl.load(next_prefill_tokens_ptr + req_state_idx)
349
350
# Shift target_input_ids by one.
351
for i in range(1, query_len, BLOCK_SIZE):
@@ -370,9 +373,9 @@ def prepare_eagle_inputs(
373
num_sampled: torch.Tensor,
374
# [num_reqs]
375
num_rejected: torch.Tensor,
373
- # [num_reqs]
376
+ # [max_num_reqs]
377
last_sampled: torch.Tensor,
375
- # [num_reqs]
378
+ # [max_num_reqs]
379
next_prefill_tokens: torch.Tensor,
380
) -> torch.Tensor:
381
num_reqs = input_batch.num_reqs
@@ -387,6 +390,7 @@ def prepare_eagle_inputs(
390
input_buffers.positions,
391
input_batch.input_ids,
392
input_batch.positions,
393
+ input_batch.idx_mapping,
394
last_sampled,
395
next_prefill_tokens,
396
num_sampled,