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

[2/N]Support DeepSeek-R1 w4a8 low latency deepep (#8464)

Browse Source 
Co-authored-by: Hank Han <hanhan7630@outlook.com>
Co-authored-by: Shangchuan Huang <2510421000@qq.com>
This commit is contained in:
Jinwu
2025-10-25 08:41:16 +08:00
committed by GitHub
parent e51046beaa
commit 13bf565d60
8 changed files with 531 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

78
sgl-kernel/csrc/moe/cutlass_moe/w4a8/w4a8_get_group_starts.cuh
View File

@@ -34,6 +34,40 @@ __global__ void int4_fp8_get_group_gemm_starts(
b_scales_offsets[expert_id] = b_scales_base_as_int + (per_out_ch ? expert_id * n * k / 128 : expert_id);
}
template <typename ElementA, typename ElementB, typename ElementC, typename ElementAccumulator>
__global__ void int4_fp8_get_group_gemm_starts_3d(
ElementA** a_offsets,
ElementB** b_offsets,
ElementC** out_offsets,
ElementAccumulator** a_scales_offsets,
cutlass::bfloat16_t** b_scales_offsets,
ElementA* a_base_as_int,
ElementB* b_base_as_int,
ElementC* out_base_as_int,
ElementAccumulator* a_scales_base_as_int,
cutlass::bfloat16_t* b_scales_base_as_int,
int64_t n,
int64_t m,
int64_t k,
bool per_act_token,
bool per_out_ch,
int num_experts) {
int expert_id = blockIdx.x * blockDim.x + threadIdx.x;
if (expert_id >= num_experts) return;
int64_t a_offset = expert_id * m * k;
int64_t b_offset = expert_id * k * n / 2;
int64_t out_offset = expert_id * m * n;
int64_t a_scales_offset = 0;
int64_t b_scales_offset = per_out_ch ? expert_id * n * 4 * k / 512 : expert_id;
a_offsets[expert_id] = a_base_as_int + a_offset;
b_offsets[expert_id] = b_base_as_int + b_offset;
out_offsets[expert_id] = out_base_as_int + out_offset;
a_scales_offsets[expert_id] = a_scales_base_as_int + a_scales_offset;
b_scales_offsets[expert_id] = b_scales_base_as_int + b_scales_offset;
}
#define __CALL_W4A8_GET_STARTS_KERNEL(TENSOR_C_TYPE, C_TYPE) \
else if (out_tensors.dtype() == TENSOR_C_TYPE) { \
int4_fp8_get_group_gemm_starts<cutlass::float_e4m3_t, cutlass::int8_t, C_TYPE, float> \
@@ -55,6 +89,28 @@ __global__ void int4_fp8_get_group_gemm_starts(
per_out_ch); \
}
#define __CALL_W4A8_GET_STARTS_KERNEL_3D(TENSOR_C_TYPE, C_TYPE) \
else if (out_tensors.dtype() == TENSOR_C_TYPE) { \
int4_fp8_get_group_gemm_starts_3d<cutlass::float_e4m3_t, cutlass::int8_t, C_TYPE, float> \
<<<1, num_experts, 0, stream>>>( \
static_cast<cutlass::float_e4m3_t**>(a_ptrs.data_ptr()), \
static_cast<cutlass::int8_t**>(b_ptrs.data_ptr()), \
static_cast<C_TYPE**>(out_ptrs.data_ptr()), \
static_cast<float**>(a_scales_ptrs.data_ptr()), \
static_cast<cutlass::bfloat16_t**>(b_scales_ptrs.data_ptr()), \
static_cast<cutlass::float_e4m3_t*>(a_tensors.data_ptr()), \
static_cast<cutlass::int8_t*>(b_tensors.data_ptr()), \
static_cast<C_TYPE*>(out_tensors.data_ptr()), \
static_cast<float*>(a_scales.data_ptr()), \
static_cast<cutlass::bfloat16_t*>(b_scales.data_ptr()), \
out_tensors.size(2), \
a_tensors.size(1), \
a_tensors.size(2), \
per_act_token, \
per_out_ch, \
num_experts); \
}
namespace {
void run_int4_fp8_get_group_gemm_starts(
@@ -80,12 +136,22 @@ void run_int4_fp8_get_group_gemm_starts(
auto stream = at::cuda::getCurrentCUDAStream(expert_offsets.device().index());
if (false) {
}
__CALL_W4A8_GET_STARTS_KERNEL(torch::kBFloat16, cutlass::bfloat16_t)
__CALL_W4A8_GET_STARTS_KERNEL(torch::kFloat16, half)
else {
TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
if (a_tensors.dim() == 3) {
if (false) {
}
__CALL_W4A8_GET_STARTS_KERNEL_3D(torch::kBFloat16, cutlass::bfloat16_t)
__CALL_W4A8_GET_STARTS_KERNEL_3D(torch::kFloat16, half)
else {
TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
}
} else {
if (false) {
}
__CALL_W4A8_GET_STARTS_KERNEL(torch::kBFloat16, cutlass::bfloat16_t)
__CALL_W4A8_GET_STARTS_KERNEL(torch::kFloat16, half)
else {
TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
}
}
}

6
sgl-kernel/csrc/moe/cutlass_moe/w4a8/w4a8_grouped_mm_c3x.cuh
View File

@@ -174,7 +174,7 @@ void cutlass_w4a8_group_gemm_caller(
bool per_out_ch = b_scales.numel() != num_experts;
// Check inputs
TORCH_CHECK(a_tensors.dim() == 2, "A tensor must be 2D");
TORCH_CHECK(a_tensors.dim() == 2 or a_tensors.dim() == 3, "A tensor must be 2D/3D");
TORCH_CHECK(b_tensors.dim() == 3, "B tensor must be 3D [E, N, K/2]");
TORCH_CHECK(b_scales.dim() == 3, "Scale tensor must be 3D [E, K//512, N*4]");
TORCH_CHECK(a_scales.dim() == 1, "A Scale tensor must be 1D [1]");
@@ -186,7 +186,9 @@ void cutlass_w4a8_group_gemm_caller(
TORCH_CHECK(problem_sizes.size(1) == 3, "problem_sizes must have 3 columns (N, M, K)");
TORCH_CHECK(b_tensors.size(0) == num_experts, "B tensor first dimension must match number of groups");
TORCH_CHECK(b_scales.size(0) == num_experts, "Scale tensor first dimension must match number of groups");
TORCH_CHECK(b_tensors.size(2) * 2 == a_tensors.size(1), "B tensor K/2 dimension must match A tensor K dimension");
TORCH_CHECK(
b_tensors.size(2) * 2 == a_tensors.size(1) or b_tensors.size(2) * 2 == a_tensors.size(2),
"B tensor K/2 dimension must match A tensor K dimension");
// Check tensor types
TORCH_CHECK(a_tensors.scalar_type() == torch::kFloat8_e4m3fn, "A tensor must be fp8 (float_e4m3_t) type");