[Bugfix] Fix UBatchWrapper CUDA graph key to sum all ubatches, not just first two (#43161)

Signed-off-by: Lanze Liu <lanzetech@gmail.com>

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()