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

Minor fix

Browse Source
This commit is contained in:
Chenggang Zhao
2025-06-19 10:38:42 +08:00
parent 24453275e3
commit 3a3398f686
1 changed files with 12 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
csrc/kernels/buffer.cuh
View File

@@ -17,12 +17,12 @@ public:
__device__ __forceinline__ Buffer(void* &gbl_ptr, int num_elems, int offset = 0) { __device__ __forceinline__ Buffer(void* &gbl_ptr, int num_elems, int offset = 0) {
total_bytes = num_elems * sizeof(dtype_t); total_bytes = num_elems * sizeof(dtype_t);
ptr = reinterpret_cast<uint8_t*>(gbl_ptr) + offset * sizeof(dtype_t); ptr = static_cast<uint8_t*>(gbl_ptr) + offset * sizeof(dtype_t);
gbl_ptr = reinterpret_cast<uint8_t*>(gbl_ptr) + total_bytes; gbl_ptr = static_cast<uint8_t*>(gbl_ptr) + total_bytes;
} }
__device__ __forceinline__ Buffer advance_also(void* &gbl_ptr) { __device__ __forceinline__ Buffer advance_also(void* &gbl_ptr) {
gbl_ptr = reinterpret_cast<uint8_t*>(gbl_ptr) + total_bytes; gbl_ptr = static_cast<uint8_t*>(gbl_ptr) + total_bytes;
return *this; return *this;
} }
@@ -51,8 +51,8 @@ public:
int per_channel_bytes = num_bytes * num_ranks; int per_channel_bytes = num_bytes * num_ranks;
total_bytes = per_channel_bytes * num_sms; total_bytes = per_channel_bytes * num_sms;
ptrs[0] = reinterpret_cast<uint8_t*>(gbl_ptr) + per_channel_bytes * sm_id + num_bytes * offset; ptrs[0] = static_cast<uint8_t*>(gbl_ptr) + per_channel_bytes * sm_id + num_bytes * offset;
gbl_ptr = reinterpret_cast<uint8_t*>(gbl_ptr) + total_bytes; gbl_ptr = static_cast<uint8_t*>(gbl_ptr) + total_bytes;
} }
__device__ __forceinline__ AsymBuffer(void** gbl_ptrs, int num_elems, int num_ranks, __device__ __forceinline__ AsymBuffer(void** gbl_ptrs, int num_elems, int num_ranks,
@@ -63,8 +63,8 @@ public:
int per_channel_bytes = num_bytes * num_ranks; int per_channel_bytes = num_bytes * num_ranks;
total_bytes = per_channel_bytes * num_sms; total_bytes = per_channel_bytes * num_sms;
for (int i = 0; i < kNumRanks; ++ i) { for (int i = 0; i < kNumRanks; ++ i) {
ptrs[i] = reinterpret_cast<uint8_t*>(gbl_ptrs[i]) + per_channel_bytes * sm_id + num_bytes * offset; ptrs[i] = static_cast<uint8_t*>(gbl_ptrs[i]) + per_channel_bytes * sm_id + num_bytes * offset;
gbl_ptrs[i] = reinterpret_cast<uint8_t*>(gbl_ptrs[i]) + total_bytes; gbl_ptrs[i] = static_cast<uint8_t*>(gbl_ptrs[i]) + total_bytes;
} }
} }
@@ -75,14 +75,14 @@ public:
} }
__device__ __forceinline__ AsymBuffer advance_also(void* &gbl_ptr) { __device__ __forceinline__ AsymBuffer advance_also(void* &gbl_ptr) {
gbl_ptr = reinterpret_cast<uint8_t*>(gbl_ptr) + total_bytes; gbl_ptr = static_cast<uint8_t*>(gbl_ptr) + total_bytes;
return *this; return *this;
} }
template<int kNumAlsoRanks> template<int kNumAlsoRanks>
__device__ __forceinline__ AsymBuffer advance_also(void** gbl_ptrs) { __device__ __forceinline__ AsymBuffer advance_also(void** gbl_ptrs) {
for (int i = 0; i < kNumAlsoRanks; ++ i) for (int i = 0; i < kNumAlsoRanks; ++ i)
gbl_ptrs[i] = reinterpret_cast<uint8_t*>(gbl_ptrs[i]) + total_bytes; gbl_ptrs[i] = static_cast<uint8_t*>(gbl_ptrs[i]) + total_bytes;
return *this; return *this;
} }
@@ -114,9 +114,9 @@ public:
int per_channel_bytes = num_bytes * num_ranks; int per_channel_bytes = num_bytes * num_ranks;
total_bytes = per_channel_bytes * num_sms * (static_cast<int>(kDecoupled) + 1); total_bytes = per_channel_bytes * num_sms * (static_cast<int>(kDecoupled) + 1);
send_ptr = reinterpret_cast<uint8_t*>(gbl_ptr) + per_channel_bytes * sm_id; send_ptr = static_cast<uint8_t*>(gbl_ptr) + per_channel_bytes * sm_id;
recv_ptr = reinterpret_cast<uint8_t*>(gbl_ptr) + per_channel_bytes * (sm_id + num_sms); recv_ptr = static_cast<uint8_t*>(gbl_ptr) + per_channel_bytes * (sm_id + num_sms);
gbl_ptr = reinterpret_cast<uint8_t*>(gbl_ptr) + total_bytes; gbl_ptr = static_cast<uint8_t*>(gbl_ptr) + total_bytes;
} }
__device__ __forceinline__ dtype_t* send_buffer(int idx = 0) { __device__ __forceinline__ dtype_t* send_buffer(int idx = 0) {