deepseek-ai / deepseek-ai/DeepEP
[Proposal] deterministic hybrid dispatch implementation
- Dominant language
- Cuda
- Stars
- 10.1k
- Forks
- 1.4k
- Avg merge
- 4d 1h
- Merged PRs (30d)
- 2
Description
## Background
commit https://github.com/deepseek-ai/DeepEP/commit/099d5f2bad488b9c534ea785062b12f2e91d1d41 implements deterministic by sorting output tensor after copy epilogue, which may take 1 milisecond.
we propose an workable implementation with less copy epilogue latency overhead:
* normal and cached overhead < 10%,
* expand overhead ~ 100%
and some GPU memory increase:
* align scaleup buffer by channel count
* add 3 field to `WorkspaceLayout` (4MB+4KB+5MB)
## Overview
the reason of non-deterministic lays in 2 merging logic:
1. dispatch forward merge channels
2. expand mode epilogue merge scaleups
they both uses atomic counter which causes non-deterministic.
to achive deterministic, we suggest to use prefix sum (like DeepEP v1).
## Memory layout changes
first increase scaleup buffer size:
```diff
- token_layout, num_scaleup_ranks, num_scaleout_ranks * num_max_tokens_per_rank);
+ token_layout, num_scaleup_ranks, num_scaleout_ranks * (num_max_tokens_per_rank + (deterministic and num_scaleout_ranks > 1 ? kNumMaxChannels : 0)));
```
then 3 counter / prefix sum to workspace layout:
```cuda
__forceinline__ __device__ __host__ int* get_rank_channel_token_count_ptr(
const int& scaleup_rank_idx, const int& channel_idx, const int& scaleout_rank_idx) const;
__forceinline__ __device__ __host__ int* get_scaleout_token_prefix_sum_ptr(
const int& scaleup_rank_idx, const int& scaleout_rank_idx) const;
__forceinline__ __device__ __host__ int* get_dispatch_epilogue_warp_psum_ptr(
const int& expert_idx, const int& warp_global_idx) const;
```
## Forward channel merge
### Forward output location
as scaleout buffer is enough, we can reserve token range for each scaleout rank x channel, and write tokens into that range:
```cuda
if (ptx::deduplicate(stored_dst_scaleup_rank_idx, lane_idx) and stored_dst_scaleup_rank_idx >= 0) {
stored_dst_slot_idx = recv_scaleout_rank_idx * kNumChannels * kNumMaxTokensPerChannel
+ channel_idx * kNumMaxTokensPerChannel
+ channel_send_value;
}
```
token count fron each scaleout rank x each channel needs to be counted, like `stored_scaleup_send_counters`.
at forward ending, the token count is written to dst scaleup's workspace `rank_channel_token_count` field.
because it's deterministic, cached mode doesn't need save `dst_slot_idx` into `EPHandle`.
### Default copy epilogue merge
Now token layout in scaleup recv buffer is a 3D array with dimension: [src scaleup][channel idx][src scaleout]
we first compute channel prefix sum of each src rank (in-place), then compute scaleout prefix sum of each scaleup (workspace 2nd new field).
with the two psum, we can locate token within the scaleup recv buffer.
## Expand mode scaleup merge
fisrt compute warp prefix sum of each local expert, by scanning token topk idx data (this is reason of its big overhead);
then use the psum to locate token write position in `recv_x`.
Note: it requires grid level sync between the prefix sum computation and lookup, so needs coopertive flag when launching kernel.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by tracing the dispatch forward channel merge and expand mode scaleup merge, then inspect WorkspaceLayout and EPHandle. Verify how the proposed counters and prefix sums locate tokens in both paths, including the required grid-level synchronization. Done means deterministic output with the stated overhead and workspace changes.
Written by the indexing model from the issue text.
Assessment
- Domain
- distributed-systems, machine-learning, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 30/100