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

fix ep_moe_reorder kernel bugs (#6858)

Browse Source 
Co-authored-by: JieXin Liang <Alcanderian@users.noreply.github.com>
This commit is contained in:
Xiaoyu Zhang
2025-06-04 19:13:59 +08:00
committed by GitHub
parent 180ff5eecc
commit bd75690f4e
3 changed files with 226 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

59
sgl-kernel/csrc/moe/ep_moe_reorder_kernel.cu
View File

@@ -7,9 +7,10 @@
#include "utils.h"
template <typename scalar_t>
__global__ void ep_pre_reorder_cuda_kernel(
const float* __restrict__ input_ptr,
float* __restrict__ gateup_input_ptr,
const scalar_t* __restrict__ input_ptr,
scalar_t* __restrict__ gateup_input_ptr,
const int* __restrict__ src2dst_ptr,
const int* __restrict__ topk_ids_ptr,
const float* __restrict__ a1_scales_ptr,
@@ -21,20 +22,20 @@ __global__ void ep_pre_reorder_cuda_kernel(
int token_idx = blockIdx.x;
int tid = threadIdx.x;
const float* src_ptr = input_ptr + int64_t(token_idx) * hidden_size;
const scalar_t* src_ptr = input_ptr + int64_t(token_idx) * hidden_size;
const int* token_src2dst = src2dst_ptr + token_idx * topk;
const int* token_topk_ids = topk_ids_ptr + token_idx * topk;
float scale = 1.0f;
if (a1_scales_ptr != nullptr and use_per_token_if_dynamic) {
scale = 1.0f / a1_scales_ptr[token_idx];
}
for (int k = 0; k < topk; ++k) {
int expert_id = token_topk_ids[k];
if (expert_id < start_expert_id || expert_id > end_expert_id) continue;
float scale = 1.0f;
if (a1_scales_ptr != nullptr and use_per_token_if_dynamic) {
scale = 1.0f / a1_scales_ptr[token_idx];
}
if (a1_scales_ptr != nullptr) {
if (!use_per_token_if_dynamic) {
scale = 1.0f / a1_scales_ptr[expert_id - start_expert_id];
@@ -42,21 +43,27 @@ __global__ void ep_pre_reorder_cuda_kernel(
}
int dst_idx = token_src2dst[k];
float* dst_ptr = gateup_input_ptr + int64_t(dst_idx) * hidden_size;
scalar_t* dst_ptr = gateup_input_ptr + int64_t(dst_idx) * hidden_size;
constexpr uint32_t vec_size = 16 / sizeof(float);
using vec_t = flashinfer::vec_t<float, vec_size>;
constexpr uint32_t vec_size = 16 / sizeof(scalar_t);
using vec_t = flashinfer::vec_t<scalar_t, vec_size>;
int vec_elements = (hidden_size / vec_size) * vec_size;
for (int idx = tid; idx < hidden_size / vec_size; idx += blockDim.x) {
vec_t input_vec, output_vec;
input_vec.cast_load(src_ptr + idx * vec_size);
#pragma unroll
for (uint32_t i = 0; i < vec_size; ++i) {
float val = static_cast<float>(input_vec[i]);
output_vec[i] = val * scale;
output_vec[i] = static_cast<scalar_t>(val * scale);
}
output_vec.cast_store(dst_ptr + idx * vec_size);
}
for (int idx = vec_elements + tid; idx < hidden_size; idx += blockDim.x) {
float val = static_cast<float>(src_ptr[idx]);
dst_ptr[idx] = static_cast<scalar_t>(val * scale);
}
}
}
@@ -75,15 +82,19 @@ void ep_moe_pre_reorder(
dim3 grid(total_blocks);
dim3 block(block_size);
int hidden_size = input.size(1);
ep_pre_reorder_cuda_kernel<<<grid, block>>>(
input.data_ptr<float>(),
gateup_input.data_ptr<float>(),
src2dst.data_ptr<int>(),
topk_ids.data_ptr<int>(),
a1_scales.defined() ? a1_scales.data_ptr<float>() : nullptr,
start_expert_id,
end_expert_id,
topk,
hidden_size,
use_per_token_if_dynamic);
DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(input.scalar_type(), scalar_t, [&] {
ep_pre_reorder_cuda_kernel<scalar_t><<<grid, block>>>(
static_cast<scalar_t*>(input.data_ptr()),
static_cast<scalar_t*>(gateup_input.data_ptr()),
src2dst.data_ptr<int>(),
topk_ids.data_ptr<int>(),
a1_scales.defined() ? a1_scales.data_ptr<float>() : nullptr,
start_expert_id,
end_expert_id,
topk,
hidden_size,
use_per_token_if_dynamic);
return true;
});
}