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 14, 05:30 UTC
[BugFix] Pad input buffers in _dummy_run (#26209)
Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Varun Sundar Rabindranath committed
Oct 6, 2025 at 16:07 UTC
f23b4c04fd099aa15ba02df779f7bf6ae66d99a5
1 file changed
+10
-8
vllm/v1/worker/gpu_model_runner.py
+10
-8
@@ -3434,26 +3434,28 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3434
with self.maybe_dummy_run_with_lora(
3435
self.lora_config, num_scheduled_tokens, remove_lora
3436
):
3437
- model_kwargs = self._init_model_kwargs(num_tokens)
3437
+ # Make sure padding doesn't exceed max_num_tokens
3438
+ assert num_tokens_after_padding <= self.max_num_tokens
3439
+ model_kwargs = self._init_model_kwargs(num_tokens_after_padding)
3440
if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
3441
input_ids = None
3440
- inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
3442
+ inputs_embeds = self.inputs_embeds.gpu[:num_tokens_after_padding]
3443
model_kwargs = {
3444
**model_kwargs,
3445
**self._dummy_mm_kwargs(num_reqs),
3446
}
3447
elif self.enable_prompt_embeds:
3448
input_ids = None
3447
- inputs_embeds = self.inputs_embeds.gpu[:num_tokens]
3448
- model_kwargs = self._init_model_kwargs(num_tokens)
3449
+ inputs_embeds = self.inputs_embeds.gpu[:num_tokens_after_padding]
3450
+ model_kwargs = self._init_model_kwargs(num_tokens_after_padding)
3451
else:
3450
- input_ids = self.input_ids.gpu[:num_tokens]
3452
+ input_ids = self.input_ids.gpu[:num_tokens_after_padding]
3453
inputs_embeds = None
3454
3455
if self.uses_mrope:
3454
- positions = self.mrope_positions.gpu[:, :num_tokens]
3456
+ positions = self.mrope_positions.gpu[:, :num_tokens_after_padding]
3457
else:
3456
- positions = self.positions.gpu[:num_tokens]
3458
+ positions = self.positions.gpu[:num_tokens_after_padding]
3459
3460
if get_pp_group().is_first_rank:
3461
intermediate_tensors = None
@@ -3468,7 +3470,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3470
)
3471
3472
intermediate_tensors = self.sync_and_slice_intermediate_tensors(
3471
- num_tokens, None, False
3473
+ num_tokens_after_padding, None, False
3474
)
3475
3476
# filter out the valid batch descriptor