diff --git a/docs/developer_guide/development_jit_kernel_guide.md b/docs/developer_guide/development_jit_kernel_guide.md index 2fb342274..ce0e9b8aa 100644 --- a/docs/developer_guide/development_jit_kernel_guide.md +++ b/docs/developer_guide/development_jit_kernel_guide.md @@ -257,3 +257,53 @@ from sglang.jit_kernel.add_constant import add_constant ``` For a complete, runnable example, refer to [test_add_constant.py](../../python/sglang/jit_kernel/tests/test_add_constant.py). + +## C++ Include Library Reference + +The JIT kernel framework provides a set of reusable C++ headers in +`python/sglang/jit_kernel/include/sgl_kernel/`. Each header is designed +to be lightweight and self-contained. Below is a summary of each header +and its key APIs. + +### Core Utilities + +| Header | Namespace | Purpose | +|--------|-----------|---------| +| `utils.h` | `host` | Host-side essentials: `RuntimeCheck`, `Panic`, `div_ceil`, `irange` | +| `utils.cuh` | `device` / `host` | Type aliases (`fp16_t`, `bf16_t`, ...), `SGL_DEVICE` macro, PDL helpers, `LaunchKernel`, `RuntimeDeviceCheck` | +| `source_location.h` | (global) | Portable `std::source_location` wrapper for error reporting | +| `runtime.cuh` | `host::runtime` | CUDA runtime queries: `get_blocks_per_sm`, `get_sm_count`, `get_cc_major`, `get_runtime_version`, `get_available_dynamic_smem_per_block` | + +### Tensor Validation + +| Header | Namespace | Purpose | +|--------|-----------|---------| +| `tensor.h` | `host` | `TensorMatcher`, `SymbolicSize`, `SymbolicDType`, `SymbolicDevice` | + +### Math & Type System + +| Header | Namespace | Purpose | +|--------|-----------|---------| +| `math.cuh` | `device::math` | `max`, `min`, `abs`, `sqrt`, `rsqrt`, `exp`, `sin`, `cos`, constants | +| `type.cuh` | (global) / `device` | `dtype_trait`, `packed_t`, `device::cast(from)` | + +### Memory Access + +| Header | Namespace | Purpose | +|--------|-----------|---------| +| `vec.cuh` | `device` | `AlignedVector` - vectorized load/store (up to 128-bit; 256-bit requires Blackwell GPUs) | +| `tile.cuh` | `device::tile` | `Memory` - cooperative tiled memory I/O (thread/warp/CTA) | + +### Parallel Primitives + +| Header | Namespace | Purpose | +|--------|-----------|---------| +| `warp.cuh` | `device::warp` | `reduce_sum`, `reduce_max` via `__shfl_xor_sync` | +| `cta.cuh` | `device::cta` | `reduce_max` across warps via shared memory | +| `atomic.cuh` | `device::atomic` | `max` - atomic float max (CUDA + ROCm fallback) | + +### Reusable Kernel Templates + +| Header | Namespace | Purpose | +|--------|-----------|---------| +| `impl/norm.cuh` | `host::norm` / `device::norm` | RMSNorm building blocks (warp & CTA paths, `StorageType`) | diff --git a/python/sglang/jit_kernel/include/sgl_kernel/atomic.cuh b/python/sglang/jit_kernel/include/sgl_kernel/atomic.cuh index 574f79b8d..c9da765f4 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/atomic.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/atomic.cuh @@ -1,8 +1,20 @@ +/// \file atomic.cuh +/// \brief Device-side atomic operations. + #pragma once #include namespace device::atomic { +/** + * \brief Atomically computes the maximum of `*addr` and `value`, storing the + * result in `*addr`. + * \param addr Pointer to the value in global/shared memory to be updated. + * \param value The value to compare against. + * \return The old value at `*addr` before the update. + * \note On CUDA, this uses `atomicMax`/`atomicMin` on the reinterpreted + * integer representation. On ROCm, a CAS loop is used as a fallback. + */ SGL_DEVICE float max(float* addr, float value) { #ifndef USE_ROCM float old; diff --git a/python/sglang/jit_kernel/include/sgl_kernel/cta.cuh b/python/sglang/jit_kernel/include/sgl_kernel/cta.cuh index 28db34f02..b47a4a27b 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/cta.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/cta.cuh @@ -1,3 +1,6 @@ +/// \file cta.cuh +/// \brief CTA (Cooperative Thread Array / thread-block) level primitives. + #pragma once #include #include @@ -5,6 +8,21 @@ namespace device::cta { +/** + * \brief Compute the maximum of `value` across all threads in the CTA. + * + * Uses a two-level reduction: first within each warp via `warp::reduce_max`, + * then across warps using shared memory. The final result is stored in + * `smem[0]`. + * + * \tparam T Numeric type (must be supported by `warp::reduce_max`). + * \param value Per-thread input value. + * \param smem Shared memory buffer (must have at least `blockDim.x / 32` + * elements). + * \param min_value Identity element for max (default 0.0f). + * \note This function does NOT issue a trailing `__syncthreads()`. + * Callers must synchronize before reading `smem[0]`. + */ template SGL_DEVICE void reduce_max(T value, float* smem, float min_value = 0.0f) { const uint32_t warp_id = threadIdx.x / kWarpThreads; diff --git a/python/sglang/jit_kernel/include/sgl_kernel/math.cuh b/python/sglang/jit_kernel/include/sgl_kernel/math.cuh index 2287b31e2..4f9ac4814 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/math.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/math.cuh @@ -1,48 +1,68 @@ +/// \file math.cuh +/// \brief Device-side math helper functions and constants. +/// +/// Provides type-generic wrappers around CUDA math intrinsics by +/// dispatching through `dtype_trait`. All functions are forced-inline +/// device functions. + #pragma once #include +#include + namespace device::math { +/// \brief Constant: log2(e) inline constexpr float log2e = 1.44269504088896340736f; +/// \brief Constant: ln(2) inline constexpr float loge2 = 0.693147180559945309417f; +/// \brief Maximum representable value for FP8 E4M3 format. inline constexpr float FP8_E4M3_MAX = 448.0f; static_assert(log2e * loge2 == 1.0f, "log2e * loge2 must be 1"); +/// \brief Returns the larger of `a` and `b`. template SGL_DEVICE T max(T a, T b) { return dtype_trait::max(a, b); } +/// \brief Returns the smaller of `a` and `b`. template SGL_DEVICE T min(T a, T b) { return dtype_trait::min(a, b); } +/// \brief Returns the absolute value of `a`. template SGL_DEVICE T abs(T a) { return dtype_trait::abs(a); } +/// \brief Returns the square root of `a`. template SGL_DEVICE T sqrt(T a) { return dtype_trait::sqrt(a); } +/// \brief Returns the reciprocal square root of `a` (i.e. 1 / sqrt(a)). template SGL_DEVICE T rsqrt(T a) { return dtype_trait::rsqrt(a); } +/// \brief Returns e^a. template SGL_DEVICE T exp(T a) { return dtype_trait::exp(a); } +/// \brief Returns sin(a). template SGL_DEVICE T sin(T a) { return dtype_trait::sin(a); } +/// \brief Returns cos(a). template SGL_DEVICE T cos(T a) { return dtype_trait::cos(a); diff --git a/python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh b/python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh index bf582999b..2812a2f8e 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh @@ -1,3 +1,9 @@ +/// \file runtime.cuh +/// \brief Host-side CUDA runtime query helpers. +/// +/// Thin wrappers around CUDA occupancy and device-property APIs with +/// automatic error checking via `RuntimeDeviceCheck`. + #pragma once #include diff --git a/python/sglang/jit_kernel/include/sgl_kernel/source_location.h b/python/sglang/jit_kernel/include/sgl_kernel/source_location.h index 9616fa7da..7c9fd5213 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/source_location.h +++ b/python/sglang/jit_kernel/include/sgl_kernel/source_location.h @@ -1,3 +1,9 @@ +/// \file source_location.h +/// \brief Portable `source_location` wrapper. +/// +/// Uses `std::source_location` when available (C++20), otherwise falls +/// back to a minimal stub that returns empty/zero values. + #pragma once #include diff --git a/python/sglang/jit_kernel/include/sgl_kernel/tensor.h b/python/sglang/jit_kernel/include/sgl_kernel/tensor.h index de7ed8a0c..484c969b5 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/tensor.h +++ b/python/sglang/jit_kernel/include/sgl_kernel/tensor.h @@ -1,3 +1,14 @@ +/// \file tensor.h +/// \brief Tensor validation and symbolic matching utilities. +/// +/// Provides the `TensorMatcher` fluent API for validating tensor shapes, +/// strides, dtypes, and devices at kernel entry points, along with +/// `SymbolicSize`, `SymbolicDType`, and `SymbolicDevice` for capturing +/// and cross-checking tensor metadata across multiple tensors. +/// +/// See the "Tensor Checking" section in the JIT kernel dev guide for +/// usage examples. + #pragma once #include @@ -151,11 +162,25 @@ inline auto& operator<<(std::ostream& os, PrintAbleSpan span) { } // namespace details +/// \brief Check whether `dtype` matches the DLDataType for C++ type `T`. template inline bool is_type(DLDataType dtype) { return dtype == details::_dtype_trait::value; } +/** + * \brief A symbolic dimension size that can be bound once and + * verified across multiple tensors. + * + * Create with an optional annotation string for error messages: + * \code + * auto N = SymbolicSize{"num_tokens"}; + * \endcode + * + * Call `verify()` during tensor matching to either bind the first + * observed value or check subsequent values match. Call `unwrap()` + * to retrieve the bound value (panics if unset). + */ struct SymbolicSize { public: SymbolicSize(std::string_view annotation = {}) : m_value(details::kNullSize), m_annotation(annotation) {} @@ -219,6 +244,12 @@ inline auto operator==(DLDevice lhs, DLDevice rhs) -> bool { return lhs.device_type == rhs.device_type && lhs.device_id == rhs.device_id; } +/** + * \brief A symbolic data type that can be constrained and verified. + * + * Optionally restrict allowed types via `set_options()`. + * Use `verify()` to bind/check the dtype, and `unwrap()` to retrieve it. + */ struct SymbolicDType { public: SymbolicDType() : m_value({details::kNullDType, 0, 0}) {} @@ -276,6 +307,12 @@ struct SymbolicDType { DLDataType m_value; }; +/** + * \brief A symbolic device that can be constrained and verified. + * + * Optionally restrict allowed device types via + * `set_options()`. The device id can be wildcarded. + */ struct SymbolicDevice { public: SymbolicDevice() : m_value({details::kNullDevice, details::kAnyDeviceID}) {} @@ -407,6 +444,24 @@ struct DeviceRef : BaseRef { } // namespace details +/** + * \brief Fluent API for validating tensor shape, strides, dtype, and device. + * + * Construct with the expected shape (using `SymbolicSize` or literal + * integers), chain `.with_strides()`, `.with_dtype<...>()`, and + * `.with_device<...>()`, then call `.verify(tensor)`. + * + * Example: + * \code + * auto N = SymbolicSize{"N"}; + * TensorMatcher({N, 128}) + * .with_dtype() + * .with_device() + * .verify(input_tensor); + * \endcode + * + * \note `TensorMatcher` is a move-only temporary. Do not store in a variable. + */ struct TensorMatcher { private: using SizeRef = details::SizeRef; diff --git a/python/sglang/jit_kernel/include/sgl_kernel/tile.cuh b/python/sglang/jit_kernel/include/sgl_kernel/tile.cuh index 1328c64f9..1adc82170 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/tile.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/tile.cuh @@ -1,3 +1,13 @@ +/// \file tile.cuh +/// \brief Tiled memory access helpers for coalesced global memory I/O. +/// +/// `tile::Memory` represents a contiguous memory region where multiple +/// threads cooperatively load/store elements. The three factory methods +/// determine the thread group: +/// - `thread()` - single thread (no tiling). +/// - `warp()` - all threads in a warp cooperate. +/// - `cta()` - all threads in the CTA cooperate. + #pragma once #include @@ -5,25 +15,41 @@ namespace device::tile { +/** + * \brief Represents a contiguous memory region for cooperative tiled access. + * + * Each instance is parameterized by an element type `T` and bound to a + * specific thread id (`tid`) within a group of `tsize` threads. + * + * \tparam T The storage element type (e.g. `AlignedVector, 4>`). + */ template struct Memory { public: SGL_DEVICE constexpr Memory(uint32_t tid, uint32_t tsize) : tid(tid), tsize(tsize) {} + /// \brief Create a Memory accessor for a single thread (no cooperation). SGL_DEVICE static constexpr Memory thread() { return Memory{0, 1}; } + /// \brief Create a Memory accessor distributed across warp threads. SGL_DEVICE static Memory warp(int warp_threads = kWarpThreads) { return Memory{static_cast(threadIdx.x % warp_threads), static_cast(warp_threads)}; } + /// \brief Create a Memory accessor distributed across all CTA threads. SGL_DEVICE static Memory cta(int cta_threads = blockDim.x) { return Memory{static_cast(threadIdx.x), static_cast(cta_threads)}; } + /// \brief Load one element from `ptr` at the position assigned to this thread. + /// \param ptr Base pointer (cast to `const T*`). + /// \param offset Optional tile offset (multiplied by `tsize`). SGL_DEVICE T load(const void* ptr, int64_t offset = 0) const { return static_cast(ptr)[tid + offset * tsize]; } + /// \brief Store one element to `ptr` at the position assigned to this thread. SGL_DEVICE void store(void* ptr, T val, int64_t offset = 0) const { static_cast(ptr)[tid + offset * tsize] = val; } + /// \brief Check whether this thread's element index is within bounds. SGL_DEVICE bool in_bound(int64_t element_count, int64_t offset = 0) const { return tid + offset * tsize < element_count; } diff --git a/python/sglang/jit_kernel/include/sgl_kernel/type.cuh b/python/sglang/jit_kernel/include/sgl_kernel/type.cuh index f06bc1407..15018fbb2 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/type.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/type.cuh @@ -1,3 +1,18 @@ +/// \file type.cuh +/// \brief Dtype trait system for CUDA scalar/packed types. +/// +/// `dtype_trait` provides per-type metadata: packed type alias, +/// conversion functions (`from`), and unary/binary math operations. +/// Use `device::cast(from_value)` for type conversion on device. +/// +/// Registered types: +/// | Scalar | Packed (x2) | Notes | +/// |-----------|-------------|-------------------------------| +/// | `fp32_t` | `fp32x2_t` | Full math ops (abs,sqrt,...) | +/// | `fp16_t` | `fp16x2_t` | Conversion only | +/// | `bf16_t` | `bf16x2_t` | Conversion only | +/// | `fp32x2_t`| `fp32x4_t` | Packed float2 <-> half2/bf162 | + #pragma once #include @@ -65,11 +80,18 @@ SGL_REGISTER_DTYPE_TRAIT( #undef SGL_REGISTER_DTYPE_TRAIT #undef SGL_REGISTER_FROM_FUNCTION +/// \brief Alias: the packed (x2) type for `T`. template using packed_t = typename dtype_trait::packed_t; namespace device { +/** + * \brief Cast a value from type `From` to type `To` on device. + * + * Dispatches through `dtype_trait::from()`, which uses the appropriate + * CUDA intrinsic (e.g. `__half2float`, `__float22half2_rn`). + */ template SGL_DEVICE To cast(const From& value) { return dtype_trait::from(value); diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh index eae04e68d..ae0577618 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh @@ -1,3 +1,18 @@ +/// \file utils.cuh +/// \brief Core CUDA/device utilities: type aliases, PDL helpers, +/// typed pointer access, kernel launch wrapper, and error checking. +/// +/// This header is included (directly or transitively) by nearly every +/// JIT kernel. It provides: +/// - Scalar/packed type aliases (`fp16_t`, `bf16_t`, `fp8_e4m3_t`, ...). +/// - `SGL_DEVICE` macro (forced-inline device function qualifier). +/// - `kWarpThreads` constant (32). +/// - PDL (Programmatic Dependent Launch) helpers for Hopper (sm_90+). +/// - Typed `load_as` / `store_as` for void-pointer access. +/// - `pointer::offset` for safe void-pointer arithmetic. +/// - `host::LaunchKernel` - kernel launcher with optional PDL. +/// - `host::RuntimeDeviceCheck` - CUDA error checking. + #pragma once #include @@ -70,11 +85,21 @@ using fp32x4_t = float4; namespace device { +/// \brief Macro: forced-inline device function qualifier. #define SGL_DEVICE __forceinline__ __device__ +/// \brief Number of threads per warp (always 32 on NVIDIA/AMD GPUs). inline constexpr auto kWarpThreads = 32u; +/// \brief Full warp active mask (all 32 lanes). inline constexpr auto kFullMask = 0xffffffffu; +/** + * \brief PDL (Programmatic Dependent Launch): wait for the primary kernel. + * + * On Hopper (sm_90+), inserts a `griddepcontrol.wait` instruction to + * synchronize with a preceding kernel in the same stream. On older + * architectures or ROCm this is a no-op. + */ template SGL_DEVICE void PDLWaitPrimary() { #if !defined(USE_ROCM) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) @@ -84,6 +109,12 @@ SGL_DEVICE void PDLWaitPrimary() { #endif } +/** + * \brief PDL: trigger dependent (secondary) kernel launch. + * + * On Hopper (sm_90+), inserts a `griddepcontrol.launch_dependents` + * instruction. On older architectures or ROCm this is a no-op. + */ template SGL_DEVICE void PDLTriggerSecondary() { #if !defined(USE_ROCM) && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) @@ -118,6 +149,7 @@ SGL_DEVICE void store_as(void* ptr, std::type_identity_t val, int64_t offset static_cast(ptr)[offset] = val; } +/// \brief Safe void-pointer arithmetic (byte-level by default). namespace pointer { // we only allow void * pointer arithmetic for safety @@ -138,6 +170,9 @@ SGL_DEVICE auto offset(const void* ptr, U... offset) -> const void* { namespace host { +/** + * \brief Check the CUDA error code and panic with location info on failure. + */ inline void RuntimeDeviceCheck(::cudaError_t error, DebugInfo location = {}) { if (error != ::cudaSuccess) { [[unlikely]]; @@ -145,10 +180,25 @@ inline void RuntimeDeviceCheck(::cudaError_t error, DebugInfo location = {}) { } } +/// \brief Check the last CUDA error (calls `cudaGetLastError`). inline void RuntimeDeviceCheck(DebugInfo location = {}) { return RuntimeDeviceCheck(::cudaGetLastError(), location); } +/** + * \brief Kernel launcher with automatic stream resolution and PDL support. + * + * Usage: + * \code + * host::LaunchKernel(grid, block, device) + * .enable_pdl(true) + * (my_kernel, arg1, arg2); + * \endcode + * + * The constructor resolves the CUDA stream from a `DLDevice` (via + * `TVMFFIEnvGetStream`) or accepts a raw `cudaStream_t`. The call + * operator launches the kernel and checks for errors. + */ struct LaunchKernel { public: explicit LaunchKernel( diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.h b/python/sglang/jit_kernel/include/sgl_kernel/utils.h index 78eae19fc..3226f79dd 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/utils.h +++ b/python/sglang/jit_kernel/include/sgl_kernel/utils.h @@ -1,3 +1,15 @@ +/// \file utils.h +/// \brief Host-side C++ utilities used by JIT kernel wrappers. +/// +/// Provides: +/// - `DebugInfo` - wraps `std::source_location` for error reporting. +/// - `RuntimeCheck` - runtime assertion with formatted error messages. +/// - `Panic` - unconditional abort with formatted error messages. +/// - `pointer::offset` - safe void-pointer arithmetic (host side). +/// - `div_ceil` - integer ceiling division. +/// - `dtype_bytes` - byte width of a `DLDataType`. +/// - `irange` - Python-style integer range for range-for loops. + #pragma once // ref: https://forums.developer.nvidia.com/t/c-20s-source-location-compilation-error-when-using-nvcc-12-1/258026/3 @@ -47,10 +59,12 @@ namespace host { template inline constexpr bool dependent_false_v = false; +/// \brief Source-location wrapper for debug/error messages. struct DebugInfo : public source_location_t { DebugInfo(source_location_t loc = source_location_t::current()) : source_location_t(loc) {} }; +/// \brief Exception type thrown by `RuntimeCheck` and `Panic`. struct PanicError : public std::runtime_error { public: explicit PanicError(std::string msg) : runtime_error(msg), m_message(std::move(msg)) {} @@ -64,6 +78,7 @@ struct PanicError : public std::runtime_error { std::string m_message; }; +/// \brief Unconditionally abort with a formatted error message. template [[noreturn]] inline auto panic(DebugInfo location, Args&&... args) -> void { @@ -78,6 +93,15 @@ inline auto panic(DebugInfo location, Args&&... args) -> void { throw PanicError(std::move(os).str()); } +/** + * \brief Runtime assertion: panics with a formatted message when `condition` + * is false. Extra `args` are streamed to the error message. + * + * Example: + * \code + * RuntimeCheck(n > 0, "n must be positive, got ", n); + * \endcode + */ template struct RuntimeCheck { template @@ -133,11 +157,13 @@ inline auto offset(const void* ptr, U... offset) -> const void* { } // namespace pointer +/// \brief Integer ceiling division: ceil(a / b). template inline constexpr auto div_ceil(T a, U b) { return (a + b - 1) / b; } +/// \brief Returns the byte width of a DLPack data type. inline auto dtype_bytes(DLDataType dtype) -> std::size_t { return static_cast(dtype.bits / 8); } @@ -145,11 +171,13 @@ inline auto dtype_bytes(DLDataType dtype) -> std::size_t { namespace stdr = std::ranges; namespace stdv = stdr::views; +/// \brief Python-style integer range: `irange(n)` -> `[0, n)`. template inline auto irange(T end) { return stdv::iota(static_cast(0), end); } +/// \brief Python-style integer range: `irange(start, end)` -> `[start, end)`. template inline auto irange(T start, T end) { return stdv::iota(start, end); diff --git a/python/sglang/jit_kernel/include/sgl_kernel/vec.cuh b/python/sglang/jit_kernel/include/sgl_kernel/vec.cuh index c1150d428..86b80c789 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/vec.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/vec.cuh @@ -1,3 +1,11 @@ +/// \file vec.cuh +/// \brief Aligned vector types for coalesced global memory access. +/// +/// `AlignedVector` wraps `N` elements of type `T` in a naturally +/// aligned struct so that the compiler emits wide (vectorized) load/store +/// instructions (e.g. `LDG.128`). The maximum supported vector width is +/// 256 bits (32 bytes), matching CUDA's widest vector load. + #pragma once #include @@ -8,6 +16,7 @@ namespace device { namespace details { +/// \brief Maps byte-width to the corresponding unsigned integer type. template struct uint_trait {}; @@ -31,35 +40,58 @@ struct uint_trait<8> { using type = uint64_t; }; +/// \brief Alias: maps `sizeof(T)` to matching unsigned int type. template using sized_int = typename uint_trait::type; } // namespace details +/// \brief Raw aligned storage for `N` elements of type `T`. template struct alignas(sizeof(T) * N) AlignedStorage { T data[N]; }; +/** + * \brief Aligned vector for vectorized memory access on GPU. + * + * Stores `N` elements of type `T` with natural alignment so that a single + * `load`/`store` call compiles to a wide memory transaction. + * + * \tparam T Element type (e.g. `fp16_t`, `bf16_t`, `float`). + * \tparam N Number of elements. Must be a power of two and + * `sizeof(T) * N <= 32` (256 bits). + * + * Example: + * \code + * AlignedVector vec; // 16 bytes, 128-bit aligned + * vec.load(input_ptr, tid); // vectorized load + * vec[0] = vec[0] + 1; + * vec.store(output_ptr, tid); // vectorized store + * \endcode + */ template 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 <= 32, "CUDA only support at most 256B vector op"); + /// NOTE: N must be a power of two and sizeof(T) * N <= 32 bytes (256 bits) + static_assert((N > 0 && (N & (N - 1)) == 0) && sizeof(T) * N <= 32, "CUDA only supports at most 256-bit vector op"); using element_t = typename details::sized_int; using storage_t = AlignedStorage; public: + /// \brief Vectorized load from `ptr` at the given element `offset`. template SGL_DEVICE void load(const U* ptr, std::size_t offset = 0) { static_assert(std::is_same_v || std::is_same_v); m_storage = reinterpret_cast(ptr)[offset]; } + /// \brief Vectorized store to `ptr` at the given element `offset`. template SGL_DEVICE void store(U* ptr, std::size_t offset = 0) const { static_assert(std::is_same_v || std::is_same_v); reinterpret_cast(ptr)[offset] = m_storage; } + /// \brief Fill all N elements with the same `value`. SGL_DEVICE void fill(T value) { const auto store_value = *reinterpret_cast(&value); #pragma unroll diff --git a/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh b/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh index d69526e97..52952ff81 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh @@ -1,11 +1,26 @@ +/// \file warp.cuh +/// \brief Warp-level reduction primitives using `__shfl_xor_sync`. + #pragma once #include -// Some warp primitives namespace device::warp { +/// \brief Full 32-thread active mask. static constexpr uint32_t kFullMask = 0xffffffffu; +/** + * \brief Warp-level sum reduction. + * + * Computes the sum of `value` across all active lanes specified by + * `active_mask` using butterfly (XOR) shuffles. The result is + * broadcast to all participating lanes. + * + * \tparam T Numeric type (e.g. float). + * \param value Per-lane input value. + * \param active_mask Bitmask of participating lanes (default: all 32). + * \return The sum across all active lanes. + */ template SGL_DEVICE T reduce_sum(T value, uint32_t active_mask = kFullMask) { #pragma unroll @@ -14,6 +29,18 @@ SGL_DEVICE T reduce_sum(T value, uint32_t active_mask = kFullMask) { return value; } +/** + * \brief Warp-level max reduction. + * + * Computes the maximum of `value` across all active lanes using + * butterfly shuffles. The result is broadcast to all participating + * lanes. + * + * \tparam T Numeric type (must be supported by `math::max`). + * \param value Per-lane input value. + * \param active_mask Bitmask of participating lanes (default: all 32). + * \return The maximum across all active lanes. + */ template SGL_DEVICE T reduce_max(T value, uint32_t active_mask = kFullMask) { #pragma unroll