deepseek-ai/DeepEP

Fork 0

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

[PR #634] [MERGED] Hybrid ep nixl restripe #675

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/634
Author: @QiaoK
Created: 5/19/2026
Status: ✅ Merged
Merged: 6/1/2026
Merged by: @jershi425

Base: hybrid-ep ← Head: hybrid-ep-nixl-restripe

📝 Commits (2)

c40e209 hybrid_ep: fix scan permute-fusion local-experts sum array sizing
da11c84 nixl: coalesce dispatch prob put with dest-major source layout

📊 Changes

4 files changed (+144 additions, -26 deletions)

View changed files

📝 csrc/hybrid_ep/backend/hybrid_ep_backend.cuh (+47 -26)
📝 csrc/hybrid_ep/buffer/internode_nixl.cu (+64 -0)
📝 csrc/hybrid_ep/buffer/internode_nixl.cuh (+16 -0)
📝 csrc/hybrid_ep/executor/executor.cu (+17 -0)

📄 Description

Baseline problem

The user-supplied prob tensor lives in [num_tokens][num_of_nodes][prob_per_token] layout: each token holds num_of_nodes consecutive per-destination prob slices. From the perspective of the dispatch N2N warp, which sends one destination's slice per token at a time, the per-destination slice of consecutive tokens is not contiguous in source memory — it is strided by num_of_nodes × prob_per_token floats.
In the baseline build the NIXL dispatch N2N warp therefore had to issue, per (chunk, destination): one nixlPut for the token (HIDDEN_DIM × sizeof(bf16) bytes, contiguous) plus NUM_OF_TOKENS_PER_CHUNK separate nixlPuts for prob (one slice per token, each prob_per_token × sizeof(float) bytes).
With NUM_OF_TOKENS_PER_CHUNK = 256, that is 1 + 256 = 257 puts per (chunk, destination), with ~256 of them being small prob slices. Each nixlPut has fixed-per-op overhead (atomic WQE reservation in the QP, descriptor lookup, WQE construction). For a 16-expert × 8-rank layout, prob_per_token = 128 floats = 512 bytes — the prob put is small enough that issue overhead dominates the actual data transfer.

Fix

The change replaces the existing cudaMemcpyAsync(prob) in dispatch_preprocess with a custom restripe_prob_for_nixl_dispatch kernel that copies prob into a destination-major staging layout:
[num_of_nodes][max_tokens_per_rank][prob_per_token]
In this layout, all prob slices for a given destination over a range of consecutive tokens are contiguous in source memory. The N2N warp can now issue one coalesced nixlPut of token_range × prob_per_token × sizeof(float) bytes per chunk per destination (and one per run in the sparse path). The per-(chunk, destination) put count drops from ~257 to 2 in the dense path.
Total HBM traffic is identical: the restripe kernel touches the same bytes the displaced cudaMemcpyAsync did. Added cost is the 2D index swap inside the kernel (one block per (token, destination), threads cooperatively copy prob_per_token floats with vector loads). That cost is negligible compared to the savings.

Why the win scales with cluster size

WQE-issue cost (atomic FA, descriptor lookup, doorbell) is roughly constant per op. Savings per chunk per destination is fixed at ~256 puts. Total puts removed per dispatch call grows like num_remote_nodes × num_chunks_per_rank × 256. The baseline pathology — issue-bound dispatch with many small WQEs queued across many QPs — gets worse than linearly with cluster size, so the relative win grows from +23% at 2 nodes to roughly +50–55% at 8 nodes.

Receive-side and DOCA path are unchanged

The G2S warp's local-node prob branch was updated to address the new dest-major layout (local slice starts at node_rank × MAX_TOKENS; per-token stride is prob_per_token). The remote-source branch was not touched because the receive-side rdma_inter_node_group_prob buffer keeps its original token-major layout — receivers pull from a per-rank buffer whose strides did not need to change. Both N2N and G2S DOCA paths are preserved via #ifdef USE_NIXL guards.

Fix scan permute-fusion local-experts sum array sizing:

A previous git rebase accidentally used older commits. This PR fix this problem.

B300 results (8 GPUs per node, 2 nics per GPU)
Test with the following variables to fully use ucx multi-rail feature.
#UCX dual port (per-GPU multi-NIC selection) :
export UCX_IB_GDA_MAX_HCA_PER_GPU=2
#decouple NIXL channel count from SM count.
export DEEPEP_NIXL_GDA_NUM_CHANNELS=16

Cluster	Kernel	RDMA base (GB/s)	RDMA Restripe (GB/s)	RDMA Δ	avg_t base (μs)	avg_t Restripe (μs)	avg_t Δ
2 nodes (16 ranks)	dispatch	132.40	163.27	+23.3%	884.1	716.9	-18.9%
2 nodes (16 ranks)	dispatch+permute	123.57	153.80	+24.5%	947.3	761.1	-19.7%
4 nodes (32 ranks)	dispatch	88.40	110.54	+25.0%	2397.9	1917.7	-20.0%
4 nodes (32 ranks)	dispatch+permute	80.94	108.96	+34.6%	2618.8	1945.5	-25.7%
8 nodes (64 ranks)	dispatch	64.37	96.44	+49.8%	4829.0	3223.2	-33.3%
8 nodes (64 ranks)	dispatch+permute	61.42	95.05	+54.8%	5060.7	3270.3	-35.4%

</html>

_{🔄 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/634 **Author:** [@QiaoK](https://github.com/QiaoK) **Created:** 5/19/2026 **Status:** ✅ Merged **Merged:** 6/1/2026 **Merged by:** [@jershi425](https://github.com/jershi425) **Base:** `hybrid-ep` ← **Head:** `hybrid-ep-nixl-restripe` --- ### 📝 Commits (2) - [`c40e209`](https://github.com/deepseek-ai/DeepEP/commit/c40e2098cbe927bfe4c90926a48495b8855ad727) hybrid_ep: fix scan permute-fusion local-experts sum array sizing - [`da11c84`](https://github.com/deepseek-ai/DeepEP/commit/da11c84e2c49658d777912b1f7147c088b23f16d) nixl: coalesce dispatch prob put with dest-major source layout ### 📊 Changes **4 files changed** (+144 additions, -26 deletions) <details> <summary>View changed files</summary> 📝 `csrc/hybrid_ep/backend/hybrid_ep_backend.cuh` (+47 -26) 📝 `csrc/hybrid_ep/buffer/internode_nixl.cu` (+64 -0) 📝 `csrc/hybrid_ep/buffer/internode_nixl.cuh` (+16 -0) 📝 `csrc/hybrid_ep/executor/executor.cu` (+17 -0) </details> ### 📄 Description ### Baseline problem The user-supplied prob tensor lives in [num_tokens][num_of_nodes][prob_per_token] layout: each token holds num_of_nodes consecutive per-destination prob slices. From the perspective of the dispatch N2N warp, which sends one destination's slice per token at a time, the per-destination slice of consecutive tokens is not contiguous in source memory — it is strided by num_of_nodes × prob_per_token floats. In the baseline build the NIXL dispatch N2N warp therefore had to issue, per (chunk, destination): one nixlPut for the token (HIDDEN_DIM × sizeof(bf16) bytes, contiguous) plus NUM_OF_TOKENS_PER_CHUNK separate nixlPuts for prob (one slice per token, each prob_per_token × sizeof(float) bytes). With NUM_OF_TOKENS_PER_CHUNK = 256, that is 1 + 256 = 257 puts per (chunk, destination), with ~256 of them being small prob slices. Each nixlPut has fixed-per-op overhead (atomic WQE reservation in the QP, descriptor lookup, WQE construction). For a 16-expert × 8-rank layout, prob_per_token = 128 floats = 512 bytes — the prob put is small enough that issue overhead dominates the actual data transfer. ### Fix The change replaces the existing cudaMemcpyAsync(prob) in dispatch_preprocess with a custom restripe_prob_for_nixl_dispatch kernel that copies prob into a destination-major staging layout: [num_of_nodes][max_tokens_per_rank][prob_per_token] In this layout, all prob slices for a given destination over a range of consecutive tokens are contiguous in source memory. The N2N warp can now issue one coalesced nixlPut of token_range × prob_per_token × sizeof(float) bytes per chunk per destination (and one per run in the sparse path). The per-(chunk, destination) put count drops from ~257 to 2 in the dense path. Total HBM traffic is identical: the restripe kernel touches the same bytes the displaced cudaMemcpyAsync did. Added cost is the 2D index swap inside the kernel (one block per (token, destination), threads cooperatively copy prob_per_token floats with vector loads). That cost is negligible compared to the savings. ### Why the win scales with cluster size WQE-issue cost (atomic FA, descriptor lookup, doorbell) is roughly constant per op. Savings per chunk per destination is fixed at ~256 puts. Total puts removed per dispatch call grows like num_remote_nodes × num_chunks_per_rank × 256. The baseline pathology — issue-bound dispatch with many small WQEs queued across many QPs — gets worse than linearly with cluster size, so the relative win grows from +23% at 2 nodes to roughly +50–55% at 8 nodes. ### Receive-side and DOCA path are unchanged The G2S warp's local-node prob branch was updated to address the new dest-major layout (local slice starts at node_rank × MAX_TOKENS; per-token stride is prob_per_token). The remote-source branch was not touched because the receive-side rdma_inter_node_group_prob buffer keeps its original token-major layout — receivers pull from a per-rank buffer whose strides did not need to change. Both N2N and G2S DOCA paths are preserved via #ifdef USE_NIXL guards. ### Fix scan permute-fusion local-experts sum array sizing: A previous git rebase accidentally used older commits. This PR fix this problem. B300 results (8 GPUs per node, 2 nics per GPU) Test with the following variables to fully use ucx multi-rail feature. \#UCX dual port (per-GPU multi-NIC selection) : export UCX_IB_GDA_MAX_HCA_PER_GPU=2 \#decouple NIXL channel count from SM count. export DEEPEP_NIXL_GDA_NUM_CHANNELS=16 <html xmlns:o="urn:schemas-microsoft-com:office:office" xmlns:w="urn:schemas-microsoft-com:office:word" xmlns:m="http://schemas.microsoft.com/office/2004/12/omml" xmlns="http://www.w3.org/TR/REC-html40"> <head> <meta name=ProgId content=Word.Document> <meta name=Generator content="Microsoft Word 15"> <meta name=Originator content="Microsoft Word 15"> <link rel=File-List href="file:///C:/Users/qkang/AppData/Local/Packages/oice_16_974fa576_32c1d314_14e5/AC/Temp/msohtmlclip1/01/clip_filelist.xml">  <link rel=themeData href="file:///C:/Users/qkang/AppData/Local/Packages/oice_16_974fa576_32c1d314_14e5/AC/Temp/msohtmlclip1/01/clip_themedata.thmx"> <link rel=colorSchemeMapping href="file:///C:/Users/qkang/AppData/Local/Packages/oice_16_974fa576_32c1d314_14e5/AC/Temp/msohtmlclip1/01/clip_colorschememapping.xml">  <style>  </style>  </head> <body lang=EN-US style='tab-interval:.5in;word-wrap:break-word'>  Cluster | Kernel | RDMA base (GB/s) | RDMA Restripe (GB/s) | RDMA Δ | avg_t base (μs) | avg_t Restripe (μs) | avg_t Δ -- | -- | -- | -- | -- | -- | -- | -- 2 nodes (16 ranks) | dispatch | 132.40 | 163.27 | +23.3% | 884.1 | 716.9 | -18.9% 2 nodes (16 ranks) | dispatch+permute | 123.57 | 153.80 | +24.5% | 947.3 | 761.1 | -19.7% 4 nodes (32 ranks) | dispatch | 88.40 | 110.54 | +25.0% | 2397.9 | 1917.7 | -20.0% 4 nodes (32 ranks) | dispatch+permute | 80.94 | 108.96 | +34.6% | 2618.8 | 1945.5 | -25.7% 8 nodes (64 ranks) | dispatch | 64.37 | 96.44 | +49.8% | 4829.0 | 3223.2 | -33.3% 8 nodes (64 ranks) | dispatch+permute | 61.42 | 95.05 | +54.8% | 5060.7 | 3270.3 | -35.4%  </body> </html> --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

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

closed this issue
added the
pull-request
label

No labels

pull-request

No milestone

No project

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#675

No description provided.

Rows
Columns