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

[perf]optimize w4afp8 kernel on deepseek-v3-0324 (#12921)

Browse Source 
Signed-off-by: bruce.xu <bruce.x@gmicloud.ai>
This commit is contained in:
Bruce-x-1997
2025-12-18 18:13:22 +08:00
committed by GitHub
parent d1f0063262
commit 793c96c3d2
3 changed files with 160 additions and 264 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/w4afp8.py
-1
View File
@@ -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(
sgl-kernel/csrc/moe/cutlass_moe/w4a8/w4a8_grouped_mm_c3x.cu
+64 -236
View File
@@ -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<SM90_PP<64, 32, 512, 2, 1, 1>>(
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<SM90_CO<128, 16, 512, 2, 1, 1>>(
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<SM90_CO<128, 16, 512, 1, 1, 1>>(
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<SM90_CO<128, 32, 512, 2, 1, 1>>(
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<SM90_CO<128, 64, 512, 1, 1, 1>>(
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<SM90_PP<64, 16, 512, 1, 1, 1>>(
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<SM90_CO<128, 32, 512, 1, 1, 1>>(
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<SM90_CO<128, 64, 512, 1, 1, 1>>(
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<SM90_PP<64, 32, 512, 2, 1, 1>>(
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<SM90_CO<128, 16, 512, 2, 1, 1>>(
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<SM90_CO<128, 16, 512, 2, 1, 1>>(
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<SM90_CO<128, 32, 512, 2, 1, 1>>(
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<SM90_CO<128, 64, 512, 1, 1, 1>>(
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<SM90_PP<64, 16, 128, 1, 1, 1>>(
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<SM90_PP<128, 32, 128, 2, 1, 1>>(
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<SM90_PP<128, 64, 128, 1, 1, 1>>(
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<SM90_CO<128, 32, 512, 1, 1, 1>>(
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<SM90_PP<128, 64, 128, 1, 1, 1>>(
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>));
}
}
}
}
sgl-kernel/csrc/moe/cutlass_moe/w4a8/w4a8_moe_data.cu
+96 -27
View File
@@ -2,29 +2,48 @@
#include <cudaTypedefs.h>
#include <torch/all.h>
#include <iostream>
constexpr uint64_t THREADS_PER_EXPERT = 512;
#include <cub/block/block_reduce.cuh>
#include <cub/block/block_scan.cuh>
template <int BLOCK_SIZE>
__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<uintptr_t>(topk_ids) % 16 == 0);
if (aligned) {
const int4* vec_ptr = reinterpret_cast<const int4*>(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<int, BLOCK_SIZE>;
__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 <int BLOCK_SIZE>
__device__ void
cumsum_block_scan(const int32_t* __restrict__ input, int32_t* __restrict__ output, int n, int input_stride) {
using BlockScan = cub::BlockScan<int32_t, BLOCK_SIZE>;
__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 <int BLOCK_SIZE>
__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<BLOCK_SIZE>(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<BLOCK_SIZE> \
<<<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<<<num_experts, num_threads, 0, stream>>>(
constexpr uint64_t BLOCK_SIZE = 512;
compute_problem_sizes_w4a8<BLOCK_SIZE><<<num_experts, BLOCK_SIZE, 0, stream>>>(
static_cast<const int32_t*>(topk_ids.data_ptr()),
static_cast<int32_t*>(problem_sizes1.data_ptr()),
static_cast<int32_t*>(problem_sizes2.data_ptr()),
static_cast<int32_t*>(atomic_buffer.data_ptr()),
topk_ids.numel(),
n,
k);
compute_expert_offsets_w4a8<<<1, 1, 0, stream>>>(
compute_expert_offsets_w4a8(
stream,
static_cast<const int32_t*>(problem_sizes1.data_ptr()),
static_cast<int32_t*>(expert_offsets.data_ptr()),
static_cast<int32_t*>(atomic_buffer.data_ptr()),
num_experts);
num_experts,
3,
1);
}