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, 09:30 UTC
[Model Runner V2] Fix draft logits not populated during cudagraph replay (#37639)
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Giancarlo Delfin committed
Mar 20, 2026 at 00:43 UTC
dcee9be95a0f7fce32ab82060733ab31f90b9154
3 files changed
+17
-25
vllm/v1/worker/gpu/model_runner.py
+2
-4
@@ -195,7 +195,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
195
num_speculative_steps=self.num_speculative_steps,
196
vocab_size=self.vocab_size,
197
device=self.device,
198
- cache_draft_logits=not use_strict_rejection_sampling,
198
)
199
self.input_buffers = InputBuffers(
200
max_num_reqs=self.max_num_reqs,
@@ -446,7 +445,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
445
next_prefill_tokens=self.req_states.next_prefill_tokens,
446
temperature=self.sampler.sampling_states.temperature.gpu,
447
seeds=self.sampler.sampling_states.seeds.gpu,
449
- draft_logits_out=self.req_states.draft_logits,
448
num_tokens_across_dp=num_tokens_across_dp,
449
dummy_run=True,
450
skip_attn_for_dummy_run=skip_attn,
@@ -815,11 +813,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
813
else:
814
# Rejection sampling for spec decoding.
815
assert self.rejection_sampler is not None
816
+ assert self.speculator is not None
817
sampler_output = self.rejection_sampler(
818
logits,
819
input_batch,
820
# Draft logits are needed for probabilistic rejection sampling.
822
- self.req_states.draft_logits,
821
+ self.speculator.draft_logits,
822
)
823
824
# Get the number of sampled and rejected tokens.
@@ -1145,7 +1144,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1144
self.req_states.next_prefill_tokens,
1145
self.sampler.sampling_states.temperature.gpu,
1146
self.sampler.sampling_states.seeds.gpu,
1148
- self.req_states.draft_logits,
1147
num_tokens_across_dp=num_tokens_across_dp,
1148
)
1149
self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+15
-8
@@ -76,6 +76,17 @@ class EagleSpeculator:
76
device=device,
77
)
78
79
+ cache_draft_logits = self.speculative_config.rejection_sample_method != "strict"
80
+ self.draft_logits: torch.Tensor | None = None
81
+ if cache_draft_logits:
82
+ self.draft_logits = torch.zeros(
83
+ self.max_num_reqs,
84
+ self.num_speculative_steps,
85
+ self.vocab_size,
86
+ dtype=torch.float32,
87
+ device=device,
88
+ )
89
+
90
# currently we don't support PIECEWISE for Eagle.
91
cudagraph_mode = vllm_config.compilation_config.cudagraph_mode
92
if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL:
@@ -158,7 +169,6 @@ class EagleSpeculator:
169
slot_mappings: dict[str, torch.Tensor] | None,
170
num_tokens_across_dp: torch.Tensor | None,
171
cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
161
- draft_logits_out: torch.Tensor | None = None,
172
) -> None:
173
pos = self.input_buffers.positions[:num_reqs]
174
query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
@@ -185,8 +195,8 @@ class EagleSpeculator:
195
self.seeds,
196
pos + 1,
197
apply_temperature=True,
188
- processed_logits_out=draft_logits_out[:, step]
189
- if draft_logits_out is not None
198
+ processed_logits_out=self.draft_logits[:, step]
199
+ if self.draft_logits is not None
200
else None,
201
)
202
self.draft_tokens[:num_reqs, step] = draft_tokens
@@ -241,8 +251,6 @@ class EagleSpeculator:
251
temperature: torch.Tensor,
252
# [max_num_reqs]
253
seeds: torch.Tensor,
244
- # [max_num_reqs, num_speculative_steps, vocab_size]
245
- draft_logits_out: torch.Tensor | None,
254
num_tokens_across_dp: torch.Tensor | None = None,
255
dummy_run: bool = False,
256
skip_attn_for_dummy_run: bool = False,
@@ -308,8 +316,8 @@ class EagleSpeculator:
316
self.seeds,
317
pos + 1,
318
apply_temperature=True,
311
- processed_logits_out=draft_logits_out[:, 0]
312
- if draft_logits_out is not None
319
+ processed_logits_out=self.draft_logits[:, 0]
320
+ if self.draft_logits is not None
321
else None,
322
)
323
@@ -394,7 +402,6 @@ class EagleSpeculator:
402
slot_mappings_updated,
403
num_tokens_across_dp=num_tokens_across_dp,
404
cudagraph_runtime_mode=batch_desc.cg_mode,
397
- draft_logits_out=draft_logits_out,
405
)
406
return self.draft_tokens[:num_reqs]
407
vllm/v1/worker/gpu/states.py
-13
@@ -15,7 +15,6 @@ class RequestState:
15
num_speculative_steps: int,
16
vocab_size: int,
17
device: torch.device,
18
- cache_draft_logits: bool,
18
):
19
self.max_num_reqs = max_num_reqs
20
self.max_model_len = max_model_len
@@ -71,18 +70,6 @@ class RequestState:
70
dtype=torch.int64,
71
device=device,
72
)
74
- # Draft token logits.
75
- # NOTE: This tensor maintains the "processed" logits after applying temperature,
76
- # top-p, etc.
77
- self.draft_logits: torch.Tensor | None = None
78
- if cache_draft_logits:
79
- self.draft_logits = torch.zeros(
80
- self.max_num_reqs,
81
- self.num_speculative_steps,
82
- self.vocab_size,
83
- dtype=torch.float32,
84
- device=device,
85
- )
73
74
self.next_prefill_tokens = torch.zeros(
75
self.max_num_reqs, dtype=torch.int32, device=device