[PR #600] Fix stream race in dispatch/combine with deterministic algorithms #655

New issue

Open

opened 2026-10-27 21:04:56 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-27 21:04:56 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/600
Author: @ccs1112
Created: 4/6/2026
Status: 🔄 Open

Base: main ← Head: fix/stream-race-nan-deterministic

📝 Commits (1)

a1577f4 Fix stream race producing NaN with deterministic algorithms

📊 Changes

2 files changed (+259 additions, -2 deletions)

View changed files

📝 csrc/deep_ep.cpp (+28 -2)
➕ tests/test_stream_race.py (+231 -0)

📄 Description

The cause of #589 so far as I can tell is as follows in deep_ep.cpp:

// The comm stream waits for everything on the compute stream.
stream_wait(comm_stream, compute_stream)

// The output tensor is allocated, but in deterministic mode a NaN-fill is enqueued on the compute stream.
torch::empty(recv_x)

// The communication kernel writes recv_x on to the comm stream.
intranode::combine(..., recv_x, ..., comm_stream)

Now this could be fixed by putting the stream wait until after the empty allocation, but we have the comment "NOTES: do not allocate tensors upfront!". I don't want to mess up the stream ownership so I just add a second sync stream wait right before the intranode combine:

// The output tensor is allocated, but in deterministic mode a NaN-fill is enqueued on the compute stream.
torch::empty(recv_x)

// Ensure comm_stream observes all compute_stream work above this point.
if (!allocate_on_comm_stream)
    stream_wait(comm_stream, compute_stream);

// The communication kernel writes recv_x on to the comm stream.
intranode::combine(..., recv_x, ..., comm_stream)

The tests should deterministically catch the bug once the fix is in. I don't have a multi-GPU setup on hand to validate myself, any recommendations?

_{🔄 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/600 **Author:** [@ccs1112](https://github.com/ccs1112) **Created:** 4/6/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/stream-race-nan-deterministic` --- ### 📝 Commits (1) - [`a1577f4`](https://github.com/deepseek-ai/DeepEP/commit/a1577f4ea3f1095a717b7ca5323e625d2ee6190d) Fix stream race producing NaN with deterministic algorithms ### 📊 Changes **2 files changed** (+259 additions, -2 deletions) <details> <summary>View changed files</summary> 📝 `csrc/deep_ep.cpp` (+28 -2) ➕ `tests/test_stream_race.py` (+231 -0) </details> ### 📄 Description The cause of #589 so far as I can tell is as follows in `deep_ep.cpp`: ```c // The comm stream waits for everything on the compute stream. stream_wait(comm_stream, compute_stream) // The output tensor is allocated, but in deterministic mode a NaN-fill is enqueued on the compute stream. torch::empty(recv_x) // The communication kernel writes recv_x on to the comm stream. intranode::combine(..., recv_x, ..., comm_stream) ``` Now this could be fixed by putting the stream wait until after the empty allocation, but we have the comment `"NOTES: do not allocate tensors upfront!"`. I don't want to mess up the stream ownership so I just add a second sync stream wait right before the intranode combine: ``` // The output tensor is allocated, but in deterministic mode a NaN-fill is enqueued on the compute stream. torch::empty(recv_x) // Ensure comm_stream observes all compute_stream work above this point. if (!allocate_on_comm_stream) stream_wait(comm_stream, compute_stream); // The communication kernel writes recv_x on to the comm stream. intranode::combine(..., recv_x, ..., comm_stream) ``` The tests should deterministically catch the bug once the fix is in. I don't have a multi-GPU setup on hand to validate myself, any recommendations? --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-27 21:04:56 +00:00

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#655

No description provided.

Rows
Columns