EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

update sgl-kernel for EP: kernel part (#8514)

Browse Source 
Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>
This commit is contained in:
Cheng Wan
2025-07-30 22:19:55 -07:00
committed by GitHub
parent 59aab76f0a
commit a5f5ab4030
7 changed files with 12 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
sgl-kernel/csrc/common_extension.cc
View File

@@ -165,7 +165,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
*/
m.def(
"moe_align_block_size(Tensor topk_ids, int num_experts, int block_size, Tensor! sorted_token_ids, Tensor! "
"experts_ids, Tensor! num_tokens_post_pad, Tensor! token_cnts_buffer, Tensor! cumsum_buffer, bool "
"experts_ids, Tensor! num_tokens_post_pad, Tensor! cumsum_buffer, bool "
"pad_sorted_token_ids) -> ()");
m.impl("moe_align_block_size", torch::kCUDA, &moe_align_block_size);

13
sgl-kernel/csrc/moe/moe_align_kernel.cu
View File

@@ -36,7 +36,7 @@ __global__ void count_and_sort_expert_tokens_kernel(
const size_t stride = blockDim.x * gridDim.x;
for (size_t i = tid; i < numel; i += stride) {
int32_t expert_id = topk_ids[i];
int32_t expert_id = topk_ids[i] + 1;
int32_t rank_post_pad = atomicAdd(&cumsum_buffer[expert_id], 1);
sorted_token_ids[rank_post_pad] = i;
}
@@ -82,7 +82,7 @@ __global__ void moe_align_block_size_kernel(
__syncthreads();
for (size_t i = tid; i < numel; i += stride) {
int expert_id = topk_ids[i];
int expert_id = topk_ids[i] + 1;
atomicAdd(&shared_counts[expert_id], 1);
}
@@ -215,7 +215,7 @@ __global__ void moe_align_block_size_kernel(
right = mid;
}
}
expert_ids[i] = left - 1;
expert_ids[i] = left - 2;
}
if (pad_sorted_token_ids) {
@@ -251,7 +251,7 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
}
for (size_t i = tid; i < numel; i += stride) {
++tokens_cnts[(threadIdx.x + 1) * num_experts + topk_ids[i]];
++tokens_cnts[(threadIdx.x + 1) * num_experts + topk_ids[i] + 1];
}
__syncthreads();
@@ -277,7 +277,7 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
if (threadIdx.x < num_experts) {
for (int i = cumsum[threadIdx.x]; i < cumsum[threadIdx.x + 1]; i += block_size) {
expert_ids[i / block_size] = threadIdx.x;
expert_ids[i / block_size] = threadIdx.x - 1;
}
}
@@ -294,7 +294,7 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
__syncthreads();
for (size_t i = tid; i < numel; i += stride) {
int32_t expert_id = topk_ids[i];
int32_t expert_id = topk_ids[i] + 1;
int32_t rank_post_pad = tokens_cnts[threadIdx.x * num_experts + expert_id] + cumsum[expert_id];
sorted_token_ids[rank_post_pad] = i;
++tokens_cnts[threadIdx.x * num_experts + expert_id];
@@ -308,7 +308,6 @@ void moe_align_block_size(
torch::Tensor sorted_token_ids,
torch::Tensor experts_ids,
torch::Tensor num_tokens_post_pad,
torch::Tensor token_cnts_buffer,
torch::Tensor cumsum_buffer,
bool pad_sorted_token_ids) {
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

2
sgl-kernel/csrc/torch_extension_rocm.cc
View File

@@ -92,7 +92,7 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
*/
m.def(
"moe_align_block_size(Tensor topk_ids, int num_experts, int block_size, Tensor! sorted_token_ids, Tensor! "
"experts_ids, Tensor! num_tokens_post_pad, Tensor! token_cnts_buffer, Tensor! cumsum_buffer, bool "
"experts_ids, Tensor! num_tokens_post_pad, Tensor! cumsum_buffer, bool "
"pad_sorted_token_ids) -> ()");
m.impl("moe_align_block_size", torch::kCUDA, &moe_align_block_size);