From 8b745527cd6ddedd3afb2c2afc9f2480f471ea70 Mon Sep 17 00:00:00 2001 From: Lanze Liu <86434077+liulanze@users.noreply.github.com> Date: Tue, 7 Jul 2026 05:42:31 -0700 Subject: [PATCH] [Bugfix] Fix UBatchWrapper CUDA graph key to sum all ubatches, not just first two (#43161) Signed-off-by: Lanze Liu --- vllm/v1/worker/gpu_ubatch_wrapper.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py index 76fa12b412..7b61999843 100644 --- a/vllm/v1/worker/gpu_ubatch_wrapper.py +++ b/vllm/v1/worker/gpu_ubatch_wrapper.py @@ -252,7 +252,7 @@ class UBatchWrapper: results: list[tuple[int, torch.Tensor]] = [] compute_stream = ubatch_metadata[0].context.compute_stream - num_tokens = ubatch_metadata[0].num_tokens + ubatch_metadata[1].num_tokens + num_tokens = sum(m.num_tokens for m in ubatch_metadata) # Ubatches will manually manage the forward context, so we override # it to None here so we can have it restored correctly later @@ -268,7 +268,7 @@ class UBatchWrapper: ) ubatch_threads.append(thread) thread.start() - self.ready_barrier.wait() # Wait for both threads to be ready + self.ready_barrier.wait() # Wait for all ubatch threads to be ready # Capture the cudagraph cudagraph_metadata = CUDAGraphMetaData( @@ -332,7 +332,7 @@ class UBatchWrapper: ) ubatch_threads.append(thread) thread.start() - self.ready_barrier.wait() # Wait for both threads to be ready + self.ready_barrier.wait() # Wait for all ubatch threads to be ready ubatch_metadata[0].context.cpu_wait_event.set() for thread in ubatch_threads: thread.join()