DeepEP: CUDA stream race between masked_fill and get_dispatch_layout causes grouped_mm assertion failure
## Bug
DeepEP + EP=2 + TP=2 crashes at step 1 with `GroupMMCommon.cuh:89: offset <= tensor_ShapeA[0]`. Fixed by `CUDA_LAUNCH_BLOCKING=1`.
**Repro** (commit 83e49042 / main, needs `moe_comm_backend="deepep"` in config):
```bash
NCCL_NVLS_ENABLE=0 NGPU=4 MODULE=llama4 CONFIG=llama4_debugmodel_ep ./run_train.sh \
--parallelism.data_parallel_shard_degree 2 --parallelism.tensor_parallel_degree 2 \
--parallelism.expert_parallel_degree 2 --debug.deterministic --debug.seed=42
```
## Root cause
In `dispatch_tokens()` (`deepep.py`):
1. `masked_fill(top_scores == 0, -1)` enqueues on the **compute stream** (not yet executed)
2. `buffer.get_dispatch_layout(topk_idx=selected_experts_indices)` reads the tensor on DeepEP's **comm stream** — sees pre-masked values, counts all tokens as valid
3. `buffer.dispatch()` runs after `masked_fill` completes — produces `dispatched_indices` with -1 entries
4. `_permute_tokens` filters -1 → fewer tokens than `num_tokens_per_expert` expects → assertion
Diagnostic output:
```
sum(num_tokens_per_expert)=8848 != permuted_tokens=8361
```
## Fix
`torch.cuda.synchronize()` before `get_dispatch_layout` ensures `masked_fill` results are visible. A targeted CUDA event wait would be more efficient.
1 条评论