[yuweih205]

Summary

Retain a completion event after a successfully executed EventOverlap hook, so
later waits on other CUDA streams also depend on the work enqueued by that hook.

The event returned by the communication operation precedes its deferred epilogue.
The first wait runs the hook on its current stream and clears the hook; another
stream currently waits only for the original communication event. It can then
read tensors before the epilogue's GPU work is complete. The same issue applies
to deterministic sorting registered as an after-wait hook.

The fix keeps the original event and its retained tensors alive, records a
separate native EventHandle after the hook succeeds, and includes it in later
waits. Explicit release_handle=True releases both event handles. Hooks still
run once and only the first wait returns their result. A failing hook remains
retryable. The ordinary no-hook path records no additional CUDA event. No
host/device synchronization is added.

Regression coverage

The current diff includes a standalone two-function CUDA regression in
tests/utils/test_event.py:

• A deliberately delayed hook writer and a second-stream reader, with a raw
  original-event negative control, repeated waits, single-hook execution and
  explicit event-handle release. The test requires the delayed write to still
  be pending when the negative control completes; an ineffective window fails.
• A no-hook wait and context-manager release, checking that no extra epilogue
  event is recorded.

Run it in an installed/built checkout:

python tests/utils/test_event.py

Historical extended harness

The earlier standalone CUDA regression harness, retained in the tested commit
e15813d, checks:

• Two subsequent reader streams after a deliberately delayed writer hook,
  with and without explicit release.
• Context-manager waits after a prior hook on a different stream.
• Multi-stream CUDA Graph capture and three replays.
• No extra event capture when no hook is registered.
• Hook failure/retry and first-wait return semantics.

Its eager race tests require the delayed write to still be pending before the
reader is queued. Otherwise they fail as an ineffective race test. They drain
both streams even if an assertion fails. This historical five-method harness
is not the current two-function regression.

Validation

Current polished revision

Head: ebaaab5f4b58d3fd96ccbc9ba66481ba41cd1703, based on public main
b5a33eefda39485f438dc764695b71ba7b9249b6.

Both current CUDA regression functions passed on H200 with PyTorch
2.13.0+cu132, CUDA 13.2 and NCCL 2.32.3. This run covers the current
two-function regression, not a rerun of the extended historical harness or
its CUDA Graph cases. The production wrapper's Python AST matches the
previously published fixed implementation.

Historical extended validation

Validated on eight H200 141GB GPUs, PyTorch 2.13.0+cu132, CUDA 13.2 and NCCL
2.32.3. Base: public main 93eb6eb238127e96c6d7a4a625a6dad158348509; pinned
DeepJIT cba802df081985141eeee36803685775c183a047. The _C extension was built
from those unmodified native sources. Only the Python event wrapper differs
between the original/fixed regression arms.

• Original wrapper: all 24 suites (three per GPU) reproduce stale reads in
  both explicit-reader release cases and the context-manager case. The CUDA
  Graph case also reproduces stale reads in all 24 suites.
• Fixed wrapper: all 24 suites pass all five methods (120 method runs), with
  no skipped tests. Each Graph test verifies three replays.
• Upstream EP8 functional suite, deterministic, CPU counts enabled: all 36
  deferred-epilogue configurations and all 72 non-deferred configurations pass.
  This covers fresh/cached, compact/expanded, BF16/FP8 dispatch and combine,
  bias variants, previous-event and allocation-stream settings, and repeated
  deterministic dispatch checks.

These results belong to the earlier extended harness and functional-suite
validation. That larger harness is not included in the current diff; the
two-function CUDA regression above is included.

Re-run the upstream EP8 functional suite in an installed/built checkout:

for deferred in 1 0; do
    python tests/ep/test_ep.py --num-processes 8 --num-tokens 128 --hidden 512 \
        --num-topk 2 --num-experts 16 --num-sms 8 --skip-perf-test \
        --deterministic --defer-epilogue "$deferred"
done

The first historical build attempt stopped before tests because the image
lacked the DeepJIT dependency elfutils/libdwfl.h. Reusing the available
libdw/libelf development headers allowed the frozen source/test payload to
build and pass.

This is a stream-ordering correctness change, not a performance claim or a
full-model/training/backward validation.