verify_done: wait not synchronize (#25465)

2026-05-19 14:57:45 -07:00
parent fab097d66d
commit 16bcc4583e
3 changed files with 75 additions and 37 deletions
@@ -2451,10 +2451,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
            )

    def maybe_wait_verify_done(self):
+        # Use event.wait() (stream-level wait) instead of .synchronize()
+        # (CPU block). Schedule-stream prep ops following this call get
+        # ordered after the forward-stream verify via the wait; CPU is not
+        # blocked. Subsequent .cpu()/.item() naturally sync the stream.
        if self.is_spec_v2:
            draft_input: EagleDraftInput = self.spec_info
            if draft_input.verify_done is not None:
-                draft_input.verify_done.synchronize()
+                draft_input.verify_done.wait()

    def filter_batch(
        self,