deepseek-ai/DeepEP
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepEP.git synced 2026-11-05 00:47:32 +00:00
Code Issues 308 Projects Packages Wiki Activity

[PR #690] DeepEP v2:Add experimental direct recv_x dispatch path. #709

New issue
Open
opened 2026-10-27 21:05:28 +00:00 by bsakaguchi · 0 comments
bsakaguchi commented 2026-10-27 21:05:28 +00:00
Owner
Copy link

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/690
Author: @kzlxd
Created: 7/14/2026
Status: 🔄 Open

Base: main ← Head: Dispacth_Direct_recv_x

📝 Commits (1)

  • 6757592 Add experimental direct recv_x dispatch path.

📊 Changes

7 files changed (+594 additions, -185 deletions)

View changed files

📝 csrc/elastic/buffer.hpp (+74 -11)
📝 csrc/kernels/elastic/dispatch.hpp (+18 -3)
📝 deep_ep/buffers/elastic.py (+26 -1)
📝 deep_ep/include/deep_ep/impls/dispatch.cuh (+179 -23)
📝 deep_ep/include/deep_ep/impls/dispatch_copy_epilogue.cuh (+27 -9)
📝 deep_ep/include/deep_ep/impls/hybrid_dispatch.cuh (+70 -4)
📝 tests/elastic/test_ep.py (+200 -134)

📄 Description

Add Direct Recv X dispatch path

Summary

This PR adds an experimental Direct Recv X path for compact EP dispatch.

The optimization removes the extra local hidden-state copy in the dispatch copy epilogue. In the original path, hidden states are first written into the NCCL symmetric staging buffer, then copied again into the final recv_x tensor by the epilogue. With Direct Recv X enabled, dispatch writes hidden states directly into the final compact recv_x window, while metadata and top-k tensors continue to use the existing staging + epilogue path.

Direct Recv X is enabled by:

EP_V2_EXPERIMENTAL_DIRECT_RECV_X=1

This PR also includes a small SM90 / H800 JIT compilation fix for an SM100-only shared-memory bulk-store instruction.

Background

The original dispatch path has two buffer layers:

                 dispatch kernel                      copy epilogue
src x/topk  ------------------------------> staging -----------------> recv_x/topk
            write hidden + metadata/top-k              read staging, write final tensors

staging is NCCL symmetric / window memory and can be accessed by the communication path directly.

recv_x is a normal PyTorch tensor. It is materialized after dispatch by dispatch_copy_epilogue.

This means hidden states always pay one extra local copy:

staging.hidden -> recv_x.hidden

The copy cost grows linearly with token count and becomes more visible with wider hidden sizes.

Direct Recv X design

The dispatch kernel already knows:

  • which destination rank a token is sent to;
  • which slot index the token occupies on the destination rank.

The receiver notify warp also knows how many tokens each source rank contributes. From those counts, it can compute the exclusive prefix base for each source rank in the compact recv_x layout.

Direct Recv X publishes this receiver-side prefix information back to the source rank:

receiver:
  rank_count[source_rank] -> exclusive prefix
  publish prefix[source_rank] back to source

source dispatch warp:
  compact_idx = prefix[dst/source group] + dst_slot_idx
  write hidden directly to recv_x[compact_idx]

The hidden-state data path becomes:

src rank x
  -> dispatch kernel
  -> target rank recv_x

The metadata / top-k path stays unchanged:

metadata/top-k
  -> staging
  -> dispatch_copy_epilogue
  -> recv_topk_idx / recv_topk_weights / recv_src_metadata / channel_linked_list

So the epilogue no longer copies hidden states, but it still materializes the metadata and top-k outputs.

Implementation

  • Reserve a Direct Recv X output window at the tail of the DeepEP symmetric communication buffer when EP_V2_EXPERIMENTAL_DIRECT_RECV_X=1.

    • Size: 
      num_ranks * num_max_tokens_per_rank * hidden * elem_size

  • Return recv_x as a view into this symmetric-window tail in Direct Recv X mode.

  • Extend dispatch runtime arguments and kernels so dispatch can write hidden states directly into the receiver-side recv_x window.

  • Let dispatch_copy_epilogue skip the hidden-state copy when Direct Recv X is enabled.

    • The epilogue still writes:
      • recv_topk_idx
      • recv_topk_weights
      • recv_src_metadata
      • channel_linked_list
      • FP8 scale-factor tensors

  • Support both compact dispatch variants:

    • single-node / scaleup dispatch;
    • hybrid scaleout + scaleup dispatch.

  • Keep the feature behind one experimental env flag:

    • EP_V2_EXPERIMENTAL_DIRECT_RECV_X=1

  • Add guardrails for unsupported modes:

    • non-cached dispatch only;
    • non-expanded dispatch only;
    • deterministic dispatch is not supported;
    • previous_event_before_epilogue is rejected because recv_x is produced during dispatch instead of during the epilogue.

H100 benchmark results

Configuration:

  • GPU: H100
  • Path: compact EP dispatch
  • Direct mode: EP_V2_EXPERIMENTAL_DIRECT_RECV_X=1
  • Metric: dispatch + copy epilogue latency

Parameter：

  • direct-recv-x
  • num-sms 12
  • num-topk 8
  • prefer-overlap-with-compute 1

EP8

Tokens Dtype Dispatch non-direct / direct Copy non-direct / direct Total non-direct / direct Total delta
4096 FP8 545.6 / 536.6 us 161.7 / 51.9 us 707.3 / 588.5 us -118.9 us, 16.8% faster
8192 FP8 1066.0 / 1042.4 us 325.7 / 92.4 us 1391.7 / 1134.8 us -256.8 us, 18.5% faster
4096 BF16 1038.3 / 1076.7 us 304.3 / 96.1 us 1342.7 / 1172.8 us -169.9 us, 12.7% faster
8192 BF16 2048.8 / 2124.8 us 606.3 / 187.5 us 2655.1 / 2312.4 us -342.8 us, 12.9% faster

EP16

Tokens Dtype Dispatch non-direct / direct Copy non-direct / direct Total non-direct / direct Total delta
4096 FP8 753.9 / 768.0 us 166.5 / 34.6 us 920.4 / 802.6 us -117.8 us, 12.8% faster
8192 FP8 1416.9 / 1374.6 us 339.3 / 35.4 us 1756.3 / 1410.1 us -346.2 us, 19.7% faster
4096 BF16 1341.8 / 1437.8 us 289.6 / 103.7 us 1631.4 / 1541.5 us -89.9 us, 5.5% faster
8192 BF16 2600.5 / 2737.4 us 588.0 / 195.4 us 3188.5 / 2932.8 us -255.7 us, 8.0% faster

Result summary

Direct Recv X substantially reduces the copy epilogue cost by removing the hidden-state copy from staging to recv_x.

End-to-end dispatch + copy improvement:

  • EP8: 12.7% - 18.5% faster
  • EP16: 5.5% - 19.7% faster

In some BF16 cases, the dispatch kernel itself becomes slightly slower because it now performs direct remote recv_x writes. The removed epilogue hidden copy still more than offsets that overhead, so total latency improves.

🔄 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/690 **Author:** [@kzlxd](https://github.com/kzlxd) **Created:** 7/14/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `Dispacth_Direct_recv_x` --- ### 📝 Commits (1) - [`6757592`](https://github.com/deepseek-ai/DeepEP/commit/6757592e03f28eba9a11c3e2170af06ebd284ebc) Add experimental direct recv_x dispatch path. ### 📊 Changes **7 files changed** (+594 additions, -185 deletions) <details> <summary>View changed files</summary> 📝 `csrc/elastic/buffer.hpp` (+74 -11) 📝 `csrc/kernels/elastic/dispatch.hpp` (+18 -3) 📝 `deep_ep/buffers/elastic.py` (+26 -1) 📝 `deep_ep/include/deep_ep/impls/dispatch.cuh` (+179 -23) 📝 `deep_ep/include/deep_ep/impls/dispatch_copy_epilogue.cuh` (+27 -9) 📝 `deep_ep/include/deep_ep/impls/hybrid_dispatch.cuh` (+70 -4) 📝 `tests/elastic/test_ep.py` (+200 -134) </details> ### 📄 Description # Add Direct Recv X dispatch path ## Summary This PR adds an experimental **Direct Recv X** path for compact EP dispatch. The optimization removes the extra local hidden-state copy in the dispatch copy epilogue. In the original path, hidden states are first written into the NCCL symmetric staging buffer, then copied again into the final `recv_x` tensor by the epilogue. With Direct Recv X enabled, dispatch writes hidden states directly into the final compact `recv_x` window, while metadata and top-k tensors continue to use the existing staging + epilogue path. Direct Recv X is enabled by: ```bash EP_V2_EXPERIMENTAL_DIRECT_RECV_X=1 ``` This PR also includes a small SM90 / H800 JIT compilation fix for an SM100-only shared-memory bulk-store instruction. ## Background The original dispatch path has two buffer layers: ```text dispatch kernel copy epilogue src x/topk ------------------------------> staging -----------------> recv_x/topk write hidden + metadata/top-k read staging, write final tensors ``` `staging` is NCCL symmetric / window memory and can be accessed by the communication path directly. `recv_x` is a normal PyTorch tensor. It is materialized after dispatch by `dispatch_copy_epilogue`. This means hidden states always pay one extra local copy: ```text staging.hidden -> recv_x.hidden ``` The copy cost grows linearly with token count and becomes more visible with wider hidden sizes. ## Direct Recv X design The dispatch kernel already knows: - which destination rank a token is sent to; - which slot index the token occupies on the destination rank. The receiver notify warp also knows how many tokens each source rank contributes. From those counts, it can compute the exclusive prefix base for each source rank in the compact `recv_x` layout. Direct Recv X publishes this receiver-side prefix information back to the source rank: ```text receiver: rank_count[source_rank] -> exclusive prefix publish prefix[source_rank] back to source source dispatch warp: compact_idx = prefix[dst/source group] + dst_slot_idx write hidden directly to recv_x[compact_idx] ``` The hidden-state data path becomes: ```text src rank x -> dispatch kernel -> target rank recv_x ``` The metadata / top-k path stays unchanged: ```text metadata/top-k -> staging -> dispatch_copy_epilogue -> recv_topk_idx / recv_topk_weights / recv_src_metadata / channel_linked_list ``` So the epilogue no longer copies hidden states, but it still materializes the metadata and top-k outputs. ## Implementation - Reserve a Direct Recv X output window at the tail of the DeepEP symmetric communication buffer when `EP_V2_EXPERIMENTAL_DIRECT_RECV_X=1`. - Size: ```text num_ranks * num_max_tokens_per_rank * hidden * elem_size ``` - Return `recv_x` as a view into this symmetric-window tail in Direct Recv X mode. - Extend dispatch runtime arguments and kernels so dispatch can write hidden states directly into the receiver-side `recv_x` window. - Let `dispatch_copy_epilogue` skip the hidden-state copy when Direct Recv X is enabled. - The epilogue still writes: - `recv_topk_idx` - `recv_topk_weights` - `recv_src_metadata` - `channel_linked_list` - FP8 scale-factor tensors - Support both compact dispatch variants: - single-node / scaleup dispatch; - hybrid scaleout + scaleup dispatch. - Keep the feature behind one experimental env flag: - `EP_V2_EXPERIMENTAL_DIRECT_RECV_X=1` - Add guardrails for unsupported modes: - non-cached dispatch only; - non-expanded dispatch only; - deterministic dispatch is not supported; - `previous_event_before_epilogue` is rejected because `recv_x` is produced during dispatch instead of during the epilogue. ## H100 benchmark results Configuration: - GPU: H100 - Path: compact EP dispatch - Direct mode: `EP_V2_EXPERIMENTAL_DIRECT_RECV_X=1` - Metric: dispatch + copy epilogue latency Parameter： - direct-recv-x - num-sms 12 - num-topk 8 - prefer-overlap-with-compute 1 ### EP8 | Tokens | Dtype | Dispatch non-direct / direct | Copy non-direct / direct | Total non-direct / direct | Total delta | |---:|:---:|---:|---:|---:|---:| | 4096 | FP8 | 545.6 / 536.6 us | 161.7 / 51.9 us | 707.3 / 588.5 us | -118.9 us, 16.8% faster | | 8192 | FP8 | 1066.0 / 1042.4 us | 325.7 / 92.4 us | 1391.7 / 1134.8 us | -256.8 us, 18.5% faster | | 4096 | BF16 | 1038.3 / 1076.7 us | 304.3 / 96.1 us | 1342.7 / 1172.8 us | -169.9 us, 12.7% faster | | 8192 | BF16 | 2048.8 / 2124.8 us | 606.3 / 187.5 us | 2655.1 / 2312.4 us | -342.8 us, 12.9% faster | ### EP16 | Tokens | Dtype | Dispatch non-direct / direct | Copy non-direct / direct | Total non-direct / direct | Total delta | |---:|:---:|---:|---:|---:|---:| | 4096 | FP8 | 753.9 / 768.0 us | 166.5 / 34.6 us | 920.4 / 802.6 us | -117.8 us, 12.8% faster | | 8192 | FP8 | 1416.9 / 1374.6 us | 339.3 / 35.4 us | 1756.3 / 1410.1 us | -346.2 us, 19.7% faster | | 4096 | BF16 | 1341.8 / 1437.8 us | 289.6 / 103.7 us | 1631.4 / 1541.5 us | -89.9 us, 5.5% faster | | 8192 | BF16 | 2600.5 / 2737.4 us | 588.0 / 195.4 us | 3188.5 / 2932.8 us | -255.7 us, 8.0% faster | ## Result summary Direct Recv X substantially reduces the copy epilogue cost by removing the hidden-state copy from staging to `recv_x`. End-to-end dispatch + copy improvement: - EP8: 12.7% - 18.5% faster - EP16: 5.5% - 19.7% faster In some BF16 cases, the dispatch kernel itself becomes slightly slower because it now performs direct remote `recv_x` writes. The removed epilogue hidden copy still more than offsets that overhead, so total latency improves. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
Sign in to join this conversation.
No Branch/Tag specified
Branches Tags
main
hybrid-ep
chore/remove-clang-format-yapf
v2.5
nvDev
try-fix-cu128
mori-ep
antgroup-opt
compat-old-torch
add-linter
eager-rdma
v1.2.1
Labels
Clear labels   
pull-request
Mirrored from GitHub Pull Request

No labels
pull-request
Milestone
Clear milestone
No items
No milestone
Projects
Clear projects
No items
No project
Assignees
Clear assignees
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
deepseek-ai/DeepEP#709
No description provided.