chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
DarkSharpness
2026-01-13 17:54:16 +08:00
committed by GitHub
parent 740d3c0b39
commit ba9f6d8f26
30 changed files with 928 additions and 513 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
python/sglang/jit_kernel/include/sgl_kernel/vec.cuh
View File

@@ -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);
}