EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Remove vllm ops scaled fp8 quant and accelerate per token quant by 20-28% (#4215)

Browse Source 
Co-authored-by: Stefan He <bhe@linkedin.com>
This commit is contained in:
Stefan He
2025-03-12 00:08:03 -07:00
committed by GitHub
parent 7130a7cea9
commit e0917e6bd0
5 changed files with 202 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -14,7 +14,6 @@ __global__ void per_token_quant_fp8_kernel(
const int64_t hidden_dim,
const int64_t num_tokens) {
const int token_idx = blockIdx.x;
if (token_idx >= num_tokens) return;
const int tid = threadIdx.x;
@@ -25,9 +24,20 @@ __global__ void per_token_quant_fp8_kernel(
float max_value = 0.0f;
for (int i = tid; i < hidden_dim; i += block_dim) {
float val = static_cast<float>(token_input[i]);
max_value = fmaxf(max_value, fabsf(val));
constexpr uint32_t vec_size = 16 / sizeof(T);
using vec_t = flashinfer::vec_t<T, vec_size>;
const int32_t num_vec_elems = hidden_dim / vec_size;
// Find max using vectorized loads
for (int32_t i = tid; i < num_vec_elems; i += block_dim) {
vec_t input_vec;
input_vec.cast_load(token_input + i * vec_size);
#pragma unroll
for (uint32_t j = 0; j < vec_size; ++j) {
float val = static_cast<float>(input_vec[j]);
max_value = fmaxf(max_value, fabsf(val));
}
}
max_value = blockReduceMax(max_value);
@@ -41,11 +51,7 @@ __global__ void per_token_quant_fp8_kernel(
const float scale_val = 1.0f / block_max;
constexpr uint32_t vec_size = 16 / sizeof(T);
using vec_t = flashinfer::vec_t<T, vec_size>;
const int32_t num_vec_elems = hidden_dim / vec_size;
// Quantize using vectorized loads
for (int32_t i = tid; i < num_vec_elems; i += block_dim) {
vec_t input_vec;
input_vec.cast_load(token_input + i * vec_size);
@@ -53,7 +59,7 @@ __global__ void per_token_quant_fp8_kernel(
FP8_TYPE output_arr[vec_size];
#pragma unroll
for (uint32_t j = 0; j < vec_size; ++j) {
float val = fmax(fmin(static_cast<float>(input_vec[j]) * scale_val, FP8_E4M3_MAX), -FP8_E4M3_MAX);
float val = fmaxf(fminf(static_cast<float>(input_vec[j]) * scale_val, FP8_E4M3_MAX), -FP8_E4M3_MAX);
#ifndef USE_ROCM
output_arr[j] = static_cast<FP8_TYPE>(val);
#else
@@ -68,18 +74,6 @@ __global__ void per_token_quant_fp8_kernel(
token_output[i * vec_size + j] = output_arr[j];
}
}
const int32_t remaining_start = num_vec_elems * vec_size;
for (int32_t idx = remaining_start + tid; idx < hidden_dim; idx += block_dim) {
float val = fmax(-FP8_E4M3_MAX, fmin(static_cast<float>(token_input[idx]) * scale_val, FP8_E4M3_MAX));
#ifndef USE_ROCM
token_output[idx] = static_cast<FP8_TYPE>(val);
#else
token_output[idx] = c10::Float8_e4m3fnuz(
__hip_cvt_float_to_fp8(val, fp8::fp8_type::__default_saturation, fp8::fp8_type::__default_interpret),
c10::Float8_e4m3fnuz::from_bits());
#endif
}
}
void sgl_per_token_quant_fp8(torch::Tensor input, torch::Tensor output_q, torch::Tensor output_s) {
@@ -91,7 +85,9 @@ void sgl_per_token_quant_fp8(torch::Tensor input, torch::Tensor output_q, torch:
const int64_t num_tokens = input_sizes[0];
const int64_t hidden_dim = input_sizes[1];
const int block_size = 128;
TORCH_CHECK(hidden_dim % 8 == 0, "Hidden dimension must be divisible by 8, but got ", hidden_dim);
const int block_size = 256;
const int num_blocks = num_tokens;
dim3 grid(num_blocks);