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 5, 14:45 UTC
Fix invalid logprobs with MTP enabled and sync scheduling (#38711)
Signed-off-by: Daniel Serebrenik <daserebrenik@nvidia.com>
danisereb committed
Apr 3, 2026 at 19:24 UTC
50cd5674b39c69e60fefd0ec1d61652d693e6a06
1 file changed
+12
-10
vllm/v1/worker/gpu_model_runner.py
+12
-10
@@ -4192,6 +4192,7 @@ class GPUModelRunner(
4192
spec_config = self.speculative_config
4193
propose_drafts_after_bookkeeping = False
4194
if spec_config is not None:
4195
+ # Decide whether to run the drafter or zero out draft tokens.
4196
input_fits_in_drafter = spec_decode_common_attn_metadata is not None and (
4197
spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens
4198
<= self.effective_drafter_max_model_len
@@ -4227,10 +4228,6 @@ class GPUModelRunner(
4228
self._copy_valid_sampled_token_count(
4229
next_token_ids, valid_sampled_tokens_count
4230
)
4230
- self._draft_token_ids = torch.zeros(
4231
- 1, device=self.device, dtype=torch.int32
4232
- ).expand(len(self.input_batch.req_ids), self.num_spec_tokens)
4233
- self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True)
4231
elif (
4232
spec_config.use_ngram_gpu()
4233
and not spec_config.disable_padded_drafter_batch
@@ -4253,15 +4250,20 @@ class GPUModelRunner(
4250
self._copy_valid_sampled_token_count(
4251
next_token_ids, valid_sampled_tokens_count
4252
)
4256
- # Since we couldn't run the drafter,
4257
- # just use zeros for the draft tokens.
4258
- self._draft_token_ids = torch.zeros(
4259
- 1, device=self.device, dtype=torch.int32
4260
- ).expand(len(self.input_batch.req_ids), self.num_spec_tokens)
4261
- self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True)
4253
else:
4254
propose_drafts_after_bookkeeping = input_fits_in_drafter
4255
4256
+ if not input_fits_in_drafter:
4257
+ # Zero out draft tokens so the scheduler doesn't schedule
4258
+ # stale drafts from the previous step.
4259
+ # For Nemotron-H: it is necessary to zero out the draft tokens,
4260
+ # otherwise the stale tokens will corrupt Mamba recurrent
4261
+ # state and logprobs for sequences near max_model_len.
4262
+ self._draft_token_ids = torch.zeros(
4263
+ 1, device=self.device, dtype=torch.int32
4264
+ ).expand(len(self.input_batch.req_ids), self.num_spec_tokens)
4265
+ self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True)
4266
+
4267
with record_function_or_nullcontext("gpu_model_runner: bookkeep"):
4268
(
4269
num_nans_in_logits,