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

[Refactor] Clean up JIT kernel utilites (#16884)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
This commit is contained in:
DarkSharpness
2026-01-13 17:54:16 +08:00
committed by GitHub
parent 740d3c0b39
commit ba9f6d8f26
30 changed files with 928 additions and 513 deletions
Download Patch File Download Diff File Expand all files Collapse all files

220
python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh
View File

@@ -1,54 +1,62 @@
#include <sgl_kernel/fp8_utils.cuh>
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/atomic.cuh>
#include <sgl_kernel/cta.cuh>
#include <sgl_kernel/math.cuh>
#include <sgl_kernel/tile.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
#include <cstddef>
#include <cstdint>
#include <cub/block/block_reduce.cuh>
#include <flashinfer/vec_dtypes.cuh>
namespace {
using device::atomicMaxFloat;
using device::blockReduceMax;
using device::FP8_E4M3_MAX;
constexpr size_t kBlockSize = 256;
// each warp will handle 512B data
template <typename T>
__global__ void
per_tensor_absmax_kernel(const T* __restrict__ input, float* __restrict__ output_s, const int64_t num_elements) {
using namespace device;
constexpr uint32_t VEC_SIZE = 16 / sizeof(T);
const int64_t gid = blockIdx.x * gridDim.x + threadIdx.x;
float max_value = 0.0f;
unsigned int tid = threadIdx.x;
unsigned int gid = blockIdx.x * blockDim.x + threadIdx.x;
const int grid_size = blockDim.x * gridDim.x;
constexpr uint32_t vec_size = 16 / sizeof(T);
using vec_t = flashinfer::vec_t<T, 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);
if (gid * VEC_SIZE + VEC_SIZE <= num_elements) {
using vec_t = AlignedVector<T, VEC_SIZE>;
const auto gmem_in = tile::Memory<vec_t>::thread();
const auto input_vec = gmem_in.load(input, gid);
#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));
for (uint32_t i = 0; i < VEC_SIZE; ++i) {
const float value = static_cast<float>(input_vec[i]);
max_value = math::max(max_value, math::abs(value));
}
} else if (gid * VEC_SIZE < num_elements) {
[[unlikely]]; // poorly aligned case, do not optimize
const auto remainder = num_elements - gid * VEC_SIZE;
for (uint32_t i = 0; i < remainder; ++i) {
const float value = static_cast<float>(input[gid * VEC_SIZE + i]);
max_value = math::max(max_value, math::abs(value));
}
}
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 = static_cast<float>(input[idx]);
max_value = fmaxf(max_value, fabsf(val));
// reduce within block and then atomic reduce between blocks
__shared__ float smem[kWarpThreads];
cta::reduce_max(max_value, smem);
if (threadIdx.x == 0) {
const auto max_value = smem[0];
atomic::max(output_s, max_value / math::FP8_E4M3_MAX);
}
}
max_value = blockReduceMax(max_value);
if (tid == 0) {
atomicMaxFloat(output_s, max_value / FP8_E4M3_MAX);
}
[[maybe_unused]]
SGL_DEVICE float fp8_e4m3_clip(float val) {
namespace math = device::math;
return math::max(math::min(val, math::FP8_E4M3_MAX), -math::FP8_E4M3_MAX);
}
template <typename T, typename DST_DTYPE>
@@ -57,123 +65,75 @@ __global__ void per_tensor_quant_fp8_kernel(
DST_DTYPE* __restrict__ output,
const float* __restrict__ scale,
const int64_t num_elements) {
const int gid = blockIdx.x * blockDim.x + threadIdx.x;
const int grid_size = blockDim.x * gridDim.x;
using namespace device;
constexpr uint32_t VEC_SIZE = 16 / sizeof(T);
const int64_t gid = blockIdx.x * blockDim.x + threadIdx.x;
const float scale_val = 1.0f / (*scale);
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;
for (int32_t i = gid; i < num_vec_elems; i += grid_size) {
vec_t input_vec;
input_vec.cast_load(input + i * VEC_SIZE);
DST_DTYPE output_arr[VEC_SIZE];
if (gid * VEC_SIZE + VEC_SIZE <= num_elements) {
using input_vec_t = AlignedVector<T, VEC_SIZE>;
using output_vec_t = AlignedVector<DST_DTYPE, VEC_SIZE>;
const auto gmem_in = tile::Memory<input_vec_t>::thread();
const auto gmem_out = tile::Memory<output_vec_t>::thread();
const auto input_vec = gmem_in.load(input, gid);
output_vec_t output_vec;
#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);
#if !defined(USE_ROCM) || defined(HIP_FP8_TYPE_E4M3)
output_arr[j] = static_cast<DST_DTYPE>(val);
#else
output_arr[j] = 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
for (uint32_t i = 0; i < VEC_SIZE; ++i) {
const float value = fp8_e4m3_clip(static_cast<float>(input_vec[i]) * scale_val);
output_vec[i] = static_cast<DST_DTYPE>(value);
}
gmem_out.store(output, output_vec, gid);
} else if (gid * VEC_SIZE < num_elements) {
[[unlikely]]; // poorly aligned case, do not optimize
const auto remainder = num_elements - gid * VEC_SIZE;
for (uint32_t i = 0; i < remainder; ++i) {
const float value = fp8_e4m3_clip(static_cast<float>(input[gid * VEC_SIZE + i]) * scale_val);
output[gid * VEC_SIZE + i] = static_cast<DST_DTYPE>(value);
}
*(uint4*)(output + i * VEC_SIZE) = *(uint4*)output_arr;
}
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));
#if !defined(USE_ROCM) || defined(HIP_FP8_TYPE_E4M3)
output[idx] = static_cast<DST_DTYPE>(val);
#else
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
}
}
constexpr size_t kBlockSize = 256;
template <bool kIsStatic>
template <bool kIsStatic, typename DType>
void per_tensor_quant_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, tvm::ffi::TensorView output_s) {
using namespace host;
const DLDevice device = input.device();
RuntimeCheck(device.device_type == kDLCUDA, "input must be on CUDA");
RuntimeCheck(input.is_contiguous(), "input must be contiguous");
const int64_t ndim = input.dim();
RuntimeCheck(ndim >= 1, "input.ndim must be >= 1, but got ", ndim);
RuntimeCheck(output_q.device() == device, "output_q must be on the same device as input");
RuntimeCheck(output_q.is_contiguous(), "output_q must be contiguous");
RuntimeCheck(output_q.dim() == ndim, "output_q.ndim must match input.ndim");
for (int64_t i = 0; i < ndim; ++i) {
RuntimeCheck(
output_q.size(i) == input.size(i),
"output_q.shape mismatch at dim ",
i,
": expected ",
input.size(i),
" but got ",
output_q.size(i));
}
auto device = SymbolicDevice{};
auto N = SymbolicSize{"num_elements"};
device.set_options<kDLCUDA>();
TensorMatcher({N}) //
.with_dtype<DType>()
.with_device(device)
.verify(input);
TensorMatcher({N}) //
.with_dtype<fp8_e4m3_t>()
.with_device(device)
.verify(output_q);
TensorMatcher({1}) //
.with_dtype<float>()
.with_device<kDLCUDA>()
.with_device(device)
.verify(output_s);
RuntimeCheck(output_s.device() == device, "output_s must be on the same device as input");
const DLDataType in_dtype = input.dtype();
const bool in_ok = (in_dtype.code == kDLFloat && in_dtype.bits == 32) ||
(in_dtype.code == kDLFloat && in_dtype.bits == 16) ||
(in_dtype.code == kDLBfloat && in_dtype.bits == 16);
RuntimeCheck(in_ok, "input dtype must be fp32/fp16/bf16, but got ", in_dtype);
const auto num_elements = N.unwrap();
const DLDataType out_dtype = output_q.dtype();
RuntimeCheck(
out_dtype.code == kDLFloat8_e4m3fn && out_dtype.bits == 8,
"output_q dtype must be fp8_e4m3fn, but got ",
out_dtype);
constexpr size_t kElementsPerBlock = kBlockSize * (16 / sizeof(DType));
const uint32_t num_blocks = div_ceil(num_elements, kElementsPerBlock);
size_t total_elements = 1;
for (const auto s : input.shape()) {
RuntimeCheck(s > 0, "Input tensor must be non-empty");
total_elements *= static_cast<size_t>(s);
if constexpr (!kIsStatic) {
LaunchKernel(num_blocks, kBlockSize, device.unwrap())(
per_tensor_absmax_kernel<DType>,
static_cast<const DType*>(input.data_ptr()),
static_cast<float*>(output_s.data_ptr()),
static_cast<int64_t>(num_elements));
}
const size_t num_blocks = std::min((total_elements + kBlockSize - 1) / kBlockSize, size_t(1024));
auto launch_kernels = [&]<typename T>() {
if constexpr (!kIsStatic) {
LaunchKernel(num_blocks, kBlockSize, device)(
per_tensor_absmax_kernel<T>,
static_cast<const T*>(input.data_ptr()),
static_cast<float*>(output_s.data_ptr()),
static_cast<int64_t>(total_elements));
}
LaunchKernel(num_blocks, kBlockSize, device)(
per_tensor_quant_fp8_kernel<T, __nv_fp8_e4m3>,
static_cast<const T*>(input.data_ptr()),
static_cast<__nv_fp8_e4m3*>(output_q.data_ptr()),
static_cast<const float*>(output_s.data_ptr()),
static_cast<int64_t>(total_elements));
};
if (in_dtype.code == kDLFloat && in_dtype.bits == 32) {
launch_kernels.template operator()<float>();
} else if (in_dtype.code == kDLBfloat && in_dtype.bits == 16) {
launch_kernels.template operator()<__nv_bfloat16>();
} else if (in_dtype.code == kDLFloat && in_dtype.bits == 16) {
launch_kernels.template operator()<__half>();
}
LaunchKernel(num_blocks, kBlockSize, device.unwrap())(
per_tensor_quant_fp8_kernel<DType, fp8_e4m3_t>,
static_cast<const DType*>(input.data_ptr()),
static_cast<fp8_e4m3_t*>(output_q.data_ptr()),
static_cast<const float*>(output_s.data_ptr()),
static_cast<int64_t>(num_elements));
}
} // namespace