diff --git a/sgl-kernel/benchmark/bench_per_token_group_quant_fp8.py b/sgl-kernel/benchmark/bench_per_token_group_quant_fp8.py index c56df30f552a..244bcc78e2cc 100644 --- a/sgl-kernel/benchmark/bench_per_token_group_quant_fp8.py +++ b/sgl-kernel/benchmark/bench_per_token_group_quant_fp8.py @@ -81,7 +81,11 @@ def triton_per_token_group_quant_fp8( assert x.is_contiguous(), "`x` is not contiguous" finfo = torch.finfo(dtype) - fp8_max = finfo.max + + if is_hip_: + fp8_max = 224 + else: + fp8_max = finfo.max fp8_min = -fp8_max @@ -127,7 +131,11 @@ def sglang_per_token_group_quant_fp8( assert x.is_contiguous(), "`x` is not contiguous" finfo = torch.finfo(dtype) - fp8_max = finfo.max + + if is_hip_: + fp8_max = 224 + else: + fp8_max = finfo.max fp8_min = -fp8_max diff --git a/sgl-kernel/setup_rocm.py b/sgl-kernel/setup_rocm.py index 6530cd7c7437..248c1d331edf 100644 --- a/sgl-kernel/setup_rocm.py +++ b/sgl-kernel/setup_rocm.py @@ -44,6 +44,7 @@ def _get_version(): sources = [ "src/sgl-kernel/torch_extension_rocm.cc", "src/sgl-kernel/csrc/moe_align_kernel.cu", + "src/sgl-kernel/csrc/per_token_group_quant_fp8.cu", ] cxx_flags = ["-O3"] diff --git a/sgl-kernel/src/sgl-kernel/csrc/per_token_group_quant_fp8.cu b/sgl-kernel/src/sgl-kernel/csrc/per_token_group_quant_fp8.cu index 9c466e0e73a9..d78b5d97d4d0 100644 --- a/sgl-kernel/src/sgl-kernel/csrc/per_token_group_quant_fp8.cu +++ b/sgl-kernel/src/sgl-kernel/csrc/per_token_group_quant_fp8.cu @@ -5,13 +5,19 @@ #include "utils.h" +#ifndef USE_ROCM +#define FULL_MASK 0xffff +#else +#define FULL_MASK 0xffffffff +#endif + using FP8_TYPE = c10::Float8_e4m3fn; __device__ __forceinline__ float GroupReduce(float val, const int tid) { - val = fmaxf(val, __shfl_xor_sync(0xffff, val, 8)); - val = fmaxf(val, __shfl_xor_sync(0xffff, val, 4)); - val = fmaxf(val, __shfl_xor_sync(0xffff, val, 2)); - val = fmaxf(val, __shfl_xor_sync(0xffff, val, 1)); + val = fmaxf(val, __shfl_xor_sync(FULL_MASK, val, 8)); + val = fmaxf(val, __shfl_xor_sync(FULL_MASK, val, 4)); + val = fmaxf(val, __shfl_xor_sync(FULL_MASK, val, 2)); + val = fmaxf(val, __shfl_xor_sync(FULL_MASK, val, 1)); return val; } @@ -32,11 +38,13 @@ __global__ void per_token_group_quant_fp8_kernel(const T* __restrict__ input, vo if (block_group_id + local_group_id < num_groups) { const T* group_input = input + (block_group_id + local_group_id) * group_size; + FP8_TYPE* group_output = static_cast(output_q) + (block_group_id + local_group_id) * group_size; float* scale_output = output_s + block_group_id + local_group_id; for (int i = local_tid; i < group_size; i += 16) { - float val = static_cast(group_input[i]); + float val = castToFloat(group_input[i]); + float abs_val = fabsf(val); local_absmax = fmaxf(local_absmax, abs_val); } @@ -56,7 +64,8 @@ __global__ void per_token_group_quant_fp8_kernel(const T* __restrict__ input, vo } for (int i = local_tid; i < group_size; i += 16) { - float val = static_cast(group_input[i]); + float val = castToFloat(group_input[i]); + float q_val = fminf(fmaxf(val / y_s, fp8_min), fp8_max); group_output[i] = FP8_TYPE(q_val); } @@ -73,7 +82,7 @@ void sgl_per_token_group_quant_fp8(torch::Tensor input, torch::Tensor output_q, CHECK_EQ(input.numel() % group_size, 0); - dim3 grid((num_groups + 15) / 16); + dim3 grid(CEILDIV(num_groups, 16)); dim3 block(256); cudaStream_t stream = at::cuda::getCurrentCUDAStream(); diff --git a/sgl-kernel/src/sgl-kernel/include/hip_math_def.h b/sgl-kernel/src/sgl-kernel/include/hip_math_def.h new file mode 100644 index 000000000000..5365e6875d9a --- /dev/null +++ b/sgl-kernel/src/sgl-kernel/include/hip_math_def.h @@ -0,0 +1,58 @@ +#pragma once + +#if defined(__HIP_PLATFORM_AMD__) + +#include +#include +#include +#include + +// Adapted from flashinfer-rocm [PR#491](https://github.com/flashinfer-ai/flashinfer/pull/491) + +namespace amdgpu { + +template +__forceinline__ __device__ T shfl_xor_sync(unsigned mask, T var, int laneMask, int width = warpSize); + +template +__forceinline__ __device__ destDtype cast(srcDtype val); + +// specialization +template <> +__forceinline__ __device__ float shfl_xor_sync(unsigned mask, float var, int laneMask, int width) { + return __shfl_xor(var, laneMask, width); +} + +template <> +__forceinline__ __device__ int shfl_xor_sync(unsigned mask, int var, int laneMask, int width) { + return __shfl_xor(var, laneMask, width); +} + +template <> +__forceinline__ __device__ float cast(float val) { + return val; +} + +template <> +__forceinline__ __device__ float cast(__half val) { + return __half2float(val); +} + +template <> +__forceinline__ __device__ float cast(__hip_bfloat16 val) { + return __bfloat162float(val); +} + +} // namespace amdgpu + +template +__forceinline__ __device__ float __shfl_xor_sync(unsigned mask, T var, int laneMask, int width = warpSize) { + return amdgpu::shfl_xor_sync(var, laneMask, width); +} + +template +__device__ __forceinline__ float castToFloat(srcDtype val) { + return amdgpu::cast(val); +} + +#endif diff --git a/sgl-kernel/src/sgl-kernel/include/utils.h b/sgl-kernel/src/sgl-kernel/include/utils.h index a342dee10140..e98d48d7185e 100644 --- a/sgl-kernel/src/sgl-kernel/include/utils.h +++ b/sgl-kernel/src/sgl-kernel/include/utils.h @@ -18,6 +18,21 @@ limitations under the License. #include #ifndef USE_ROCM #include +#else + +// Adapted from flashinfer-rocm [PR#491](https://github.com/flashinfer-ai/flashinfer/pull/491) +#define _DISPATCH_CASE_F16(c_type, ...) \ + case at::ScalarType::Half: { \ + using c_type = __half; \ + return __VA_ARGS__(); \ + } + +#define _DISPATCH_CASE_BF16(c_type, ...) \ + case at::ScalarType::BFloat16: { \ + using c_type = __hip_bfloat16; \ + return __VA_ARGS__(); \ + } + #endif #include @@ -55,6 +70,8 @@ struct cuda_error : public std::runtime_error { CHECK_IS_CUDA(x); \ CHECK_IS_CONTIGUOUS(x) +#define CHECK_CUDA_EQ(x, y) TORCH_CHECK(x == y, #x " must be" #y) + inline int getSMVersion() { int device{-1}; CHECK_CUDA_SUCCESS(cudaGetDevice(&device)); @@ -65,7 +82,14 @@ inline int getSMVersion() { return sm_major * 10 + sm_minor; } -#ifndef USE_ROCM +#ifdef USE_ROCM + +#define CHECK_INPUT(x) CHECK_CUDA_INPUT(x); + +#define CHECK_EQ(x, y) CHECK_CUDA_EQ(x, y); + +#endif + #define DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(pytorch_dtype, c_type, ...) \ [&]() -> bool { \ switch (pytorch_dtype) { \ @@ -82,7 +106,6 @@ inline int getSMVersion() { return false; \ } \ }() -#endif #define DISPATCH_CASE_INTEGRAL_TYPES(...) \ AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) \ @@ -95,3 +118,22 @@ inline int getSMVersion() { AT_DISPATCH_SWITCH(TYPE, NAME, DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__)) #define CEILDIV(x, y) (((x) + (y)-1) / (y)) + +#ifndef USE_ROCM +#define WARP_SIZE 32 +#else +#define WARP_SIZE warpSize // 64 +#endif + +#if defined(__HIP_PLATFORM_AMD__) + +#include "hip_math_def.h" + +#else + +template +__device__ __forceinline__ float castToFloat(srcDtype val) { + return static_cast(val); +} + +#endif diff --git a/sgl-kernel/src/sgl-kernel/torch_extension_rocm.cc b/sgl-kernel/src/sgl-kernel/torch_extension_rocm.cc index 22f40da10911..913de0fc0399 100644 --- a/sgl-kernel/src/sgl-kernel/torch_extension_rocm.cc +++ b/sgl-kernel/src/sgl-kernel/torch_extension_rocm.cc @@ -24,6 +24,12 @@ TORCH_LIBRARY_EXPAND(sgl_kernels, m) { "moe_align_block_size(Tensor topk_ids, int num_experts, int block_size, Tensor! sorted_token_ids, Tensor! " "experts_ids, Tensor! num_tokens_post_pad, Tensor! token_cnts_buffer, Tensor! cumsum_buffer) -> ()"); m.impl("moe_align_block_size", torch::kCUDA, &moe_align_block_size); + + // per_token_group_quant_fp8 + m.def( + "sgl_per_token_group_quant_fp8(Tensor input, Tensor output_q, Tensor output_s, int group_size," + " float eps, float fp8_min, float fp8_max) -> ()"); + m.impl("sgl_per_token_group_quant_fp8", torch::kCUDA, &sgl_per_token_group_quant_fp8); } REGISTER_EXTENSION(_kernels)