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, 02:00 UTC
Make microbatch optimization (DBO) work with general models (#37926)
Signed-off-by: Junhao Li <junhao@ubicloud.com>
Junhao committed
Mar 24, 2026 at 17:40 UTC
b73b5b06290c0d1439b09db71eef15fe59bc1fbb
1 file changed
+9
-5
vllm/v1/worker/gpu_ubatch_wrapper.py
+9
-5
@@ -389,16 +389,20 @@ class UBatchWrapper:
389
inputs_embeds,
390
intermediate_tensors,
391
):
392
- sliced_input_ids = input_ids[tokens_slice]
392
+ sliced_input_ids = input_ids[tokens_slice] if input_ids is not None else None
393
# if we are using mrope. Mrope adds an additional dimension to the
394
# positions tensor
395
if positions.ndim == 2:
396
sliced_positions = positions[:, tokens_slice]
397
else:
398
sliced_positions = positions[tokens_slice]
399
- sliced_inputs_embeds = inputs_embeds[tokens_slice] if inputs_embeds else None
399
+ sliced_inputs_embeds = (
400
+ inputs_embeds[tokens_slice] if inputs_embeds is not None else None
401
+ )
402
sliced_intermediate_tensors = (
401
- intermediate_tensors[tokens_slice] if intermediate_tensors else None
403
+ intermediate_tensors[tokens_slice]
404
+ if intermediate_tensors is not None
405
+ else None
406
)
407
408
return (
@@ -478,7 +482,7 @@ class UBatchWrapper:
482
cudagraph_runtime_mode=CUDAGraphMode.NONE,
483
)
484
with self.sm_control:
481
- return self._capture_ubatches(ubatch_metadata, self.model)
485
+ return self._capture_ubatches(ubatch_metadata, self.runnable)
486
elif (
487
num_tokens in self.cudagraphs
488
and cudagraph_runtime_mode is CUDAGraphMode.FULL
@@ -504,4 +508,4 @@ class UBatchWrapper:
508
cudagraph_runtime_mode=CUDAGraphMode.NONE,
509
)
510
with self.sm_control:
507
- return self._run_ubatches(ubatch_metadata, self.model)
511
+ return self._run_ubatches(ubatch_metadata, self.runnable)