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

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 {
-  DebugInfo(std::source_location loc = std::source_location::current()) : std::source_location(loc) {}
+struct DebugInfo : public source_location_t {
+  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