diff --git a/python/sglang/srt/layers/quantization/w4afp8.py b/python/sglang/srt/layers/quantization/w4afp8.py index 6ab8f4a69..1dc592957 100644 --- a/python/sglang/srt/layers/quantization/w4afp8.py +++ b/python/sglang/srt/layers/quantization/w4afp8.py @@ -300,7 +300,6 @@ class W4AFp8MoEMethod(FusedMoEMethodBase): x = dispatch_output.hidden_states topk_output = dispatch_output.topk_output - topk_weights, topk_ids, _ = topk_output output = cutlass_w4a8_moe( diff --git a/sgl-kernel/csrc/moe/cutlass_moe/w4a8/w4a8_grouped_mm_c3x.cu b/sgl-kernel/csrc/moe/cutlass_moe/w4a8/w4a8_grouped_mm_c3x.cu index bd63d2ee1..94a6003a4 100644 --- a/sgl-kernel/csrc/moe/cutlass_moe/w4a8/w4a8_grouped_mm_c3x.cu +++ b/sgl-kernel/csrc/moe/cutlass_moe/w4a8/w4a8_grouped_mm_c3x.cu @@ -66,6 +66,25 @@ inline void invoke_gemm( chunk_size); } +// Helper macro to reduce code duplication +// Note: Config must be wrapped in parentheses when it contains commas (e.g., template parameters) +// This uses a helper macro to strip the parentheses from the template parameter +#define INVOKE_GEMM_WITH_CONFIG_HELPER(...) \ + invoke_gemm<__VA_ARGS__>( \ + d_tensors, \ + a_tensors, \ + b_tensors, \ + a_scales, \ + b_scales, \ + expert_offsets, \ + problem_sizes, \ + a_strides, \ + b_strides, \ + d_strides, \ + s_strides, \ + chunk_size) +#define INVOKE_GEMM_WITH_CONFIG(Config) INVOKE_GEMM_WITH_CONFIG_HELPER Config + void dispatch_w4a8_moe_mm_sm90( torch::Tensor& d_tensors, torch::Tensor const& a_tensors, @@ -87,268 +106,77 @@ void dispatch_w4a8_moe_mm_sm90( if (n == 4096 && k == 7168) { // group gemm 1 if (m <= 4) { - invoke_gemm>( - d_tensors, - a_tensors, - b_tensors, - a_scales, - b_scales, - expert_offsets, - problem_sizes, - a_strides, - b_strides, - d_strides, - s_strides, - chunk_size); - } else if (m <= 16) { - invoke_gemm>( - d_tensors, - a_tensors, - b_tensors, - a_scales, - b_scales, - expert_offsets, - problem_sizes, - a_strides, - b_strides, - d_strides, - s_strides, - chunk_size); + INVOKE_GEMM_WITH_CONFIG((SM90_PP<64, 32, 512, 2, 1, 1>)); + } else if (m <= 32) { + INVOKE_GEMM_WITH_CONFIG((SM90_CO<128, 16, 512, 2, 1, 1>)); } else if (m <= 256) { - invoke_gemm>( - d_tensors, - a_tensors, - b_tensors, - a_scales, - b_scales, - expert_offsets, - problem_sizes, - a_strides, - b_strides, - d_strides, - s_strides, - chunk_size); + INVOKE_GEMM_WITH_CONFIG((SM90_CO<128, 16, 512, 1, 1, 1>)); } else if (m <= 1024) { - invoke_gemm>( - d_tensors, - a_tensors, - b_tensors, - a_scales, - b_scales, - expert_offsets, - problem_sizes, - a_strides, - b_strides, - d_strides, - s_strides, - chunk_size); + INVOKE_GEMM_WITH_CONFIG((SM90_CO<128, 32, 512, 2, 1, 1>)); + } else if (m <= 4096) { + // Optimized for prefill: seq_len up to 4096 (m=4096 with topk=1) + INVOKE_GEMM_WITH_CONFIG((SM90_CO<128, 64, 512, 2, 1, 1>)); } else { - invoke_gemm>( - d_tensors, - a_tensors, - b_tensors, - a_scales, - b_scales, - expert_offsets, - problem_sizes, - a_strides, - b_strides, - d_strides, - s_strides, - chunk_size); + // Optimized for prefill: seq_len up to 8192 (m=8192 with topk=1) + INVOKE_GEMM_WITH_CONFIG((SM90_CO<128, 64, 512, 1, 1, 1>)); } } else if (n == 7168 && k == 2048) { // group gemm 2 if (m <= 8) { - invoke_gemm>( - d_tensors, - a_tensors, - b_tensors, - a_scales, - b_scales, - expert_offsets, - problem_sizes, - a_strides, - b_strides, - d_strides, - s_strides, - chunk_size); + INVOKE_GEMM_WITH_CONFIG((SM90_PP<64, 16, 512, 1, 1, 1>)); } else if (m <= 512) { - invoke_gemm>( - d_tensors, - a_tensors, - b_tensors, - a_scales, - b_scales, - expert_offsets, - problem_sizes, - a_strides, - b_strides, - d_strides, - s_strides, - chunk_size); + INVOKE_GEMM_WITH_CONFIG((SM90_CO<128, 32, 512, 1, 1, 1>)); + } else if (m <= 4096) { + // Optimized for prefill: larger cluster for better throughput + INVOKE_GEMM_WITH_CONFIG((SM90_CO<128, 64, 512, 2, 1, 1>)); } else { - invoke_gemm>( - d_tensors, - a_tensors, - b_tensors, - a_scales, - b_scales, - expert_offsets, - problem_sizes, - a_strides, - b_strides, - d_strides, - s_strides, - chunk_size); + INVOKE_GEMM_WITH_CONFIG((SM90_CO<128, 64, 512, 1, 1, 1>)); } } else if (n == 512 && k == 7168) { // group gemm 1 for tp if (m <= 4) { - invoke_gemm>( - d_tensors, - a_tensors, - b_tensors, - a_scales, - b_scales, - expert_offsets, - problem_sizes, - a_strides, - b_strides, - d_strides, - s_strides, - chunk_size); - } else if (m <= 16) { - invoke_gemm>( - d_tensors, - a_tensors, - b_tensors, - a_scales, - b_scales, - expert_offsets, - problem_sizes, - a_strides, - b_strides, - d_strides, - s_strides, - chunk_size); + INVOKE_GEMM_WITH_CONFIG((SM90_PP<64, 32, 512, 2, 1, 1>)); + } else if (m <= 32) { + INVOKE_GEMM_WITH_CONFIG((SM90_CO<128, 16, 512, 2, 1, 1>)); } else if (m <= 256) { - invoke_gemm>( - d_tensors, - a_tensors, - b_tensors, - a_scales, - b_scales, - expert_offsets, - problem_sizes, - a_strides, - b_strides, - d_strides, - s_strides, - chunk_size); + INVOKE_GEMM_WITH_CONFIG((SM90_CO<128, 16, 512, 1, 1, 1>)); } else if (m <= 1024) { - invoke_gemm>( - d_tensors, - a_tensors, - b_tensors, - a_scales, - b_scales, - expert_offsets, - problem_sizes, - a_strides, - b_strides, - d_strides, - s_strides, - chunk_size); + INVOKE_GEMM_WITH_CONFIG((SM90_CO<128, 32, 512, 2, 1, 1>)); } else { - invoke_gemm>( - d_tensors, - a_tensors, - b_tensors, - a_scales, - b_scales, - expert_offsets, - problem_sizes, - a_strides, - b_strides, - d_strides, - s_strides, - chunk_size); + INVOKE_GEMM_WITH_CONFIG((SM90_CO<128, 64, 512, 1, 1, 1>)); } } else if (n == 7168 && k == 256) { // group gemm 2 for tp if (m <= 8) { - invoke_gemm>( - d_tensors, - a_tensors, - b_tensors, - a_scales, - b_scales, - expert_offsets, - problem_sizes, - a_strides, - b_strides, - d_strides, - s_strides, - chunk_size); + INVOKE_GEMM_WITH_CONFIG((SM90_PP<64, 16, 128, 1, 1, 1>)); + } else if (m <= 32) { + INVOKE_GEMM_WITH_CONFIG((SM90_PP<128, 32, 128, 1, 1, 1>)); } else if (m <= 512) { - invoke_gemm>( - d_tensors, - a_tensors, - b_tensors, - a_scales, - b_scales, - expert_offsets, - problem_sizes, - a_strides, - b_strides, - d_strides, - s_strides, - chunk_size); + INVOKE_GEMM_WITH_CONFIG((SM90_PP<128, 32, 128, 2, 1, 1>)); } else { - invoke_gemm>( - d_tensors, - a_tensors, - b_tensors, - a_scales, - b_scales, - expert_offsets, - problem_sizes, - a_strides, - b_strides, - d_strides, - s_strides, - chunk_size); + INVOKE_GEMM_WITH_CONFIG((SM90_PP<128, 64, 128, 1, 1, 1>)); } } else { if (k % 512 == 0) { - invoke_gemm>( - d_tensors, - a_tensors, - b_tensors, - a_scales, - b_scales, - expert_offsets, - problem_sizes, - a_strides, - b_strides, - d_strides, - s_strides, - chunk_size); + // For large m (prefill), prefer larger cluster + if (m <= 32) { + // Decode: target batch size (16-32) - use cluster size 1 for better latency + INVOKE_GEMM_WITH_CONFIG((SM90_CO<128, 16, 512, 1, 1, 1>)); + } else if (m <= 1024) { + // Decode: large batch or small prefill + INVOKE_GEMM_WITH_CONFIG((SM90_CO<128, 32, 512, 1, 1, 1>)); + } else { + // Prefill: large sequence length - prefer larger cluster + INVOKE_GEMM_WITH_CONFIG((SM90_CO<128, 64, 512, 1, 1, 1>)); + } } else { - invoke_gemm>( - d_tensors, - a_tensors, - b_tensors, - a_scales, - b_scales, - expert_offsets, - problem_sizes, - a_strides, - b_strides, - d_strides, - s_strides, - chunk_size); + if (m <= 32) { + // Decode: target batch size (16-32) - use larger tile for better throughput + INVOKE_GEMM_WITH_CONFIG((SM90_PP<128, 32, 128, 1, 1, 1>)); + } else { + // Prefill: larger sequence length + INVOKE_GEMM_WITH_CONFIG((SM90_PP<128, 64, 128, 1, 1, 1>)); + } } } } diff --git a/sgl-kernel/csrc/moe/cutlass_moe/w4a8/w4a8_moe_data.cu b/sgl-kernel/csrc/moe/cutlass_moe/w4a8/w4a8_moe_data.cu index f30f7b025..6316a913a 100644 --- a/sgl-kernel/csrc/moe/cutlass_moe/w4a8/w4a8_moe_data.cu +++ b/sgl-kernel/csrc/moe/cutlass_moe/w4a8/w4a8_moe_data.cu @@ -2,29 +2,48 @@ #include #include -#include - -constexpr uint64_t THREADS_PER_EXPERT = 512; +#include +#include +template __global__ void compute_problem_sizes_w4a8( const int32_t* __restrict__ topk_ids, int32_t* problem_sizes1, int32_t* problem_sizes2, - int32_t* atomic_buffer, const int topk_length, const int n, const int k) { int expert_id = blockIdx.x; int occurrences = 0; - for (int i = threadIdx.x; i < topk_length; i += THREADS_PER_EXPERT) { - occurrences += (topk_ids[i] == expert_id); + // Optimized: vectorized memory access using int4 for better memory bandwidth + // Process vectorized chunks first + bool aligned = (reinterpret_cast(topk_ids) % 16 == 0); + + if (aligned) { + const int4* vec_ptr = reinterpret_cast(topk_ids); + int vec_length = topk_length / 4; + + for (int i = threadIdx.x; i < vec_length; i += BLOCK_SIZE) { + int4 vec_data = vec_ptr[i]; + occurrences += + (vec_data.x == expert_id) + (vec_data.y == expert_id) + (vec_data.z == expert_id) + (vec_data.w == expert_id); + } + + for (int i = vec_length * 4 + threadIdx.x; i < topk_length; i += BLOCK_SIZE) { + occurrences += (topk_ids[i] == expert_id); + } + } else { + for (int i = threadIdx.x; i < topk_length; i += BLOCK_SIZE) { + occurrences += (topk_ids[i] == expert_id); + } } - atomicAdd(&atomic_buffer[expert_id], occurrences); - __syncthreads(); + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + int final_occurrences = BlockReduce(temp_storage).Sum(occurrences); if (threadIdx.x == 0) { - int final_occurrences = atomic_buffer[expert_id]; problem_sizes1[expert_id * 3] = 2 * n; problem_sizes1[expert_id * 3 + 1] = final_occurrences; problem_sizes1[expert_id * 3 + 2] = k; @@ -34,20 +53,69 @@ __global__ void compute_problem_sizes_w4a8( } } -__global__ void compute_expert_offsets_w4a8( - const int32_t* __restrict__ problem_sizes1, - int32_t* expert_offsets, - int32_t* atomic_buffer, - const int num_experts) { - int32_t tot_offset = 0; - expert_offsets[0] = 0; - for (int i = 0; i < num_experts; ++i) { - atomic_buffer[i] = tot_offset; - tot_offset += problem_sizes1[i * 3 + 1]; - expert_offsets[i + 1] = tot_offset; +template +__device__ void +cumsum_block_scan(const int32_t* __restrict__ input, int32_t* __restrict__ output, int n, int input_stride) { + using BlockScan = cub::BlockScan; + __shared__ typename BlockScan::TempStorage temp_scan_storage; + __shared__ int32_t s_broadcast_val; + + int tid = threadIdx.x; + int32_t base_prefix_sum = 0; + const int num_chunks = (n + BLOCK_SIZE - 1) / BLOCK_SIZE; + + for (int chunk = 0; chunk < num_chunks; chunk++) { + const int base_idx = chunk * BLOCK_SIZE; + const int index = base_idx + tid; + + const int32_t val = (index < n) ? input[index * input_stride] : 0; + int32_t local_prefix_sum; + BlockScan(temp_scan_storage).InclusiveSum(val, local_prefix_sum); + const int32_t prefix_sum = local_prefix_sum + base_prefix_sum; + if (index < n) { + output[index] = prefix_sum; + } + if (chunk < num_chunks - 1) { + if (tid == BLOCK_SIZE - 1) { + s_broadcast_val = prefix_sum; + } + __syncthreads(); + base_prefix_sum = s_broadcast_val; + } } } +template +__global__ void compute_expert_offsets_w4a8_kernel( + const int32_t* __restrict__ problem_sizes1, int32_t* __restrict__ expert_offsets, int n, int stride) { + if (threadIdx.x == 0) { + expert_offsets[0] = 0; + } + cumsum_block_scan(problem_sizes1, expert_offsets + 1, n, stride); +} + +void compute_expert_offsets_w4a8( + cudaStream_t stream, const int32_t* problem_sizes1, int32_t* expert_offsets, int n, int stride = 1, int off = 0) { +#define compute_expert_offsets_w4a8_call(BLOCK_SIZE) \ + compute_expert_offsets_w4a8_kernel \ + <<<1, BLOCK_SIZE, 0, stream>>>(problem_sizes1 + off, expert_offsets, n, stride); + + if (n <= 32) { + compute_expert_offsets_w4a8_call(32); + } else if (n <= 64) { + compute_expert_offsets_w4a8_call(64); + } else if (n <= 128) { + compute_expert_offsets_w4a8_call(128); + } else if (n <= 256) { + compute_expert_offsets_w4a8_call(256); + } else if (n <= 512) { + compute_expert_offsets_w4a8_call(512); + } else { + compute_expert_offsets_w4a8_call(1024); + } +#undef compute_expert_offsets_w4a8_call +} + void get_cutlass_w4a8_moe_mm_data_caller( const torch::Tensor& topk_ids, torch::Tensor& expert_offsets, @@ -60,20 +128,21 @@ void get_cutlass_w4a8_moe_mm_data_caller( const int64_t k) { auto stream = at::cuda::getCurrentCUDAStream(topk_ids.device().index()); auto options_int32 = torch::TensorOptions().dtype(torch::kInt32).device(topk_ids.device()); - torch::Tensor atomic_buffer = torch::zeros(num_experts, options_int32); - int num_threads = min(THREADS_PER_EXPERT, topk_ids.numel()); - compute_problem_sizes_w4a8<<>>( + constexpr uint64_t BLOCK_SIZE = 512; + compute_problem_sizes_w4a8<<>>( static_cast(topk_ids.data_ptr()), static_cast(problem_sizes1.data_ptr()), static_cast(problem_sizes2.data_ptr()), - static_cast(atomic_buffer.data_ptr()), topk_ids.numel(), n, k); - compute_expert_offsets_w4a8<<<1, 1, 0, stream>>>( + + compute_expert_offsets_w4a8( + stream, static_cast(problem_sizes1.data_ptr()), static_cast(expert_offsets.data_ptr()), - static_cast(atomic_buffer.data_ptr()), - num_experts); + num_experts, + 3, + 1); }