chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Opt moe align block size kernel (#14133)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2025-12-02 19:13:55 +08:00
committed by GitHub
parent 9530b76630
commit c5947ecd85
3 changed files with 96 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
sgl-kernel/benchmark/bench_moe_align_block_size.py
View File

@@ -162,14 +162,17 @@ def calculate_diff(num_tokens, num_experts=256, block_size=128, topk=8):
]
)
max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
# SGL kernel uses dynamic padding optimization
max_num_tokens_padded_sgl = topk_ids.numel() + num_experts * (block_size - 1)
if topk_ids.numel() < num_experts + 1:
max_num_tokens_padded_sgl = topk_ids.numel() * block_size
sorted_ids_cuda = torch.empty(
(max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device
(max_num_tokens_padded_sgl,), dtype=torch.int32, device=topk_ids.device
)
sorted_ids_cuda.fill_(topk_ids.numel())
max_num_m_blocks = max_num_tokens_padded // block_size
max_num_m_blocks_sgl = max_num_tokens_padded_sgl // block_size
expert_ids_cuda = torch.zeros(
(max_num_m_blocks,), dtype=torch.int32, device=topk_ids.device
(max_num_m_blocks_sgl,), dtype=torch.int32, device=topk_ids.device
)
num_tokens_post_pad_cuda = torch.empty(
(1), dtype=torch.int32, device=topk_ids.device
@@ -178,14 +181,21 @@ def calculate_diff(num_tokens, num_experts=256, block_size=128, topk=8):
num_experts + 1, dtype=torch.int32, device=topk_ids.device
)
sorted_ids_triton = torch.empty_like(sorted_ids_cuda)
# Triton and vLLM use original padding calculation
max_num_tokens_padded_triton = topk_ids.numel() + num_experts * (block_size - 1)
max_num_m_blocks_triton = max_num_tokens_padded_triton // block_size
sorted_ids_triton = torch.empty(
(max_num_tokens_padded_triton,), dtype=torch.int32, device=topk_ids.device
)
sorted_ids_triton.fill_(topk_ids.numel())
expert_ids_triton = torch.zeros_like(expert_ids_cuda)
expert_ids_triton = torch.zeros(
(max_num_m_blocks_triton,), dtype=torch.int32, device=topk_ids.device
)
num_tokens_post_pad_triton = torch.empty_like(num_tokens_post_pad_cuda)
sorted_ids_vllm = torch.empty_like(sorted_ids_cuda)
sorted_ids_vllm = torch.empty_like(sorted_ids_triton)
sorted_ids_vllm.fill_(topk_ids.numel())
expert_ids_vllm = torch.zeros_like(expert_ids_cuda)
expert_ids_vllm = torch.zeros_like(expert_ids_triton)
num_tokens_post_pad_vllm = torch.empty_like(num_tokens_post_pad_cuda)
# compare the performance of cuda, triton and vllm implementation
@@ -326,7 +336,17 @@ def benchmark(num_tokens, num_experts, topk, provider):
device="cuda",
)
max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
# Calculate max_num_tokens_padded based on provider
if provider == "sgl" or provider == "sgl_fusion":
# Apply dynamic padding optimization for SGL kernel
max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
if topk_ids.numel() < num_experts:
max_num_tokens_padded = topk_ids.numel() * block_size
else: # triton
# Use original padding calculation for Triton
max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
# Create tensors
sorted_ids = torch.empty(
(max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device
)

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

@@ -63,7 +63,22 @@ __global__ void moe_align_block_size_kernel(
size_t numel,
int32_t* __restrict__ cumsum,
bool pad_sorted_token_ids,
const int32_t scan_size) {
const int32_t scan_size,
int32_t max_num_tokens_padded) {
// Use a separate thread block to populate sorted_token_ids
if (blockIdx.x == 1) {
if (pad_sorted_token_ids) {
Vec fill_vec;
fill_vec.x = fill_vec.y = fill_vec.z = fill_vec.w = numel;
int32_t total_vecs = (max_num_tokens_padded + VEC_SIZE - 1) / VEC_SIZE;
Vec* out_ptr = reinterpret_cast<Vec*>(sorted_token_ids);
for (int32_t i = threadIdx.x; i < total_vecs; i += blockDim.x) {
out_ptr[i] = fill_vec;
}
}
return;
}
extern __shared__ int32_t smem[];
int32_t* shared_counts = smem; // [num_experts]
int32_t* prefix = shared_counts + num_experts; // [num_experts + 1]
@@ -215,19 +230,9 @@ __global__ void moe_align_block_size_kernel(
}
expert_ids[i] = left - 2;
}
if (pad_sorted_token_ids) {
Vec fill_vec;
fill_vec.x = fill_vec.y = fill_vec.z = fill_vec.w = numel;
int32_t total_vecs = (s_total_tokens_post_pad + VEC_SIZE - 1) / VEC_SIZE;
Vec* out_ptr = reinterpret_cast<Vec*>(sorted_token_ids);
for (int32_t i = tid; i < total_vecs; i += stride) {
out_ptr[i] = fill_vec;
}
}
}
template <typename scalar_t>
template <typename scalar_t, int32_t fill_threads>
__global__ void moe_align_block_size_small_batch_expert_kernel(
const scalar_t* __restrict__ topk_ids,
int32_t* __restrict__ sorted_token_ids,
@@ -236,66 +241,76 @@ __global__ void moe_align_block_size_small_batch_expert_kernel(
int32_t num_experts,
int32_t block_size,
size_t numel,
bool pad_sorted_token_ids) {
const size_t tid = threadIdx.x;
const size_t stride = blockDim.x;
bool pad_sorted_token_ids,
int32_t max_num_tokens_padded) {
// Adapted from
// https://github.com/vllm-project/vllm/pull/29642/files#diff-5647b1413f4ae9aacba904eca8f8a8aee9079321eadff4c10101a2c6962dcc53R226
// Use an additional group of threads to fill sorted_token_ids.
// Since the kernel will use sorted_token_ids afterward,
// we fill sorted_token_ids within the same threadblock to make
// synchronization easier.
if (threadIdx.x < fill_threads) {
// Initialize sorted_token_ids with numel
if (pad_sorted_token_ids) {
for (int32_t it = threadIdx.x; it < max_num_tokens_padded; it += fill_threads) {
sorted_token_ids[it] = numel;
}
}
// Three __syncthreads() corresponding to the other threads
__syncthreads();
__syncthreads();
__syncthreads();
return;
}
const size_t tid = threadIdx.x - fill_threads;
const size_t stride = blockDim.x - fill_threads;
extern __shared__ int32_t shared_mem[];
int32_t* cumsum = shared_mem;
int32_t* tokens_cnts = (int32_t*)(shared_mem + num_experts + 1);
for (int i = 0; i < num_experts; ++i) {
tokens_cnts[(threadIdx.x + 1) * num_experts + i] = 0;
tokens_cnts[(tid + 1) * num_experts + i] = 0;
}
for (size_t i = tid; i < numel; i += stride) {
++tokens_cnts[(threadIdx.x + 1) * num_experts + topk_ids[i] + 1];
int32_t expert_id = topk_ids[i] + 1;
++tokens_cnts[(tid + 1) * num_experts + expert_id];
}
__syncthreads();
if (threadIdx.x < num_experts) {
tokens_cnts[threadIdx.x] = 0;
for (int i = 1; i <= blockDim.x; ++i) {
tokens_cnts[i * num_experts + threadIdx.x] += tokens_cnts[(i - 1) * num_experts + threadIdx.x];
if (tid < num_experts) {
tokens_cnts[tid] = 0;
for (int i = 1; i <= stride; ++i) {
tokens_cnts[i * num_experts + tid] += tokens_cnts[(i - 1) * num_experts + tid];
}
}
__syncthreads();
if (threadIdx.x == 0) {
if (tid == 0) {
cumsum[0] = 0;
for (int i = 1; i <= num_experts; ++i) {
cumsum[i] = cumsum[i - 1] + CEILDIV(tokens_cnts[blockDim.x * num_experts + i - 1], block_size) * block_size;
cumsum[i] = cumsum[i - 1] + CEILDIV(tokens_cnts[stride * num_experts + i - 1], block_size) * block_size;
}
*total_tokens_post_pad = static_cast<int32_t>(cumsum[num_experts]);
}
__syncthreads();
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 - 1;
if (tid < num_experts) {
for (int i = cumsum[tid]; i < cumsum[tid + 1]; i += block_size) {
expert_ids[i / block_size] = tid - 1;
}
}
if (pad_sorted_token_ids) {
Vec fill_vec;
fill_vec.x = fill_vec.y = fill_vec.z = fill_vec.w = numel;
int32_t total_vecs = (*total_tokens_post_pad + VEC_SIZE - 1) / VEC_SIZE;
Vec* out_ptr = reinterpret_cast<Vec*>(sorted_token_ids);
for (int32_t i = tid; i < total_vecs; i += stride) {
out_ptr[i] = fill_vec;
}
}
__syncthreads();
for (size_t i = tid; i < numel; i += stride) {
int32_t expert_id = topk_ids[i] + 1;
int32_t rank_post_pad = tokens_cnts[threadIdx.x * num_experts + expert_id] + cumsum[expert_id];
int32_t rank_post_pad = tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id];
sorted_token_ids[rank_post_pad] = i;
++tokens_cnts[threadIdx.x * num_experts + expert_id];
++tokens_cnts[tid * num_experts + expert_id];
}
}
@@ -311,18 +326,20 @@ void moe_align_block_size(
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
int threads = 1024;
threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
int64_t max_num_tokens_padded = sorted_token_ids.size(0);
DISPATCH_INTEGRAL_TYPES(topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] {
bool small_batch_expert_mode = (topk_ids.numel() < 1024) && (num_experts <= 64);
if (small_batch_expert_mode) {
const int32_t threads = max((int32_t)num_experts, WARP_SIZE);
constexpr int32_t fill_threads = 256;
const int32_t shared_mem_size = ((threads + 1) * num_experts + (num_experts + 1)) * sizeof(int32_t);
auto small_batch_expert_kernel = moe_align_block_size_small_batch_expert_kernel<scalar_t>;
small_batch_expert_kernel<<<1, threads, shared_mem_size, stream>>>(
auto small_batch_expert_kernel = moe_align_block_size_small_batch_expert_kernel<scalar_t, fill_threads>;
small_batch_expert_kernel<<<1, fill_threads + threads, shared_mem_size, stream>>>(
topk_ids.data_ptr<scalar_t>(),
sorted_token_ids.data_ptr<int32_t>(),
experts_ids.data_ptr<int32_t>(),
@@ -330,13 +347,14 @@ void moe_align_block_size(
num_experts,
block_size,
topk_ids.numel(),
pad_sorted_token_ids);
pad_sorted_token_ids,
max_num_tokens_padded);
} else {
auto align_kernel = moe_align_block_size_kernel<scalar_t>;
const size_t scan_size = next_pow2(num_experts);
const size_t shared_mem_size = (num_experts + (num_experts + 1) + scan_size + WARP_SIZE) * sizeof(int32_t);
align_kernel<<<1, threads, shared_mem_size, stream>>>(
align_kernel<<<2, threads, shared_mem_size, stream>>>(
topk_ids.data_ptr<scalar_t>(),
sorted_token_ids.data_ptr<int32_t>(),
experts_ids.data_ptr<int32_t>(),
@@ -346,7 +364,8 @@ void moe_align_block_size(
topk_ids.numel(),
cumsum_buffer.data_ptr<int32_t>(),
pad_sorted_token_ids,
scan_size);
scan_size,
max_num_tokens_padded);
const int block_threads = std::min(256, (int)threads);
const int num_blocks = (topk_ids.numel() + block_threads - 1) / block_threads;

2
sgl-kernel/tests/test_moe_align.py
View File

@@ -165,6 +165,8 @@ def test_moe_align_block_size_compare_implementations(
]
max_num_tokens_padded = topk_ids.numel() + (num_experts + 1) * (block_size - 1)
if topk_ids.numel() < num_experts + 1:
max_num_tokens_padded = topk_ids.numel() * block_size
sorted_ids_cuda = torch.empty(
(max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device