deepseek-ai/DeepEP

Fork 0

mirror of https://github.com/deepseek-ai/DeepEP.git synced 2026-11-05 00:47:32 +00:00

[PR #753] DeepEP v2: Kernel-side deterministic dispatch for the non-expand layout #757

New issue

Open

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

bsakaguchi commented

2026-10-27 21:05:59 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/753
Author: @EricWang008
Created: 9/17/2026
Status: 🔄 Open

Base: main ← Head: deterministic-hybrid-dispatch

📝 Commits (4)

3586de8 feat(elastic): add kernel-side deterministic dispatch path
2d69687 feat(elastic): wire the deterministic prologue and row-form epilogue launchers
35f4c8d feat(elastic): orchestrate the deterministic dispatch on the host side
4cd6431 feat(elastic): keep the sorting epilogue only for the expand layout

📊 Changes

7 files changed (+472 additions, -45 deletions)

View changed files

📝 csrc/elastic/buffer.hpp (+75 -8)
📝 csrc/kernels/elastic/dispatch.hpp (+98 -6)
📝 deep_ep/buffers/elastic.py (+14 -4)
📝 deep_ep/include/deep_ep/common/layout.cuh (+10 -0)
📝 deep_ep/include/deep_ep/impls/dispatch_copy_epilogue.cuh (+27 -6)
➕ deep_ep/include/deep_ep/impls/dispatch_deterministic_prologue.cuh (+151 -0)
📝 deep_ep/include/deep_ep/impls/hybrid_dispatch.cuh (+97 -21)

📄 Description

Motivation

Deterministic dispatch is currently implemented as a post-hoc sort in EPHandle.deterministic_sort: after the dispatch kernels finish, the received tokens are sorted by src_token_global_idx and every order-dependent tensor (recv_x, recv_sf, recv_topk_weights, recv_topk_idx, recv_src_metadata, channel_linked_list) is permuted in place.

This is correct but costly. The sort plus the permutations add a full extra round trip over the received activations, and the cost grows with num_recv_tokens * hidden.

The root cause of the non-determinism is that a sender picks its destination slot with an atomicAdd on a shared counter, so the placement depends on the inter-warp race:

if (ptx::deduplicate(stored_dst_rank_idx, lane_idx) and stored_dst_rank_idx >= 0)
    stored_dst_slot_idx = atomicAdd(workspace_layout.get_scaleup_atomic_sender_counter() + stored_dst_rank_idx, 1);

This PR removes that race for the non-expand layout instead of repairing its consequences afterwards.

Approach

A lightweight prologue kernel (dispatch_deterministic_prologue_impl) pre-computes the destination slot of every (token, top-k) entry, making the placement a pure function of topk_idx. The dispatch kernels then replay those slots rather than racing for them, so the output is bitwise reproducible without any sorting.

The emitted value is row_base_idx * num_max_tokens_per_rank + row-local slot, where the row-local slot is the sender's running count of deduplicated entries to that destination. A row is private to a single source rank, which can never send more than num_max_tokens_per_rank deduplicated entries to one destination, so rows cannot overflow.

The prologue is a cooperative kernel doing a two-pass scan: the first pass accumulates per-warp / per-SM destination counts, a grid-wide sync exchanges them, and the second pass turns the prefix sums into per-entry slots.

Direct vs hybrid dispatch

The two paths need different row bases:

	`row_base_idx`	receive-buffer rows
direct (`num_scaleout_ranks == 1`)	`scaleup_rank_idx`	already one row per source rank
hybrid (`num_scaleout_ranks > 1`)	`scaleout_rank_idx`	scale-up segment is selected by the forwarder

The direct path needs no layout change: its receive buffer is [num_ranks][num_max_tokens_per_rank], i.e. it is already row-form by construction.

The hybrid path is different. Its scale-up buffer has only num_scaleup_ranks segments while there are num_ranks possible sources, so locating a token requires per-source-rank counts. This PR therefore adds a row-form interpretation of the hybrid scale-up receive buffer, enabled by the new kRowForm template argument of dispatch_copy_epilogue_impl: every global source rank owns a private (partially filled) row, and a token sits at segment = src_scaleup_rank_idx, in-segment offset = src_scaleout_rank_idx * num_max_tokens_per_rank + row-local slot.

To feed that traversal, the hybrid notify phase now also publishes per-source-rank received-token counts, and psum_num_recv_tokens_per_scaleup_rank is extended from [num_scaleup_ranks] to [num_scaleup_ranks | num_ranks] in this mode.

Compatibility

No existing behaviour is removed. The kernel path is an addition, not a replacement:

non-expand → kernel-side path (no sort)
expand → keeps using deterministic_sort

The expanded layout assigns its output rows with atomicAdd on psum_num_recv_tokens_per_expert, which the kernel path cannot make reproducible, so those cases must keep the sorting epilogue. The two are strictly complementary: C++ gates on deterministic and not do_expand, Python gates on self.deterministic and do_expand.

kDoZeroPadding, kExpertAlignment, the cached-expand path and deterministic_sort itself are all preserved.

Incidental fix

combine asserted num_scaleup_ranks == nccl_context->num_scaleup_ranks on the psum tensor. Since the row-form dispatch appends num_ranks extra prefix sums, that check is relaxed to accept either length.

New API

dispatch() gains an optional num_prologue_sms (0 = automatic, capped at 64 SMs). It only affects the deterministic prologue.

Testing

Verified on H800 2×8 with tests/elastic/test_ep.py --deterministic, covering both the direct and hybrid dispatch paths as well as the expand layout (which still falls back to the sorting epilogue).

Performance

Configuration:


GPU	H800 2×8
NIC	CX7 (2×200 Gbps) 2×8
experts	8 / 256
tokens	4096
hidden	7168
SM	24
QPs	129 / 129
other	`do_expand=false`, BF16, `expert_alignment=128`, `num_bias=2`

	deterministic_prologue	dispatch	copy_epilogue	deterministic_sort	whole time
This PR	6 us	2590 us	300 us	–	2896 us
Original	–	2598 us	300 us	635 us	3533 us

The 635 us sorting epilogue is replaced by a 6 us prologue, a ~18% reduction in end-to-end dispatch time (1.22× speedup). dispatch and copy_epilogue themselves are unchanged, confirming the row-form traversal and the slot replay add no measurable overhead.

_{🔄 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/753 **Author:** [@EricWang008](https://github.com/EricWang008) **Created:** 9/17/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `deterministic-hybrid-dispatch` --- ### 📝 Commits (4) - [`3586de8`](https://github.com/deepseek-ai/DeepEP/commit/3586de8977b9938c78f10b92af7b6070d682e519) feat(elastic): add kernel-side deterministic dispatch path - [`2d69687`](https://github.com/deepseek-ai/DeepEP/commit/2d69687b1c1264a32c0625d5aa74e8a0f6b90df3) feat(elastic): wire the deterministic prologue and row-form epilogue launchers - [`35f4c8d`](https://github.com/deepseek-ai/DeepEP/commit/35f4c8dbff911d081c19e6de4ee904de2760d8f6) feat(elastic): orchestrate the deterministic dispatch on the host side - [`4cd6431`](https://github.com/deepseek-ai/DeepEP/commit/4cd64318c9c22572256dd6fa176d33766f999c98) feat(elastic): keep the sorting epilogue only for the expand layout ### 📊 Changes **7 files changed** (+472 additions, -45 deletions) <details> <summary>View changed files</summary> 📝 `csrc/elastic/buffer.hpp` (+75 -8) 📝 `csrc/kernels/elastic/dispatch.hpp` (+98 -6) 📝 `deep_ep/buffers/elastic.py` (+14 -4) 📝 `deep_ep/include/deep_ep/common/layout.cuh` (+10 -0) 📝 `deep_ep/include/deep_ep/impls/dispatch_copy_epilogue.cuh` (+27 -6) ➕ `deep_ep/include/deep_ep/impls/dispatch_deterministic_prologue.cuh` (+151 -0) 📝 `deep_ep/include/deep_ep/impls/hybrid_dispatch.cuh` (+97 -21) </details> ### 📄 Description ## Motivation Deterministic dispatch is currently implemented as a post-hoc sort in `EPHandle.deterministic_sort`: after the dispatch kernels finish, the received tokens are sorted by `src_token_global_idx` and every order-dependent tensor (`recv_x`, `recv_sf`, `recv_topk_weights`, `recv_topk_idx`, `recv_src_metadata`, `channel_linked_list`) is permuted in place. This is correct but costly. The sort plus the permutations add a full extra round trip over the received activations, and the cost grows with `num_recv_tokens * hidden`. The root cause of the non-determinism is that a sender picks its destination slot with an `atomicAdd` on a shared counter, so the placement depends on the inter-warp race: ```cpp if (ptx::deduplicate(stored_dst_rank_idx, lane_idx) and stored_dst_rank_idx >= 0) stored_dst_slot_idx = atomicAdd(workspace_layout.get_scaleup_atomic_sender_counter() + stored_dst_rank_idx, 1); ``` This PR removes that race for the non-expand layout instead of repairing its consequences afterwards. ## Approach A lightweight prologue kernel (`dispatch_deterministic_prologue_impl`) pre-computes the destination slot of every `(token, top-k)` entry, making the placement a **pure function of `topk_idx`**. The dispatch kernels then replay those slots rather than racing for them, so the output is bitwise reproducible without any sorting. The emitted value is `row_base_idx * num_max_tokens_per_rank + row-local slot`, where the row-local slot is the sender's running count of deduplicated entries to that destination. A row is private to a single source rank, which can never send more than `num_max_tokens_per_rank` deduplicated entries to one destination, so rows cannot overflow. The prologue is a cooperative kernel doing a two-pass scan: the first pass accumulates per-warp / per-SM destination counts, a grid-wide sync exchanges them, and the second pass turns the prefix sums into per-entry slots. ### Direct vs hybrid dispatch The two paths need different row bases: | | `row_base_idx` | receive-buffer rows | |---|---|---| | direct (`num_scaleout_ranks == 1`) | `scaleup_rank_idx` | already one row per source rank | | hybrid (`num_scaleout_ranks > 1`) | `scaleout_rank_idx` | scale-up segment is selected by the forwarder | The direct path needs no layout change: its receive buffer is `[num_ranks][num_max_tokens_per_rank]`, i.e. it is already row-form by construction. The hybrid path is different. Its scale-up buffer has only `num_scaleup_ranks` segments while there are `num_ranks` possible sources, so locating a token requires per-source-rank counts. This PR therefore adds a **row-form** interpretation of the hybrid scale-up receive buffer, enabled by the new `kRowForm` template argument of `dispatch_copy_epilogue_impl`: every global source rank owns a private (partially filled) row, and a token sits at `segment = src_scaleup_rank_idx`, `in-segment offset = src_scaleout_rank_idx * num_max_tokens_per_rank + row-local slot`. To feed that traversal, the hybrid notify phase now also publishes per-source-rank received-token counts, and `psum_num_recv_tokens_per_scaleup_rank` is extended from `[num_scaleup_ranks]` to `[num_scaleup_ranks | num_ranks]` in this mode. ## Compatibility **No existing behaviour is removed.** The kernel path is an addition, not a replacement: - **non-expand** → kernel-side path (no sort) - **expand** → keeps using `deterministic_sort` The expanded layout assigns its output rows with `atomicAdd` on `psum_num_recv_tokens_per_expert`, which the kernel path cannot make reproducible, so those cases must keep the sorting epilogue. The two are strictly complementary: C++ gates on `deterministic and not do_expand`, Python gates on `self.deterministic and do_expand`. `kDoZeroPadding`, `kExpertAlignment`, the cached-expand path and `deterministic_sort` itself are all preserved. ### Incidental fix `combine` asserted `num_scaleup_ranks == nccl_context->num_scaleup_ranks` on the psum tensor. Since the row-form dispatch appends `num_ranks` extra prefix sums, that check is relaxed to accept either length. ### New API `dispatch()` gains an optional `num_prologue_sms` (0 = automatic, capped at 64 SMs). It only affects the deterministic prologue. ## Testing Verified on H800 2×8 with `tests/elastic/test_ep.py --deterministic`, covering both the direct and hybrid dispatch paths as well as the expand layout (which still falls back to the sorting epilogue). ## Performance Configuration: | | | |---|---| | GPU | H800 2×8 | | NIC | CX7 (2×200 Gbps) 2×8 | | experts | 8 / 256 | | tokens | 4096 | | hidden | 7168 | | SM | 24 | | QPs | 129 / 129 | | other | `do_expand=false`, BF16, `expert_alignment=128`, `num_bias=2` | | | deterministic_prologue | dispatch | copy_epilogue | deterministic_sort | whole time | |---|---|---|---|---|---| | **This PR** | 6 us | 2590 us | 300 us | – | **2896 us** | | **Original** | – | 2598 us | 300 us | 635 us | **3533 us** | The 635 us sorting epilogue is replaced by a 6 us prologue, a **~18% reduction in end-to-end dispatch time (1.22× speedup)**. `dispatch` and `copy_epilogue` themselves are unchanged, confirming the row-form traversal and the slot replay add no measurable overhead. --- <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:06:00 +00:00

bsakaguchi referenced this issue

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

[PR #757] [CLOSED] Fix/issue756 unpermute scalar #761