[Docs] Add docstrings to JIT kernel include headers (#19770)

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<T>`, `packed_t<T>`, `device::cast<To>(from)` |
+
+### Memory Access
+
+| Header | Namespace | Purpose |
+|--------|-----------|---------|
+| `vec.cuh` | `device` | `AlignedVector<T, N>` - vectorized load/store (up to 128-bit; 256-bit requires Blackwell GPUs) |
+| `tile.cuh` | `device::tile` | `Memory<T>` - 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`) |

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

@@ -1,8 +1,20 @@
+/// \file atomic.cuh
+/// \brief Device-side atomic operations.
+
 #pragma once
 #include <sgl_kernel/utils.cuh>
 
 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;

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 <sgl_kernel/math.cuh>
 #include <sgl_kernel/utils.cuh>
@@ -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 <typename T>
 SGL_DEVICE void reduce_max(T value, float* smem, float min_value = 0.0f) {
   const uint32_t warp_id = threadIdx.x / kWarpThreads;

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<T>`. All functions are forced-inline
+/// device functions.
+
 #pragma once
 #include <sgl_kernel/type.cuh>
 
+#include <cmath>
+
 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 <typename T>
 SGL_DEVICE T max(T a, T b) {
   return dtype_trait<T>::max(a, b);
 }
 
+/// \brief Returns the smaller of `a` and `b`.
 template <typename T>
 SGL_DEVICE T min(T a, T b) {
   return dtype_trait<T>::min(a, b);
 }
 
+/// \brief Returns the absolute value of `a`.
 template <typename T>
 SGL_DEVICE T abs(T a) {
   return dtype_trait<T>::abs(a);
 }
 
+/// \brief Returns the square root of `a`.
 template <typename T>
 SGL_DEVICE T sqrt(T a) {
   return dtype_trait<T>::sqrt(a);
 }
 
+/// \brief Returns the reciprocal square root of `a` (i.e. 1 / sqrt(a)).
 template <typename T>
 SGL_DEVICE T rsqrt(T a) {
   return dtype_trait<T>::rsqrt(a);
 }
 
+/// \brief Returns e^a.
 template <typename T>
 SGL_DEVICE T exp(T a) {
   return dtype_trait<T>::exp(a);
 }
 
+/// \brief Returns sin(a).
 template <typename T>
 SGL_DEVICE T sin(T a) {
   return dtype_trait<T>::sin(a);
 }
 
+/// \brief Returns cos(a).
 template <typename T>
 SGL_DEVICE T cos(T a) {
   return dtype_trait<T>::cos(a);

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 <sgl_kernel/utils.cuh>

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 <version>
 

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 <sgl_kernel/utils.h>
 
@@ -151,11 +162,25 @@ inline auto& operator<<(std::ostream& os, PrintAbleSpan<T> span) {
 
 }  // namespace details
 
+/// \brief Check whether `dtype` matches the DLDataType for C++ type `T`.
 template <typename T>
 inline bool is_type(DLDataType dtype) {
   return dtype == details::_dtype_trait<T>::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<fp16_t, bf16_t>()`.
+ * 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<kDLCUDA, kDLCPU>()`. The device id can be wildcarded.
+ */
 struct SymbolicDevice {
  public:
   SymbolicDevice() : m_value({details::kNullDevice, details::kAnyDeviceID}) {}
@@ -407,6 +444,24 @@ struct DeviceRef : BaseRef<SymbolicDevice> {
 
 }  // 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<fp16_t, bf16_t>()
+ *       .with_device<kDLCUDA>()
+ *       .verify(input_tensor);
+ * \endcode
+ *
+ * \note `TensorMatcher` is a move-only temporary. Do not store in a variable.
+ */
 struct TensorMatcher {
  private:
   using SizeRef = details::SizeRef;

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<T>` 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 <sgl_kernel/utils.cuh>
 
@@ -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<packed_t<float>, 4>`).
+ */
 template <typename T>
 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<uint32_t>(threadIdx.x % warp_threads), static_cast<uint32_t>(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<uint32_t>(threadIdx.x), static_cast<uint32_t>(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<const T*>(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<T*>(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;
   }

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<T>` provides per-type metadata: packed type alias,
+/// conversion functions (`from`), and unary/binary math operations.
+/// Use `device::cast<To>(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 <sgl_kernel/utils.cuh>
 
@@ -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 <typename T>
 using packed_t = typename dtype_trait<T>::packed_t;
 
 namespace device {
 
+/**
+ * \brief Cast a value from type `From` to type `To` on device.
+ *
+ * Dispatches through `dtype_trait<To>::from()`, which uses the appropriate
+ * CUDA intrinsic (e.g. `__half2float`, `__float22half2_rn`).
+ */
 template <typename To, typename From>
 SGL_DEVICE To cast(const From& value) {
   return dtype_trait<To>::from(value);

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 <sgl_kernel/utils.h>
@@ -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 <bool kUsePDL>
 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 <bool kUsePDL>
 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<T> val, int64_t offset
   static_cast<T*>(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(

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 <typename>
 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 <typename... Args>
 [[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 <typename... Args>
 struct RuntimeCheck {
   template <typename Cond>
@@ -133,11 +157,13 @@ inline auto offset(const void* ptr, U... offset) -> const void* {
 
 }  // namespace pointer
 
+/// \brief Integer ceiling division: ceil(a / b).
 template <std::integral T, std::integral U>
 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<std::size_t>(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 <std::integral T>
 inline auto irange(T end) {
   return stdv::iota(static_cast<T>(0), end);
 }
 
+/// \brief Python-style integer range: `irange(start, end)` -> `[start, end)`.
 template <std::integral T>
 inline auto irange(T start, T end) {
   return stdv::iota(start, end);

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<T, N>` 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 <sgl_kernel/utils.cuh>
 
@@ -8,6 +16,7 @@ namespace device {
 
 namespace details {
 
+/// \brief Maps byte-width to the corresponding unsigned integer type.
 template <std::size_t N>
 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 <typename T>
 using sized_int = typename uint_trait<sizeof(T)>::type;
 
 }  // namespace details
 
+/// \brief Raw aligned storage for `N` elements of type `T`.
 template <typename T, std::size_t N>
 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<fp16_t, 8> 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 <typename T, std::size_t N>
 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<T>;
   using storage_t = AlignedStorage<element_t, N>;
 
  public:
+  /// \brief Vectorized load from `ptr` at the given element `offset`.
   template <typename U>
   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];
   }
+  /// \brief Vectorized store to `ptr` at the given element `offset`.
   template <typename U>
   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;
   }
+  /// \brief Fill all N elements with the same `value`.
   SGL_DEVICE void fill(T value) {
     const auto store_value = *reinterpret_cast<element_t*>(&value);
 #pragma unroll

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 <sgl_kernel/math.cuh>
 
-// 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 <typename T>
 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 <typename T>
 SGL_DEVICE T reduce_max(T value, uint32_t active_mask = kFullMask) {
 #pragma unroll