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

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>

python/sglang/jit_kernel/include/sgl_kernel/atomic.cuh

@@ -0,0 +1,23 @@
+#pragma once
+#include <sgl_kernel/utils.cuh>
+
+namespace device::atomic {
+
+SGL_DEVICE float max(float* addr, float value) {
+#ifndef USE_ROCM
+  float old;
+  old = (value >= 0) ? __int_as_float(atomicMax((int*)addr, __float_as_int(value)))
+                     : __uint_as_float(atomicMin((unsigned int*)addr, __float_as_uint(value)));
+  return old;
+#else
+  int* addr_as_i = (int*)addr;
+  int old = *addr_as_i, assumed;
+  do {
+    assumed = old;
+    old = atomicCAS(addr_as_i, assumed, __float_as_int(fmaxf(value, __int_as_float(assumed))));
+  } while (assumed != old);
+  return __int_as_float(old);
+#endif
+}
+
+}  // namespace device::atomic

python/sglang/jit_kernel/include/sgl_kernel/cta.cuh

@@ -0,0 +1,22 @@
+#pragma once
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/warp.cuh>
+
+namespace device::cta {
+
+template <typename T>
+SGL_DEVICE void reduce_max(T value, float* smem, float min_value = 0.0f) {
+  const uint32_t warp_id = threadIdx.x / kWarpThreads;
+  smem[warp_id] = warp::reduce_max(value);
+  __syncthreads();
+  if (warp_id == 0) {
+    const auto tx = threadIdx.x;
+    const auto local_value = tx * kWarpThreads < blockDim.x ? smem[tx] : min_value;
+    const auto max_value = warp::reduce_max(local_value);
+    smem[0] = max_value;
+  }
+  // no extra sync; it is caller's responsibility to sync if needed
+}
+
+}  // namespace device::cta

python/sglang/jit_kernel/include/sgl_kernel/fp8_utils.cuh

@@ -1,13 +0,0 @@
-#pragma once
-
-#ifdef __CUDACC__
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
-#include <cuda_fp8.h>
-#endif
-
-namespace device {
-
-inline constexpr float FP8_E4M3_MAX = 448.0f;
-
-}  // namespace device

python/sglang/jit_kernel/include/sgl_kernel/impl/norm.cuh

@@ -0,0 +1,168 @@
+#pragma once
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <cstdint>
+#include <type_traits>
+
+namespace host::norm {
+
+/**
+ * \brief Check if the given configuration is supported.
+ * \tparam T Element type (only fp16_t/bf16_t is supported)
+ * \tparam kDim Dimension size (usually hidden size)
+ */
+template <typename T, int64_t kDim>
+inline constexpr bool is_config_supported() {
+  if (!std::is_same_v<T, fp16_t> && !std::is_same_v<T, bf16_t>) return false;
+  if (kDim <= 256) {
+    return (kDim == 64 || kDim == 128 || kDim == 256);
+  } else {
+    return (kDim % 256 == 0 && kDim <= 8192);
+  }
+}
+
+/**
+ * \brief Determine whether to use cta norm based on dimension size.
+ * TL;DR: use warp norm for dim <= 256, cta norm otherwise.
+ * \tparam T Element type (fp16_t or bf16_t)
+ * \tparam kDim Dimension size (usually hidden size)
+ * \note This function assumes that the configuration is supported.
+ * \see `is_config_supported`
+ */
+template <typename T, int64_t kDim>
+inline constexpr bool should_use_cta() {
+  static_assert(is_config_supported<T, kDim>(), "Unsupported norm configuration");
+  return kDim > 256;
+}
+
+/**
+ * \brief Get the number of threads per CTA for cta norm.
+ * \tparam T Element type (fp16_t or bf16_t)
+ * \tparam kDim Dimension size (usually hidden size)
+ * \return Number of threads per CTA
+ */
+template <typename T, int64_t kDim>
+inline constexpr uint32_t get_cta_threads() {
+  static_assert(should_use_cta<T, kDim>());
+  return (kDim / 256) * device::kWarpThreads;
+}
+
+}  // namespace host::norm
+
+namespace device::norm {
+
+namespace details {
+
+template <int64_t kDim, bool kUseCTA, typename PackedFloat, std::size_t N>
+SGL_DEVICE AlignedVector<PackedFloat, N> apply_norm_impl(
+    const AlignedVector<PackedFloat, N> input,
+    const AlignedVector<PackedFloat, N> weight,
+    const float eps,
+    [[maybe_unused]] float* smem_buffer,
+    [[maybe_unused]] uint32_t num_warps) {
+  float sum_of_squares = 0.0f;
+
+#pragma unroll
+  for (auto i = 0u; i < N; ++i) {
+    const auto fp32_input = cast<fp32x2_t>(input[i]);
+    sum_of_squares += fp32_input.x * fp32_input.x;
+    sum_of_squares += fp32_input.y * fp32_input.y;
+  }
+
+  sum_of_squares = warp::reduce_sum(sum_of_squares);
+  float norm_factor;
+  if constexpr (kUseCTA) {
+    // need to synchronize across the cta
+    const auto warp_id = threadIdx.x / kWarpThreads;
+    smem_buffer[warp_id] = sum_of_squares;
+    __syncthreads();
+    // use the first warp to reduce
+    if (warp_id == 0) {
+      const auto tx = threadIdx.x;
+      const auto local_sum = tx < num_warps ? smem_buffer[tx] : 0.0f;
+      sum_of_squares = warp::reduce_sum(local_sum);
+      smem_buffer[32] = math::rsqrt(sum_of_squares / kDim + eps);
+    }
+    __syncthreads();
+    norm_factor = smem_buffer[32];
+  } else {
+    norm_factor = math::rsqrt(sum_of_squares / kDim + eps);
+  }
+
+  AlignedVector<PackedFloat, N> output;
+
+#pragma unroll
+  for (auto i = 0u; i < N; ++i) {
+    const auto fp32_input = cast<fp32x2_t>(input[i]);
+    const auto fp32_weight = cast<fp32x2_t>(weight[i]);
+    output[i] = cast<PackedFloat, fp32x2_t>({
+        fp32_input.x * norm_factor * fp32_weight.x,
+        fp32_input.y * norm_factor * fp32_weight.y,
+    });
+  }
+
+  return output;
+}
+
+}  // namespace details
+
+/**
+ * \brief Apply norm using warp-level implementation.
+ * \tparam kDim Dimension size
+ * \tparam T Element type (fp16_t or bf16_t)
+ * \param input Input vector
+ * \param weight Weight vector
+ * \param eps Epsilon value for numerical stability
+ * \return Normalized output vector
+ */
+template <int64_t kDim, typename T>
+SGL_DEVICE T apply_norm_warp(const T& input, const T& weight, float eps) {
+  static_assert(kDim <= 256, "Warp norm only supports dim <= 256");
+  return details::apply_norm_impl<kDim, false>(input, weight, eps, nullptr, 0);
+}
+
+/**
+ * \brief Apply norm using CTA-level implementation.
+ * \tparam kDim Dimension size
+ * \tparam T Element type (fp16_t or bf16_t)
+ * \param input Input vector
+ * \param weight Weight vector
+ * \param eps Epsilon value for numerical stability
+ * \param smem Shared memory buffer
+ * \param num_warps Number of warps in the CTA
+ * \return Normalized output vector
+ */
+template <int64_t kDim, typename T>
+SGL_DEVICE T apply_norm_cta(
+    const T& input, const T& weight, float eps, float* smem, uint32_t num_warps = blockDim.x / kWarpThreads) {
+  static_assert(kDim > 256, "CTA norm only supports dim > 256");
+  return details::apply_norm_impl<kDim, true>(input, weight, eps, smem, num_warps);
+}
+
+/**
+ * \brief Storage type for norm operation.
+ * For warp norm, the storage size depends on kDim.
+ * For cta norm, the storage size is fixed to 16B.
+ * We will also pack the input 16-bit floats into 32-bit types
+ * for faster CUDA core operations.
+ *
+ * \tparam T Element type (fp16_t or bf16_t)
+ * \tparam kDim Dimension size
+ */
+template <typename T, int64_t kDim>
+using StorageType = std::conditional_t<                    // storage type
+    (kDim > 256),                                          // whether to use cta norm
+    AlignedVector<packed_t<T>, 4>,                         // cta norm storage, fixed to 16B
+    AlignedVector<packed_t<T>, kDim / (2 * kWarpThreads)>  // warp norm storage
+    >;
+
+/**
+ * \brief Minimum shared memory size (in bytes) required for cta norm.
+ */
+inline constexpr uint32_t kSmemBufferSize = 33;
+
+}  // namespace device::norm

python/sglang/jit_kernel/include/sgl_kernel/math.cuh

@@ -0,0 +1,36 @@
+#pragma once
+#include <sgl_kernel/type.cuh>
+
+namespace device::math {
+
+inline constexpr float log2e = 1.44269504088896340736f;
+inline constexpr float loge2 = 0.693147180559945309417f;
+inline constexpr float FP8_E4M3_MAX = 448.0f;
+static_assert(log2e * loge2 == 1.0f, "log2e * loge2 must be 1");
+
+template <typename T>
+SGL_DEVICE T max(T a, T b) {
+  return dtype_trait<T>::max(a, b);
+}
+
+template <typename T>
+SGL_DEVICE T min(T a, T b) {
+  return dtype_trait<T>::min(a, b);
+}
+
+template <typename T>
+SGL_DEVICE T abs(T a) {
+  return dtype_trait<T>::abs(a);
+}
+
+template <typename T>
+SGL_DEVICE T sqrt(T a) {
+  return dtype_trait<T>::sqrt(a);
+}
+
+template <typename T>
+SGL_DEVICE T rsqrt(T a) {
+  return dtype_trait<T>::rsqrt(a);
+}
+
+}  // namespace device::math

python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh

@@ -4,6 +4,7 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <cuda_runtime.h>
 
 namespace host::runtime {
 

python/sglang/jit_kernel/include/sgl_kernel/tensor.h

@@ -21,9 +21,7 @@
 #include <utility>
 
 #ifdef __CUDACC__
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
-#include <cuda_fp8.h>
+#include <sgl_kernel/utils.cuh>
 #endif
 
 namespace host {
@@ -41,10 +39,10 @@ struct DTypeRef;
 struct DeviceRef;
 
 template <typename T>
-struct dtype_trait {};
+struct _dtype_trait {};
 
 template <std::integral T>
-struct dtype_trait<T> {
+struct _dtype_trait<T> {
   inline static constexpr DLDataType value = {
       .code = std::is_signed_v<T> ? DLDataTypeCode::kDLInt : DLDataTypeCode::kDLUInt,
       .bits = static_cast<std::uint8_t>(sizeof(T) * 8),
@@ -52,36 +50,36 @@ struct dtype_trait<T> {
 };
 
 template <std::floating_point T>
-struct dtype_trait<T> {
+struct _dtype_trait<T> {
   inline static constexpr DLDataType value = {
       .code = DLDataTypeCode::kDLFloat, .bits = static_cast<std::uint8_t>(sizeof(T) * 8), .lanes = 1};
 };
 
 #ifdef __CUDACC__
 template <>
-struct dtype_trait<__half> {
+struct _dtype_trait<fp16_t> {
   inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLFloat, .bits = 16, .lanes = 1};
 };
 template <>
-struct dtype_trait<__nv_bfloat16> {
+struct _dtype_trait<bf16_t> {
   inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLBfloat, .bits = 16, .lanes = 1};
 };
 template <>
-struct dtype_trait<__nv_fp8_e4m3> {
+struct _dtype_trait<fp8_e4m3_t> {
   inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLFloat8_e4m3fn, .bits = 8, .lanes = 1};
 };
 #endif
 
 template <DLDeviceType Code>
-struct device_trait {
+struct _device_trait {
   inline static constexpr DLDevice value = {.device_type = Code, .device_id = kAnyDeviceID};
 };
 
 template <typename... Ts>
-inline constexpr auto kDTypeList = std::array<DLDataType, sizeof...(Ts)>{dtype_trait<Ts>::value...};
+inline constexpr auto kDTypeList = std::array<DLDataType, sizeof...(Ts)>{_dtype_trait<Ts>::value...};
 
 template <DLDeviceType... Codes>
-inline constexpr auto kDeviceList = std::array<DLDevice, sizeof...(Codes)>{device_trait<Codes>::value...};
+inline constexpr auto kDeviceList = std::array<DLDevice, sizeof...(Codes)>{_device_trait<Codes>::value...};
 
 template <typename T>
 struct PrintAbleSpan {
@@ -155,7 +153,7 @@ inline auto& operator<<(std::ostream& os, PrintAbleSpan<T> span) {
 
 template <typename T>
 inline bool is_type(DLDataType dtype) {
-  return dtype == details::dtype_trait<T>::value;
+  return dtype == details::_dtype_trait<T>::value;
 }
 
 struct SymbolicSize {

python/sglang/jit_kernel/include/sgl_kernel/tile.cuh

@@ -0,0 +1,36 @@
+#pragma once
+#include <sgl_kernel/utils.cuh>
+
+#include <cstdint>
+
+namespace device::tile {
+
+template <typename T>
+struct Memory {
+ public:
+  SGL_DEVICE constexpr Memory(uint32_t tid, uint32_t tsize) : tid(tid), tsize(tsize) {}
+  SGL_DEVICE static constexpr Memory thread() {
+    return Memory{0, 1};
+  }
+  SGL_DEVICE static Memory warp(int warp_threads = kWarpThreads) {
+    return Memory{threadIdx.x % warp_threads, warp_threads};
+  }
+  SGL_DEVICE static Memory cta(int cta_threads = blockDim.x) {
+    return Memory{threadIdx.x, cta_threads};
+  }
+  SGL_DEVICE T load(const void* ptr, int64_t offset = 0) const {
+    return static_cast<const T*>(ptr)[tid + offset * tsize];
+  }
+  SGL_DEVICE void store(void* ptr, T val, int64_t offset = 0) const {
+    static_cast<T*>(ptr)[tid + offset * tsize] = val;
+  }
+  SGL_DEVICE bool in_bound(int64_t element_count, int64_t offset = 0) const {
+    return tid + offset * tsize < element_count;
+  }
+
+ private:
+  uint32_t tid;
+  uint32_t tsize;
+};
+
+}  // namespace device::tile

python/sglang/jit_kernel/include/sgl_kernel/type.cuh

@@ -0,0 +1,72 @@
+#pragma once
+#include <sgl_kernel/utils.cuh>
+
+template <typename T>
+struct dtype_trait {};
+
+#define SGL_REGISTER_DTYPE_TRAIT(TYPE, PACK2, ...)  \
+  template <>                                       \
+  struct dtype_trait<TYPE> {                        \
+    using self_t = TYPE;                            \
+    using packed_t = PACK2;                         \
+    template <typename S>                           \
+    SGL_DEVICE static self_t from(const S& value) { \
+      return static_cast<TYPE>(value);              \
+    }                                               \
+    __VA_ARGS__                                     \
+  }
+
+#define SGL_REGISTER_TYPE_END static_assert(true)
+
+#define SGL_REGISTER_FROM_FUNCTION(FROM, FN)     \
+  SGL_DEVICE static self_t from(const FROM& x) { \
+    return FN(x);                                \
+  }                                              \
+  static_assert(true)
+
+#define SGL_REGISTER_UNARY_FUNCTION(NAME, FN)      \
+  SGL_DEVICE static self_t NAME(const self_t& x) { \
+    return FN(x);                                  \
+  }                                                \
+  static_assert(true)
+
+#define SGL_REGISTER_BINARY_FUNCTION(NAME, FN)                      \
+  SGL_DEVICE static self_t NAME(const self_t& x, const self_t& y) { \
+    return FN(x, y);                                                \
+  }                                                                 \
+  static_assert(true)
+
+SGL_REGISTER_DTYPE_TRAIT(fp32_t, fp32x2_t, SGL_REGISTER_TYPE_END;  //
+                         SGL_REGISTER_UNARY_FUNCTION(abs, fabsf);
+                         SGL_REGISTER_UNARY_FUNCTION(sqrt, sqrtf);
+                         SGL_REGISTER_UNARY_FUNCTION(rsqrt, rsqrtf);
+                         SGL_REGISTER_BINARY_FUNCTION(max, fmaxf);
+                         SGL_REGISTER_BINARY_FUNCTION(min, fminf););
+SGL_REGISTER_DTYPE_TRAIT(fp16_t, fp16x2_t);
+SGL_REGISTER_DTYPE_TRAIT(bf16_t, bf16x2_t);
+
+/// TODO: Add ROCM implementation
+SGL_REGISTER_DTYPE_TRAIT(fp32x2_t, fp32x4_t, SGL_REGISTER_TYPE_END;
+                         SGL_REGISTER_FROM_FUNCTION(fp16x2_t, __half22float2);
+                         SGL_REGISTER_FROM_FUNCTION(bf16x2_t, __bfloat1622float2););
+
+SGL_REGISTER_DTYPE_TRAIT(fp16x2_t, void, SGL_REGISTER_TYPE_END;
+                         SGL_REGISTER_FROM_FUNCTION(fp32x2_t, __float22half2_rn););
+
+SGL_REGISTER_DTYPE_TRAIT(bf16x2_t, void, SGL_REGISTER_TYPE_END;
+                         SGL_REGISTER_FROM_FUNCTION(fp32x2_t, __float22bfloat162_rn););
+
+#undef SGL_REGISTER_DTYPE_TRAIT
+#undef SGL_REGISTER_FROM_FUNCTION
+
+template <typename T>
+using packed_t = typename dtype_trait<T>::packed_t;
+
+namespace device {
+
+template <typename To, typename From>
+SGL_DEVICE To cast(const From& value) {
+  return dtype_trait<To>::from(value);
+}
+
+}  // namespace device

python/sglang/jit_kernel/include/sgl_kernel/utils.cuh

@@ -2,146 +2,77 @@
 
 #include <sgl_kernel/utils.h>
 
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
 #include <dlpack/dlpack.h>
 #include <tvm/ffi/extra/c_env_api.h>
 
 #include <concepts>
 #include <cstddef>
-#include <type_traits>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+
+#ifndef USE_ROCM
+using fp32_t = float;
+using fp16_t = __half;
+using bf16_t = __nv_bfloat16;
+using fp8_e4m3_t = __nv_fp8_e4m3;
+using fp8_e5m2_t = __nv_fp8_e5m2;
+
+using fp32x2_t = float2;
+using fp16x2_t = __half2;
+using bf16x2_t = __nv_bfloat162;
+using fp8x2_e4m3_t = __nv_fp8x2_e4m3;
+using fp8x2_e5m2_t = __nv_fp8x2_e5m2;
+
+using fp32x4_t = float4;
+#endif
 
 namespace device {
 
+#define SGL_DEVICE __forceinline__ __device__
+
 inline constexpr auto kWarpThreads = 32u;
 inline constexpr auto kFullMask = 0xffffffffu;
 
-__device__ __forceinline__ float atomicMaxFloat(float* addr, float value) {
+template <bool kUsePDL>
+SGL_DEVICE void PDLWaitPrimary() {
 #ifndef USE_ROCM
-  float old;
-  old = (value >= 0) ? __int_as_float(atomicMax((int*)addr, __float_as_int(value)))
-                     : __uint_as_float(atomicMin((unsigned int*)addr, __float_as_uint(value)));
-  return old;
-#else
-  int* addr_as_i = (int*)addr;
-  int old = *addr_as_i, assumed;
-  do {
-    assumed = old;
-    old = atomicCAS(addr_as_i, assumed, __float_as_int(fmaxf(value, __int_as_float(assumed))));
-  } while (assumed != old);
-  return __int_as_float(old);
+  if constexpr (kUsePDL) {
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+  }
 #endif
 }
 
-__device__ __forceinline__ float warpReduceMax(float value) {
-  value = fmaxf(value, __shfl_xor_sync(kFullMask, value, 16));
-  value = fmaxf(value, __shfl_xor_sync(kFullMask, value, 8));
-  value = fmaxf(value, __shfl_xor_sync(kFullMask, value, 4));
-  value = fmaxf(value, __shfl_xor_sync(kFullMask, value, 2));
-  value = fmaxf(value, __shfl_xor_sync(kFullMask, value, 1));
-  return value;
-}
-
-__device__ __forceinline__ float blockReduceMax(float value) {
-  static __shared__ float warpLevelMaxs[kWarpThreads];
-  const int laneId = threadIdx.x % kWarpThreads;
-  const int warpId = threadIdx.x / kWarpThreads;
-
-  value = warpReduceMax(value);
-
-  if (laneId == 0) warpLevelMaxs[warpId] = value;
-  __syncthreads();
-
-  value = (threadIdx.x < blockDim.x / kWarpThreads) ? warpLevelMaxs[laneId] : 0;
-  if (warpId == 0) value = warpReduceMax(value);
-
-  return value;
+template <bool kUsePDL>
+SGL_DEVICE void PDLTriggerSecondary() {
+#ifndef USE_ROCM
+  if constexpr (kUsePDL) {
+    asm volatile("griddepcontrol.launch_dependents;" :::);
+  }
+#endif
 }
 
 namespace pointer {
 
 // we only allow void * pointer arithmetic for safety
 
-template <typename T, std::integral... U>
-__always_inline __device__ auto offset(T* ptr, U... offset) -> void* {
-  static_assert(std::is_same_v<T, void>, "Pointer arithmetic is only allowed for void* pointers");
-  return static_cast<char*>(ptr) + (... + offset);
+template <typename T = char, std::integral... U>
+SGL_DEVICE auto offset(void* ptr, U... offset) -> void* {
+  return static_cast<T*>(ptr) + (... + offset);
 }
 
-template <typename T, std::integral... U>
-__always_inline __device__ auto offset(const T* ptr, U... offset) -> const void* {
-  static_assert(std::is_same_v<T, void>, "Pointer arithmetic is only allowed for void* pointers");
-  return static_cast<const char*>(ptr) + (... + offset);
+template <typename T = char, std::integral... U>
+SGL_DEVICE auto offset(const void* ptr, U... offset) -> const void* {
+  return static_cast<const T*>(ptr) + (... + offset);
 }
 
 }  // namespace pointer
 
-template <bool kUsePDL>
-__forceinline__ __device__ void PDLWaitPrimary() {
-#ifndef USE_ROCM
-  if constexpr (kUsePDL) {
-    asm volatile("griddepcontrol.wait;");
-  }
-#endif
-}
-
-template <bool kUsePDL>
-__forceinline__ __device__ void PDLTriggerSecondary() {
-#ifndef USE_ROCM
-  if constexpr (kUsePDL) {
-    asm volatile("griddepcontrol.launch_dependents;");
-  }
-#endif
-}
-
 }  // namespace device
 
 namespace host {
 
-// DType
-template <typename DType, int Pack>
-struct PackedDType {
-  static_assert(dependent_false_v<DType>, "Unsupported dtype for Packed");
-};
-
-template <>
-struct PackedDType<float, 2> {
-  using type = float2;
-};
-
-template <>
-struct PackedDType<float, 4> {
-  using type = float4;
-};
-
-template <>
-struct PackedDType<__half, 2> {
-  using type = __half2;
-};
-
-struct alignas(8) half4 {
-  __half x, y, z, w;
-};
-
-template <>
-struct PackedDType<__half, 4> {
-  using type = half4;
-};
-
-template <>
-struct PackedDType<nv_bfloat16, 2> {
-  using type = nv_bfloat162;
-};
-
-struct alignas(8) bf16_4 {
-  nv_bfloat16 x, y, z, w;
-};
-
-template <>
-struct PackedDType<nv_bfloat16, 4> {
-  using type = bf16_4;
-};
-
 inline void RuntimeDeviceCheck(::cudaError_t error, DebugInfo location = {}) {
   if (error != ::cudaSuccess) {
     [[unlikely]];

python/sglang/jit_kernel/include/sgl_kernel/utils.h

@@ -121,16 +121,14 @@ namespace pointer {
 
 // we only allow void * pointer arithmetic for safety
 
-template <typename T, std::integral... U>
-inline auto offset(T* ptr, U... offset) -> void* {
-  static_assert(std::is_same_v<T, void>, "Pointer arithmetic is only allowed for void* pointers");
-  return static_cast<char*>(ptr) + (... + offset);
+template <typename T = char, std::integral... U>
+inline auto offset(void* ptr, U... offset) -> void* {
+  return static_cast<T*>(ptr) + (... + offset);
 }
 
-template <typename T, std::integral... U>
-inline auto offset(const T* ptr, U... offset) -> const void* {
-  static_assert(std::is_same_v<T, void>, "Pointer arithmetic is only allowed for void* pointers");
-  return static_cast<const char*>(ptr) + (... + offset);
+template <typename T = char, std::integral... U>
+inline auto offset(const void* ptr, U... offset) -> const void* {
+  return static_cast<const T*>(ptr) + (... + offset);
 }
 
 }  // namespace pointer

python/sglang/jit_kernel/include/sgl_kernel/vec.cuh

@@ -1,6 +1,5 @@
 #pragma once
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
+#include <sgl_kernel/utils.cuh>
 
 #include <cstddef>
 #include <cstdint>
@@ -38,30 +37,30 @@ using sized_int = typename uint_trait<sizeof(T)>::type;
 }  // namespace details
 
 template <typename T, std::size_t N>
-struct alignas(sizeof(T) * N) aligned_storage {
+struct alignas(sizeof(T) * N) AlignedStorage {
   T data[N];
 };
 
 template <typename T, std::size_t N>
-struct aligned_vector {
+struct AlignedVector {
  private:
   /// NOTE: 1. must be pow of two 2. 16 * 8 = 128 byte, which is the max vector size supported by most devices
   static_assert((N > 0 && (N & (N - 1)) == 0) && sizeof(T) * N <= 16, "CUDA only support at most 128B vector op");
   using element_t = typename details::sized_int<T>;
-  using storage_t = aligned_storage<element_t, N>;
+  using storage_t = AlignedStorage<element_t, N>;
 
  public:
   template <typename U>
-  __forceinline__ __device__ void load(const U* ptr, std::size_t offset = 0) {
+  SGL_DEVICE void load(const U* ptr, std::size_t offset = 0) {
     static_assert(std::is_same_v<U, T> || std::is_same_v<U, void>);
     m_storage = reinterpret_cast<const storage_t*>(ptr)[offset];
   }
   template <typename U>
-  __forceinline__ __device__ void store(U* ptr, std::size_t offset = 0) const {
+  SGL_DEVICE void store(U* ptr, std::size_t offset = 0) const {
     static_assert(std::is_same_v<U, T> || std::is_same_v<U, void>);
     reinterpret_cast<storage_t*>(ptr)[offset] = m_storage;
   }
-  __forceinline__ __device__ void fill(T value) {
+  SGL_DEVICE void fill(T value) {
     const auto store_value = *reinterpret_cast<element_t*>(&value);
 #pragma unroll
     for (std::size_t i = 0; i < N; ++i) {
@@ -69,16 +68,16 @@ struct aligned_vector {
     }
   }
 
-  __forceinline__ __device__ auto operator[](std::size_t idx) -> T& {
+  SGL_DEVICE auto operator[](std::size_t idx) -> T& {
     return reinterpret_cast<T*>(&m_storage)[idx];
   }
-  __forceinline__ __device__ auto operator[](std::size_t idx) const -> T {
+  SGL_DEVICE auto operator[](std::size_t idx) const -> T {
     return reinterpret_cast<const T*>(&m_storage)[idx];
   }
-  __forceinline__ __device__ auto data() -> T* {
+  SGL_DEVICE auto data() -> T* {
     return reinterpret_cast<T*>(&m_storage);
   }
-  __forceinline__ __device__ auto data() const -> const T* {
+  SGL_DEVICE auto data() const -> const T* {
     return reinterpret_cast<const T*>(&m_storage);
   }
 

python/sglang/jit_kernel/include/sgl_kernel/warp.cuh

@@ -1,33 +1,25 @@
 #pragma once
-#include <sgl_kernel/utils.cuh>
-#include <sgl_kernel/vec.cuh>
-
-#include <cstddef>
+#include <sgl_kernel/math.cuh>
 
 // Some warp primitives
 namespace device::warp {
 
+static constexpr uint32_t kFullMask = 0xffffffffu;
+
 template <typename T>
-__forceinline__ __device__ T reduce_sum(T val, uint32_t active_mask = 0xffffffff) {
+SGL_DEVICE T reduce_sum(T value, uint32_t active_mask = kFullMask) {
 #pragma unroll
   for (int mask = 16; mask > 0; mask >>= 1)
-    val += __shfl_xor_sync(active_mask, val, mask, 32);
-  return val;
+    value = value + __shfl_xor_sync(active_mask, value, mask, 32);
+  return value;
 }
 
-template <typename T, std::size_t kThreads = kWarpThreads>
-__forceinline__ __device__ T load(const void* ptr) {
-  return static_cast<const T*>(ptr)[threadIdx.x % kWarpThreads];
-}
-
-template <std::size_t kThreads = kWarpThreads, typename T>
-__forceinline__ __device__ T load(const T* ptr) {
-  return static_cast<const T*>(ptr)[threadIdx.x % kWarpThreads];
-}
-
-template <std::size_t kThreads = kWarpThreads, typename T>
-__forceinline__ __device__ void store(void* ptr, T val) {
-  static_cast<T*>(ptr)[threadIdx.x % kWarpThreads] = val;
+template <typename T>
+SGL_DEVICE T reduce_max(T value, uint32_t active_mask = kFullMask) {
+#pragma unroll
+  for (int mask = 16; mask > 0; mask >>= 1)
+    value = math::max(value, __shfl_xor_sync(active_mask, value, mask, 32));
+  return value;
 }
 
 }  // namespace device::warp