`buffer.combine()` produces NaN when `torch.use_deterministic_algorithms(True)` due to stream race on output tensor

# Summary

When using DeepEP with `torch.use_deterministic_algorithms(True)`, `buffer.combine()` intermittently produces NaN outputs. The root cause is a **write-write CUDA stream race condition** between PyTorch's NaN-fill kernel (on `compute_stream`) and DeepEP's combine kernel (on `comm_stream`) writing to the same output tensor `recv_x`.

# Root Cause

In `deep_ep.cpp`, all dispatch/combine functions follow this pattern:

```cpp
// Step 1: stream_wait synchronizes comm_stream with compute_stream
// comm_stream will wait for everything *currently* enqueued on compute_stream
stream_wait(comm_stream, compute_stream);

// Step 2: torch::empty() allocates tensors AFTER the sync point
// When fill_uninitialized_memory=True, this launches a NaN-fill kernel
// on compute_stream — which comm_stream does NOT wait for
auto recv_x = torch::empty({num_recv_tokens, hidden}, x.options());

// Step 3: communication kernel runs on comm_stream, writing to recv_x
intranode::combine(..., recv_x.data_ptr(), ..., comm_stream, ...);
```

The NaN-fill kernel (Step 2, `compute_stream`) and the combine kernel (Step 3, `comm_stream`) execute concurrently with no ordering guarantee. If the NaN-fill kernel writes to a memory location **after** the combine kernel has already written the correct value there, the correct value is overwritten with NaN.

### Timeline

```
compute_stream: ··· work ··· |← sync point | NaN-fill(recv_x) ···
↕ RACE!
comm_stream: ── wait ──→ | combine(recv_x) ··················
```

## Why `torch.use_deterministic_algorithms` triggers this

`torch.use_deterministic_algorithms(True)` sets [`fill_uninitialized_memory = True`](https://docs.pytorch.org/docs/stable/deterministic.html), which causes every `torch::empty()` call to launch a NaN-fill kernel on the current stream. Without this flag, `torch::empty()` returns raw memory from the caching allocator without launching any kernel, so there is nothing to race with.

## Affected code paths

The `stream_wait`-before-`torch::empty` pattern exists in all four core functions in `csrc/deep_ep.cpp`:

| Function | `stream_wait` | `torch::empty` (output tensor) |
| -------------------- | ------------- | ------------------------------ |
| `intranode_dispatch` | L569–574 | L656 |
| `intranode_combine` | L813–818 | L859 |
| `internode_dispatch` | L1048–1054 | L1170 |
| `internode_combine` | L1365–1371 | (similar pattern) |

# Reproduction

## Environment

- PyTorch 2.8 with CUDA
- DeepEP commit `567632dd59810d77b3cc05553df953cc0f779799`
- 8× NVIDIA H20

## Minimal reproduction script

```python
"""
repro_deepep_nan.py — Reproduce stream-race NaN with DeepEP dispatch/combine.

Usage:
# Reproduce NaN
torchrun --nproc_per_node 8 repro_deepep_nan.py

# Control 1: disable fill → no race
NO_FILL_UNINIT=1 torchrun --nproc_per_node 8 repro_deepep_nan.py

# Control 2: serialize kernels → no race
CUDA_LAUNCH_BLOCKING=1 torchrun --nproc_per_node 8 repro_deepep_nan.py
"""

import os
import torch
import torch.distributed as dist
import deep_ep

NUM_ITERS = int(os.environ.get("NUM_ITERS", "200"))

def main():
dist.init_process_group(backend="nccl")
rank = dist.get_rank()
world_size = dist.get_world_size()
local_rank = int(os.environ.get("LOCAL_RANK", rank % torch.cuda.device_count()))
torch.cuda.set_device(local_rank)
torch.set_default_dtype(torch.bfloat16)
torch.set_default_device('cuda')

os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
torch.use_deterministic_algorithms(True, warn_only=True)

if os.environ.get("NO_FILL_UNINIT", "0") == "1":
torch.utils.deterministic.fill_uninitialized_memory = False

test = torch.empty(10, device="cuda")
fill_active = torch.isnan(test).any().item()

num_ranks = world_size
num_tokens = 4096
hidden = 4096
num_topk = 8
num_experts = 256
assert num_experts % num_ranks == 0

if rank == 0:
print(f"fill_uninitialized_memory active: {fill_active}", flush=True)
print(f"num_ranks={num_ranks}, tokens={num_tokens}, hidden={hidden}", flush=True)
print(f"experts={num_experts}, top_k={num_topk}, iters={NUM_ITERS}", flush=True)
print("", flush=True)

x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
scores = torch.randn((num_tokens, num_experts), dtype=torch.float32, device='cuda').abs() + 1
topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False)[1]
topk_idx = topk_idx.to(deep_ep.topk_idx_t)
topk_weights = torch.randn((num_tokens, num_topk), dtype=torch.float32, device='cuda')

group = dist.new_group(list(range(world_size)))
config = deep_ep.Config(24, 8, 256)
buffer = deep_ep.Buffer(
group=group,
num_nvl_bytes=int(2e9),
num_rdma_bytes=0,
allow_mnnvl=False,
)

num_tokens_per_rank, _, num_tokens_per_expert, is_token_in_rank, _ = \
buffer.get_dispatch_layout(topk_idx, num_experts)

# Initial dispatch to establish handle
recv_x, recv_topk_idx, recv_topk_weights, _, handle, event = buffer.dispatch(
x=x, topk_idx=topk_idx, topk_weights=topk_weights,
num_tokens_per_rank=num_tokens_per_rank,
is_token_in_rank=is_token_in_rank,
num_tokens_per_expert=num_tokens_per_expert,
config=config, async_finish=True,
)
event.current_stream_wait()

# Baseline check
combined_x, _, _ = buffer.combine(
x=recv_x, handle=handle, topk_weights=recv_topk_weights,
config=config, async_finish=False,
)
torch.cuda.synchronize()
if rank == 0:
print(f"Baseline combine NaN: {torch.isnan(combined_x).any().item()}", flush=True)
print("", flush=True)

# Tight loop: cached dispatch + combine
nan_counter = torch.zeros(1, dtype=torch.int64, device="cuda")
group.barrier()

if rank == 0:
print(f"Running {NUM_ITERS} cached dispatch + combine ...", flush=True)

for it in range(NUM_ITERS):
recv_x, _, _, _, _, event = buffer.dispatch(
x=x, handle=handle, config=config, async_finish=True,
)
event.current_stream_wait()

combined_x, _, _ = buffer.combine(
x=recv_x, handle=handle, config=config, async_finish=False,
)

nan_counter += torch.isnan(combined_x).any().to(torch.int64).unsqueeze(0)

torch.cuda.synchronize()
local_nans = nan_counter.item()

tc = torch.tensor([local_nans], dtype=torch.int64, device="cuda")
dist.all_reduce(tc)

if rank == 0:
print(f"\n{'='*60}", flush=True)
print(f"COMBINE NaN: {tc.item()} (all ranks, {NUM_ITERS} iterations)", flush=True)
if tc.item() > 0:
print(f" >>> BUG REPRODUCED <<<", flush=True)
else:
print(f" Not triggered this run", flush=True)
print(f"{'='*60}", flush=True)

dist.destroy_process_group()

if __name__ == "__main__":
main()
```

## Expected vs actual output

**Expected:** `COMBINE NaN: 0`

**Actual (typical):**

```
fill_uninitialized_memory active: True
num_ranks=8, tokens=4096, hidden=4096
experts=256, top_k=8, iters=200

Baseline combine NaN: False

Running 200 cached dispatch + combine ...

============================================================
COMBINE NaN: 37 (all ranks, 200 iterations)
>>> BUG REPRODUCED <<<
============================================================
```

## Control experiments

| Experiment | NaN count | Explanation |
| ------------------------------ | --------- | ------------------------------------------ |
| Default (`deterministic=True`) | >0 | Bug reproduced |
| `NO_FILL_UNINIT=1` | 0 | No fill kernel → no race |
| `CUDA_LAUNCH_BLOCKING=1` | 0 | All kernels serialized → no race |
| `deterministic=False` | 0 | `torch::empty` does not launch fill kernel |

# Suggested Fix

## Option A (recommended): Move `stream_wait` after all `torch::empty` allocations

```cpp
// 1. Allocate FIRST — NaN-fill kernels are enqueued on compute_stream
auto recv_x = torch::empty({num_recv_tokens, hidden}, x.options());

// 2. THEN synchronize — comm_stream now also waits for the NaN-fill kernels
if (previous_event.has_value()) {
stream_wait(comm_stream, previous_event.value());
} else {
stream_wait(comm_stream, compute_stream);
}

// 3. Communication kernel runs after NaN-fill is guaranteed to complete
intranode::combine(..., recv_x.data_ptr(), ..., comm_stream, ...);
```

This ensures `comm_stream` observes (and waits for) the NaN-fill kernel before launching the combine kernel. The combine kernel then overwrites all NaN values with correct results.

## Option B: Allocate output tensors on `comm_stream`

Use the existing `allocate_on_comm_stream=True` parameter. When tensors are allocated on `comm_stream`, the NaN-fill kernel runs on `comm_stream` as well, serialized before the combine kernel. No cross-stream race.

## Option C (user-side workaround): Disable `fill_uninitialized_memory`

```python
torch.use_deterministic_algorithms(True, warn_only=True)
torch.utils.deterministic.fill_uninitialized_memory = False # workaround
```

This is safe — `fill_uninitialized_memory` is a debug aid and does not affect numerical determinism. However, it requires every user to know about this workaround.