From 601e1c7c51a10d6fd884d39f079f38bb942e436f Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Sat, 19 Sep 2026 22:58:30 -0600 Subject: [PATCH] [Perf] Give a prefill-only batch the model state's number of logit rows (#57416) Signed-off-by: Matt Mastracci Co-authored-by: Claude Fable 5.1 --- vllm/v1/worker/gpu/model_runner.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index e76c447ea5..8c736394cc 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1291,7 +1291,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): # Get the number of draft tokens for each request. draft_tokens = scheduler_output.scheduled_spec_decode_tokens num_draft_tokens_per_req = None - if not draft_tokens: + if not draft_tokens and self.model_state.num_new_sampled_tokens_per_step == 1: # No draft token scheduled (common case). total_num_draft_tokens = 0 total_num_logits = num_reqs @@ -1322,9 +1322,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): np.cumsum(num_logits, out=cu_num_logits_np[1:]) cu_num_logits = async_tensor_h2d(cu_num_logits_np, device=self.device) - adaptive_verification = ( - self.adaptive_verification if num_draft_tokens_per_req is not None else None - ) + # The general branch also serves a no-draft batch with k != 1, and + # compact_batch's budget is only primed when drafts are scheduled. + adaptive_verification = self.adaptive_verification if draft_tokens else None num_scheduled_tokens_upper_bound = num_scheduled_tokens_np if adaptive_verification is not None: # num_scheduled_tokens represents the draft budget evenly distributed across @@ -1354,7 +1354,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): adaptive_verification.reallocate_drafts(req_ids, idx_mapping) ) total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens - if draft_tokens: + if num_draft_tokens_per_req is not None: expanded_idx_mapping, expanded_local_pos = expand_idx_mapping( idx_mapping, total_num_logits, cu_num_logits, self.decode_query_len ) @@ -1534,7 +1534,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): sample_hidden_states = hidden_states[input_batch.logits_indices] logits = self.model.compute_logits(sample_hidden_states) - if grammar_output is not None: + # A diffusion prefill has no logit rows even when a bitmask row + # arrived for it. + if grammar_output is not None and logits.shape[0] > 0: # Apply grammar bitmask to the logits in-place. assert self.structured_outputs_worker is not None self.structured_outputs_worker.apply_grammar_bitmask(