From 2a7f25fbe2286f57a93984d0b54e15c6430f7039 Mon Sep 17 00:00:00 2001
From: Sage Moore <sage@neuralmagic.com>
Date: Wed, 21 May 2025 21:29:41 +0000
Subject: [PATCH] fix hang

---
 vllm/v1/worker/ubatching.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/vllm/v1/worker/ubatching.py b/vllm/v1/worker/ubatching.py
index 7449f196636f0..2dc84dc579092 100644
--- a/vllm/v1/worker/ubatching.py
+++ b/vllm/v1/worker/ubatching.py
@@ -64,9 +64,9 @@ class UBatchContext:
         self.cpu_signal_event.set()
 
     def _wait(self):
-        self.stream.wait_event(self.gpu_wait_event)
         self.cpu_wait_event.wait()
         self.cpu_wait_event.clear()
+        self.stream.wait_event(self.gpu_wait_event)
         self._restore_context()
 
 _CURRENT_CONTEXT: dict = {}