DeepSeek/DeepEP
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepEP synced 2025-06-26 18:28:11 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add low-latency kernel PCIe usage flag (#195)

Browse Source 
* Add low-latency kernel usage flag

* Update comments
This commit is contained in:
Chenggang Zhao
2025-06-09 14:37:13 +08:00
committed by GitHub
parent 564e375234
commit 0d1a855d81
6 changed files with 57 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
csrc/kernels/api.cuh
View File

@@ -138,7 +138,8 @@ void dispatch(void* packed_recv_x, float* packed_recv_x_scales,
int* next_clean, int num_next_clean_int,
int num_tokens, int hidden, int num_max_dispatch_tokens_per_rank,
int num_topk, int num_experts, int rank, int num_ranks, bool use_fp8,
void* workspace, cudaStream_t stream, int phases);
void* workspace, int* usage_flag,
cudaStream_t stream, int phases);
void combine(void* combined_x,
void* rdma_recv_x, int* rdma_recv_flag, void* rdma_send_x,
@@ -147,8 +148,8 @@ void combine(void* combined_x,
int* next_clean, int num_next_clean_int,
int num_combined_tokens, int hidden, int num_max_dispatch_tokens_per_rank,
int num_topk, int num_experts, int rank, int num_ranks,
void* workspace, cudaStream_t stream,
int phases, bool zero_copy);
void* workspace, int* usage_flag,
cudaStream_t stream, int phases, bool zero_copy);
} // namespace internode_ll

26
csrc/kernels/internode_ll.cu
View File

@@ -47,7 +47,7 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales,
int* next_clean, int num_next_clean_int,
int num_tokens, int num_max_dispatch_tokens_per_rank,
int num_topk, int num_experts, int rank, int num_ranks,
int phases) {
int* usage_flag, int phases) {
const auto sm_id = static_cast<int>(blockIdx.x);
const auto thread_id = static_cast<int>(threadIdx.x);
const auto warp_id = thread_id / 32, lane_id = get_lane_id();
@@ -180,6 +180,10 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales,
#pragma unroll
for (int i = lane_id; i < num_experts; i += 32)
atomic_add_release_global(atomic_finish_counter_per_expert + i, FINISHED_SUM_TAG);
} else if (sm_id == 1) {
// The second SM is also responsible for notifying PCIe usage
if (lane_id == 0)
atomicAdd_system(usage_flag, 1);
}
// This SM should be responsible for some destination experts, read `topk_idx` for them
@@ -311,7 +315,8 @@ void dispatch(void* packed_recv_x, float* packed_recv_x_scales,
int* next_clean, int num_next_clean_int,
int num_tokens, int hidden, int num_max_dispatch_tokens_per_rank,
int num_topk, int num_experts, int rank, int num_ranks, bool use_fp8,
void* workspace, cudaStream_t stream, int phases) {
void* workspace, int* usage_flag,
cudaStream_t stream, int phases) {
constexpr int kNumMaxTopK = 9;
constexpr int kNumWarpsPerGroup = 10;
constexpr int kNumWarpGroups = 3;
@@ -338,7 +343,8 @@ LAUNCH_KERNEL(&cfg, dispatch_func, \
atomic_counter_per_expert, atomic_finish_counter_per_expert, \
next_clean, num_next_clean_int, \
num_tokens, num_max_dispatch_tokens_per_rank, \
num_topk, num_experts, rank, num_ranks, phases); } break
num_topk, num_experts, rank, num_ranks, \
usage_flag, phases); } break
SETUP_LAUNCH_CONFIG(num_sms, num_warps * 32, stream);
SWITCH_HIDDEN(DISPATCH_LAUNCH_CASE);
@@ -356,7 +362,7 @@ combine(void* combined_x,
int num_combined_tokens, int hidden, int num_topk,
int num_max_dispatch_tokens_per_rank,
int num_experts, int rank, int num_ranks,
int phases, bool zero_copy) {
int* usage_flag, int phases, bool zero_copy) {
const auto sm_id = static_cast<int>(blockIdx.x);
const auto num_sms = static_cast<int>(gridDim.x);
const auto thread_id = static_cast<int>(threadIdx.x);
@@ -451,11 +457,14 @@ combine(void* combined_x,
if ((phases & LOW_LATENCY_RECV_PHASE) == 0)
return;
// Wait all ranks to arrive
// Wait all ranks to arrive and notify usages
if (responsible_expert_idx < num_experts) {
EP_STATIC_ASSERT(kNumWarpsPerGroup > 1, "Invalid number of warps per group");
if (sub_warp_id == 0 and lane_id == 0)
if (sub_warp_id == 0 and lane_id == 0) {
while (ld_acquire_sys_global(rdma_recv_flag + responsible_expert_idx) == 0);
} else if (sm_id == 0 and sub_warp_id == 1 and lane_id == 0) {
atomicAdd_system(usage_flag, 1);
}
}
cg::this_grid().sync();
@@ -506,8 +515,8 @@ void combine(void* combined_x,
int* next_clean, int num_next_clean_int,
int num_combined_tokens, int hidden, int num_max_dispatch_tokens_per_rank,
int num_topk, int num_experts, int rank, int num_ranks,
void* workspace, cudaStream_t stream,
int phases, bool zero_copy) {
void* workspace, int* usage_flag,
cudaStream_t stream, int phases, bool zero_copy) {
constexpr int kNumWarpsPerGroup = 10;
constexpr int kNumWarpGroups = 3;
constexpr int kNumMaxTopk = 9;
@@ -531,6 +540,7 @@ LAUNCH_KERNEL(&cfg, combine_func, \
num_combined_tokens, hidden, num_topk, \
num_max_dispatch_tokens_per_rank, \
num_experts, rank, num_ranks, \
usage_flag, \
phases, zero_copy); } break
SETUP_LAUNCH_CONFIG(num_sms, num_warps * 32, stream);