diff --git a/sgl-kernel/benchmark/bench_moe_align_block_size.py b/sgl-kernel/benchmark/bench_moe_align_block_size.py index 2156c5cd4..5dc4a0e4d 100644 --- a/sgl-kernel/benchmark/bench_moe_align_block_size.py +++ b/sgl-kernel/benchmark/bench_moe_align_block_size.py @@ -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 ) diff --git a/sgl-kernel/csrc/moe/moe_align_kernel.cu b/sgl-kernel/csrc/moe/moe_align_kernel.cu index 92fd34270..b7bad43da 100644 --- a/sgl-kernel/csrc/moe/moe_align_kernel.cu +++ b/sgl-kernel/csrc/moe/moe_align_kernel.cu @@ -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(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(sorted_token_ids); - for (int32_t i = tid; i < total_vecs; i += stride) { - out_ptr[i] = fill_vec; - } - } } -template +template __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(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(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; - small_batch_expert_kernel<<<1, threads, shared_mem_size, stream>>>( + auto small_batch_expert_kernel = moe_align_block_size_small_batch_expert_kernel; + small_batch_expert_kernel<<<1, fill_threads + threads, shared_mem_size, stream>>>( topk_ids.data_ptr(), sorted_token_ids.data_ptr(), experts_ids.data_ptr(), @@ -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; 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(), sorted_token_ids.data_ptr(), experts_ids.data_ptr(), @@ -346,7 +364,8 @@ void moe_align_block_size( topk_ids.numel(), cumsum_buffer.data_ptr(), 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; diff --git a/sgl-kernel/tests/test_moe_align.py b/sgl-kernel/tests/test_moe_align.py index 40a37f563..6a27f6e72 100644 --- a/sgl-kernel/tests/test_moe_align.py +++ b/sgl-kernel/tests/test_moe_align.py @@ -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