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

Support hidden_dim % 4 == 0 in per_token_quant_fp8 (#12883)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2025-11-10 17:13:14 +08:00
committed by GitHub
parent 7bffc5dc25
commit 05559a4a90
2 changed files with 19 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
sgl-kernel/csrc/gemm/per_token_quant_fp8.cu
View File

@@ -170,13 +170,14 @@ void sgl_per_token_quant_fp8(torch::Tensor input, torch::Tensor output_q, torch:
const auto input_sizes = input.sizes();
const int64_t num_tokens = input_sizes[0];
const int64_t hidden_dim = input_sizes[1];
TORCH_CHECK(hidden_dim % 8 == 0, "Hidden dimension must be divisible by 8, but got ", hidden_dim);
TORCH_CHECK(hidden_dim % 4 == 0, "Hidden dimension must be divisible by 4, but got ", hidden_dim);
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const int sm_count = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
const int TOKENS_PER_CTA = 8;
const bool use_warp_kernel = (num_tokens >= sm_count * 2 * TOKENS_PER_CTA);
const bool use_vec16 = (hidden_dim % 16 == 0);
const bool use_vec8 = (hidden_dim % 8 == 0);
DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(input.scalar_type(), scalar_t, [&] {
if (use_warp_kernel) {
@@ -192,13 +193,20 @@ void sgl_per_token_quant_fp8(torch::Tensor input, torch::Tensor output_q, torch:
static_cast<float*>(output_s.data_ptr()),
hidden_dim,
num_tokens);
} else {
} else if (use_vec8) {
per_token_quant_fp8_kernel<scalar_t, __nv_fp8_e4m3, TOKENS_PER_CTA, 8><<<grid, block, 0, stream>>>(
static_cast<const scalar_t*>(input.data_ptr()),
static_cast<__nv_fp8_e4m3*>(output_q.data_ptr()),
static_cast<float*>(output_s.data_ptr()),
hidden_dim,
num_tokens);
} else {
per_token_quant_fp8_kernel<scalar_t, __nv_fp8_e4m3, TOKENS_PER_CTA, 4><<<grid, block, 0, stream>>>(
static_cast<const scalar_t*>(input.data_ptr()),
static_cast<__nv_fp8_e4m3*>(output_q.data_ptr()),
static_cast<float*>(output_s.data_ptr()),
hidden_dim,
num_tokens);
}
} else {
// -------- baseline -----------------------------------------------------
@@ -213,13 +221,20 @@ void sgl_per_token_quant_fp8(torch::Tensor input, torch::Tensor output_q, torch:
static_cast<float*>(output_s.data_ptr()),
hidden_dim,
num_tokens);
} else {
} else if (use_vec8) {
per_token_quant_fp8_small_batch_kernel<scalar_t, __nv_fp8_e4m3, 8><<<grid, block, 0, stream>>>(
static_cast<const scalar_t*>(input.data_ptr()),
static_cast<__nv_fp8_e4m3*>(output_q.data_ptr()),
static_cast<float*>(output_s.data_ptr()),
hidden_dim,
num_tokens);
} else {
per_token_quant_fp8_small_batch_kernel<scalar_t, __nv_fp8_e4m3, 4><<<grid, block, 0, stream>>>(
static_cast<const scalar_t*>(input.data_ptr()),
static_cast<__nv_fp8_e4m3*>(output_q.data_ptr()),
static_cast<float*>(output_s.data_ptr()),
hidden_dim,
num_tokens);
}
}
return true;

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

@@ -36,7 +36,7 @@ def sglang_per_token_quant_fp8(
@pytest.mark.parametrize(
"num_tokens,hidden_dim",
list(itertools.product([128, 256, 512], [512, 1368, 2048, 4096])),
list(itertools.product([128, 256, 512], [512, 1076, 1368, 2048, 4096])),
)
def test_per_token_quant_compare_implementations(
num_tokens: int,