[Perf] Give a prefill-only batch the model state's number of logit rows (#57416)
Signed-off-by: Matt Mastracci <matthew@mastracci.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Matt Mastracci committed
Sep 19, 2026 at 22:58 UTC
601e1c7c51a10d6fd884d39f079f38bb942e436f
1 file changed
+8
-6
vllm/v1/worker/gpu/model_runner.py
+8
-6
@@ -1291,7 +1291,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1291
# Get the number of draft tokens for each request.
1292
draft_tokens = scheduler_output.scheduled_spec_decode_tokens
1293
num_draft_tokens_per_req = None
1294
- if not draft_tokens:
1294
+ if not draft_tokens and self.model_state.num_new_sampled_tokens_per_step == 1:
1295
# No draft token scheduled (common case).
1296
total_num_draft_tokens = 0
1297
total_num_logits = num_reqs
@@ -1322,9 +1322,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1322
np.cumsum(num_logits, out=cu_num_logits_np[1:])
1323
cu_num_logits = async_tensor_h2d(cu_num_logits_np, device=self.device)
1324
1325
- adaptive_verification = (
1326
- self.adaptive_verification if num_draft_tokens_per_req is not None else None
1327
- )
1325
+ # The general branch also serves a no-draft batch with k != 1, and
1326
+ # compact_batch's budget is only primed when drafts are scheduled.
1327
+ adaptive_verification = self.adaptive_verification if draft_tokens else None
1328
num_scheduled_tokens_upper_bound = num_scheduled_tokens_np
1329
if adaptive_verification is not None:
1330
# num_scheduled_tokens represents the draft budget evenly distributed across
@@ -1354,7 +1354,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1354
adaptive_verification.reallocate_drafts(req_ids, idx_mapping)
1355
)
1356
total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens
1357
- if draft_tokens:
1357
+ if num_draft_tokens_per_req is not None:
1358
expanded_idx_mapping, expanded_local_pos = expand_idx_mapping(
1359
idx_mapping, total_num_logits, cu_num_logits, self.decode_query_len
1360
)
@@ -1534,7 +1534,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1534
sample_hidden_states = hidden_states[input_batch.logits_indices]
1535
logits = self.model.compute_logits(sample_hidden_states)
1536
1537
- if grammar_output is not None:
1537
+ # A diffusion prefill has no logit rows even when a bitmask row
1538
+ # arrived for it.
1539
+ if grammar_output is not None and logits.shape[0] > 0:
1540
# Apply grammar bitmask to the logits in-place.
1541
assert self.structured_outputs_worker is not None
1542
self.structured_outputs_worker.apply_grammar_bitmask(