Fix cross-stream waits after deferred epilogues

deep_ep/utils/event.py

@@ -36,6 +36,9 @@ def __init__(self, event: Optional[EventHandle] = None, extra_tensors: Optional[
         # Useful for deterministic dispatch, which requires a sort (on the current stream) after `self.current_stream_wait()` is invoked
         self.hook_after_wait: Optional[Callable] = None
 
+        # Later waits must include GPU work enqueued by the hook.
+        self._epilogue_event: Optional[EventHandle] = None
+
     def current_stream_wait(self, release_handle: bool = False) -> Any:
         """
         The current stream `torch.cuda.current_stream()` waits for the event to be finished.
@@ -45,18 +48,22 @@ def current_stream_wait(self, release_handle: bool = False) -> Any:
         """
         assert self.event is not None
         self.event.current_stream_wait()
+        if self._epilogue_event is not None:
+            self._epilogue_event.current_stream_wait()
 
         # Call epilogue hook
         result = None
         if self.hook_after_wait is not None:
             result = self.hook_after_wait()
+            self._epilogue_event = EventHandle()
             self.hook_after_wait = None
 
         # In `self.event`, we also have some V2 APIs storing tensors to record in it,
         # So, after waiting the current stream, those tensors can be released by deleting `self.event`
         # However, you better do it by yourself (to be compatible with multi-stream waits)
         if release_handle:
             self.event = None
+            self._epilogue_event = None
         return result
 
     def wait(self) -> Any:

tests/utils/test_event.py

@@ -0,0 +1,72 @@
+import torch
+
+from deep_ep import EventHandle, EventOverlap
+
+
+def test_multistream_hook():
+    hook_stream, reader_stream = torch.cuda.Stream(), torch.cuda.Stream()
+    payload = torch.zeros(1, device='cuda', dtype=torch.int32)
+    control, received = torch.empty_like(payload), torch.empty_like(payload)
+    control_done, hook_done = torch.cuda.Event(), torch.cuda.Event()
+
+    # Warm up the kernels and event API before constructing the delayed dependency.
+    control.copy_(payload)
+    received.copy_(payload)
+    torch.cuda._sleep(1)
+    EventHandle().current_stream_wait()
+    torch.cuda.synchronize()
+
+    event = EventHandle()
+    overlap = EventOverlap(event)
+    calls = 0
+
+    def hook():
+        nonlocal calls
+        calls += 1
+        torch.cuda._sleep(200_000_000)
+        payload.fill_(123)
+        hook_done.record()
+        return payload
+
+    overlap.register_hook_after_wait(hook)
+    with torch.cuda.stream(hook_stream):
+        assert overlap.current_stream_wait() is payload
+    epilogue_event = overlap._epilogue_event
+    assert epilogue_event is not None
+
+    with torch.cuda.stream(reader_stream):
+        # The original event precedes the hook: this negative control must see zero.
+        event.current_stream_wait()
+        control.copy_(payload)
+        control_done.record()
+        # A second wait must additionally wait for the first stream's hook writes.
+        assert overlap.wait() is None
+        received.copy_(payload)
+
+    control_done.synchronize()
+    assert not hook_done.query(), 'The hook finished before the cross-stream dependency was exercised'
+    assert control.item() == 0, 'The original event unexpectedly covered the hook'
+    reader_stream.synchronize()
+    assert received.item() == 123
+    assert calls == 1 and overlap.hook_after_wait is None
+    assert overlap._epilogue_event is epilogue_event
+
+    with torch.cuda.stream(reader_stream):
+        assert overlap.current_stream_wait(release_handle=True) is None
+    assert overlap.event is None and overlap._epilogue_event is None
+    assert calls == 1
+
+
+def test_no_hook():
+    overlap = EventOverlap(EventHandle())
+    assert overlap.wait() is None
+    assert overlap._epilogue_event is None
+    with overlap(release_handle=True):
+        pass
+    assert overlap.event is None and overlap._epilogue_event is None
+
+
+if __name__ == '__main__':
+    test_multistream_hook()
+    test_no_hook()
+    print('EventOverlap tests passed')