ITADN

DeepEP: CUDA stream race between masked_fill and get_dispatch_layout causes grouped_mm assertion failure

#3413OpenacisseJZhong 创建于 2026-05-21
A
acisseJZhongcommented
## 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 条评论