[PR #347] Feature/sm free normal kernel #529

New issue

Open

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

bsakaguchi commented

2026-10-27 21:03:41 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepEP/pull/347
Author: @ZhiyiHu1999
Created: 7/31/2025
Status: 🔄 Open

Base: antgroup-opt ← Head: feature/sm_free_normal_kernel

📝 Commits (9)

971ee33 First version of SM Free Normal Kernel for Internode Dispatch/Combine
26f06ac address conflicts due to TMA in Combine kNVLAndRDMAForwarder
2f93b35 remove duplicate codes due to faults in addressing conflicts
dc79091 minor modifications and add an annotation
83e27a2 minor modification
a5782dc minor modifications
afa77d3 minor modifications
5e4093b minor modifications
0b1ccde add arg return_recv_hook for get_dispatch_layout, so the kernel will be on compute stream in hook mode

📊 Changes

12 files changed (+1515 additions, -535 deletions)

View changed files

➖ .github/workflows/format.yml (+0 -28)
📝 csrc/config.hpp (+61 -10)
📝 csrc/deep_ep.cpp (+185 -117)
📝 csrc/deep_ep.hpp (+15 -9)
📝 csrc/kernels/api.cuh (+21 -8)
📝 csrc/kernels/buffer.cuh (+10 -10)
📝 csrc/kernels/configs.cuh (+2 -0)
📝 csrc/kernels/internode.cu (+353 -321)
📝 csrc/kernels/utils.cuh (+279 -0)
📝 deep_ep/buffer.py (+52 -27)
📝 tests/test_internode.py (+5 -5)
➕ tests/test_internode_hook.py (+532 -0)

📄 Description

1. Motivation

In MOE training and the prefilling phase of inference, the current ring-based RDMA buffer implementation for normal kernels significantly wastes SM resources. Due to the small size of the ring buffer and its frequent reuse for token transmission, SMs are often stalled, continuously polling for RDMA buffer availability instead of performing useful computation. This inefficient resource usage severely limits overall system throughput.

To address this issue, this PR implements a SM-friendly buffer design that frees SMs from buffer polling duties. This design is inspired by the large RDMA buffer approach discussed in #39. By allocating a larger RDMA buffer in HBM, tokens can be moved from SMs to the RDMA buffer in one go. While the NIC handles transmission asynchronously in the background, SMs can immediately resume computation tasks. Once data transfer completes, SMs can process the received tokens without blocking.

2. Design

2.1. Feature

The SM Free mode decouples the execution phases of the native mode: when the user first launches internode dispatch/combine, only the send phase is executed and a recv hook is returned. The user can wait for the network transmission to complete and then launch the receive phase of the internode dispatch/combine using the recv hook.
Implementation

2.2.1. Principles

Full Compatibility: Seamlessly integrates with existing native communication modes, requiring minimal changes to existing codebases

2.2.2. Highlights

Mode Control via return_recv_hook: Introduces a user-controllable argument return_recv_hook to switch between native mode and hook mode.
RDMA Buffer Management:
- Retains the native buffer structure for ease of integration.
- Enlarges token capacity per RDMA recv chunk to fit all tokens in one batch.
- Provides a utility function get_normal_hook_rdma_size_hint() to help users estimate the minimum required RDMA buffer size.
Improved Compute Stream Utilization:
- In hook mode, the kernel runs on the compute stream, enabling more efficient offloading of data transfer tasks to the NIC and the network while freeing up SMs.
- Compared to native mode (which allocates 2 SMs per channel), hook mode maps one channel per SM, improving scalability and SM utilization.

3. Performance Evaluation

3.1. Experiment Setup

2/4 nodes, with 8 × H20 GPUs per node.
4096 tokens per batch, 7168 hidden, top-8 experts.

3.2. Effect

3.2.1. Estimated Performance (native mode → hook mode)

A GEMM operation of shape [4096, 4096] is used to overlap the network phase. It typically takes 4-5 μs in my system, sufficient for completing the token transmission.
For hook mode, the kernel execution time is the sum of the send kernel time and the receive kernel time, and the RDMA bandwidth is calculated as: RDMA Bandwidth = RDMA Bytes / Kernel Execution Time

Dispatch Kernel Execution Time & Bandwidth

#EP	FP8 Dispatch Kernel Execution Time	FP8 Dispatch RDMA Bandwidth	BF16 Dispatch Kernel Execution Time	BF16 Dispatch RDMA Bandwidth
16	1358 → 765 μs	44 → 79 GB/s	2530 → 1381 μs	46 → 85 GB/s
32	3535 → 884 μs	30 → 124 GB/s	6780 → 2468 μs	32 → 86 GB/s

Combine Kernel Execution Time & Bandwidth

#EP	BF16 Combine Kernel Execution Time	BF16 Combine RDMA Bandwidth
16	2590 → 1425 μs	45 → 82 GB/s
32	6850 → 2250 μs	31 → 81 GB/s

3.3. Cost

3.3.1. HBM Cost

The main HBM cost in hook mode comes from two sides:

RDMA Buffer:
rdma_buffer_size = num_max_dispatch_tokens_per_rank × hidden_size × size_of(element) × num_nodes × 2
The HBM cost from using the large RDMA buffer increases with num_max_dispatch_tokens_per_rank and num_nodes; for our experiment setup, the HBM cost for the RDMA buffer on each rank is about 270 MB.
NVLink Buffer:
nvl_buffer_size = num_max_nvl_chunked_recv_tokens × hidden_size × size_of(element) × num_nvl_peers × num_channels
HBM used for NVLink buffer increased with the number of SMs (channels) in the hook mode, but the number of tokens assigned to each channel is smaller, so the nvl_recv_chunk could be shrunk and the overall HBM cost for this part changed slightly.

3.3.2. Bandwidth Cost

For SM Free mode, the data movement in the recv phase of dispatch and the send phase of combine are inter-GPU through NVLink. So the respective kernel execution time is also limited by the NVLink bandwidth.

4. RoadMap

Code Refactor
Optimize HBM Usage
Implement Dynamic Resizing in RDMA Buffer
More TMA Optimization in Normal Hook mode

_{🔄 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/347 **Author:** [@ZhiyiHu1999](https://github.com/ZhiyiHu1999) **Created:** 7/31/2025 **Status:** 🔄 Open **Base:** `antgroup-opt` ← **Head:** `feature/sm_free_normal_kernel` --- ### 📝 Commits (9) - [`971ee33`](https://github.com/deepseek-ai/DeepEP/commit/971ee332b916a5231c511c8656230af6e5a478ad) First version of SM Free Normal Kernel for Internode Dispatch/Combine - [`26f06ac`](https://github.com/deepseek-ai/DeepEP/commit/26f06acfd0d73db94cdf7177694f2ae36c7b5c17) address conflicts due to TMA in Combine kNVLAndRDMAForwarder - [`2f93b35`](https://github.com/deepseek-ai/DeepEP/commit/2f93b35a2fbf84be22e8154ecd6b0d54eb379fb4) remove duplicate codes due to faults in addressing conflicts - [`dc79091`](https://github.com/deepseek-ai/DeepEP/commit/dc7909194702656ba22a7899017aa912bc87569c) minor modifications and add an annotation - [`83e27a2`](https://github.com/deepseek-ai/DeepEP/commit/83e27a2a93f216b21c387dcfa41f45562f6ddbac) minor modification - [`a5782dc`](https://github.com/deepseek-ai/DeepEP/commit/a5782dcf8217bae4efbc5a0834c350ada6321f3d) minor modifications - [`afa77d3`](https://github.com/deepseek-ai/DeepEP/commit/afa77d379407ffa6273b3c4ff57c1c42fc0c9dfb) minor modifications - [`5e4093b`](https://github.com/deepseek-ai/DeepEP/commit/5e4093bc454a1f91f9ad42c49d7462341f5a8053) minor modifications - [`0b1ccde`](https://github.com/deepseek-ai/DeepEP/commit/0b1ccde7b8b4678341dfee5d06cabb54b6f477f9) add arg return_recv_hook for get_dispatch_layout, so the kernel will be on compute stream in hook mode ### 📊 Changes **12 files changed** (+1515 additions, -535 deletions) <details> <summary>View changed files</summary> ➖ `.github/workflows/format.yml` (+0 -28) 📝 `csrc/config.hpp` (+61 -10) 📝 `csrc/deep_ep.cpp` (+185 -117) 📝 `csrc/deep_ep.hpp` (+15 -9) 📝 `csrc/kernels/api.cuh` (+21 -8) 📝 `csrc/kernels/buffer.cuh` (+10 -10) 📝 `csrc/kernels/configs.cuh` (+2 -0) 📝 `csrc/kernels/internode.cu` (+353 -321) 📝 `csrc/kernels/utils.cuh` (+279 -0) 📝 `deep_ep/buffer.py` (+52 -27) 📝 `tests/test_internode.py` (+5 -5) ➕ `tests/test_internode_hook.py` (+532 -0) </details> ### 📄 Description ## 1. Motivation In MOE training and the prefilling phase of inference, the current ring-based RDMA buffer implementation for normal kernels significantly wastes SM resources. Due to the small size of the ring buffer and its frequent reuse for token transmission, SMs are often stalled, continuously polling for RDMA buffer availability instead of performing useful computation. This inefficient resource usage severely limits overall system throughput. To address this issue, this PR implements a SM-friendly buffer design that frees SMs from buffer polling duties. This design is inspired by the large RDMA buffer approach discussed in #39. By allocating a larger RDMA buffer in HBM, tokens can be moved from SMs to the RDMA buffer in one go. While the NIC handles transmission asynchronously in the background, SMs can immediately resume computation tasks. Once data transfer completes, SMs can process the received tokens without blocking. ## 2. Design ### 2.1. Feature The SM Free mode decouples the execution phases of the native mode: when the user first launches internode dispatch/combine, only the send phase is executed and a recv hook is returned. The user can wait for the network transmission to complete and then launch the receive phase of the internode dispatch/combine using the recv hook. Implementation ### 2.2.1. Principles - Full Compatibility: Seamlessly integrates with existing native communication modes, requiring minimal changes to existing codebases ### 2.2.2. Highlights - Mode Control via `return_recv_hook`: Introduces a user-controllable argument `return_recv_hook` to switch between native mode and hook mode. - RDMA Buffer Management: - Retains the native buffer structure for ease of integration. - Enlarges token capacity per RDMA recv chunk to fit all tokens in one batch. - Provides a utility function `get_normal_hook_rdma_size_hint()` to help users estimate the minimum required RDMA buffer size. - Improved Compute Stream Utilization: - In hook mode, the kernel runs on the compute stream, enabling more efficient offloading of data transfer tasks to the NIC and the network while freeing up SMs. - Compared to native mode (which allocates 2 SMs per channel), hook mode maps one channel per SM, improving scalability and SM utilization. <img width="3054" height="1862" alt="image" src="https://github.com/user-attachments/assets/20868256-04bd-4a18-b5d2-fbc1fcf94c10" /> ## 3. Performance Evaluation ### 3.1. Experiment Setup - 2/4 nodes, with 8 × H20 GPUs per node. - 4096 tokens per batch, 7168 hidden, top-8 experts. --- ### 3.2. Effect #### 3.2.1. Estimated Performance (native mode → hook mode) - A GEMM operation of shape [4096, 4096] is used to overlap the network phase. It typically takes 4-5 μs in my system, sufficient for completing the token transmission. - For hook mode, the kernel execution time is the sum of the send kernel time and the receive kernel time, and the RDMA bandwidth is calculated as: RDMA Bandwidth = RDMA Bytes / Kernel Execution Time #### Dispatch Kernel Execution Time & Bandwidth | #EP | FP8 Dispatch Kernel Execution Time | FP8 Dispatch RDMA Bandwidth | BF16 Dispatch Kernel Execution Time | BF16 Dispatch RDMA Bandwidth | |:---:|:-------------------------------------:|:-------------------------------:|:--------------------------------------:|:-------------------------------:| | 16 | 1358 → 765 μs | 44 → 79 GB/s | 2530 → 1381 μs | 46 → 85 GB/s | | 32 | 3535 → 884 μs | 30 → 124 GB/s | 6780 → 2468 μs | 32 → 86 GB/s | #### Combine Kernel Execution Time & Bandwidth | #EP | BF16 Combine Kernel Execution Time | BF16 Combine RDMA Bandwidth | |:---:|:-------------------------------------:|:------------------------------:| | 16 | 2590 → 1425 μs | 45 → 82 GB/s | | 32 | 6850 → 2250 μs | 31 → 81 GB/s | --- ### 3.3. Cost #### 3.3.1. HBM Cost The main HBM cost in hook mode comes from two sides: - **RDMA Buffer:** rdma_buffer_size = num_max_dispatch_tokens_per_rank × hidden_size × size_of(element) × num_nodes × 2 The HBM cost from using the large RDMA buffer increases with `num_max_dispatch_tokens_per_rank` and `num_nodes`; for our experiment setup, the HBM cost for the RDMA buffer on each rank is about 270 MB. - **NVLink Buffer:** nvl_buffer_size = num_max_nvl_chunked_recv_tokens × hidden_size × size_of(element) × num_nvl_peers × num_channels HBM used for NVLink buffer increased with the number of SMs (channels) in the hook mode, but the number of tokens assigned to each channel is smaller, so the `nvl_recv_chunk` could be shrunk and the overall HBM cost for this part changed slightly. #### 3.3.2. Bandwidth Cost For SM Free mode, the data movement in the recv phase of dispatch and the send phase of combine are inter-GPU through NVLink. So the respective kernel execution time is also limited by the NVLink bandwidth. --- ## 4. RoadMap - [ ] Code Refactor - [ ] Optimize HBM Usage - [ ] Implement Dynamic Resizing in RDMA Buffer - [ ] More TMA Optimization in Normal Hook mode --- 🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.

bsakaguchi added the

pull-request

label

2026-10-27 21:03:41 +00:00

bsakaguchi referenced this issue

2026-10-27 21:04:32 +00:00

[PR #529] [MERGED] Add UCCL-EP to Community Forks section in README #611