[Docs] Add docstrings to JIT kernel include headers

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,10 +1,28 @@
+/// \file cta.cuh
+/// \brief CTA (Cooperative Thread Array / thread-block) level primitives.
+
 #pragma once
 #include <sgl_kernel/math.cuh>
 #include <sgl_kernel/utils.cuh>
 #include <sgl_kernel/warp.cuh>
 
 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,69 @@
+/// \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,29 +1,55 @@
+/// \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>
 
 #include <cstdint>
 
 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);