[Bugfix] Fix UBatchWrapper CUDA graph key to sum all ubatches, not just first two (#43161)
Signed-off-by: Lanze Liu <lanzetech@gmail.com>
Lanze Liu committed
Jul 7, 2026 at 05:42 UTC
8b745527cd6ddedd3afb2c2afc9f2480f471ea70
1 file changed
+3
-3
vllm/v1/worker/gpu_ubatch_wrapper.py
+3
-3
@@ -252,7 +252,7 @@ class UBatchWrapper:
252
253
results: list[tuple[int, torch.Tensor]] = []
254
compute_stream = ubatch_metadata[0].context.compute_stream
255
- num_tokens = ubatch_metadata[0].num_tokens + ubatch_metadata[1].num_tokens
255
+ num_tokens = sum(m.num_tokens for m in ubatch_metadata)
256
257
# Ubatches will manually manage the forward context, so we override
258
# it to None here so we can have it restored correctly later
@@ -268,7 +268,7 @@ class UBatchWrapper:
268
)
269
ubatch_threads.append(thread)
270
thread.start()
271
- self.ready_barrier.wait() # Wait for both threads to be ready
271
+ self.ready_barrier.wait() # Wait for all ubatch threads to be ready
272
273
# Capture the cudagraph
274
cudagraph_metadata = CUDAGraphMetaData(
@@ -332,7 +332,7 @@ class UBatchWrapper:
332
)
333
ubatch_threads.append(thread)
334
thread.start()
335
- self.ready_barrier.wait() # Wait for both threads to be ready
335
+ self.ready_barrier.wait() # Wait for all ubatch threads to be ready
336
ubatch_metadata[0].context.cpu_wait_event.set()
337
for thread in ubatch_threads:
338
thread.join()