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

Speed up per token and per tensor quant by 15% (#4639)

Browse Source
This commit is contained in:
Chunan Zeng
2025-03-22 00:37:57 -07:00
committed by GitHub
parent f69e0696bc
commit 6a384d5c01
2 changed files with 29 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
sgl-kernel/csrc/gemm/per_tensor_quant_fp8.cu
View File

@@ -54,42 +54,40 @@ __global__ void per_tensor_quant_fp8_kernel(
const int grid_size = blockDim.x * gridDim.x;
const float scale_val = 1.0f / (*scale);
constexpr uint32_t vec_size = 16 / sizeof(T);
using vec_t = flashinfer::vec_t<T, vec_size>;
// We want to store 128 bits of data at a time. 16 = 128 / 8 bits
// Load is already vectorized, so 16 elements work for T.
const uint32_t VEC_SIZE = 16;
using vec_t = flashinfer::vec_t<T, VEC_SIZE>;
const int32_t num_vec_elems = num_elements / vec_size;
const int32_t num_vec_elems = num_elements / VEC_SIZE;
for (int32_t i = gid; i < num_vec_elems; i += grid_size) {
vec_t input_vec;
input_vec.cast_load(input + i * vec_size);
input_vec.cast_load(input + i * VEC_SIZE);
FP8_TYPE output_arr[vec_size];
FP8_TYPE output_arr[VEC_SIZE];
#pragma unroll
for (uint32_t j = 0; j < vec_size; ++j) {
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);
#ifndef USE_ROCM
output_arr[j] = static_cast<FP8_TYPE>(val);
#else
output_arr[j] = c10::Float8_e4m3fnuz(
__hip_cvt_float_to_fp8(value, fp8::fp8_type::__default_saturation, fp8::fp8_type::__default_interpret),
__hip_cvt_float_to_fp8(val, fp8::fp8_type::__default_saturation, fp8::fp8_type::__default_interpret),
c10::Float8_e4m3fnuz::from_bits());
#endif
}
#pragma unroll
for (uint32_t j = 0; j < vec_size; ++j) {
output[i * vec_size + j] = output_arr[j];
}
*(uint4*)(output + i * VEC_SIZE) = *(uint4*)output_arr;
}
const int32_t remaining_start = num_vec_elems * vec_size;
const int32_t remaining_start = num_vec_elems * VEC_SIZE;
for (int32_t idx = remaining_start + gid; idx < num_elements; idx += grid_size) {
float val = fmax(-FP8_E4M3_MAX, fmin(static_cast<float>(input[idx]) * scale_val, FP8_E4M3_MAX));
#ifndef USE_ROCM
output[idx] = static_cast<FP8_TYPE>(val);
#else
output[idx] = c10::Float8_e4m3fnuz(
__hip_cvt_float_to_fp8(value, fp8::fp8_type::__default_saturation, fp8::fp8_type::__default_interpret),
__hip_cvt_float_to_fp8(val, fp8::fp8_type::__default_saturation, fp8::fp8_type::__default_interpret),
c10::Float8_e4m3fnuz::from_bits());
#endif
}