[PR #776] [CLOSED] Fix cross-stream waits after deferred epilogues #775

Closed
opened 2026-10-27 21:06:13 +00:00 by bsakaguchi · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/776
Author: @yuweih205
Created: 10/9/2026
Status: ❌ Closed

Base: main ← Head: fix/deferred-epilogue-multistream-wait


📝 Commits (1)

  • ebaaab5 Fix repeated cross-stream waits after deferred epilogues

📊 Changes

2 files changed (+79 additions, -0 deletions)

View changed files

📝 deep_ep/utils/event.py (+7 -0)
➕ tests/utils/test_event.py (+72 -0)

📄 Description

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.


🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepEP/pull/776 **Author:** [@yuweih205](https://github.com/yuweih205) **Created:** 10/9/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `fix/deferred-epilogue-multistream-wait` --- ### 📝 Commits (1) - [`ebaaab5`](https://github.com/deepseek-ai/DeepEP/commit/ebaaab5f4b58d3fd96ccbc9ba66481ba41cd1703) Fix repeated cross-stream waits after deferred epilogues ### 📊 Changes **2 files changed** (+79 additions, -0 deletions) <details> <summary>View changed files</summary> 📝 `deep_ep/utils/event.py` (+7 -0) ➕ `tests/utils/test_event.py` (+72 -0) </details> ### 📄 Description ## 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: ```bash python tests/utils/test_event.py ``` ### Historical extended harness The earlier standalone CUDA regression harness, retained in the tested commit [`e15813d`](https://github.com/yuweih205/DeepEP/blob/e15813df91be6a4288b280b8cb26a5e84b96c0e0/tests/utils/test_event.py), 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: ```bash 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. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
bsakaguchi 2026-10-27 21:06:13 +00:00
Sign in to join this conversation.
No labels
pull-request
No milestone
No project
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
deepseek-ai/DeepEP#775
No description provided.