From 595d18051de85d8900290a73e63f3992a48987a2 Mon Sep 17 00:00:00 2001 From: yiakwy-xpu-ml-framework-team <961186938@qq.com> Date: Tue, 18 Feb 2025 17:21:49 +0800 Subject: [PATCH 1/5] facilitate compiling in AMD platform --- .../bench_per_token_group_quant_fp8.py | 12 +++- .../csrc/per_token_group_quant_fp8.cu | 23 +++++--- .../src/sgl-kernel/include/hip_math_def.h | 58 +++++++++++++++++++ sgl-kernel/src/sgl-kernel/include/utils.h | 43 +++++++++++++- sgl-kernel/src/sgl-kernel/ops/__init__.py | 8 ++- .../src/sgl-kernel/torch_extension_rocm.cc | 6 ++ 6 files changed, 136 insertions(+), 14 deletions(-) create mode 100644 sgl-kernel/src/sgl-kernel/include/hip_math_def.h 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/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..6fd3751d1b1e 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 = static_cast(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 = static_cast(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..22d582d3a113 --- /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 + +// this is adoption 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..ec83fb41893b 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 + +// adding dispatch utils adpated from flashinfer +#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,18 @@ 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 +110,6 @@ inline int getSMVersion() { return false; \ } \ }() -#endif #define DISPATCH_CASE_INTEGRAL_TYPES(...) \ AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) \ @@ -95,3 +122,15 @@ 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" + +#endif diff --git a/sgl-kernel/src/sgl-kernel/ops/__init__.py b/sgl-kernel/src/sgl-kernel/ops/__init__.py index 0a3ea9aca928..51ce13bf4721 100644 --- a/sgl-kernel/src/sgl-kernel/ops/__init__.py +++ b/sgl-kernel/src/sgl-kernel/ops/__init__.py @@ -600,6 +600,8 @@ def sgl_per_token_group_quant_fp8( fp8_min: float, fp8_max: float, ) -> None: - torch.ops.sgl_kernels.sgl_per_token_group_quant_fp8( - input, output_q, output_s, group_size, eps, fp8_min, fp8_max - ) + if hasattr(torch.ops.sgl_kernels, "sgl_per_token_group_quant_fp8"): + _impl = torch.ops.sgl_kernels.sgl_per_token_group_quant_fp8 + else: + _impl = sgl_kernel.ops._kernels.sgl_per_token_group_quant_fp8 + _impl(input, output_q, output_s, group_size, eps, fp8_min, fp8_max) 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) From d996df5132f4dd00fee660fc14189ff2b3a14f4b Mon Sep 17 00:00:00 2001 From: yiakwy-xpu-ml-framework-team <961186938@qq.com> Date: Tue, 18 Feb 2025 17:40:24 +0800 Subject: [PATCH 2/5] remove redundant static_cast --- sgl-kernel/src/sgl-kernel/csrc/per_token_group_quant_fp8.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 6fd3751d1b1e..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 @@ -43,7 +43,7 @@ __global__ void per_token_group_quant_fp8_kernel(const T* __restrict__ input, vo 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(castToFloat(group_input[i])); + float val = castToFloat(group_input[i]); float abs_val = fabsf(val); local_absmax = fmaxf(local_absmax, abs_val); @@ -64,7 +64,7 @@ __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(castToFloat(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); From e05389853840a074776b771ba56a5cf0fdef109f Mon Sep 17 00:00:00 2001 From: yiakwy-xpu-ml-framework-team <961186938@qq.com> Date: Tue, 18 Feb 2025 17:57:34 +0800 Subject: [PATCH 3/5] add castToFloat symbol for CUDA platform --- .../src/sgl-kernel/include/hip_math_def.h | 2 +- sgl-kernel/src/sgl-kernel/include/utils.h | 17 ++++++++++------- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/sgl-kernel/src/sgl-kernel/include/hip_math_def.h b/sgl-kernel/src/sgl-kernel/include/hip_math_def.h index 22d582d3a113..5365e6875d9a 100644 --- a/sgl-kernel/src/sgl-kernel/include/hip_math_def.h +++ b/sgl-kernel/src/sgl-kernel/include/hip_math_def.h @@ -7,7 +7,7 @@ #include #include -// this is adoption from flashinfer-rocm [PR#491](https://github.com/flashinfer-ai/flashinfer/pull/491) +// Adapted from flashinfer-rocm [PR#491](https://github.com/flashinfer-ai/flashinfer/pull/491) namespace amdgpu { diff --git a/sgl-kernel/src/sgl-kernel/include/utils.h b/sgl-kernel/src/sgl-kernel/include/utils.h index ec83fb41893b..e98d48d7185e 100644 --- a/sgl-kernel/src/sgl-kernel/include/utils.h +++ b/sgl-kernel/src/sgl-kernel/include/utils.h @@ -20,7 +20,7 @@ limitations under the License. #include #else -// adding dispatch utils adpated from flashinfer +// 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; \ @@ -84,13 +84,9 @@ inline int getSMVersion() { #ifdef USE_ROCM -#define CHECK_INPUT(x) \ - \ - CHECK_CUDA_INPUT(x); +#define CHECK_INPUT(x) CHECK_CUDA_INPUT(x); -#define CHECK_EQ(x, y) \ - \ - CHECK_CUDA_EQ(x, y); +#define CHECK_EQ(x, y) CHECK_CUDA_EQ(x, y); #endif @@ -133,4 +129,11 @@ inline int getSMVersion() { #include "hip_math_def.h" +#else + +template +__device__ __forceinline__ float castToFloat(srcDtype val) { + return static_cast(val); +} + #endif From 1d75eaf0189dfa8b0ecbae6daa480587214b8a69 Mon Sep 17 00:00:00 2001 From: yiakwy-xpu-ml-framework-team <961186938@qq.com> Date: Tue, 18 Feb 2025 18:06:36 +0800 Subject: [PATCH 4/5] remove guard --- sgl-kernel/src/sgl-kernel/ops/__init__.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/sgl-kernel/src/sgl-kernel/ops/__init__.py b/sgl-kernel/src/sgl-kernel/ops/__init__.py index 51ce13bf4721..0a3ea9aca928 100644 --- a/sgl-kernel/src/sgl-kernel/ops/__init__.py +++ b/sgl-kernel/src/sgl-kernel/ops/__init__.py @@ -600,8 +600,6 @@ def sgl_per_token_group_quant_fp8( fp8_min: float, fp8_max: float, ) -> None: - if hasattr(torch.ops.sgl_kernels, "sgl_per_token_group_quant_fp8"): - _impl = torch.ops.sgl_kernels.sgl_per_token_group_quant_fp8 - else: - _impl = sgl_kernel.ops._kernels.sgl_per_token_group_quant_fp8 - _impl(input, output_q, output_s, group_size, eps, fp8_min, fp8_max) + torch.ops.sgl_kernels.sgl_per_token_group_quant_fp8( + input, output_q, output_s, group_size, eps, fp8_min, fp8_max + ) From f5514fcedf1b6730ab52a2e68a84437140f66c96 Mon Sep 17 00:00:00 2001 From: yiakwy-xpu-ml-framework-team <961186938@qq.com> Date: Wed, 19 Feb 2025 08:31:22 +0800 Subject: [PATCH 5/5] update setup_rocm.py --- sgl-kernel/setup_rocm.py | 1 + 1 file changed, 1 insertion(+) 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"]