[PR #766] [MERGED] [HybridEP] Fix the fused combine on a non-blocking dispatch #767

New issue

Closed

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

bsakaguchi commented

2026-10-27 21:06:10 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/766
Author: @VesVlad
Created: 9/26/2026
Status: ✅ Merged
Merged: 10/16/2026
Merged by: @jershi425

Base: hybrid-ep ← Head: hybridep-fix-token-drop-in-fused-combine

📝 Commits (1)

aa3122c [HybridEP] Fix the fused combine on a non-blocking dispatch

📊 Changes

8 files changed (+177 additions, -5 deletions)

View changed files

📝 csrc/hybrid_ep/backend/hybrid_ep_backend.cuh (+19 -4)
📝 csrc/hybrid_ep/hybrid_ep.cu (+2 -0)
📝 csrc/hybrid_ep/hybrid_ep.cuh (+1 -0)
📝 csrc/hybrid_ep/pybind_hybrid_ep.cu (+1 -0)
📝 deep_ep/hybrid_ep_buffer.py (+2 -0)
📝 docs/Hybrid-EP_Implementation.md (+2 -1)
📝 tests/test_graphed_hybrid_ep.py (+2 -0)
➕ tests/test_hybrid_ep_token_drop.py (+148 -0)

📄 Description

A non-blocking dispatch drops the tokens that overflow its fixed capacity. The fused combine hangs on them, and once it stops hanging it returns NaN.

Why

The scan kernel clamps a dropped token's entries in dense_to_expert_map to -1, which leaves dst tokens with no src token at all. Two things break:

Hang. The dummy entry that serves such a dst token exists only under HYBRID_EP_BUILD_TOKEN_DROP_ENABLE, which is set when fuse_unpermute_combine && non_blocking. But combine_with_unpermute never took a non_blocking argument, so the combine was always built without it: the producer pushes nothing, the consumer waits for the end-of-group flag forever.
NaN. With that branch present, the dummy entry's payload is still never written, so the reduction sums whatever the previous dst token left in that shared-memory stage or uninitialized memory, which decodes to NaN in bf16.

What changed

non_blocking plumbed through combine_with_unpermute, mirroring dispatch_with_permute, and assigned to args.non_blocking.
The dummy entry is zeroed – payload, plus the prob and its local expert id under BACKWARD_COMBINE – with a __threadfence_block() before the mbarrier arrive.

A fully dropped dst token now contributes zero, which is what the standalone unpermute path in extension/permute.cu already did.

Testing

tests/test_hybrid_ep_token_drop.py (new): top-1 routing at half the capacity a balanced dispatch needs, so tokens really are dropped. Asserts the overflow flag fires, the run does not hang, surviving tokens return bitwise unchanged and fully dropped ones return zeroed. tests/test_graphed_hybrid_ep.py now passes non_blocking=True to the combine, matching its dispatch.

_{🔄 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/766 **Author:** [@VesVlad](https://github.com/VesVlad) **Created:** 9/26/2026 **Status:** ✅ Merged **Merged:** 10/16/2026 **Merged by:** [@jershi425](https://github.com/jershi425) **Base:** `hybrid-ep` ← **Head:** `hybridep-fix-token-drop-in-fused-combine` --- ### 📝 Commits (1) - [`aa3122c`](https://github.com/deepseek-ai/DeepEP/commit/aa3122c18fc7c16300505264b59b758fc1ee2c46) [HybridEP] Fix the fused combine on a non-blocking dispatch ### 📊 Changes **8 files changed** (+177 additions, -5 deletions) <details> <summary>View changed files</summary> 📝 `csrc/hybrid_ep/backend/hybrid_ep_backend.cuh` (+19 -4) 📝 `csrc/hybrid_ep/hybrid_ep.cu` (+2 -0) 📝 `csrc/hybrid_ep/hybrid_ep.cuh` (+1 -0) 📝 `csrc/hybrid_ep/pybind_hybrid_ep.cu` (+1 -0) 📝 `deep_ep/hybrid_ep_buffer.py` (+2 -0) 📝 `docs/Hybrid-EP_Implementation.md` (+2 -1) 📝 `tests/test_graphed_hybrid_ep.py` (+2 -0) ➕ `tests/test_hybrid_ep_token_drop.py` (+148 -0) </details> ### 📄 Description A non-blocking dispatch drops the tokens that overflow its fixed capacity. The fused combine hangs on them, and once it stops hanging it returns NaN. **Why** The scan kernel clamps a dropped token's entries in dense_to_expert_map to -1, which leaves dst tokens with no src token at all. Two things break: 1. **Hang.** The dummy entry that serves such a dst token exists only under `HYBRID_EP_BUILD_TOKEN_DROP_ENABLE`, which is set when `fuse_unpermute_combine && non_blocking`. But `combine_with_unpermute` never took a `non_blocking` argument, so the combine was always built without it: the producer pushes nothing, the consumer waits for the end-of-group flag forever. 2. **NaN.** With that branch present, the dummy entry's payload is still never written, so the reduction sums whatever the previous dst token left in that shared-memory stage or uninitialized memory, which decodes to NaN in bf16. **What changed** - `non_blocking` plumbed through `combine_with_unpermute`, mirroring `dispatch_with_permute`, and assigned to `args.non_blocking`. - The dummy entry is zeroed – payload, plus the prob and its local expert id under `BACKWARD_COMBINE` – with a `__threadfence_block()` before the mbarrier arrive. A fully dropped dst token now contributes zero, which is what the standalone unpermute path in extension/permute.cu already did. **Testing** `tests/test_hybrid_ep_token_drop.py` (new): top-1 routing at half the capacity a balanced dispatch needs, so tokens really are dropped. Asserts the overflow flag fires, the run does not hang, surviving tokens return bitwise unchanged and fully dropped ones return zeroed. `tests/test_graphed_hybrid_ep.py` now passes `non_blocking=True` to the combine, matching its dispatch. --- <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:10 +00:00