[ROCm] Facilitate shfl_xor usage in AMD platform

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
 

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<FP8_TYPE*>(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<float>(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<float>(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();

sgl-kernel/src/sgl-kernel/include/hip_math_def.h

@@ -0,0 +1,58 @@
+#pragma once
+
+#if defined(__HIP_PLATFORM_AMD__)
+
+#include <hip/hip_bf16.h>
+#include <hip/hip_common.h>
+#include <hip/hip_fp16.h>
+#include <hip/hip_fp8.h>
+
+// this is adoption from flashinfer-rocm [PR#491](https://github.com/flashinfer-ai/flashinfer/pull/491)
+
+namespace amdgpu {
+
+template <typename T>
+__forceinline__ __device__ T shfl_xor_sync(unsigned mask, T var, int laneMask, int width = warpSize);
+
+template <typename srcDtype, typename destDtype>
+__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 <typename T>
+__forceinline__ __device__ float __shfl_xor_sync(unsigned mask, T var, int laneMask, int width = warpSize) {
+  return amdgpu::shfl_xor_sync(var, laneMask, width);
+}
+
+template <typename srcDtype>
+__device__ __forceinline__ float castToFloat(srcDtype val) {
+  return amdgpu::cast<srcDtype, float>(val);
+}
+
+#endif

sgl-kernel/src/sgl-kernel/include/utils.h

@@ -18,6 +18,21 @@ limitations under the License.
 #include <cuda_runtime.h>
 #ifndef USE_ROCM
 #include <pytorch_extension_utils.h>
+#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 <torch/extension.h>
 
@@ -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

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)

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)