[Feature] add aligned_vector type for JIT kernel (#16162)

python/sglang/jit_kernel/csrc/hicache.cuh

@@ -12,6 +12,11 @@
 
 namespace device::warp {
 
+template <typename T, std::size_t N>
+struct device_vec {
+  T data[N];
+};
+
 namespace details {
 
 template <std::size_t kUnit>

python/sglang/jit_kernel/csrc/norm.cuh

@@ -2,6 +2,7 @@
 #include <sgl_kernel/tensor.h>
 #include <sgl_kernel/utils.cuh>
 #include <sgl_kernel/utils.h>
+#include <sgl_kernel/vec.cuh>
 #include <sgl_kernel/warp.cuh>
 
 #include <cuda_bf16.h>
@@ -9,6 +10,7 @@
 #include <dlpack/dlpack.h>
 #include <tvm/ffi/container/tensor.h>
 
+#include <cstddef>
 #include <cstdint>
 #include <type_traits>
 
@@ -52,18 +54,17 @@ template <int64_t kHeadDim, typename PackedFloat>
 __always_inline __device__ void apply_norm(void* __restrict__ input, const void* __restrict__ weight, float eps) {
   using namespace device;
 
-  constexpr auto kLoopCount = kHeadDim / (kWarpThreads * 2);
+  constexpr std::size_t kLoopCount = kHeadDim / (kWarpThreads * 2);
   static_assert(kHeadDim % (kWarpThreads * 2) == 0);
 
-  const auto lane_id = threadIdx.x % kWarpThreads;
   float sum_of_squares = 0.0f;
 
-  using vec_t = device_vec<PackedFloat, kLoopCount>;
+  using vec_t = aligned_vector<PackedFloat, kLoopCount>;
-  auto input_vec = static_cast<const vec_t*>(input)[lane_id];
+  const auto input_vec = warp::load<vec_t>(input);
 
 #pragma unroll
   for (auto i = 0u; i < kLoopCount; ++i) {
-    const auto fp16_input = input_vec.data[i];
+    const auto fp16_input = input_vec[i];
     const auto fp32_input = to_float2(fp16_input);
     sum_of_squares += fp32_input.x * fp32_input.x;
     sum_of_squares += fp32_input.y * fp32_input.y;
@@ -71,20 +72,20 @@ __always_inline __device__ void apply_norm(void* __restrict__ input, const void*
 
   sum_of_squares = warp::reduce_sum(sum_of_squares);
   const auto norm_factor = rsqrtf(sum_of_squares / kHeadDim + eps);
-  const auto weight_vec = static_cast<const vec_t*>(weight)[lane_id];
+  const auto weight_vec = warp::load<vec_t>(weight);
-
   vec_t output_vec;
+
 #pragma unroll
   for (auto i = 0u; i < kLoopCount; ++i) {
-    const auto fp32_weight = to_float2(weight_vec.data[i]);
+    const auto fp32_input = to_float2(input_vec[i]);
-    const auto fp32_input = to_float2(input_vec.data[i]);
+    const auto fp32_weight = to_float2(weight_vec[i]);
-    output_vec.data[i] = from_float2<PackedFloat>({
+    output_vec[i] = from_float2<PackedFloat>({
         fp32_input.x * norm_factor * fp32_weight.x,
         fp32_input.y * norm_factor * fp32_weight.y,
     });
   }
 
-  static_cast<vec_t*>(input)[lane_id] = output_vec;
+  warp::store(input, output_vec);
 }
 
 constexpr uint32_t kWarpsPerBlock = 4;

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

@@ -0,0 +1,34 @@
+#pragma once
+#include <version>
+
+/// NOTE: fallback to a minimal source_location implementation
+#if defined(__cpp_lib_source_location)
+#include <source_location>
+
+using source_location_t = std::source_location;
+
+#else
+
+struct source_location_fallback {
+ public:
+  static constexpr source_location_fallback current() noexcept {
+    return source_location_fallback{};
+  }
+  constexpr source_location_fallback() noexcept = default;
+  constexpr unsigned line() const noexcept {
+    return 0;
+  }
+  constexpr unsigned column() const noexcept {
+    return 0;
+  }
+  constexpr const char* file_name() const noexcept {
+    return "";
+  }
+  constexpr const char* function_name() const noexcept {
+    return "";
+  }
+};
+
+using source_location_t = source_location_fallback;
+
+#endif

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

@@ -74,11 +74,6 @@ __always_inline __device__ auto offset(const T* ptr, U... offset) -> const void*
 
 }  // namespace pointer
 
-template <typename T, std::size_t N>
-struct device_vec {
-  T data[N];
-};
-
 template <bool kUsePDL>
 __forceinline__ __device__ void PDLWaitPrimary() {
 #ifndef USE_ROCM

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

@@ -21,16 +21,16 @@
 
 #define consteval constexpr
 
-#include <source_location>
+#include "source_location.h"
 
 #undef consteval
 #pragma pop_macro("__cpp_consteval")
 #pragma pop_macro("_NODISCARD")
 #else  // __CUDACC__ && CUDA_VERSION > 12010
-#include <source_location>
+#include "source_location.h"
 #endif
 #else  // no __CUDACC__
-#include <source_location>
+#include "source_location.h"
 #endif
 
 #include <dlpack/dlpack.h>
@@ -44,8 +44,8 @@
 
 namespace host {
 
-struct DebugInfo : public std::source_location {
+struct DebugInfo : public source_location_t {
-  DebugInfo(std::source_location loc = std::source_location::current()) : std::source_location(loc) {}
+  DebugInfo(source_location_t loc = source_location_t::current()) : source_location_t(loc) {}
 };
 
 struct PanicError : public std::runtime_error {

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

@@ -0,0 +1,89 @@
+#pragma once
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
+#include <cstddef>
+#include <cstdint>
+
+namespace device {
+
+namespace details {
+
+template <std::size_t N>
+struct uint_trait {};
+
+template <>
+struct uint_trait<1> {
+  using type = uint8_t;
+};
+
+template <>
+struct uint_trait<2> {
+  using type = uint16_t;
+};
+
+template <>
+struct uint_trait<4> {
+  using type = uint32_t;
+};
+
+template <>
+struct uint_trait<8> {
+  using type = uint64_t;
+};
+
+template <typename T>
+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 {
+  T data[N];
+};
+
+template <typename T, std::size_t N>
+struct aligned_vector {
+ 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>;
+
+ public:
+  template <typename U>
+  __forceinline__ __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 {
+    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) {
+    const auto store_value = *reinterpret_cast<element_t*>(&value);
+#pragma unroll
+    for (std::size_t i = 0; i < N; ++i) {
+      m_storage.data[i] = store_value;
+    }
+  }
+
+  __forceinline__ __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 {
+    return reinterpret_cast<const T*>(&m_storage)[idx];
+  }
+  __forceinline__ __device__ auto data() -> T* {
+    return reinterpret_cast<T*>(&m_storage);
+  }
+  __forceinline__ __device__ auto data() const -> const T* {
+    return reinterpret_cast<const T*>(&m_storage);
+  }
+
+ private:
+  storage_t m_storage;
+};
+
+}  // namespace device

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

@@ -1,14 +1,33 @@
 #pragma once
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <cstddef>
 
 // Some warp primitives
 namespace device::warp {
 
 template <typename T>
-__always_inline __device__ T reduce_sum(T val, uint32_t active_mask = 0xffffffff) {
+__forceinline__ __device__ T reduce_sum(T val, uint32_t active_mask = 0xffffffff) {
 #pragma unroll
   for (int mask = 16; mask > 0; mask >>= 1)
     val += __shfl_xor_sync(active_mask, val, mask, 32);
   return val;
 }
 
+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;
+}
+
 }  // namespace device::warp