chore: import upstream snapshot with attribution

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

@@ -0,0 +1,35 @@
+/// \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;
+  old = (value >= 0) ? __int_as_float(atomicMax((int*)addr, __float_as_int(value)))
+                     : __uint_as_float(atomicMin((unsigned int*)addr, __float_as_uint(value)));
+  return old;
+#else
+  int* addr_as_i = (int*)addr;
+  int old = *addr_as_i, assumed;
+  do {
+    assumed = old;
+    old = atomicCAS(addr_as_i, assumed, __float_as_int(fmaxf(value, __int_as_float(assumed))));
+  } while (assumed != old);
+  return __int_as_float(old);
+#endif
+}
+
+}  // namespace device::atomic

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

@@ -0,0 +1,40 @@
+/// \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;
+  smem[warp_id] = warp::reduce_max(value);
+  __syncthreads();
+  if (warp_id == 0) {
+    const auto tx = threadIdx.x;
+    const auto local_value = tx * kWarpThreads < blockDim.x ? smem[tx] : min_value;
+    const auto max_value = warp::reduce_max(local_value);
+    smem[0] = max_value;
+  }
+  // no extra sync; it is caller's responsibility to sync if needed
+}
+
+}  // namespace device::cta

python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/compress.cuh

@@ -0,0 +1,37 @@
+#pragma once
+
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tuple.h>
+
+#include <cstdint>
+
+namespace device::compress {
+
+struct alignas(16) PrefillPlan {
+  uint32_t ragged_id;
+  uint32_t batch_id;
+  uint32_t position;
+  uint32_t window_len;  // must be in `[0, compress_ratio * (1 + is_overlap))`
+
+  bool is_valid(const uint32_t ratio, const bool is_overlap) const {
+    const uint32_t max_window_len = ratio * (1 + is_overlap);
+    return window_len < max_window_len;
+  }
+};
+
+}  // namespace device::compress
+
+namespace host::compress {
+
+using device::compress::PrefillPlan;
+using PrefillPlanTensorDtype = uint8_t;
+inline constexpr int64_t kPrefillPlanDim = 16;
+
+static_assert(alignof(PrefillPlan) == sizeof(PrefillPlan));
+static_assert(sizeof(PrefillPlan) == kPrefillPlanDim * sizeof(PrefillPlanTensorDtype));
+
+}  // namespace host::compress

python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/compress_v2.cuh

@@ -0,0 +1,99 @@
+#pragma once
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace device::compress {
+
+/// \brief Per-batch decode plan. Layout: 16 bytes.
+struct alignas(16) DecodePlan {
+  uint32_t seq_len;
+  int32_t write_loc;
+  int32_t read_page_0;
+  int32_t read_page_1;
+};
+
+/// \brief Per-token compress plan (used by c4/c128 prefill). Layout: 16 bytes.
+struct alignas(16) CompressPlan {
+  uint32_t seq_len;
+  uint16_t ragged_id;
+  uint16_t buffer_len;
+  int32_t read_page_0;
+  /// \brief Stage 0 (CPU): batch_id (used to look up page table).
+  /// \brief Stage 1 (GPU): final state-pool write location.
+  int32_t read_page_1;
+
+  static SGL_DEVICE __host__ CompressPlan invalid() {
+    return CompressPlan{-1u, 0, 0, -1, -1};
+  }
+
+  SGL_DEVICE __host__ bool is_invalid() const {
+    return seq_len == -1u;
+  }
+};
+
+/// \brief Per-token write plan (used by c4/c128 prefill). Layout: 8 bytes.
+struct alignas(8) WritePlan {
+  /// \brief Stage 0 (CPU): packed `(batch_id << 16) | ragged_id`.
+  /// \brief Stage 1 (GPU): just `ragged_id`.
+  uint32_t ragged_id;
+  /// \brief Stage 0 (CPU): position + 1 (used to look up state slot).
+  /// \brief Stage 1 (GPU): final state-pool write location.
+  int32_t write_loc;
+
+  static SGL_DEVICE __host__ WritePlan invalid() {
+    return WritePlan{-1u, -1};
+  }
+
+  SGL_DEVICE __host__ bool is_invalid() const {
+    return ragged_id == -1u;
+  }
+};
+
+}  // namespace device::compress
+
+namespace host::compress {
+
+using device::compress::CompressPlan;
+using device::compress::DecodePlan;
+using device::compress::WritePlan;
+
+static_assert(alignof(DecodePlan) == sizeof(DecodePlan));
+static_assert(sizeof(DecodePlan) == 16);
+static_assert(alignof(CompressPlan) == sizeof(CompressPlan));
+static_assert(sizeof(CompressPlan) == 16);
+static_assert(alignof(WritePlan) == sizeof(WritePlan));
+static_assert(sizeof(WritePlan) == 8);
+
+inline auto verify_plan_d(tvm::ffi::TensorView t, SymbolicSize& N, SymbolicDevice& device) -> const DecodePlan* {
+  TensorMatcher({N, sizeof(DecodePlan)})  //
+      .with_dtype<uint8_t>()
+      .with_device(device)
+      .verify(t);
+  return static_cast<const DecodePlan*>(t.data_ptr());
+}
+
+inline auto verify_plan_c(tvm::ffi::TensorView t, SymbolicSize& N, SymbolicDevice& device) -> const CompressPlan* {
+  TensorMatcher({N, sizeof(CompressPlan)})  //
+      .with_dtype<uint8_t>()
+      .with_device(device)
+      .verify(t);
+  return static_cast<const CompressPlan*>(t.data_ptr());
+}
+
+inline auto verify_plan_w(tvm::ffi::TensorView t, SymbolicSize& N, SymbolicDevice& device) -> const WritePlan* {
+  TensorMatcher({N, sizeof(WritePlan)})  //
+      .with_dtype<uint8_t>()
+      .with_device(device)
+      .verify(t);
+  return static_cast<const WritePlan*>(t.data_ptr());
+}
+
+}  // namespace host::compress

python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/fp8_utils.cuh

@@ -0,0 +1,120 @@
+#pragma once
+
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+
+#include <cstdint>
+#ifndef USE_ROCM
+#include <cuda_fp8.h>
+#endif
+
+// Small helpers shared by the DeepSeek-V4 FP8/UE8M0 quantization kernels
+// (silu_and_mul_masked_post_quant, store, mega_moe_pre_dispatch, ...).
+// All functions are `SGL_DEVICE` (= `__forceinline__ __device__`) so
+// including this header in multiple translation units is ODR-safe.
+
+namespace deepseek_v4::fp8 {
+
+// Round `x` to the nearest representable UE8M0 value. Returns the raw
+// 8-bit biased exponent; the actual fp32 scale is `2^(exp - 127)`
+// (i.e. `__uint_as_float(exp << 23)`).
+SGL_DEVICE int32_t cast_to_ue8m0(float x) {
+  uint32_t u = __float_as_uint(x);
+  int32_t exp = int32_t((u >> 23) & 0xFF);
+  uint32_t mant = u & 0x7FFFFF;
+  return exp + (mant != 0);
+}
+
+// 1 / 2^(exp - 127) as fp32. Equivalent to `1.0f / __uint_as_float(exp << 23)`.
+SGL_DEVICE float inv_scale_ue8m0(int32_t exp) {
+  return __uint_as_float((127 + 127 - exp) << 23);
+}
+
+// Clamp to [-FP8_E4M3_MAX, FP8_E4M3_MAX].
+// Uses platform-specific max from type.cuh (448 for E4M3FN, 224 for E4M3FNUZ).
+SGL_DEVICE float fp8_e4m3_clip(float val) {
+  return fmaxf(fminf(val, kFP8E4M3Max), -kFP8E4M3Max);
+}
+
+#ifndef USE_ROCM
+// Pack two fp32 values into a single fp8x2_e4m3 with clamping.
+SGL_DEVICE fp8x2_e4m3_t pack_fp8(float x, float y) {
+  return fp8x2_e4m3_t{fp32x2_t{fp8_e4m3_clip(x), fp8_e4m3_clip(y)}};
+}
+#else
+// Software float -> FP8 E4M3 conversion for ROCm/HIP.
+// Supports both E4M3FN (MI350X, gfx950) and E4M3FNUZ (MI300X, gfx942).
+SGL_DEVICE uint8_t cvt_float_to_fp8_e4m3(float val) {
+  val = fp8_e4m3_clip(val);
+  if (val == 0.0f) return 0;
+
+  uint32_t f32 = __float_as_uint(val);
+  uint8_t sign = static_cast<uint8_t>((f32 >> 31) << 7);
+  int32_t exp32 = static_cast<int32_t>((f32 >> 23) & 0xFF) - 127;
+  uint32_t mant23 = f32 & 0x7FFFFF;
+
+#if HIP_FP8_TYPE_FNUZ
+  // E4M3FNUZ: bias=8, max=240, no negative zero, NaN=0x80
+  constexpr int32_t kBias = 8;
+  constexpr int32_t kMaxExp = 15;
+  constexpr int32_t kMinSubnormExp = -10;  // min subnormal exponent
+  constexpr int32_t kMinNormExp = -7;      // min normal exponent
+  constexpr uint8_t kSaturate = 0x7Fu;     // max normal = 0_1111_111 = 240.0
+#else
+  // E4M3FN: bias=7, max=448, NaN=0x7F
+  constexpr int32_t kBias = 7;
+  constexpr int32_t kMaxExp = 15;
+  constexpr int32_t kMinSubnormExp = -9;
+  constexpr int32_t kMinNormExp = -6;
+  constexpr uint8_t kSaturate = 0x7Eu;  // max normal = 0_1111_110 = 448.0
+#endif
+
+  int32_t exp8;
+  uint8_t mant3;
+
+  if (exp32 < kMinSubnormExp) {
+#if HIP_FP8_TYPE_FNUZ
+    // E4M3FNUZ (gfx942) has no negative zero: byte 0x80 is NaN, not -0.0.
+    // Returning `sign` (0x80) for an underflowing negative injects NaN into the
+    // fp8 KV cache -> NaN attention/logits. Flush underflow to +0 instead.
+    return 0;
+#else
+    // E4M3FN (gfx950): 0x80 == -0.0, harmless.
+    return sign;
+#endif
+  } else if (exp32 < kMinNormExp) {
+    // Subnormal range
+    int32_t shift = -(kBias - 1) - exp32;  // 1..3
+    uint32_t subnorm_mant = (0x800000 | mant23) >> (shift + 20);
+    uint32_t round_bit = ((0x800000 | mant23) >> (shift + 19)) & 1;
+    subnorm_mant += round_bit;
+    mant3 = static_cast<uint8_t>(subnorm_mant & 0x07);
+    exp8 = 0;
+    if (subnorm_mant > 7) {
+      exp8 = 1;
+      mant3 = 0;
+    }
+  } else {
+    exp8 = exp32 + kBias;
+    mant3 = static_cast<uint8_t>(mant23 >> 20);
+    uint32_t round_bit = (mant23 >> 19) & 1;
+    mant3 += round_bit;
+    if (mant3 > 7) {
+      mant3 = 0;
+      exp8++;
+    }
+    if (exp8 >= kMaxExp) return sign | kSaturate;
+  }
+  return sign | (static_cast<uint8_t>(exp8) << 3) | mant3;
+}
+
+// Pack two fp32 values into a single fp8x2_e4m3 (uint16_t on HIP).
+SGL_DEVICE fp8x2_e4m3_t pack_fp8(float x, float y) {
+  uint8_t x8 = cvt_float_to_fp8_e4m3(x);
+  uint8_t y8 = cvt_float_to_fp8_e4m3(y);
+  return static_cast<uint16_t>(x8) | (static_cast<uint16_t>(y8) << 8);
+}
+#endif
+
+}  // namespace deepseek_v4::fp8

python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/kvcacheio.cuh

@@ -0,0 +1,76 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+namespace device::hisparse {
+
+/// NOTE: We call nope+rope as a "value" here.
+/// Paged C4 cache layout:
+/// VALUE 0, VALUE 1, ..., VALUE 63,
+/// SCALE 0, SCALE 1, ..., SCALE 63,
+/// [Padding to align to 576 bytes]
+inline constexpr int64_t kPageSize = 64;
+inline constexpr int64_t kPageBits = 6;  // log2(kPageSize)
+inline constexpr int64_t kValueBytes = 576;
+inline constexpr int64_t kScaleBytes = 8;
+/// NOTE: FlashMLA requires each page to be aligned to 576 bytes
+inline constexpr int64_t kItemBytes = kValueBytes + kScaleBytes;
+inline constexpr int64_t kPageBytes = host::div_ceil(kItemBytes * kPageSize, 576) * 576;
+inline constexpr int64_t kScaleOffset = kValueBytes * kPageSize;
+
+struct PointerInfo {
+  int64_t* value_ptr;
+  int64_t* scale_ptr;
+};
+
+SGL_DEVICE PointerInfo get_pointer_paged(void* cache, int32_t index) {
+  using namespace device;
+  static_assert(1 << kPageBits == kPageSize);
+  const int32_t page_num = index >> kPageBits;
+  const int32_t page_offset = index & (kPageSize - 1);
+  const auto page_ptr = pointer::offset(cache, page_num * kPageBytes);
+  const auto value_ptr = pointer::offset(page_ptr, page_offset * kValueBytes);
+  const auto scale_ptr = pointer::offset(page_ptr, kScaleOffset + page_offset * kScaleBytes);
+  return {static_cast<int64_t*>(value_ptr), static_cast<int64_t*>(scale_ptr)};
+}
+
+SGL_DEVICE void transfer_item(void* dst_cache, void* src_cache, const int32_t dst_index, const int32_t src_index) {
+  const auto [dst_value_ptr, dst_scale_ptr] = get_pointer_paged(dst_cache, dst_index);
+  const auto [src_value_ptr, src_scale_ptr] = get_pointer_paged(src_cache, src_index);
+
+  int64_t local_items[2];
+  const int64_t* tail_src_ptr;
+  int64_t* tail_dst_ptr;
+
+  const int32_t lane_id = threadIdx.x % 32;
+
+  for (int i = 0; i < 2; ++i) {
+    const auto j = lane_id + i * 32;
+    local_items[i] = src_value_ptr[j];
+  }
+
+  if (lane_id < 8) {  // handle the tail element safely
+    const auto last_id = 64 + lane_id;
+    tail_src_ptr = src_value_ptr + last_id;
+    tail_dst_ptr = dst_value_ptr + last_id;
+  } else {  // broadcast load/store is safe
+    tail_src_ptr = src_scale_ptr;
+    tail_dst_ptr = dst_scale_ptr;
+  }
+
+  const auto tail_item = *tail_src_ptr;
+
+  // store first 512 bytes of value
+  for (int i = 0; i < 2; ++i) {
+    const auto j = lane_id + i * 32;
+    dst_value_ptr[j] = local_items[i];
+  }
+
+  // store the tail element
+  *tail_dst_ptr = tail_item;
+}
+
+}  // namespace device::hisparse

python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk_impl.cuh

@@ -0,0 +1,842 @@
+/// \file topk_impl.cuh
+/// \brief DeepSeek-V4 (DSA indexer) top-k implementation classes.
+///
+/// This header holds ONLY the device-side implementation classes + helpers; the
+/// `__global__` kernels and the host dispatcher live in csrc/deepseek_v4/topk_v2.cuh.
+///
+/// Design notes:
+///  - top-k (`topk`) is a *runtime* value (<= kMaxTopK = 2048), never a
+///    compile-time constant.
+///  - the output is the page-table transform of the selected raw indices
+///    (`TopKProblem::emit` then `transform_output`).
+///  - each block reads its own `seq_len` (per-batch ragged lengths) -- the host
+///    launches one universal kernel and dispatches per block.
+///  - the cluster size is fixed at 8 (dynamic persistent clusters are hard).
+///
+/// Algorithm: fp16 coarse histogram -> threshold bin -> fp32-boundary collect ->
+/// exact radix tie-break.
+
+#pragma once
+
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <cfloat>
+#include <cooperative_groups.h>
+#include <cstdint>
+#include <limits>
+
+namespace device::topk {
+
+namespace cg = cooperative_groups;
+
+/// sgl_kernel names the warp size `kWarpThreads`; alias it locally as `kWarpSize`.
+inline constexpr uint32_t kWarpSize = kWarpThreads;
+
+// ---------------------------------------------------------------------------
+// Shared-memory storage sized/aligned for several impl `Smem` types
+// ---------------------------------------------------------------------------
+
+/// Compile-time max over a non-empty pack (avoids an <algorithm> dependency).
+template <typename T>
+constexpr T ct_max(T a) {
+  return a;
+}
+template <typename T, typename... Ts>
+constexpr T ct_max(T a, Ts... rest) {
+  const T m = ct_max(rest...);
+  return a > m ? a : m;
+}
+
+/// Static shared-memory buffer sized + aligned to hold any one of the given
+/// impl `Smem` types. A kernel that dispatches across several paths (e.g. the
+/// fused small-batch kernel runs either Streaming or Cluster; the main kernel
+/// runs any of Register2/Register4/Streaming) declares one
+/// `__shared__ MaxSmem<...> smem` and hands `&smem` to whichever forward() it
+/// calls -- instead of hand-picking "the largest" type and relying on it
+/// staying the largest. `&smem` converts to the `void*` the forwards expect;
+/// the buffer is aligned to the strictest member, so the cast is well-aligned.
+template <typename... Smems>
+struct MaxSmem {
+  static constexpr size_t kSize = ct_max(sizeof(Smems)...);
+  static constexpr size_t kAlign = ct_max(alignof(Smems)...);
+  alignas(kAlign) uint8_t storage[kSize];
+};
+
+// ---------------------------------------------------------------------------
+// Order-preserving float -> integer key extraction
+// ---------------------------------------------------------------------------
+
+SGL_DEVICE uint32_t extract_exact_bin(float x) {
+  uint32_t bits = __float_as_uint(x);
+  return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
+}
+
+template <uint32_t kBits>
+SGL_DEVICE uint32_t extract_coarse_bin(float x) {
+  static_assert(0 < kBits && kBits < 15);
+  const auto hx = cast<fp16_t>(x);
+  const uint16_t bits = *reinterpret_cast<const uint16_t*>(&hx);
+  const uint16_t key = (bits & 0x8000) ? ~bits : bits | 0x8000;
+  return key >> (16 - kBits);
+}
+
+// Smallest fp32 value `v` for which `extract_coarse_bin<kBits>(v) >= bin`, i.e. the
+// lower fp32 boundary of coarse bin `bin`. Because `extract_coarse_bin` is monotonic
+// non-decreasing in its argument, the collect pass can classify an element with two
+// fp32 comparisons against these boundaries instead of recomputing the fp16 bin --
+// removing the F2F conversion and bit-twiddle from the (compute-bound) second pass.
+// Returns -inf for bin 0 (everything qualifies) and +inf for bins past the top.
+template <uint32_t kBits>
+SGL_DEVICE float coarse_bin_lower_bound(uint32_t bin) {
+  constexpr uint32_t kShift = 16 - kBits;
+  const uint32_t key = bin << kShift;  // ordered16 key at the low edge of `bin`
+  // ordered16 -> fp16 value (inverse of the transform in extract_coarse_bin);
+  // finite keys only.
+  const auto to_finite_val = [](uint32_t okey) -> float {
+    const uint16_t ob = static_cast<uint16_t>(okey);
+    const uint16_t hb = (ob & 0x8000) ? static_cast<uint16_t>(ob ^ 0x8000) : static_cast<uint16_t>(~ob);
+    return cast<float>(*reinterpret_cast<const fp16_t*>(&hb));
+  };
+  // Fast path, hoisted above the per-key special cases so both keys are
+  // range-checked at once: `key` and `key - 1` both land in the finite band
+  // [0x0401, 0xFBFF] -- every boundary a finite-score threshold produces.
+  // fp16 rounds to nearest, so the fp32 boundary is the midpoint between the
+  // fp16 values at `key` and `key - 1`. (Verified bit-exact against the slow
+  // path for every bin of kBits 10 and 12, and measured faster than either
+  // per-key dispatch or an ordered-bit decrement trick -- the two conversions
+  // are independent and issue in parallel.)
+  if (key - 0x0401u <= 0xFBFFu - 0x0401u && bin < (1u << kBits)) {
+    return 0.5f * (to_finite_val(key) + to_finite_val(key - 1));
+  }
+  // Slow path: an edge of `bin` touches the +/-inf keys or NaN key space.
+  // The ordered-key line is: [0, 0x03FF) negative-NaN space, 0x03FF = -inf,
+  // [0x0400, 0xFC00) finite, 0xFC00 = +inf, (0xFC00, 0xFFFF] positive-NaN
+  // space. Treat the +/-inf keys as +/-65536 (one ideal step past fp16 max,
+  // so the midpoint lands exactly on +/-65520 -- the fp32->fp16
+  // round-to-nearest overflow threshold) and saturate NaN-space keys, keeping
+  // the returned boundaries finite-or-inf and monotone. Otherwise a threshold
+  // bin at/next to the inf bin gets NaN boundaries, the collect pass matches
+  // nothing, and rows whose scores contain >= topk (+/-)inf or >65504 values
+  // come back short -- the padded slots then illegal-address downstream.
+  if (bin == 0) return -FLT_MAX;
+  if (bin >= (1u << kBits)) return FLT_MAX;
+  const auto to_val = [&](uint32_t okey) -> float {
+    constexpr float k_Inf = std::numeric_limits<float>::infinity();
+    if (okey < 0x03FFu) return -k_Inf;
+    if (okey == 0x03FFu) return -65536.0f;
+    if (okey == 0xFC00u) return 65536.0f;
+    if (okey > 0xFC00u) return FLT_MAX;
+    return to_finite_val(okey);
+  };
+  return 0.5f * (to_val(key) + to_val(key - 1));
+}
+
+SGL_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t val) {
+#pragma unroll
+  for (uint32_t offset = 1; offset < 32; offset *= 2) {
+    uint32_t n = __shfl_up_sync(0xFFFFFFFF, val, offset);
+    if (lane_id >= offset) val += n;
+  }
+  return val;
+}
+
+SGL_DEVICE uint32_t warp_sum_bool(bool pred, uint32_t mask = 0xFFFFFFFF) {
+  return __popc(__ballot_sync(mask, pred));
+}
+
+struct alignas(8) TieValue {
+  float value;
+  uint32_t idx;
+  inline static constexpr TieValue invalid() {
+    return TieValue{-FLT_MAX, 0xFFFFFFFFu};
+  }
+};
+
+// ---------------------------------------------------------------------------
+// Per-batch problem description + page-table transform sink
+// ---------------------------------------------------------------------------
+
+SGL_DEVICE int32_t page_to_indices(const int32_t* __restrict__ page_table, uint32_t i, uint32_t page_bits) {
+  const uint32_t mask = (1u << page_bits) - 1u;
+  return (page_table[i >> page_bits] << page_bits) | (i & mask);
+}
+
+/// One batch element's worth of work. `emit(pos, raw_idx)` writes the selected raw
+/// index to output slot `pos`; `transform_output` then applies the page-table
+/// transform in a separate pass (and records the raw index in `raw_out` if set).
+struct TopKProblem {
+  const float* __restrict__ in;
+  int32_t* __restrict__ out;      // page_indices [topk]
+  int32_t* __restrict__ raw_out;  // optional raw (pre-transform) indices [topk]; nullptr if unused
+  const int32_t* __restrict__ page_table;
+  uint32_t topk;
+  uint32_t seq_len;
+  uint32_t page_bits;
+
+  // Write the raw selected index; the page-table transform is applied afterwards
+  // by transform_output() in a separate, pipelined pass. Keeping the per-element
+  // page_table gather off the atomic-serialized scatter loop is measurably faster
+  // for both short and long context.
+  SGL_DEVICE void emit(uint32_t pos, uint32_t raw_idx) const {
+    out[pos] = static_cast<int32_t>(raw_idx);
+  }
+  SGL_DEVICE void transform_output(uint32_t t, int32_t raw) const {
+    if (raw_out != nullptr) raw_out[t] = raw;
+    out[t] = raw < 0 ? -1 : page_to_indices(page_table, raw, page_bits);
+  }
+};
+
+// ---------------------------------------------------------------------------
+// Shared configuration + tie handling (exact radix select on the threshold bin)
+// ---------------------------------------------------------------------------
+
+struct TopKConfig {
+  static constexpr uint32_t kMaxTopK = 2048;
+  static constexpr uint32_t kBlockSize = 1024;
+  static constexpr uint32_t kOccupancy = 2;
+  static constexpr uint32_t kNumWarps = kBlockSize / kWarpSize;
+  // kMaxNumTie must be >= kMaxTopK: the collect pass keeps at most kMaxNumTie
+  // threshold-bin candidates, and up to `topk` output slots may have to be
+  // filled from them (above_count can be 0, e.g. heavily tied or all-inf
+  // scores). A smaller cap leaves slots that handle_tie can only pad, and
+  // padded slots inside the first min(seq_len, topk) entries are dereferenced
+  // by downstream sparse attention.
+  static constexpr uint32_t kMaxNumTie = 2048;
+  static constexpr uint32_t kRadixSize = 1 << 8;
+  static constexpr uint32_t kTopKItems = (kMaxTopK + kBlockSize - 1) / kBlockSize;
+  // tie candidates owned per thread in the strided handle_tie loops
+  static constexpr uint32_t kTieItems = kMaxNumTie / kBlockSize;
+  static_assert(kMaxNumTie >= kMaxTopK && kMaxNumTie % kBlockSize == 0 && kBlockSize % kNumWarps == 0);
+
+  struct TieHandleSmem {
+    struct alignas(16) MatchBin {
+      uint32_t bin;
+      uint32_t above_count;
+      uint32_t equal_count;
+      uint32_t _pad = 0;
+    };
+    alignas(128) uint32_t counter;
+    alignas(128) uint32_t counter_final;
+    MatchBin match;
+    uint32_t warp_sum[kNumWarps];
+    uint32_t histogram[2][kRadixSize];
+  };
+
+  /// Resolve the threshold bin's ties exactly. `base` is the number of strictly
+  /// "above" elements already emitted (final output starts at slot `base`);
+  /// `topk` here is the number of remaining slots to fill (== global_topk - base).
+  SGL_DEVICE static void handle_tie(  //
+      const TieValue* tie_buffer,
+      const TopKProblem& problem,
+      const uint32_t base,
+      const uint32_t num_ties,
+      const uint32_t topk,
+      TieHandleSmem* smem) {
+    constexpr auto is_greater = [](const TieValue& a, const TieValue& b) {
+      return (a.value > b.value) || (a.value == b.value && a.idx < b.idx);
+    };
+    const auto tx = threadIdx.x;
+    const auto lane_id = tx % kWarpSize;
+    const auto warp_id = tx / kWarpSize;
+    static_assert(kNumWarps == kWarpSize);
+
+    if (num_ties <= topk) {
+      for (uint32_t t = tx; t < num_ties; t += kBlockSize) {
+        problem.emit(base + t, tie_buffer[t].idx);
+      }
+      for (uint32_t t = num_ties + tx; t < topk; t += kBlockSize) {
+        problem.emit(base + t, base + t);
+      }
+    } else if (num_ties <= kWarpSize) {
+      if (lane_id >= num_ties || warp_id >= num_ties) return;  // some threads are idle
+      /// NOTE: use long long to avoid mask overflow when num_tie == 32
+      const uint32_t mask = (1ull << num_ties) - 1u;
+      const auto tie = tie_buffer[lane_id];
+      const auto target = tie_buffer[warp_id];
+      const auto rank = warp_sum_bool(is_greater(tie, target), mask);
+      if (lane_id == 0 && rank < topk) problem.emit(base + rank, target.idx);
+    } else if (num_ties <= kWarpSize * 2) {
+      // 64 x 64 topk implementation: each thread takes 2 elements
+      const auto warp_id_0 = warp_id;
+      const auto warp_id_1 = warp_id + kWarpSize;
+      const auto lane_id_1 = lane_id + kWarpSize;
+      const auto invalid = TieValue::invalid();
+      const auto tie_0 = tie_buffer[lane_id];
+      const auto tie_1 = lane_id_1 < num_ties ? tie_buffer[lane_id_1] : invalid;
+      const auto target_0 = tie_buffer[warp_id_0];
+      const auto target_1 = tie_buffer[warp_id_1];
+      if (true) {  // NOTE: warp_id_0 <= kNumWarps < num_ties
+        const auto rank_0 = warp_sum_bool(is_greater(tie_0, target_0));
+        const auto rank_1 = warp_sum_bool(is_greater(tie_1, target_0));
+        const auto rank = rank_0 + rank_1;
+        if (lane_id == 0 && rank < topk) problem.emit(base + rank, target_0.idx);
+      }
+      if (warp_id_1 < num_ties) {
+        const auto rank_0 = warp_sum_bool(is_greater(tie_0, target_1));
+        const auto rank_1 = warp_sum_bool(is_greater(tie_1, target_1));
+        const auto rank = rank_0 + rank_1;
+        if (lane_id == 0 && rank < topk) problem.emit(base + rank, target_1.idx);
+      }
+    } else if (num_ties <= kWarpSize * 4) {
+      // 128 x 128 topk implementation: each thread takes 4 elements and does local sort + merge
+      const auto invalid = TieValue::invalid();
+      const TieValue tie[] = {
+          tie_buffer[lane_id + 0 * kWarpSize],
+          tie_buffer[lane_id + 1 * kWarpSize],
+          lane_id + 2 * kWarpSize < num_ties ? tie_buffer[lane_id + 2 * kWarpSize] : invalid,
+          lane_id + 3 * kWarpSize < num_ties ? tie_buffer[lane_id + 3 * kWarpSize] : invalid,
+      };
+      const TieValue target[] = {
+          tie_buffer[warp_id + 0 * kWarpSize],
+          tie_buffer[warp_id + 1 * kWarpSize],
+          tie_buffer[warp_id + 2 * kWarpSize],
+          tie_buffer[warp_id + 3 * kWarpSize],
+      };
+#pragma unroll
+      for (int i = 0; i < 4; ++i) {
+        if (i >= 2 && warp_id + i * kWarpSize >= num_ties) break;
+        uint32_t rank = 0;
+#pragma unroll
+        for (int j = 0; j < 4; ++j) {
+          rank += warp_sum_bool(is_greater(tie[j], target[i]));
+        }
+        if (lane_id == 0 && rank < topk) problem.emit(base + rank, target[i].idx);
+      }
+    } else if (num_ties <= kBlockSize) {
+      // Common case: one candidate per thread.
+      radix_tie_select<1>(tie_buffer, problem, base, num_ties, topk, smem);
+    } else {
+      // Rare overflow case (kBlockSize < num_ties <= kMaxNumTie), kept out of
+      // the common path so it alone pays the multi-item register cost.
+      radix_tie_select<kTieItems>(tie_buffer, problem, base, num_ties, topk, smem);
+    }
+  }
+
+  /// Exact radix select over the tie candidates: each thread owns kItems
+  /// strided elements (inactive beyond num_ties). Requires
+  /// num_ties <= kItems * kBlockSize.
+  template <uint32_t kItems>
+  SGL_DEVICE static void radix_tie_select(  //
+      const TieValue* tie_buffer,
+      const TopKProblem& problem,
+      const uint32_t base,
+      const uint32_t num_ties,
+      const uint32_t topk,
+      TieHandleSmem* smem) {
+    const auto tx = threadIdx.x;
+    const auto lane_id = tx % kWarpSize;
+    const auto warp_id = tx / kWarpSize;
+
+    bool active[kItems];
+    uint32_t key[kItems];
+    uint32_t idx[kItems];
+    uint32_t write_pos[kItems];
+#pragma unroll
+    for (uint32_t i = 0; i < kItems; ++i) {
+      const auto t = tx + i * kBlockSize;
+      active[i] = t < num_ties;
+      const auto tie = active[i] ? tie_buffer[t] : TieValue::invalid();
+      key[i] = extract_exact_bin(tie.value);
+      idx[i] = tie.idx;
+      write_pos[i] = topk;
+    }
+    uint32_t topk_remain = topk;
+    if (tx < kRadixSize) smem->histogram[0][tx] = 0;
+    if (tx == kRadixSize) smem->counter = smem->counter_final = 0;
+    __syncthreads();
+    uint32_t total_active = num_ties;
+
+#pragma unroll
+    for (int round = 0; round < 4; round++) {
+      const uint32_t shift = 24 - round * 8;
+      const auto hist_idx = round % 2;
+      const auto histogram = smem->histogram[hist_idx];
+
+#pragma unroll
+      for (uint32_t i = 0; i < kItems; ++i) {
+        if (active[i]) atomicAdd(&histogram[(key[i] >> shift) & 0xFFu], 1);
+      }
+      if (round < 3 && tx < kRadixSize) {
+        smem->histogram[hist_idx ^ 1][tx] = 0;
+      }
+      __syncthreads();
+
+      uint32_t hist_val = 0;
+      uint32_t warp_inc = 0;
+      if (tx < kRadixSize) {
+        hist_val = histogram[tx];
+        warp_inc = warp_inclusive_sum(lane_id, hist_val);
+        if (lane_id == kWarpSize - 1) smem->warp_sum[warp_id] = warp_inc;
+      }
+      __syncthreads();
+      if (tx < kRadixSize) {
+        const auto inter = warp::reduce_sum(lane_id < warp_id ? smem->warp_sum[lane_id] : 0);
+        const auto prefix = inter + warp_inc;      // inclusive prefix through this bin
+        const auto above = total_active - prefix;  // elements in bins ABOVE this one
+        // 3. Find threshold bin
+        if (above < topk_remain && above + hist_val >= topk_remain) {
+          smem->match = {tx, above, hist_val};
+        }
+      }
+      __syncthreads();
+
+      const auto [threshold_bin, above_count, equal_count, __] = smem->match;
+      if (round < 3) total_active = equal_count;
+      topk_remain -= above_count;
+
+      // 4. Scatter
+#pragma unroll
+      for (uint32_t i = 0; i < kItems; ++i) {
+        if (!active[i]) continue;
+        const uint32_t bin = (key[i] >> shift) & 0xFFu;
+        if (bin > threshold_bin) {
+          write_pos[i] = atomicAdd(&smem->counter, 1);
+          active[i] = false;
+        } else if (bin < threshold_bin) {
+          active[i] = false;
+        } else if (round == 3) {
+          write_pos[i] = topk - topk_remain + atomicAdd(&smem->counter_final, 1);
+        }
+        // my_bin == thr && round < 3: stay active for next round
+      }
+
+      if (round == 3 || topk_remain == 0) break;
+    }
+
+#pragma unroll
+    for (uint32_t i = 0; i < kItems; ++i) {
+      if (write_pos[i] < topk) problem.emit(base + write_pos[i], idx[i]);
+    }
+  }
+};
+
+// ---------------------------------------------------------------------------
+// Radix base: histogram storage + input iteration + threshold-bin search
+// ---------------------------------------------------------------------------
+
+template <uint32_t kHistBits_>
+struct TopKRadixBase : TopKConfig {
+  static constexpr uint32_t kVecSize = 4;
+  static constexpr uint32_t kHistBits = kHistBits_;
+  static constexpr uint32_t kHistSize = 1 << kHistBits;
+  using vec_t = AlignedVector<float, kVecSize>;
+
+  struct Smem {
+    using kHistVec = AlignedVector<uint32_t, kHistSize / kBlockSize>;
+    alignas(128) uint32_t count_eq;
+    alignas(128) uint32_t count_gt;
+    uint32_t threshold_bin;
+    uint32_t warp_sum[kNumWarps];
+    // The coarse histogram is dead once find_threshold() has published
+    // threshold_bin, and the tie machinery only comes alive after that: the
+    // collect pass fills tie.values, then handle_tie works over them with
+    // tie.handle as scratch. Overlaying the two phases keeps the
+    // kMaxNumTie-candidate buffer from growing the block's shared-memory
+    // footprint. tie.handle and tie.values are live TOGETHER, so they sit
+    // side by side inside the overlay, not in a union with each other.
+    union {
+      uint32_t histogram[kHistSize];
+      kHistVec hist_vecs[kBlockSize];
+      struct {
+        TieHandleSmem handle;
+        TieValue values[kMaxNumTie];
+      } tie;
+    };
+  };
+
+ protected:
+  template <typename F>
+  SGL_DEVICE static void for_each_input(const float* __restrict__ in, uint32_t seq_len, F&& fn) {
+    const auto tx = threadIdx.x;
+    const uint32_t num_full = seq_len / kVecSize;  // fully-in-bounds vectors
+
+    vec_t next_vec;
+    uint32_t vi = tx;
+    if (vi < num_full) next_vec.load(in, vi);
+    while (vi < num_full) {
+      const auto cur = next_vec;
+      const auto base = vi * kVecSize;
+      vi += kBlockSize;
+      if (vi < num_full) next_vec.load(in, vi);
+#pragma unroll
+      for (uint32_t j = 0; j < kVecSize; ++j) {
+        fn(cur[j], base + j);
+      }
+    }
+
+    // Tail: at most one partial vector, `rem` in [0, kVecSize).
+    static_assert(kVecSize <= kBlockSize);  // ensure tail correctness
+    const uint32_t tail_start = num_full * kVecSize;
+    if (tx < seq_len - tail_start) {
+      const auto idx = tail_start + tx;
+      fn(in[idx], idx);
+    }
+  }
+
+  SGL_DEVICE static void find_threshold(const uint32_t topk, const uint32_t seq_len, Smem* smem) {
+    const auto tx = threadIdx.x;
+    constexpr uint32_t kItems = kHistSize / kBlockSize;
+    uint32_t orig[kItems];
+    const auto hist_vec = smem->hist_vecs[tx];
+    uint32_t tmp_local_sum = 0;
+
+#pragma unroll
+    for (uint32_t i = 0; i < kItems; ++i) {
+      orig[i] = hist_vec[i];
+      tmp_local_sum += orig[i];
+    }
+
+    const auto lane_id = tx % kWarpSize;
+    const auto warp_id = tx / kWarpSize;
+    const auto warp_inc = warp_inclusive_sum(lane_id, tmp_local_sum);
+    const auto warp_exc = warp_inc - tmp_local_sum;
+    if (lane_id == kWarpSize - 1) smem->warp_sum[warp_id] = warp_inc;
+
+    __syncthreads();
+
+    const auto tmp = smem->warp_sum[lane_id];
+    // Exactly one bin satisfies: above < K && above + count >= K
+    uint32_t prefix_sum = warp::reduce_sum(lane_id < warp_id ? tmp : 0);
+    prefix_sum += warp_exc;
+#pragma unroll
+    for (uint32_t i = 0; i < kItems; ++i) {
+      prefix_sum += orig[i];
+      const auto above = seq_len - prefix_sum;
+      if (above < topk && above + orig[i] >= topk) {
+        smem->threshold_bin = tx * kItems + i;
+      }
+    }
+    __syncthreads();
+  }
+};
+
+// ---------------------------------------------------------------------------
+// Register path: scores stay resident in registers across both passes (read
+// once). Templated on kLocalVecs so the caller picks the smallest covering
+// kernel -- a larger kLocalVecs raises kMaxSeqLen but its fixed-unrolled loop
+// wastes work on shorter sequences.
+// ---------------------------------------------------------------------------
+
+template <uint32_t kLocalVecs_>
+struct TopKRegister : TopKRadixBase<12> {
+  static constexpr uint32_t kLocalVecs = kLocalVecs_;
+  static constexpr uint32_t kMaxSeqLen = kBlockSize * kVecSize * kLocalVecs;
+  using Smem = typename TopKRadixBase<12>::Smem;
+
+  template <bool kUsePDL>
+  SGL_DEVICE static void forward(const TopKProblem problem, void* _smem) {
+    const auto tx = threadIdx.x;
+    const auto smem = static_cast<Smem*>(_smem);
+
+    {
+      Smem::kHistVec hist_vec;
+      hist_vec.fill(0);
+      smem->hist_vecs[tx] = hist_vec;
+    }
+    if (tx == 0) {
+      smem->count_eq = 0;
+      smem->count_gt = 0;
+    }
+
+    __syncthreads();
+    PDLWaitPrimary<kUsePDL>();
+
+    // A vector `vi` is fully in bounds iff vi < num_full; only full vectors are
+    // vector-loaded (16B aligned, never straddling seq_len). The <kVecSize tail is
+    // a scalar remainder on the LAST lanes (which own the fewest full vectors, so
+    // it overlaps the busy lanes' extra vector). The full path has no per-element
+    // bounds check, keeping register pressure low enough to hold all vectors.
+    const uint32_t num_full = problem.seq_len / kVecSize;
+    const uint32_t tail_start = num_full * kVecSize;
+    const uint32_t tail = problem.seq_len - tail_start;
+
+    // Phase 1: load full vectors + build histogram
+    vec_t local_vecs[kLocalVecs];
+#pragma unroll
+    for (uint32_t i = 0; i < kLocalVecs; ++i) {
+      const auto vi = tx + kBlockSize * i;
+      if (vi >= num_full) break;
+      local_vecs[i].load(problem.in, vi);
+    }
+#pragma unroll
+    for (uint32_t i = 0; i < kLocalVecs; ++i) {
+      const auto vi = tx + kBlockSize * i;
+      if (vi >= num_full) break;
+#pragma unroll
+      for (uint32_t j = 0; j < kVecSize; ++j)
+        atomicAdd(&smem->histogram[extract_coarse_bin<kHistBits>(local_vecs[i][j])], 1);
+    }
+    if (tx >= kBlockSize - tail) {
+      const uint32_t idx = tail_start + tx - (kBlockSize - tail);
+      atomicAdd(&smem->histogram[extract_coarse_bin<kHistBits>(problem.in[idx])], 1);
+    }
+    __syncthreads();
+
+    // Phase 2: Find the threshold bin
+    find_threshold(problem.topk, problem.seq_len, smem);
+
+    // Phase 3: collect by two fp32 boundaries (raw indices; transform applied later)
+    const auto topk = problem.topk;
+    const auto threshold_bin = smem->threshold_bin;
+    const auto v_hi = coarse_bin_lower_bound<kHistBits>(threshold_bin + 1);
+    const auto v_lo = coarse_bin_lower_bound<kHistBits>(threshold_bin);
+    const auto collect = [&](float val, uint32_t idx) {
+      if (val >= v_hi) {
+        const auto pos = atomicAdd(&smem->count_gt, 1);
+        if (pos < topk) [[likely]]
+          problem.emit(pos, idx);
+      } else if (val >= v_lo) {
+        const auto count_eq = atomicAdd(&smem->count_eq, 1);
+        if (count_eq < kMaxNumTie) [[likely]]
+          smem->tie.values[count_eq] = {val, idx};
+      }
+    };
+#pragma unroll
+    for (uint32_t i = 0; i < kLocalVecs; ++i) {
+      const auto vi = tx + kBlockSize * i;
+      const auto base = vi * kVecSize;
+      if (vi >= num_full) break;
+#pragma unroll
+      for (uint32_t j = 0; j < kVecSize; ++j)
+        collect(local_vecs[i][j], base + j);
+    }
+    if (tx >= kBlockSize - tail) {
+      const uint32_t idx = tail_start + tx - (kBlockSize - tail);
+      collect(problem.in[idx], idx);
+    }
+
+    // Phase 4: Handle ties.
+    __syncthreads();
+    const auto above_count = smem->count_gt;
+    const auto equal_count = smem->count_eq;
+    const auto remain_topk = above_count < topk ? topk - above_count : 0;
+    const auto tie_count = min(equal_count, kMaxNumTie);
+    handle_tie(smem->tie.values, problem, above_count, tie_count, remain_topk, &smem->tie.handle);
+  }
+};
+
+// ---------------------------------------------------------------------------
+// Streaming path: seq_len > 8192 -- two vectorized passes over global memory
+// ---------------------------------------------------------------------------
+
+struct TopKStreaming : TopKRegister<2> {
+ public:
+  static constexpr uint32_t kMaxSeqLen = std::numeric_limits<uint32_t>::max();
+
+  template <bool kUsePDL>
+  SGL_DEVICE static void forward(const TopKProblem problem, void* _smem) {
+    const auto tx = threadIdx.x;
+    const auto smem = static_cast<Smem*>(_smem);
+
+    {
+      Smem::kHistVec hist_vec;
+      hist_vec.fill(0);
+      smem->hist_vecs[tx] = hist_vec;
+    }
+    if (tx == 0) {
+      smem->count_eq = 0;
+      smem->count_gt = 0;
+    }
+    __syncthreads();
+    PDLWaitPrimary<kUsePDL>();
+
+    // Phase 1: Load and build histogram
+    for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t) {
+      const auto bin = extract_coarse_bin<kHistBits>(val);
+      atomicAdd(&smem->histogram[bin], 1);
+    });
+    __syncthreads();
+
+    // Phase 2: Find the threshold bin
+    find_threshold(problem.topk, problem.seq_len, smem);
+
+    // Phase 3: Collect candidates and sort. Classify by two fp32 boundaries derived
+    // from the threshold bin instead of recomputing the fp16 bin per element: an
+    // element is "above" iff val >= v_hi (bin > threshold) and a "tie" iff
+    // v_lo <= val < v_hi (bin == threshold). This drops the F2F + bit-twiddle from
+    // the second full pass over the input.
+    const auto threshold_bin = smem->threshold_bin;
+    const float v_hi = coarse_bin_lower_bound<kHistBits>(threshold_bin + 1);
+    const float v_lo = coarse_bin_lower_bound<kHistBits>(threshold_bin);
+    const auto topk = problem.topk;
+    for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t idx) {
+      if (val >= v_hi) {
+        const auto pos = atomicAdd(&smem->count_gt, 1);
+        if (pos < topk) [[likely]] {
+          problem.emit(pos, idx);
+        }
+      } else if (val >= v_lo) {
+        const auto count_eq = atomicAdd(&smem->count_eq, 1);
+        if (count_eq < kMaxNumTie) [[likely]] {
+          smem->tie.values[count_eq] = {val, idx};
+        }
+      }
+    });
+
+    // Phase 4: Handle ties. Drive the output layout from the *collect* counts so it
+    // is self-consistent with the fp32 classification above (rather than the fp16
+    // histogram counts), even if rounding moves a boundary element between the
+    // "above" and "tie" sets. above_count is < topk by the threshold-bin invariant,
+    // so the count_gt guard above effectively never triggers.
+    __syncthreads();
+    const auto above_count = smem->count_gt;
+    const auto equal_count = smem->count_eq;
+    const auto remain_topk = above_count < topk ? topk - above_count : 0;
+    const auto tie_count = min(equal_count, kMaxNumTie);
+    handle_tie(smem->tie.values, problem, above_count, tie_count, remain_topk, &smem->tie.handle);
+  }
+};
+
+// ---------------------------------------------------------------------------
+// Cluster path: very long seq_len, small batch. `kClusterSize` blocks cooperate
+// on one batch element via distributed shared memory (one cluster per element).
+// ---------------------------------------------------------------------------
+
+template <uint32_t kClusterSize_>
+struct TopKCluster : TopKRadixBase<10> {
+ public:
+  static constexpr uint32_t kClusterSize = kClusterSize_;
+  static constexpr uint32_t kMaxSeqLen = std::numeric_limits<uint32_t>::max();
+  using Base = TopKRadixBase<10>;
+  struct Smem : Base::Smem {
+    using kHistVec = Base::Smem::kHistVec;
+    uint32_t start_eq_local, start_gt_local;
+    int32_t tmp_out[kMaxTopK];
+  };
+
+  // Process ONE batch element (one cluster). NO PDL and NO trailing barrier --
+  // the persistent kernel does PDLWaitPrimary once before its item loop and a
+  // cluster.sync() after each forward(). Writes raw indices to out; the kernel's
+  // transform pass applies the page-table transform.
+  template <bool kUsePDL>
+  SGL_DEVICE static void forward(TopKProblem problem, void* _smem) {
+    const auto tx = threadIdx.x;
+    const auto smem = static_cast<Smem*>(_smem);
+    const auto cluster = cg::this_cluster();
+    const auto this_rank = blockIdx.y;
+    const bool is_primary = (this_rank == 0);
+
+    constexpr uint32_t kAlignElems = kWarpSize * kVecSize;
+    const uint32_t chunk_size = div_ceil(problem.seq_len, kClusterSize * kAlignElems) * kAlignElems;
+    const uint32_t chunk_start = min(this_rank * chunk_size, problem.seq_len);
+    const uint32_t chunk_finish = min(chunk_start + chunk_size, problem.seq_len);
+    const uint32_t local_seq_len = chunk_finish - chunk_start;
+    problem.in += chunk_start;
+
+    {
+      typename Smem::kHistVec hist_vec;
+      hist_vec.fill(0);
+      smem->hist_vecs[tx] = hist_vec;
+    }
+    if (tx == 0) {
+      smem->count_eq = 0;
+      smem->count_gt = 0;
+    }
+    __syncthreads();
+    PDLWaitPrimary<kUsePDL>();
+
+    // Phase 1: Load and build histogram over this rank's contiguous chunk.
+    for_each_input(problem.in, local_seq_len, [&](float val, uint32_t) {
+      const auto bin = extract_coarse_bin<kHistBits>(val);
+      atomicAdd(&smem->histogram[bin], 1);
+    });
+    __syncthreads();
+
+    // Phase 1.5: reduce the histogram across the cluster
+    {
+      // 1-shot all-reduce: each rank owns kPartition consecutive bins;
+      // for each owned bin, gather the kClusterSize peer values (one per
+      // consecutive lane) via DSMEM, sum across the lanes, then scatter back.
+      cluster.sync();
+      static_assert(kHistSize == kBlockSize);  // we optimize on top of this
+      constexpr uint32_t kPartition = kHistSize / kClusterSize;
+      const auto start = this_rank * kPartition;
+      const auto which = start + tx / kClusterSize;
+      const auto peer_rank = tx % kClusterSize;
+      const auto addr = cluster.map_shared_rank(&smem->histogram[which], peer_rank);
+      const auto value = *addr;
+      *addr = warp::reduce_sum<kClusterSize>(value);
+      cluster.sync();
+    }
+
+    // Phase 2: Find the threshold bin (uses global seq_len)
+    find_threshold(problem.topk, problem.seq_len, smem);
+
+    // Phase 3: Collect candidates over this rank's chunk; convert local indices
+    // back to global by adding chunk_start. Classify by two fp32 boundaries derived
+    // from the (global) threshold bin instead of recomputing the fp16 bin per
+    // element -- see TopKStreaming for the rationale. threshold_bin is identical
+    // across ranks, so v_hi/v_lo are too.
+    const auto topk = problem.topk;
+    const auto threshold_bin = smem->threshold_bin;
+    const float v_hi = coarse_bin_lower_bound<kHistBits>(threshold_bin + 1);
+    const float v_lo = coarse_bin_lower_bound<kHistBits>(threshold_bin);
+    const auto cur_out = is_primary ? problem.out : smem->tmp_out;
+    for_each_input(problem.in, local_seq_len, [&](float val, uint32_t local_idx) {
+      const auto idx = chunk_start + local_idx;
+      if (val >= v_hi) {
+        const auto pos = atomicAdd(&smem->count_gt, 1);
+        if (pos < topk) [[likely]] {
+          // rank 0's slots [0, a0) are final; other ranks stage raw indices and
+          // page-translate them after the cross-rank prefix sum is known.
+          cur_out[pos] = idx;
+        }
+      } else if (val >= v_lo) {
+        const auto count_eq = atomicAdd(&smem->count_eq, 1);
+        if (count_eq < kMaxNumTie) [[likely]] {
+          smem->tie.values[count_eq] = {val, idx};
+        }
+      }
+    });
+
+    // Phase 3.5: write tmp out and exit for non-primary blocks
+    uint32_t start_write = 0;
+    uint32_t num_write = 0;
+    if (!is_primary) {
+      __syncthreads();
+      const auto local_above_count = smem->count_gt;
+      const auto local_equal_count = min(smem->count_eq, kMaxNumTie);
+      const auto smem_0 = cluster.map_shared_rank(smem, 0);
+      if (tx == 0) {
+        const auto gt = atomicAdd(&smem_0->count_gt, local_above_count);
+        const auto eq = atomicAdd(&smem_0->count_eq, local_equal_count);
+        smem->start_gt_local = gt;
+        smem->start_eq_local = eq;
+      }
+      __syncthreads();
+      const auto start_gt_local = smem->start_gt_local;
+      const auto start_eq_local = smem->start_eq_local;
+#pragma unroll
+      for (uint32_t i = 0; i < kTieItems; ++i) {
+        const auto t = tx + i * kBlockSize;
+        if (t < local_equal_count && start_eq_local + t < kMaxNumTie) {
+          smem_0->tie.values[start_eq_local + t] = smem->tie.values[t];
+        }
+      }
+      start_write = start_gt_local;
+      num_write = local_above_count;
+    }
+
+    cluster.sync();
+    if (!is_primary) {
+#pragma unroll
+      for (uint32_t i = 0; i < kTopKItems; ++i) {
+        if (const auto t = tx + i * kBlockSize; t < num_write && start_write + t < topk) {
+          problem.emit(start_write + t, smem->tmp_out[t]);
+        }
+      }
+    } else {
+      // Phase 4: Handle ties.
+      const auto above_count = smem->count_gt;
+      const auto equal_count = smem->count_eq;
+      const auto remain_topk = above_count < topk ? topk - above_count : 0;
+      const auto tie_count = min(equal_count, kMaxNumTie);
+      handle_tie(smem->tie.values, problem, above_count, tie_count, remain_topk, &smem->tie.handle);
+    }
+  }
+};
+
+}  // namespace device::topk

python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh

@@ -0,0 +1,120 @@
+#pragma once
+#include <sgl_kernel/utils.cuh>
+
+namespace device::distributed {
+
+inline constexpr uint32_t kMaxNumGPU = 8;
+
+struct alignas(128) Semaphore {
+ public:
+  constexpr Semaphore() : m_flag(0), m_counter(0) {}
+
+  template <bool kFence>
+  SGL_DEVICE uint32_t get() const {
+    uint32_t val;
+    if constexpr (kFence) {
+      asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(val) : "l"(&m_flag));
+    } else {
+      asm volatile("ld.volatile.global.u32 %0, [%1];" : "=r"(val) : "l"(&m_flag));
+    }
+    return val;
+  }
+
+  template <bool kFence>
+  SGL_DEVICE uint32_t add(uint32_t val) {
+    uint32_t old_val;
+    if constexpr (kFence) {
+      asm volatile("atom.release.sys.global.add.u32 %0, [%1], %2;" : "=r"(old_val) : "l"(&m_flag), "r"(val));
+    } else {
+      asm volatile("atom.global.add.u32 %0, [%1], %2;" : "=r"(old_val) : "l"(&m_flag), "r"(val));
+    }
+    return old_val;
+  }
+
+  // Only called by the owning GPU - plain load is sufficient
+  SGL_DEVICE uint32_t get_counter() const {
+    return m_counter;
+  }
+
+  // Only called by the owning GPU - plain store is sufficient
+  SGL_DEVICE void set_counter(uint32_t val) {
+    m_counter = val;
+  }
+
+ private:
+  uint32_t m_flag;
+  uint32_t m_counter;
+};
+
+struct PullController {
+ public:
+  using SignalType = Semaphore;
+
+  PullController(void** signals, uint32_t num_gpu) {
+    for (uint32_t i = 0; i < num_gpu; ++i) {
+      m_signals[i] = static_cast<Semaphore*>(signals[i]);
+    }
+  }
+
+  /// Synchronize all GPUs.
+  /// When kFence is true, establishes happens-before across GPUs using
+  /// release/acquire semantics, ensuring prior writes are visible system-wide.
+  template <bool kFence, bool kStart>
+  SGL_DEVICE void sync(uint32_t rank, uint32_t num_gpu) const {
+    // For fenced sync: ensure all threads in this block have completed their writes,
+    // so the signaling thread's release carries them transitively.
+    static_assert(!(kFence && kStart), "Start stage does not need to wait fence");
+    if constexpr (kFence || !kStart) __syncthreads();
+    constexpr auto kStage = kStart ? 1 : 2;
+    const auto warp_id = threadIdx.x / kWarpThreads;
+    const auto lane_id = threadIdx.x % kWarpThreads;
+    if (lane_id == 0 && warp_id < num_gpu) {
+      auto& signal = m_signals[warp_id][blockIdx.x];
+      signal.add<kFence>(1);
+      if (warp_id == rank) {
+        const auto target = num_gpu * kStage;
+        /// NOTE: correctness here:
+        /// - base is only read/updated locally by the owning GPU
+        const auto base = signal.get_counter();
+        while (signal.get<kFence>() - base < target)
+          ;
+        if constexpr (!kStart) {
+          signal.set_counter(base + target);
+        }
+      }
+    }
+    if constexpr (kStart) __syncthreads();
+  }
+
+ private:
+  Semaphore* __restrict__ m_signals[kMaxNumGPU];
+};
+
+struct PushController {
+ public:
+  using SignalType = uint32_t;
+  static constexpr int64_t kNumStages = 2;
+
+  PushController(void* ptr) : m_local_signal(static_cast<SignalType*>(ptr)) {}
+
+  SGL_DEVICE SignalType epoch() const {
+    return m_local_signal[blockIdx.x];
+  }
+
+  SGL_DEVICE void exit() const {
+    __syncthreads();
+    if (threadIdx.x == 0) {
+      this->exit_unsafe(blockIdx.x);
+    }
+  }
+
+  SGL_DEVICE void exit_unsafe(uint32_t which) const {
+    auto& signal = m_local_signal[which];
+    signal = (signal + 1) % kNumStages;
+  }
+
+ private:
+  SignalType* m_local_signal;
+};
+
+}  // namespace device::distributed

python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh

@@ -0,0 +1,446 @@
+#pragma once
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <sgl_kernel/distributed/common.cuh>
+
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/tuple.h>
+#include <tvm/ffi/reflection/registry.h>
+
+#include <algorithm>
+#include <array>
+#include <cstdint>
+#include <cstring>
+#include <functional>
+#include <numeric>
+#include <optional>
+#include <span>
+#include <unordered_map>
+#include <vector>
+
+namespace host::distributed {
+
+using device::distributed::PullController, device::distributed::PushController;
+
+struct AllReduceData {
+  constexpr AllReduceData() {}
+  void* __restrict__ input[device::distributed::kMaxNumGPU];
+};
+
+using ExternHandle = tvm::ffi::Array<char>;
+
+inline ExternHandle to_extern_handle(void* ptr) {
+  ExternHandle array;
+  cudaIpcMemHandle_t handle;
+  RuntimeDeviceCheck(cudaIpcGetMemHandle(&handle, ptr));
+  for (size_t i = 0; i < sizeof(handle); ++i) {
+    array.push_back(handle.reserved[i]);
+  }
+  return array;
+}
+
+inline void* from_extern_handle(const ExternHandle& array) {
+  cudaIpcMemHandle_t handle;
+  RuntimeCheck(array.size() == sizeof(handle), "Invalid IPC handle size: ", array.size());
+  for (size_t i = 0; i < sizeof(handle); ++i) {
+    handle.reserved[i] = array[i];
+  }
+  void* ptr;
+  RuntimeDeviceCheck(cudaIpcOpenMemHandle(&ptr, handle, cudaIpcMemLazyEnablePeerAccess));
+  return ptr;
+}
+
+struct HandleHash {
+  std::size_t operator()(const cudaIpcMemHandle_t& handle) const {
+    return std::hash<std::string_view>{}({handle.reserved, sizeof(handle.reserved)});
+  }
+};
+
+struct HandleEqual {
+  bool operator()(const cudaIpcMemHandle_t& a, const cudaIpcMemHandle_t& b) const {
+    return std::memcmp(a.reserved, b.reserved, sizeof(a.reserved)) == 0;
+  }
+};
+
+/**
+ * \brief The control plane of the custom all-reduce implementation.
+ * It manages the internal state and synchronization of the participating GPUs.
+ */
+struct CustomAllReduceBase : public tvm::ffi::Object {
+ public:
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sgl.CustomAllReduce", CustomAllReduceBase, tvm::ffi::Object);
+
+  static constexpr bool _type_mutable = true;
+  using InputPair = tvm::ffi::Tuple<int64_t, ExternHandle>;  // (offset, ipc handle)
+
+  CustomAllReduceBase(
+      uint32_t rank,
+      uint32_t num_gpu,
+      uint32_t max_num_cta_pull,
+      uint32_t max_num_cta_push,
+      int64_t pull_buffer_size,
+      int64_t push_buffer_size,
+      int64_t graph_buffer_count)
+      : m_pull_buffer_bytes(pull_buffer_size),
+        m_push_buffer_bytes(push_buffer_size),
+        m_graph_buffer_count(graph_buffer_count),
+        m_rank(rank),
+        m_num_gpu(num_gpu),
+        m_max_num_cta_pull(max_num_cta_pull),
+        m_max_num_cta_push(max_num_cta_push),
+        // default config for pull kernel, can be updated by `configure()`
+        m_num_cta(max_num_cta_pull),
+        m_cta_size(256) {
+    RuntimeCheck(pull_buffer_size % 128 == 0, "Pull buffer size should be aligned to 128 bytes");
+    RuntimeCheck(push_buffer_size % 128 == 0, "Push buffer size should be aligned to 128 bytes");
+    RuntimeCheck(rank < num_gpu, "Invalid rank: ", rank);
+    const int64_t kU32Max = static_cast<int64_t>(std::numeric_limits<uint32_t>::max());
+    const int64_t push_buffer_size_all = push_all_ranks_bytes();
+    RuntimeCheck(pull_buffer_size <= kU32Max, "Pull buffer size is too large: ", pull_buffer_size);
+    RuntimeCheck(push_buffer_size_all <= kU32Max, "Push buffer size is too large: ", push_buffer_size_all);
+    RuntimeDeviceCheck(cudaMalloc(&m_storage, storage_bytes()));
+  }
+
+  ExternHandle share_storage() {
+    return to_extern_handle(m_storage);
+  }
+
+  tvm::ffi::Array<InputPair> share_graph_inputs() {
+    tvm::ffi::Array<InputPair> result;
+    const auto new_inputs_count = registered_count() - m_cum_registered_count;
+    RuntimeCheck(new_inputs_count >= 0, "Invalid new count: ", new_inputs_count);
+    result.reserve(new_inputs_count);
+    std::unordered_map<void*, ExternHandle> ipc_cache;
+    const auto get_handle = [&](void* ptr) -> ExternHandle {
+      const auto it = ipc_cache.find(ptr);
+      if (it != ipc_cache.end()) return it->second;
+      const auto handle = to_extern_handle(ptr);
+      ipc_cache.try_emplace(ptr, handle);
+      return handle;
+    };
+    for (const auto ptr : std::span(m_graph_capture_inputs).subspan(m_cum_registered_count)) {
+      // note: must share the base address of each allocation, or we get wrong address
+      void* base_ptr;
+      const auto cu_result = cuPointerGetAttribute(&base_ptr, CU_POINTER_ATTRIBUTE_RANGE_START_ADDR, (CUdeviceptr)ptr);
+      RuntimeCheck(cu_result == CUDA_SUCCESS, "failed to get pointer attr");
+      const auto offset = reinterpret_cast<char*>(ptr) - reinterpret_cast<char*>(base_ptr);
+      result.push_back(InputPair{offset, get_handle(base_ptr)});
+    }
+    return result;
+  }
+
+  void post_init(tvm::ffi::Array<ExternHandle> ipc_storages) {
+    RuntimeCheck(ipc_storages.size() == m_num_gpu, "Invalid array size: ", ipc_storages.size());
+    m_peer_storage.resize(m_num_gpu);
+    for (const auto i : irange(m_num_gpu)) {
+      if (i == m_rank) {
+        m_peer_storage[i] = m_storage;
+      } else {
+        m_peer_storage[i] = from_extern_handle(ipc_storages[i]);
+      }
+    }
+
+    // set signal buffer to zero
+    const auto pull_signal = get_pull_signal(m_storage);
+    RuntimeDeviceCheck(cudaMemset(pull_signal, 0, pull_signal_bytes()));
+
+    // update the pull controller and data pointer
+    RuntimeCheck(!m_pull_ctrl.has_value(), "Controller is already initialized");
+    m_pull_ctrl.emplace(m_peer_storage.data(), m_num_gpu);
+    AllReduceData data;
+    for (const auto i : irange(m_num_gpu)) {
+      data.input[i] = get_pull_buffer(m_peer_storage[i]);
+    }
+    const auto default_data_ptr = get_data_ptr();
+    RuntimeDeviceCheck(cudaMemcpy(default_data_ptr, &data, sizeof(AllReduceData), cudaMemcpyHostToDevice));
+
+    // update the push controller and data pointer
+    RuntimeCheck(!m_push_ctrl.has_value(), "Controller is already initialized");
+    const auto push_signal = get_push_signal(m_storage);
+    RuntimeDeviceCheck(cudaMemset(push_signal, 0, push_signal_bytes()));
+    m_push_ctrl.emplace(push_signal);
+    const auto push_buffer = get_push_buffer(m_storage);
+    RuntimeDeviceCheck(cudaMemset(push_buffer, 0, push_all_ranks_bytes()));
+  }
+
+  void register_inputs(tvm::ffi::Array<tvm::ffi::Array<InputPair>> ipc_graph_inputs) {
+    RuntimeCheck(ipc_graph_inputs.size() == m_num_gpu);
+    const auto new_registered_count = registered_count() - m_cum_registered_count;
+    RuntimeCheck(new_registered_count >= 0, "Invalid registered count: ", new_registered_count);
+    if (new_registered_count == 0) return;  // avoid `m_get_data_ptr()` out-of-bounds
+    std::vector<AllReduceData> data;
+    data.resize(new_registered_count);
+    const auto open_cached = [&](const ExternHandle& h) -> void* {
+      RuntimeCheck(h.size() == sizeof(cudaIpcMemHandle_t), "Invalid IPC handle size: ", h.size());
+      cudaIpcMemHandle_t handle;
+      for (size_t i = 0; i < sizeof(handle); ++i)
+        handle.reserved[i] = h[i];
+      const auto [it, success] = m_ipc_cache.try_emplace(handle, nullptr);
+      if (success) {
+        void* ptr;
+        RuntimeDeviceCheck(cudaIpcOpenMemHandle(&ptr, handle, cudaIpcMemLazyEnablePeerAccess));
+        it->second = ptr;
+      }
+      return it->second;
+    };
+    for (const auto i : irange(ipc_graph_inputs.size())) {
+      const auto& array = ipc_graph_inputs[i];
+      RuntimeCheck(int64_t(array.size()) == new_registered_count);
+      if (i == m_rank) {
+        for (const auto j : irange(new_registered_count)) {
+          data[j].input[i] = m_graph_capture_inputs[m_cum_registered_count + j];
+        }
+      } else {
+        for (const auto j : irange(new_registered_count)) {
+          /// NOTE: structural binding will cause intern compiler error...
+          const auto elem = array[j];
+          const auto offset = elem.get<0>();
+          const auto ipc_handle = elem.get<1>();
+          data[j].input[i] = pointer::offset(open_cached(ipc_handle), offset);
+        }
+      }
+    }
+
+    const auto new_registered_bytes = sizeof(AllReduceData) * new_registered_count;
+    const auto dst_ptr = get_data_ptr(m_cum_registered_count);
+    m_cum_registered_count += new_registered_count;
+    RuntimeDeviceCheck(cudaMemcpy(dst_ptr, data.data(), new_registered_bytes, cudaMemcpyHostToDevice));
+  }
+
+  void set_cuda_graph_capture(bool enabled) {
+    m_is_graph_capturing = enabled;
+  }
+
+  tvm::ffi::Array<int64_t> get_graph_capture_ptrs() {
+    tvm::ffi::Array<int64_t> result;
+    const auto new_count = registered_count() - m_cum_registered_count;
+    result.reserve(new_count);
+    for (const auto ptr : std::span(m_graph_capture_inputs).subspan(m_cum_registered_count)) {
+      result.push_back(reinterpret_cast<int64_t>(ptr));
+    }
+    return result;
+  }
+
+  using BaseInfo = tvm::ffi::Tuple<int64_t, int64_t>;  // (base_ptr, size)
+
+  /// Returns (unique_bases, per_input_base_indices, per_input_offset).
+  /// unique_bases[i] = (base_ptr, alloc_size) for each unique allocation.
+  /// per_input_base_indices[j] = indices of VMM allocations covering input j.
+  /// per_input_offset[j] = byte offset from the first allocation base for input j.
+  tvm::ffi::Tuple<tvm::ffi::Array<BaseInfo>, tvm::ffi::Array<tvm::ffi::Array<int64_t>>, tvm::ffi::Array<int64_t>>
+  get_graph_capture_bases() {
+    const auto new_inputs = std::span(m_graph_capture_inputs).subspan(m_cum_registered_count);
+    const auto new_input_bytes = std::span(m_graph_capture_input_bytes).subspan(m_cum_registered_count);
+    std::unordered_map<uintptr_t, int64_t> base_to_idx;
+    tvm::ffi::Array<BaseInfo> bases;
+    tvm::ffi::Array<tvm::ffi::Array<int64_t>> input_indices;
+    tvm::ffi::Array<int64_t> offsets;
+    input_indices.reserve(new_inputs.size());
+    offsets.reserve(new_inputs.size());
+    RuntimeCheck(new_inputs.size() == new_input_bytes.size(), "graph input metadata mismatch");
+    for (const auto input_idx : irange(new_inputs.size())) {
+      const auto ptr = new_inputs[input_idx];
+      auto remaining = new_input_bytes[input_idx];
+      RuntimeCheck(remaining > 0, "Invalid graph capture input size: ", remaining);
+
+      auto cursor = reinterpret_cast<CUdeviceptr>(ptr);
+      CUdeviceptr first_base = 0;
+      tvm::ffi::Array<int64_t> chunks;
+      while (remaining > 0) {
+        CUdeviceptr base = 0;
+        size_t size = 0;
+        const auto r = cuMemGetAddressRange(&base, &size, cursor);
+        RuntimeCheck(r == CUDA_SUCCESS, "cuMemGetAddressRange failed: ", r);
+        if (first_base == 0) first_base = base;
+        const auto byte_offset = static_cast<int64_t>(cursor - base);
+        RuntimeCheck(
+            byte_offset >= 0 && static_cast<size_t>(byte_offset) < size,
+            "graph capture input at ",
+            reinterpret_cast<uintptr_t>(ptr),
+            " is outside VMM allocation [base=",
+            base,
+            ", size=",
+            size,
+            "]");
+
+        auto [it, inserted] = base_to_idx.try_emplace(base, bases.size());
+        if (inserted) {
+          bases.push_back(BaseInfo{static_cast<int64_t>(base), static_cast<int64_t>(size)});
+        }
+        chunks.push_back(it->second);
+
+        const auto available = static_cast<int64_t>(size) - byte_offset;
+        const auto advance = std::min(remaining, available);
+        RuntimeCheck(advance > 0, "Failed to advance VMM graph capture span");
+        remaining -= advance;
+        cursor += advance;
+      }
+      input_indices.push_back(chunks);
+      offsets.push_back(reinterpret_cast<CUdeviceptr>(ptr) - first_base);
+    }
+    using Result =
+        tvm::ffi::Tuple<tvm::ffi::Array<BaseInfo>, tvm::ffi::Array<tvm::ffi::Array<int64_t>>, tvm::ffi::Array<int64_t>>;
+    return Result(bases, input_indices, offsets);
+  }
+
+  void register_peer_mapped_inputs(tvm::ffi::Array<tvm::ffi::Array<int64_t>> peer_ptrs_per_input) {
+    const auto new_count = registered_count() - m_cum_registered_count;
+    RuntimeCheck(int64_t(peer_ptrs_per_input.size()) == new_count, "peer_ptrs count mismatch");
+    if (new_count == 0) return;
+    std::vector<AllReduceData> data(new_count);
+    for (const auto j : irange(new_count)) {
+      const auto& ptrs = peer_ptrs_per_input[j];
+      RuntimeCheck(ptrs.size() == m_num_gpu, "peer count mismatch");
+      for (const auto i : irange(m_num_gpu)) {
+        data[j].input[i] = reinterpret_cast<void*>(static_cast<int64_t>(ptrs[i]));
+      }
+    }
+    const auto dst_ptr = get_data_ptr(m_cum_registered_count);
+    m_cum_registered_count += new_count;
+    RuntimeDeviceCheck(cudaMemcpy(dst_ptr, data.data(), sizeof(AllReduceData) * new_count, cudaMemcpyHostToDevice));
+  }
+
+  void free_ipc_handles() {
+    for (const auto& pair : m_ipc_cache) {
+      host::RuntimeDeviceCheck(cudaIpcCloseMemHandle(pair.second));
+    }
+    m_ipc_cache.clear();
+  }
+
+  void free_storage() {
+    host::RuntimeDeviceCheck(cudaFree(m_storage));
+    m_storage = nullptr;
+  }
+
+  tvm::ffi::Tuple<uint32_t, uint32_t> configure_pull(uint32_t num_cta, uint32_t cta_size) {
+    using host::RuntimeCheck;
+    const auto min_cta_size = m_num_gpu * device::kWarpThreads;
+    RuntimeCheck(num_cta > 0 && num_cta <= m_max_num_cta_pull, "Invalid number of CTAs: ", num_cta);
+    RuntimeCheck(cta_size >= min_cta_size, "Block size must be at least ", min_cta_size);
+    const auto old_num_cta = m_num_cta;
+    const auto old_block_size = m_cta_size;
+    m_num_cta = num_cta;
+    m_cta_size = cta_size;
+    return tvm::ffi::Tuple<uint32_t, uint32_t>{old_num_cta, old_block_size};
+  }
+
+ protected:
+  AllReduceData* allocate_graph_capture_input(void* data_ptr, int64_t input_bytes) {
+    const auto count = registered_count();
+    RuntimeCheck(count < m_graph_buffer_count, "Graph buffer overflow, increase `graph_buffer_count`!");
+    m_graph_capture_inputs.push_back(data_ptr);
+    m_graph_capture_input_bytes.push_back(input_bytes);
+    return get_data_ptr(count);
+  }
+  AllReduceData* get_data_ptr(int64_t which = -1) {
+    const auto count = registered_count();
+    RuntimeCheck(which >= -1 && which < count, "Invalid graph buffer index: ", which, ", count: ", count);
+    const auto start = get_pull_params(m_storage);
+    return static_cast<AllReduceData*>(start) + (1 + which);
+  }
+  int64_t registered_count() const {
+    return static_cast<int64_t>(m_graph_capture_inputs.size());
+  }
+  int64_t pull_signal_bytes() const {
+    return _align_bytes(sizeof(PullController::SignalType) * m_max_num_cta_pull);
+  }
+  int64_t push_signal_bytes() const {
+    return _align_bytes(sizeof(PushController::SignalType) * m_max_num_cta_push);
+  }
+  int64_t graph_param_bytes() const {
+    return _align_bytes(sizeof(AllReduceData) * (1 + m_graph_buffer_count));  // 1 for default
+  }
+  int64_t push_all_ranks_bytes() const {
+    return _align_bytes(PushController::kNumStages * m_num_gpu * m_push_buffer_bytes);
+  }
+  int64_t storage_bytes() const {
+    return _get_offset_impl(5);
+  }
+  void* get_pull_signal(void* ptr) const {
+    return pointer::offset(ptr, _get_offset_impl(0));
+  }
+  void* get_push_signal(void* ptr) const {
+    return pointer::offset(ptr, _get_offset_impl(1));
+  }
+  void* get_pull_params(void* ptr) const {
+    return pointer::offset(ptr, _get_offset_impl(2));
+  }
+  void* get_pull_buffer(void* ptr) const {
+    return pointer::offset(ptr, _get_offset_impl(3));
+  }
+  void* get_push_buffer(void* ptr) const {
+    return pointer::offset(ptr, _get_offset_impl(4));
+  }
+  int64_t _get_offset_impl(int64_t which) const {
+    // | SignalArray (pull + push) | GraphBuffers (pull params) | Buffers (pull + push) |
+    const int64_t offset_map[5] = {
+        /*[0]=*/pull_signal_bytes(),
+        /*[1]=*/push_signal_bytes(),
+        /*[2]=*/graph_param_bytes(),
+        /*[3]=*/m_pull_buffer_bytes,
+        /*[4]=*/push_all_ranks_bytes(),
+    };
+    RuntimeCheck(which >= 0 && which <= 5, "Invalid offset index: ", which);
+    return std::accumulate(offset_map, offset_map + which, int64_t(0));
+  }
+  static int64_t _align_bytes(int64_t size) {
+    return div_ceil(size, 128) * 128;
+  }
+
+  const int64_t m_pull_buffer_bytes;
+  const int64_t m_push_buffer_bytes;
+  const int64_t m_graph_buffer_count;
+  const uint32_t m_rank;
+  const uint32_t m_num_gpu;
+  const uint32_t m_max_num_cta_pull;
+  const uint32_t m_max_num_cta_push;
+  // these 2 config should only affect pull kernel
+  uint32_t m_num_cta;
+  uint32_t m_cta_size;
+  // other states
+  bool m_is_graph_capturing = false;
+  int64_t m_cum_registered_count = 0;
+  std::optional<PullController> m_pull_ctrl;
+  std::optional<PushController> m_push_ctrl;
+  void* m_storage = nullptr;
+  std::vector<void*> m_graph_capture_inputs;
+  std::vector<int64_t> m_graph_capture_input_bytes;
+  std::vector<void*> m_peer_storage;
+  std::unordered_map<cudaIpcMemHandle_t, void*, HandleHash, HandleEqual> m_ipc_cache;
+};
+
+struct CustomAllReduceRef : public tvm::ffi::ObjectRef {
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(CustomAllReduceRef, tvm::ffi::ObjectRef, CustomAllReduceBase);
+};
+
+}  // namespace host::distributed
+
+namespace device::distributed {
+
+template <typename DType2, size_t N, uint32_t M>
+SGL_DEVICE auto reduce_impl(AlignedVector<DType2, N> (&storage)[M]) -> AlignedVector<DType2, N> {
+  fp32x2_t acc[N] = {};
+#pragma unroll  // unroll num gpu
+  for (uint32_t i = 0; i < M; ++i) {
+#pragma unroll  // unroll vec
+    for (uint32_t j = 0; j < N; ++j) {
+      const auto [x, y] = cast<fp32x2_t>(storage[i][j]);
+      auto& [x_acc, y_acc] = acc[j];
+      x_acc += x;
+      y_acc += y;
+    }
+  }
+
+  AlignedVector<DType2, N> result;
+#pragma unroll
+  for (uint32_t j = 0; j < N; ++j) {
+    result[j] = cast<DType2>(acc[j]);
+  }
+
+  return result;
+}
+
+}  // namespace device::distributed

python/sglang/jit_kernel/include/sgl_kernel/ffi.h

@@ -0,0 +1,104 @@
+#pragma once
+#include <sgl_kernel/utils.h>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/shape.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/extra/c_env_api.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <cstdlib>
+#include <memory>
+#include <optional>
+
+namespace host::ffi {
+
+using tvm::ffi::Tensor, tvm::ffi::TensorView, tvm::ffi::ShapeView;
+
+inline Tensor empty(ShapeView shape, DLDataType dtype, DLDevice device) {
+  return Tensor::FromEnvAlloc(::TVMFFIEnvTensorAlloc, shape, dtype, device);
+}
+
+inline Tensor empty_like(TensorView tensor) {
+  return empty(tensor.shape(), tensor.dtype(), tensor.device());
+}
+
+struct _dummy_deleter {
+  void operator()(void*) const {}
+};
+
+// template <typename Fn = _dummy_deleter>
+
+template <typename Fn>
+struct FromBlobContext {
+  [[no_unique_address]] Fn deleter;
+  int64_t dimension;
+  int64_t* get_shape() {
+    return reinterpret_cast<int64_t*>(this + 1);
+  }
+  int64_t* get_stride() {
+    return this->get_shape() + dimension;
+  }
+};
+
+template <typename Fn = _dummy_deleter>
+inline Tensor from_blob(
+    void* data,
+    ShapeView shape,
+    DLDataType dtype,
+    DLDevice device,
+    Fn&& deleter = {},
+    std::optional<ShapeView> stride = {},
+    uint64_t byte_offset = 0) {
+  using Context = FromBlobContext<std::decay_t<Fn>>;
+  const auto ndim = shape.size();
+  const auto ctx = [&] {
+    auto ptr = std::malloc(sizeof(Context) + sizeof(int64_t) * ndim * 2);
+    auto ctx = static_cast<Context*>(ptr);
+    std::construct_at(ctx, std::forward<Fn>(deleter), static_cast<int64_t>(ndim));
+    stdr::copy_n(shape.data(), ndim, ctx->get_shape());
+    if (stride.has_value()) {
+      RuntimeCheck(stride->size() == ndim, "Stride ndim mismatch!");
+      stdr::copy_n(stride->data(), ndim, ctx->get_stride());
+    } else {
+      int64_t stride_val = 1;
+      for (const auto i : irange(ndim)) {
+        const auto j = ndim - 1 - i;
+        ctx->get_stride()[j] = stride_val;
+        stride_val *= shape[j];
+      }
+    }
+    return ctx;
+  }();
+  const auto tensor = DLTensor{
+      .data = data,
+      .device = device,
+      .ndim = static_cast<int32_t>(ndim),
+      .dtype = dtype,
+      .shape = ctx->get_shape(),
+      .strides = ctx->get_stride(),
+      .byte_offset = byte_offset,
+  };
+  const auto blob_deleter = [](DLManagedTensor* self) {
+    auto ctx = static_cast<Context*>(self->manager_ctx);
+    ctx->deleter(self->dl_tensor.data);
+    std::destroy_at(ctx);
+    std::free(ctx);
+  };
+  auto managed_tensor = DLManagedTensor{tensor, ctx, blob_deleter};
+  return Tensor::FromDLPack(&managed_tensor);
+}
+
+template <typename Fn = _dummy_deleter>
+inline Tensor from_blob_like(
+    void* data,
+    TensorView t,
+    Fn&& deleter = {},
+    bool is_contiguous = false,  // if override to true, the stride will be ignored
+    uint64_t byte_offset = 0) {
+  const auto stride = is_contiguous ? std::nullopt : std::optional{t.strides()};
+  return from_blob(data, t.shape(), t.dtype(), t.device(), std::forward<Fn>(deleter), stride, byte_offset);
+}
+
+}  // namespace host::ffi

python/sglang/jit_kernel/include/sgl_kernel/impl/norm.cuh

@@ -0,0 +1,168 @@
+#pragma once
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <cstdint>
+#include <type_traits>
+
+namespace host::norm {
+
+/**
+ * \brief Check if the given configuration is supported.
+ * \tparam T Element type (only fp16_t/bf16_t is supported)
+ * \tparam kDim Dimension size (usually hidden size)
+ */
+template <typename T, int64_t kDim>
+inline constexpr bool is_config_supported() {
+  if (!std::is_same_v<T, fp16_t> && !std::is_same_v<T, bf16_t>) return false;
+  if (kDim <= 256) {
+    return (kDim == 64 || kDim == 128 || kDim == 256);
+  } else {
+    return (kDim % 256 == 0 && kDim <= 8192);
+  }
+}
+
+/**
+ * \brief Determine whether to use cta norm based on dimension size.
+ * TL;DR: use warp norm for dim <= 256, cta norm otherwise.
+ * \tparam T Element type (fp16_t or bf16_t)
+ * \tparam kDim Dimension size (usually hidden size)
+ * \note This function assumes that the configuration is supported.
+ * \see `is_config_supported`
+ */
+template <typename T, int64_t kDim>
+inline constexpr bool should_use_cta() {
+  static_assert(is_config_supported<T, kDim>(), "Unsupported norm configuration");
+  return kDim > 256;
+}
+
+/**
+ * \brief Get the number of threads per CTA for cta norm.
+ * \tparam T Element type (fp16_t or bf16_t)
+ * \tparam kDim Dimension size (usually hidden size)
+ * \return Number of threads per CTA
+ */
+template <typename T, int64_t kDim>
+inline constexpr uint32_t get_cta_threads() {
+  static_assert(should_use_cta<T, kDim>());
+  return (kDim / 256) * device::kWarpThreads;
+}
+
+}  // namespace host::norm
+
+namespace device::norm {
+
+namespace details {
+
+template <int64_t kDim, bool kUseCTA, typename PackedFloat, std::size_t N>
+SGL_DEVICE AlignedVector<PackedFloat, N> apply_norm_impl(
+    const AlignedVector<PackedFloat, N> input,
+    const AlignedVector<PackedFloat, N> weight,
+    const float eps,
+    [[maybe_unused]] float* smem_buffer,
+    [[maybe_unused]] uint32_t num_warps) {
+  float sum_of_squares = 0.0f;
+
+#pragma unroll
+  for (auto i = 0u; i < N; ++i) {
+    const auto fp32_input = cast<fp32x2_t>(input[i]);
+    sum_of_squares += fp32_input.x * fp32_input.x;
+    sum_of_squares += fp32_input.y * fp32_input.y;
+  }
+
+  sum_of_squares = warp::reduce_sum(sum_of_squares);
+  float norm_factor;
+  if constexpr (kUseCTA) {
+    // need to synchronize across the cta
+    const auto warp_id = threadIdx.x / kWarpThreads;
+    smem_buffer[warp_id] = sum_of_squares;
+    __syncthreads();
+    // use the first warp to reduce
+    if (warp_id == 0) {
+      const auto tx = threadIdx.x;
+      const auto local_sum = tx < num_warps ? smem_buffer[tx] : 0.0f;
+      sum_of_squares = warp::reduce_sum(local_sum);
+      smem_buffer[32] = math::rsqrt(sum_of_squares / kDim + eps);
+    }
+    __syncthreads();
+    norm_factor = smem_buffer[32];
+  } else {
+    norm_factor = math::rsqrt(sum_of_squares / kDim + eps);
+  }
+
+  AlignedVector<PackedFloat, N> output;
+
+#pragma unroll
+  for (auto i = 0u; i < N; ++i) {
+    const auto fp32_input = cast<fp32x2_t>(input[i]);
+    const auto fp32_weight = cast<fp32x2_t>(weight[i]);
+    output[i] = cast<PackedFloat, fp32x2_t>({
+        fp32_input.x * norm_factor * fp32_weight.x,
+        fp32_input.y * norm_factor * fp32_weight.y,
+    });
+  }
+
+  return output;
+}
+
+}  // namespace details
+
+/**
+ * \brief Apply norm using warp-level implementation.
+ * \tparam kDim Dimension size
+ * \tparam T Element type (fp16_t or bf16_t)
+ * \param input Input vector
+ * \param weight Weight vector
+ * \param eps Epsilon value for numerical stability
+ * \return Normalized output vector
+ */
+template <int64_t kDim, typename T>
+SGL_DEVICE T apply_norm_warp(const T& input, const T& weight, float eps) {
+  static_assert(kDim <= 256, "Warp norm only supports dim <= 256");
+  return details::apply_norm_impl<kDim, false>(input, weight, eps, nullptr, 0);
+}
+
+/**
+ * \brief Apply norm using CTA-level implementation.
+ * \tparam kDim Dimension size
+ * \tparam T Element type (fp16_t or bf16_t)
+ * \param input Input vector
+ * \param weight Weight vector
+ * \param eps Epsilon value for numerical stability
+ * \param smem Shared memory buffer
+ * \param num_warps Number of warps in the CTA
+ * \return Normalized output vector
+ */
+template <int64_t kDim, typename T>
+SGL_DEVICE T apply_norm_cta(
+    const T& input, const T& weight, float eps, float* smem, uint32_t num_warps = blockDim.x / kWarpThreads) {
+  static_assert(kDim > 256, "CTA norm only supports dim > 256");
+  return details::apply_norm_impl<kDim, true>(input, weight, eps, smem, num_warps);
+}
+
+/**
+ * \brief Storage type for norm operation.
+ * For warp norm, the storage size depends on kDim.
+ * For cta norm, the storage size is fixed to 16B.
+ * We will also pack the input 16-bit floats into 32-bit types
+ * for faster CUDA core operations.
+ *
+ * \tparam T Element type (fp16_t or bf16_t)
+ * \tparam kDim Dimension size
+ */
+template <typename T, int64_t kDim>
+using StorageType = std::conditional_t<                    // storage type
+    (kDim > 256),                                          // whether to use cta norm
+    AlignedVector<packed_t<T>, 4>,                         // cta norm storage, fixed to 16B
+    AlignedVector<packed_t<T>, kDim / (2 * kWarpThreads)>  // warp norm storage
+    >;
+
+/**
+ * \brief Minimum shared memory size (in bytes) required for cta norm.
+ */
+inline constexpr uint32_t kSmemBufferSize = 33;
+
+}  // namespace device::norm

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

@@ -0,0 +1,75 @@
+/// \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.
+/// Arch-aware: 448 on CUDA / AMD OCP e4m3fn (gfx950), 224 on AMD e4m3fnuz
+/// (gfx942). Mirrors kFP8E4M3Max so fp8 quant scale divisors and clamps in
+/// the dsv4 compute path (indexer Q-quant, MoE silu+mul / dispatch quant,
+/// GEMM per-tensor quant) do not over-saturate fnuz hardware.
+inline constexpr float FP8_E4M3_MAX = ::kFP8E4M3Max;
+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);
+}
+
+}  // namespace device::math

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

@@ -0,0 +1,86 @@
+/// \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>
+
+#include <cstddef>
+#include <cstdint>
+#ifndef USE_ROCM
+#include <cuda_runtime.h>
+#else
+#include <hip/hip_runtime.h>
+#ifndef cudaOccupancyMaxActiveBlocksPerMultiprocessor
+#define cudaOccupancyMaxActiveBlocksPerMultiprocessor hipOccupancyMaxActiveBlocksPerMultiprocessor
+#endif
+#ifndef cudaDeviceGetAttribute
+#define cudaDeviceGetAttribute hipDeviceGetAttribute
+#endif
+#ifndef cudaDevAttrMultiProcessorCount
+#define cudaDevAttrMultiProcessorCount hipDeviceAttributeMultiprocessorCount
+#endif
+#ifndef cudaDevAttrComputeCapabilityMajor
+#define cudaDevAttrComputeCapabilityMajor hipDeviceAttributeComputeCapabilityMajor
+#endif
+#ifndef cudaRuntimeGetVersion
+#define cudaRuntimeGetVersion hipRuntimeGetVersion
+#endif
+#ifndef cudaOccupancyAvailableDynamicSMemPerBlock
+inline hipError_t
+cudaOccupancyAvailableDynamicSMemPerBlock(std::size_t* smem, const void* func, int num_blocks, int block_size) {
+  // HIP does not expose this directly; return max shared mem as conservative estimate
+  hipDeviceProp_t prop;
+  int device;
+  hipGetDevice(&device);
+  hipGetDeviceProperties(&prop, device);
+  *smem = prop.sharedMemPerBlock;
+  return hipSuccess;
+}
+#endif
+#endif
+
+namespace host::runtime {
+
+// Return the maximum number of active blocks per SM for the given kernel
+template <typename T>
+inline auto get_blocks_per_sm(T&& kernel, int32_t block_dim, std::size_t dynamic_smem = 0) -> uint32_t {
+  int num_blocks_per_sm = 0;
+  RuntimeDeviceCheck(
+      cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks_per_sm, kernel, block_dim, dynamic_smem));
+  return static_cast<uint32_t>(num_blocks_per_sm);
+}
+
+// Return the number of SMs for the given device
+inline auto get_sm_count(int device_id) -> uint32_t {
+  int sm_count;
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device_id));
+  return static_cast<uint32_t>(sm_count);
+}
+
+// Return the Major compute capability for the given device
+inline auto get_cc_major(int device_id) -> int {
+  int cc_major;
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&cc_major, cudaDevAttrComputeCapabilityMajor, device_id));
+  return cc_major;
+}
+
+// Return the runtime version
+inline auto get_runtime_version() -> int {
+  int runtime_version;
+  RuntimeDeviceCheck(cudaRuntimeGetVersion(&runtime_version));
+  return runtime_version;
+}
+
+// Return the maximum dynamic shared memory per block for the given kernel
+template <typename T>
+inline auto get_available_dynamic_smem_per_block(T&& kernel, int num_blocks, int block_size) -> std::size_t {
+  std::size_t smem_size;
+  RuntimeDeviceCheck(cudaOccupancyAvailableDynamicSMemPerBlock(&smem_size, kernel, num_blocks, block_size));
+  return smem_size;
+}
+
+}  // namespace host::runtime

python/sglang/jit_kernel/include/sgl_kernel/scalar_type.hpp

@@ -0,0 +1,334 @@
+#pragma once
+
+#include <cassert>
+#include <stdexcept>
+#ifndef __CUDACC__
+#include <variant>
+#endif
+
+namespace host {
+
+//
+//  ScalarType can represent a wide range of floating point and integer types,
+//  in particular it can be used to represent sub-byte data types (something
+//  that torch.dtype currently does not support).
+//
+//  The type definitions on the Python side can be found in: vllm/scalar_type.py
+//  these type definitions should be kept up to date with any Python API changes
+//  here.
+//
+class ScalarType {
+ public:
+  enum NanRepr : uint8_t {
+    NAN_NONE = 0,                // nans are not supported
+    NAN_IEEE_754 = 1,            // nans are: exp all 1s, mantissa not all 0s
+    NAN_EXTD_RANGE_MAX_MIN = 2,  // nans are: exp all 1s, mantissa all 1s
+
+    NAN_REPR_ID_MAX
+  };
+
+  constexpr ScalarType(
+      uint8_t exponent,
+      uint8_t mantissa,
+      bool signed_,
+      int32_t bias,
+      bool finite_values_only = false,
+      NanRepr nan_repr = NAN_IEEE_754)
+      : exponent(exponent),
+        mantissa(mantissa),
+        signed_(signed_),
+        bias(bias),
+        finite_values_only(finite_values_only),
+        nan_repr(nan_repr) {};
+
+  static constexpr ScalarType int_(uint8_t size_bits, int32_t bias = 0) {
+    return ScalarType(0, size_bits - 1, true, bias);
+  }
+
+  static constexpr ScalarType uint(uint8_t size_bits, int32_t bias = 0) {
+    return ScalarType(0, size_bits, false, bias);
+  }
+
+  // IEEE 754 compliant floating point type
+  static constexpr ScalarType float_IEEE754(uint8_t exponent, uint8_t mantissa) {
+    assert(mantissa > 0 && exponent > 0);
+    return ScalarType(exponent, mantissa, true, 0, false, NAN_IEEE_754);
+  }
+
+  // IEEE 754 non-compliant floating point type
+  static constexpr ScalarType float_(uint8_t exponent, uint8_t mantissa, bool finite_values_only, NanRepr nan_repr) {
+    assert(nan_repr < NAN_REPR_ID_MAX);
+    assert(mantissa > 0 && exponent > 0);
+    assert(nan_repr != NAN_IEEE_754);
+    return ScalarType(exponent, mantissa, true, 0, finite_values_only, nan_repr);
+  }
+
+  uint8_t const exponent;  // size of the exponent field (0 for integer types)
+  uint8_t const mantissa;  // size of the mantissa field (size of the integer
+                           // excluding the sign bit for integer types)
+  bool const signed_;      // flag if the type supports negative numbers (i.e. has a
+                           // sign bit)
+  int32_t const bias;      // stored values equal value + bias,
+                           // used for quantized type
+
+  // Extra Floating point info
+  bool const finite_values_only;  // i.e. no +/-inf if true
+  NanRepr const nan_repr;         // how NaNs are represented
+                                  // (not applicable for integer types)
+
+  using Id = int64_t;
+
+ private:
+  // Field size in id
+  template <typename T_>
+  static constexpr size_t member_id_field_width() {
+    using T = std::decay_t<T_>;
+    return std::is_same_v<T, bool> ? 1 : sizeof(T) * 8;
+  }
+
+  template <typename Fn, typename Init, typename Member, typename... Rest>
+  static constexpr auto reduce_members_helper(Fn f, Init val, Member member, Rest... rest) {
+    auto new_val = f(val, member);
+    if constexpr (sizeof...(rest) > 0) {
+      return reduce_members_helper(f, new_val, rest...);
+    } else {
+      return new_val;
+    };
+  }
+
+  template <typename Fn, typename Init>
+  constexpr auto reduce_members(Fn f, Init init) const {
+    // Should be in constructor order for `from_id`
+    return reduce_members_helper(f, init, exponent, mantissa, signed_, bias, finite_values_only, nan_repr);
+  };
+
+  template <typename Fn, typename Init>
+  static constexpr auto reduce_member_types(Fn f, Init init) {
+    constexpr auto dummy_type = ScalarType(0, 0, false, 0, false, NAN_NONE);
+    return dummy_type.reduce_members(f, init);
+  };
+
+  static constexpr auto id_size_bits() {
+    return reduce_member_types(
+        [](int acc, auto member) -> int { return acc + member_id_field_width<decltype(member)>(); }, 0);
+  }
+
+ public:
+  // unique id for this scalar type that can be computed at compile time for
+  //  c++17 template specialization this is not needed once we migrate to
+  //  c++20 and can pass literal classes as template parameters
+  constexpr Id id() const {
+    static_assert(id_size_bits() <= sizeof(Id) * 8, "ScalarType id is too large to be stored");
+
+    auto or_and_advance = [](std::pair<Id, uint32_t> result, auto member) -> std::pair<Id, uint32_t> {
+      auto [id, bit_offset] = result;
+      auto constexpr bits = member_id_field_width<decltype(member)>();
+      return {id | (int64_t(member) & ((uint64_t(1) << bits) - 1)) << bit_offset, bit_offset + bits};
+    };
+    return reduce_members(or_and_advance, std::pair<Id, uint32_t>{}).first;
+  }
+
+  // create a ScalarType from an id, for c++17 template specialization,
+  //  this is not needed once we migrate to c++20 and can pass literal
+  //  classes as template parameters
+  static constexpr ScalarType from_id(Id id) {
+    auto extract_and_advance = [id](auto result, auto member) {
+      using T = decltype(member);
+      auto [tuple, bit_offset] = result;
+      auto constexpr bits = member_id_field_width<T>();
+      auto extracted_val = static_cast<T>((int64_t(id) >> bit_offset) & ((uint64_t(1) << bits) - 1));
+      auto new_tuple = std::tuple_cat(tuple, std::make_tuple(extracted_val));
+      return std::pair<decltype(new_tuple), int>{new_tuple, bit_offset + bits};
+    };
+
+    auto [tuple_args, _] = reduce_member_types(extract_and_advance, std::pair<std::tuple<>, int>{});
+    return std::apply([](auto... args) { return ScalarType(args...); }, tuple_args);
+  }
+
+  constexpr int64_t size_bits() const {
+    return mantissa + exponent + is_signed();
+  }
+  constexpr bool is_signed() const {
+    return signed_;
+  }
+  constexpr bool is_integer() const {
+    return exponent == 0;
+  }
+  constexpr bool is_floating_point() const {
+    return exponent > 0;
+  }
+  constexpr bool is_ieee_754() const {
+    return is_floating_point() && finite_values_only == false && nan_repr == NAN_IEEE_754;
+  }
+  constexpr bool has_nans() const {
+    return is_floating_point() && nan_repr != NAN_NONE;
+  }
+  constexpr bool has_infs() const {
+    return is_floating_point() && finite_values_only == false;
+  }
+  constexpr bool has_bias() const {
+    return bias != 0;
+  }
+
+#ifndef __CUDACC__
+ private:
+  double _floating_point_max() const {
+    assert(mantissa <= 52 && exponent <= 11);
+
+    uint64_t max_mantissa = (uint64_t(1) << mantissa) - 1;
+    if (nan_repr == NAN_EXTD_RANGE_MAX_MIN) {
+      max_mantissa -= 1;
+    }
+
+    uint64_t max_exponent = (uint64_t(1) << exponent) - 2;
+    if (nan_repr == NAN_EXTD_RANGE_MAX_MIN || nan_repr == NAN_NONE) {
+      assert(exponent < 11);
+      max_exponent += 1;
+    }
+
+    // adjust the exponent to match that of a double
+    //  for now we assume the exponent bias is the standard 2^(e-1) -1, (where e
+    //  is the exponent bits), there is some precedent for non-standard biases,
+    //  example `float8_e4m3b11fnuz` here: https://github.com/jax-ml/ml_dtypes
+    //  but to avoid premature over complication we are just assuming the
+    //  standard exponent bias until there is a need to support non-standard
+    //  biases
+    uint64_t exponent_bias = (uint64_t(1) << (exponent - 1)) - 1;
+    uint64_t exponent_bias_double = (uint64_t(1) << 10) - 1;  // double e = 11
+
+    uint64_t max_exponent_double = max_exponent - exponent_bias + exponent_bias_double;
+
+    // shift the mantissa into the position for a double and
+    // the exponent
+    uint64_t double_raw = (max_mantissa << (52 - mantissa)) | (max_exponent_double << 52);
+
+    return *reinterpret_cast<double*>(&double_raw);
+  }
+
+  constexpr std::variant<int64_t, double> _raw_max() const {
+    if (is_floating_point()) {
+      return {_floating_point_max()};
+    } else {
+      assert(size_bits() < 64 || (size_bits() == 64 && is_signed()));
+      return {(int64_t(1) << mantissa) - 1};
+    }
+  }
+
+  constexpr std::variant<int64_t, double> _raw_min() const {
+    if (is_floating_point()) {
+      assert(is_signed());
+      constexpr uint64_t sign_bit_double = (uint64_t(1) << 63);
+
+      double max = _floating_point_max();
+      uint64_t max_raw = *reinterpret_cast<uint64_t*>(&max);
+      uint64_t min_raw = max_raw | sign_bit_double;
+      return {*reinterpret_cast<double*>(&min_raw)};
+    } else {
+      assert(!is_signed() || size_bits() <= 64);
+      if (is_signed()) {
+        // set the top bit to 1 (i.e. INT64_MIN) and the rest to 0
+        // then perform an arithmetic shift right to set all the bits above
+        // (size_bits() - 1) to 1
+        return {INT64_MIN >> (64 - size_bits())};
+      } else {
+        return {int64_t(0)};
+      }
+    }
+  }
+
+ public:
+  // Max representable value for this scalar type.
+  // (accounting for bias if there is one)
+  constexpr std::variant<int64_t, double> max() const {
+    return std::visit([this](auto x) -> std::variant<int64_t, double> { return {x - bias}; }, _raw_max());
+  }
+
+  // Min representable value for this scalar type.
+  // (accounting for bias if there is one)
+  constexpr std::variant<int64_t, double> min() const {
+    return std::visit([this](auto x) -> std::variant<int64_t, double> { return {x - bias}; }, _raw_min());
+  }
+#endif  // __CUDACC__
+
+ public:
+  std::string str() const {
+    /* naming generally follows: https://github.com/jax-ml/ml_dtypes
+     * for floating point types (leading f) the scheme is:
+     *  `float<size_bits>_e<exponent_bits>m<mantissa_bits>[flags]`
+     *  flags:
+     *  - no-flags: means it follows IEEE 754 conventions
+     *  - f: means finite values only (no infinities)
+     *  - n: means nans are supported (non-standard encoding)
+     * for integer types the scheme is:
+     *  `[u]int<size_bits>[b<bias>]`
+     *  - if bias is not present it means its zero
+     */
+    if (is_floating_point()) {
+      auto ret =
+          "float" + std::to_string(size_bits()) + "_e" + std::to_string(exponent) + "m" + std::to_string(mantissa);
+      if (!is_ieee_754()) {
+        if (finite_values_only) {
+          ret += "f";
+        }
+        if (nan_repr != NAN_NONE) {
+          ret += "n";
+        }
+      }
+      return ret;
+    } else {
+      auto ret = ((is_signed()) ? "int" : "uint") + std::to_string(size_bits());
+      if (has_bias()) {
+        ret += "b" + std::to_string(bias);
+      }
+      return ret;
+    }
+  }
+
+  constexpr bool operator==(ScalarType const& other) const {
+    return mantissa == other.mantissa && exponent == other.exponent && bias == other.bias && signed_ == other.signed_ &&
+           finite_values_only == other.finite_values_only && nan_repr == other.nan_repr;
+  }
+};
+
+using ScalarTypeId = ScalarType::Id;
+
+// "rust style" names generally following:
+//   https://github.com/pytorch/pytorch/blob/6d9f74f0af54751311f0dd71f7e5c01a93260ab3/torch/csrc/api/include/torch/types.h#L60-L70
+static inline constexpr auto kS4 = ScalarType::int_(4);
+static inline constexpr auto kU4 = ScalarType::uint(4);
+static inline constexpr auto kU4B8 = ScalarType::uint(4, 8);
+static inline constexpr auto kS8 = ScalarType::int_(8);
+static inline constexpr auto kU8 = ScalarType::uint(8);
+static inline constexpr auto kU8B128 = ScalarType::uint(8, 128);
+
+static inline constexpr auto kFE2M1f = ScalarType::float_(2, 1, true, ScalarType::NAN_NONE);
+static inline constexpr auto kFE3M2f = ScalarType::float_(3, 2, true, ScalarType::NAN_NONE);
+static inline constexpr auto kFE4M3fn = ScalarType::float_(4, 3, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN);
+static inline constexpr auto kFE8M0fnu = ScalarType(8, 0, false, 0, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN);
+static inline constexpr auto kFE5M2 = ScalarType::float_IEEE754(5, 2);
+static inline constexpr auto kFE8M7 = ScalarType::float_IEEE754(8, 7);
+static inline constexpr auto kFE5M10 = ScalarType::float_IEEE754(5, 10);
+
+// Fixed width style names, generally following:
+//  https://github.com/pytorch/pytorch/blob/6d9f74f0af54751311f0dd71f7e5c01a93260ab3/torch/csrc/api/include/torch/types.h#L47-L57
+static inline constexpr auto kInt4 = kS4;
+static inline constexpr auto kUint4 = kU4;
+static inline constexpr auto kUint4b8 = kU4B8;
+static inline constexpr auto kInt8 = kS8;
+static inline constexpr auto kUint8 = kU8;
+static inline constexpr auto kUint8b128 = kU8B128;
+
+static inline constexpr auto kFloat4_e2m1f = kFE2M1f;
+static inline constexpr auto kFloat6_e3m2f = kFE3M2f;
+static inline constexpr auto kFloat8_e4m3fn = kFE4M3fn;
+static inline constexpr auto kFloat8_e5m2 = kFE5M2;
+static inline constexpr auto kFloat16_e8m7 = kFE8M7;
+static inline constexpr auto kFloat16_e5m10 = kFE5M10;
+
+// colloquial names
+static inline constexpr auto kHalf = kFE5M10;
+static inline constexpr auto kFloat16 = kHalf;
+static inline constexpr auto kBFloat16 = kFE8M7;
+
+static inline constexpr auto kFloat16Id = kFloat16.id();
+}  // namespace host

python/sglang/jit_kernel/include/sgl_kernel/source_location.h

@@ -0,0 +1,40 @@
+/// \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>
+
+/// 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/tensor.h

@@ -0,0 +1,605 @@
+/// \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>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/dtype.h>
+
+#include <algorithm>
+#include <array>
+#include <concepts>
+#include <cstddef>
+#include <cstdint>
+#include <initializer_list>
+#include <optional>
+#include <ranges>
+#include <span>
+#include <sstream>
+#include <string>
+#include <string_view>
+#include <type_traits>
+#include <utility>
+
+#ifdef __CUDACC__
+#include <sgl_kernel/utils.cuh>
+#elif defined(__HIPCC__)
+#include <sgl_kernel/utils.cuh>
+#endif
+
+namespace host {
+
+namespace details {
+
+inline constexpr auto kAnyDeviceID = -1;
+inline constexpr auto kAnySize = static_cast<int64_t>(-1);
+inline constexpr auto kNullSize = static_cast<int64_t>(-1);
+inline constexpr auto kNullDType = static_cast<DLDataTypeCode>(18u);
+inline constexpr auto kNullDevice = static_cast<DLDeviceType>(-1);
+
+struct SizeRef;
+struct DTypeRef;
+struct DeviceRef;
+
+template <typename T>
+struct _dtype_trait {};
+
+template <std::integral T>
+struct _dtype_trait<T> {
+  inline static constexpr DLDataType value = {
+      .code = std::is_signed_v<T> ? DLDataTypeCode::kDLInt : DLDataTypeCode::kDLUInt,
+      .bits = static_cast<std::uint8_t>(sizeof(T) * 8),
+      .lanes = 1};
+};
+
+template <std::floating_point T>
+struct _dtype_trait<T> {
+  inline static constexpr DLDataType value = {
+      .code = DLDataTypeCode::kDLFloat, .bits = static_cast<std::uint8_t>(sizeof(T) * 8), .lanes = 1};
+};
+
+#ifdef __CUDACC__
+template <>
+struct _dtype_trait<fp16_t> {
+  inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLFloat, .bits = 16, .lanes = 1};
+};
+template <>
+struct _dtype_trait<bf16_t> {
+  inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLBfloat, .bits = 16, .lanes = 1};
+};
+template <>
+struct _dtype_trait<fp8_e4m3_t> {
+  inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLFloat8_e4m3fn, .bits = 8, .lanes = 1};
+};
+#elif defined(__HIPCC__)
+template <>
+struct _dtype_trait<fp16_t> {
+  inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLFloat, .bits = 16, .lanes = 1};
+};
+template <>
+struct _dtype_trait<bf16_t> {
+  inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLBfloat, .bits = 16, .lanes = 1};
+};
+#endif
+
+template <DLDeviceType Code>
+struct _device_trait {
+  inline static constexpr DLDevice value = {.device_type = Code, .device_id = kAnyDeviceID};
+};
+
+template <typename... Ts>
+inline constexpr auto kDTypeList = std::array<DLDataType, sizeof...(Ts)>{_dtype_trait<Ts>::value...};
+
+template <DLDeviceType... Codes>
+inline constexpr auto kDeviceList = std::array<DLDevice, sizeof...(Codes)>{_device_trait<Codes>::value...};
+
+template <typename T>
+struct PrintAbleSpan {
+  explicit PrintAbleSpan(std::span<const T> data) : data(data) {}
+  std::span<const T> data;
+};
+
+// define DLDataType comparison and printing in root namespace
+inline constexpr auto kDeviceStringMap = [] {
+  constexpr auto map = std::array<std::pair<DLDeviceType, const char*>, 16>{
+      std::pair{DLDeviceType::kDLCPU, "cpu"},
+      std::pair{DLDeviceType::kDLCUDA, "cuda"},
+      std::pair{DLDeviceType::kDLCUDAHost, "cuda_host"},
+      std::pair{DLDeviceType::kDLOpenCL, "opencl"},
+      std::pair{DLDeviceType::kDLVulkan, "vulkan"},
+      std::pair{DLDeviceType::kDLMetal, "metal"},
+      std::pair{DLDeviceType::kDLVPI, "vpi"},
+      std::pair{DLDeviceType::kDLROCM, "rocm"},
+      std::pair{DLDeviceType::kDLROCMHost, "rocm_host"},
+      std::pair{DLDeviceType::kDLExtDev, "ext_dev"},
+      std::pair{DLDeviceType::kDLCUDAManaged, "cuda_managed"},
+      std::pair{DLDeviceType::kDLOneAPI, "oneapi"},
+      std::pair{DLDeviceType::kDLWebGPU, "webgpu"},
+      std::pair{DLDeviceType::kDLHexagon, "hexagon"},
+      std::pair{DLDeviceType::kDLMAIA, "maia"},
+      std::pair{DLDeviceType::kDLTrn, "trn"},
+  };
+  constexpr auto max_type = stdr::max(map | stdv::keys);
+  auto result = std::array<std::string_view, max_type + 1>{};
+  for (const auto& [code, name] : map) {
+    result[static_cast<std::size_t>(code)] = name;
+  }
+  return result;
+}();
+
+struct PrintableDevice {
+  DLDevice device;
+};
+
+inline auto& operator<<(std::ostream& os, DLDevice device) {
+  const auto& mapping = kDeviceStringMap;
+  const auto entry = static_cast<std::size_t>(device.device_type);
+  RuntimeCheck(entry < mapping.size());
+  const auto name = mapping[entry];
+  RuntimeCheck(!name.empty(), "Unknown device: ", int(device.device_type));
+  os << name;
+  if (device.device_id != kAnyDeviceID && device.device_type != DLDeviceType::kDLCPU) {
+    os << ":" << device.device_id;
+  }
+  return os;
+}
+
+inline auto& operator<<(std::ostream& os, PrintableDevice pd) {
+  return os << pd.device;
+}
+
+template <typename T>
+inline auto& operator<<(std::ostream& os, PrintAbleSpan<T> span) {
+  os << "[";
+  for (const auto i : irange(span.data.size())) {
+    if (i > 0) {
+      os << ", ";
+    }
+    os << span.data[i];
+  }
+  os << "]";
+  return os;
+}
+
+}  // 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) {}
+  SymbolicSize(const SymbolicSize&) = delete;
+  SymbolicSize& operator=(const SymbolicSize&) = delete;
+
+  auto get_name() const -> std::string_view {
+    return m_annotation;
+  }
+
+  auto set_value(int64_t value) -> void {
+    RuntimeCheck(!this->has_value(), "Size value already set");
+    m_value = value;
+  }
+
+  auto has_value() const -> bool {
+    return m_value != details::kNullSize;
+  }
+
+  auto get_value() const -> std::optional<int64_t> {
+    return this->has_value() ? std::optional{m_value} : std::nullopt;
+  }
+
+  auto unwrap(DebugInfo info = {}) const -> int64_t {
+    RuntimeCheck(info, this->has_value(), "Size value is not set");
+    return m_value;
+  }
+
+  auto verify(int64_t value, const char* prefix, int64_t dim) -> void {
+    if (this->has_value()) {
+      if (m_value != value) {
+        [[unlikely]];
+        Panic("Size mismatch for ", m_name_str(prefix, dim), ": expected ", m_value, " but got ", value);
+      }
+    } else {
+      this->set_value(value);
+    }
+  }
+
+  auto value_or_name(const char* prefix, int64_t dim) const -> std::string {
+    if (const auto value = this->get_value()) {
+      return std::to_string(*value);
+    } else {
+      return m_name_str(prefix, dim);
+    }
+  }
+
+ private:
+  auto m_name_str(const char* prefix, int64_t dim) const -> std::string {
+    std::ostringstream os;
+    os << prefix << '#' << dim;
+    if (!m_annotation.empty()) os << "('" << m_annotation << "')";
+    return std::move(os).str();
+  }
+
+  std::int64_t m_value;
+  std::string_view m_annotation;
+};
+
+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}) {}
+  SymbolicDType(const SymbolicDType&) = delete;
+  SymbolicDType& operator=(const SymbolicDType&) = delete;
+
+  auto set_value(DLDataType value) -> void {
+    RuntimeCheck(!this->has_value(), "Dtype value already set");
+    RuntimeCheck(
+        m_check(value), "Dtype value [", value, "] not in the allowed options: ", details::PrintAbleSpan{m_options});
+    m_value = value;
+  }
+
+  auto has_value() const -> bool {
+    return m_value.code != details::kNullDType;
+  }
+
+  auto get_value() const -> std::optional<DLDataType> {
+    return this->has_value() ? std::optional{m_value} : std::nullopt;
+  }
+
+  auto unwrap(DebugInfo info = {}) const -> DLDataType {
+    RuntimeCheck(info, this->has_value(), "Dtype value is not set");
+    return m_value;
+  }
+
+  auto set_options(std::span<const DLDataType> options) -> void {
+    m_options = options;
+  }
+
+  template <typename... Ts>
+  auto set_options() -> void {
+    m_options = details::kDTypeList<Ts...>;
+  }
+
+  auto verify(DLDataType dtype) -> void {
+    if (this->has_value()) {
+      RuntimeCheck(m_value == dtype, "DType mismatch: expected ", m_value, " but got ", dtype);
+    } else {
+      this->set_value(dtype);
+    }
+  }
+
+  template <typename T>
+  auto is_type() const -> bool {
+    return ::host::is_type<T>(m_value);
+  }
+
+ private:
+  auto m_check(DLDataType value) const -> bool {
+    return stdr::empty(m_options) || (stdr::find(m_options, value) != stdr::end(m_options));
+  }
+
+  std::span<const DLDataType> m_options;
+  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}) {}
+  SymbolicDevice(const SymbolicDevice&) = delete;
+  SymbolicDevice& operator=(const SymbolicDevice&) = delete;
+
+  auto set_value(DLDevice value) -> void {
+    RuntimeCheck(!this->has_value(), "Device value already set");
+    RuntimeCheck(
+        m_check(value),
+        "Device value [",
+        details::PrintableDevice{value},
+        "] not in the allowed options: ",
+        details::PrintAbleSpan{m_options});
+    m_value = value;
+  }
+
+  auto has_value() const -> bool {
+    return m_value.device_type != details::kNullDevice;
+  }
+
+  auto get_value() const -> std::optional<DLDevice> {
+    return this->has_value() ? std::optional{m_value} : std::nullopt;
+  }
+
+  auto unwrap(DebugInfo info = {}) const -> DLDevice {
+    RuntimeCheck(info, this->has_value(), "Device value is not set");
+    return m_value;
+  }
+
+  auto set_options(std::span<const DLDevice> options) -> void {
+    m_options = options;
+  }
+
+  template <DLDeviceType... Codes>
+  auto set_options() -> void {
+    m_options = details::kDeviceList<Codes...>;
+  }
+
+  auto verify(DLDevice device) -> void {
+    if (this->has_value()) {
+      RuntimeCheck(
+          m_value == device,
+          "Device mismatch: expected ",
+          details::PrintableDevice{m_value},
+          " but got ",
+          details::PrintableDevice{device});
+    } else {
+      this->set_value(device);
+    }
+  }
+
+ private:
+  auto m_check(DLDevice value) const -> bool {
+    return stdr::empty(m_options) || (stdr::any_of(m_options, [value](const DLDevice& opt) {
+             // device type must exactly match
+             if (opt.device_type != value.device_type) return false;
+             // device id can be wildcarded
+             return opt.device_id == details::kAnyDeviceID || opt.device_id == value.device_id;
+           }));
+  }
+
+  std::span<const DLDevice> m_options;
+  DLDevice m_value;
+};
+
+namespace details {
+
+template <typename T>
+struct BaseRef {
+ public:
+  BaseRef(const BaseRef&) = delete;
+  BaseRef& operator=(const BaseRef&) = delete;
+
+  auto operator->() const -> T* {
+    return m_ref;
+  }
+  auto operator*() const -> T& {
+    return *m_ref;
+  }
+  auto rebind(T& other) -> void {
+    m_ref = &other;
+  }
+
+  explicit BaseRef() : m_ref(&m_cache), m_cache() {}
+  BaseRef(T& size) : m_ref(&size), m_cache() {}
+
+ private:
+  T* m_ref;
+  T m_cache;
+};
+
+struct SizeRef : BaseRef<SymbolicSize> {
+  using BaseRef::BaseRef;
+  SizeRef(int64_t value) {
+    if (value != kAnySize) {
+      (**this).set_value(value);
+    } else {
+      // otherwise, we can match any size
+    }
+  }
+};
+
+struct DTypeRef : BaseRef<SymbolicDType> {
+  using BaseRef::BaseRef;
+  DTypeRef(DLDataType options) {
+    (**this).set_value(options);
+  }
+  DTypeRef(std::initializer_list<DLDataType> options) {
+    (**this).set_options(options);
+  }
+  DTypeRef(std::span<const DLDataType> options) {
+    (**this).set_options(options);
+  }
+};
+
+struct DeviceRef : BaseRef<SymbolicDevice> {
+  using BaseRef::BaseRef;
+  DeviceRef(DLDevice options) {
+    (**this).set_value(options);
+  }
+  DeviceRef(std::initializer_list<DLDevice> options) {
+    (**this).set_options(options);
+  }
+  DeviceRef(std::span<const DLDevice> options) {
+    (**this).set_options(options);
+  }
+};
+
+}  // 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;
+  using DTypeRef = details::DTypeRef;
+  using DeviceRef = details::DeviceRef;
+
+ public:
+  TensorMatcher(const TensorMatcher&) = delete;
+  TensorMatcher& operator=(const TensorMatcher&) = delete;
+
+  explicit TensorMatcher(std::initializer_list<SizeRef> shape) : m_shape(shape), m_strides(), m_dtype() {}
+
+  auto with_strides(std::initializer_list<SizeRef> strides) && -> TensorMatcher&& {
+    // no partial update allowed
+    RuntimeCheck(m_strides.size() == 0, "Strides already specified");
+    RuntimeCheck(m_shape.size() == strides.size(), "Strides size must match shape size");
+    m_strides = strides;
+    return std::move(*this);
+  }
+
+  template <typename... Ts>
+  auto with_dtype(DTypeRef&& dtype) && -> TensorMatcher&& {
+    m_init_dtype();
+    m_dtype.rebind(*dtype);
+    m_dtype->set_options<Ts...>();
+    return std::move(*this);
+  }
+
+  template <typename... Ts>
+  auto with_dtype() && -> TensorMatcher&& {
+    static_assert(sizeof...(Ts) > 0, "At least one dtype option must be specified");
+    m_init_dtype();
+    m_dtype->set_options<Ts...>();
+    return std::move(*this);
+  }
+
+  template <DLDeviceType... Codes>
+  auto with_device(DeviceRef&& device) && -> TensorMatcher&& {
+    m_init_device();
+    m_device.rebind(*device);
+    m_device->set_options<Codes...>();
+    return std::move(*this);
+  }
+
+  template <DLDeviceType... Codes>
+  auto with_device() && -> TensorMatcher&& {
+    static_assert(sizeof...(Codes) > 0, "At least one device option must be specified");
+    m_init_device();
+    m_device->set_options<Codes...>();
+    return std::move(*this);
+  }
+
+  // once we start verification, we cannot modify anymore
+  auto verify(tvm::ffi::TensorView view, DebugInfo info = {}) const&& -> const TensorMatcher&& {
+    try {
+      m_verify_impl(view);
+    } catch (PanicError& e) {
+      auto oss = std::ostringstream{};
+      oss << "Tensor match failed for ";
+      s_print_tensor(oss, view);
+      oss << " at " << info.file_name() << ":" << info.line() << "\n- Root cause: " << e.root_cause();
+      throw PanicError(std::move(oss).str());
+    }
+    return std::move(*this);
+  }
+
+ private:
+  static auto s_print_tensor(std::ostringstream& oss, tvm::ffi::TensorView view) -> void {
+    oss << "Tensor<";
+    int64_t dim = 0;
+    for (const auto& size : view.shape()) {
+      if (dim++ > 0) oss << ", ";
+      oss << size;
+    }
+    oss << ">[strides=<";
+    dim = 0;
+    for (const auto& stride : view.strides()) {
+      if (dim++ > 0) {
+        oss << ", ";
+      }
+      oss << stride;
+    }
+    oss << ">, dtype=" << view.dtype();
+    oss << ", device=" << details::PrintableDevice{view.device()} << "]";
+  }
+
+  auto m_verify_impl(tvm::ffi::TensorView view) const -> void {
+    const auto dim = static_cast<std::size_t>(view.dim());
+    RuntimeCheck(dim == m_shape.size(), "Tensor dimension mismatch: expected ", m_shape.size(), " but got ", dim);
+    for (const auto i : irange(dim)) {
+      m_shape[i]->verify(view.size(i), "shape", i);
+    }
+    if (m_has_strides()) {
+      for (const auto i : irange(dim)) {
+        if (view.size(i) != 1 || !m_strides[i]->has_value()) {
+          // skip stride check for size 1 dimension
+          m_strides[i]->verify(view.stride(i), "stride", i);
+        }
+      }
+    } else {
+      RuntimeCheck(view.is_contiguous(), "Tensor is not contiguous as expected");
+    }
+    // since we may double verify, we will force to check
+    m_dtype->verify(view.dtype());
+    m_device->verify(view.device());
+  }
+
+  auto m_init_dtype() -> void {
+    RuntimeCheck(!m_has_dtype, "DType already specified");
+    m_has_dtype = true;
+  }
+
+  auto m_init_device() -> void {
+    RuntimeCheck(!m_has_device, "Device already specified");
+    m_has_device = true;
+  }
+
+  auto m_has_strides() const -> bool {
+    return !m_strides.empty();
+  }
+
+  std::span<const SizeRef> m_shape;
+  std::span<const SizeRef> m_strides;
+  DTypeRef m_dtype;
+  DeviceRef m_device;
+  bool m_has_dtype = false;
+  bool m_has_device = false;
+};
+
+}  // namespace host

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

@@ -0,0 +1,62 @@
+/// \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;
+  }
+
+ private:
+  uint32_t tid;
+  uint32_t tsize;
+};
+
+}  // namespace device::tile

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

@@ -0,0 +1,120 @@
+/// \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>
+
+template <typename T>
+struct dtype_trait {};
+
+#define SGL_REGISTER_DTYPE_TRAIT(TYPE, PACK2, ...)  \
+  template <>                                       \
+  struct dtype_trait<TYPE> {                        \
+    using self_t = TYPE;                            \
+    using packed_t = PACK2;                         \
+    template <typename S>                           \
+    SGL_DEVICE static self_t from(const S& value) { \
+      return static_cast<TYPE>(value);              \
+    }                                               \
+    __VA_ARGS__                                     \
+  }
+
+#define SGL_REGISTER_TYPE_END static_assert(true)
+
+#define SGL_REGISTER_FROM_FUNCTION(FROM, FN)     \
+  SGL_DEVICE static self_t from(const FROM& x) { \
+    return FN(x);                                \
+  }                                              \
+  static_assert(true)
+
+#define SGL_REGISTER_UNARY_FUNCTION(NAME, FN)      \
+  SGL_DEVICE static self_t NAME(const self_t& x) { \
+    return FN(x);                                  \
+  }                                                \
+  static_assert(true)
+
+#define SGL_REGISTER_BINARY_FUNCTION(NAME, FN)                      \
+  SGL_DEVICE static self_t NAME(const self_t& x, const self_t& y) { \
+    return FN(x, y);                                                \
+  }                                                                 \
+  static_assert(true)
+
+SGL_REGISTER_DTYPE_TRAIT(
+    fp32_t, fp32x2_t, SGL_REGISTER_TYPE_END;  //
+    SGL_REGISTER_FROM_FUNCTION(fp16_t, __half2float);
+    SGL_REGISTER_FROM_FUNCTION(bf16_t, __bfloat162float);
+    SGL_REGISTER_UNARY_FUNCTION(abs, fabsf);
+    SGL_REGISTER_UNARY_FUNCTION(sqrt, sqrtf);
+    SGL_REGISTER_UNARY_FUNCTION(rsqrt, rsqrtf);
+    SGL_REGISTER_UNARY_FUNCTION(exp, expf);
+    SGL_REGISTER_UNARY_FUNCTION(sin, sinf);
+    SGL_REGISTER_UNARY_FUNCTION(cos, cosf);
+    SGL_REGISTER_BINARY_FUNCTION(max, fmaxf);
+    SGL_REGISTER_BINARY_FUNCTION(min, fminf););
+SGL_REGISTER_DTYPE_TRAIT(fp16_t, fp16x2_t);
+SGL_REGISTER_DTYPE_TRAIT(bf16_t, bf16x2_t);
+
+/// TODO: Add ROCM implementation
+SGL_REGISTER_DTYPE_TRAIT(
+    fp32x2_t, fp32x4_t, SGL_REGISTER_TYPE_END; SGL_REGISTER_FROM_FUNCTION(fp16x2_t, __half22float2);
+    SGL_REGISTER_FROM_FUNCTION(bf16x2_t, __bfloat1622float2););
+
+SGL_REGISTER_DTYPE_TRAIT(
+    fp16x2_t, void, SGL_REGISTER_TYPE_END; SGL_REGISTER_FROM_FUNCTION(fp32x2_t, __float22half2_rn););
+
+SGL_REGISTER_DTYPE_TRAIT(
+    bf16x2_t, void, SGL_REGISTER_TYPE_END; SGL_REGISTER_FROM_FUNCTION(fp32x2_t, __float22bfloat162_rn););
+
+#ifndef USE_ROCM
+SGL_REGISTER_DTYPE_TRAIT(fp8_e4m3_t, fp8x2_e4m3_t);
+#endif
+
+#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);
+}
+
+}  // namespace device
+
+// ---------------------------------------------------------------------------
+// FP8 max clamp value — platform-dependent
+//   CUDA (e4m3fn):      448.0f
+//   AMD FNUZ (e4m3fnuz): 224.0f
+//   AMD E4M3 (e4m3fn):  448.0f
+// ---------------------------------------------------------------------------
+#ifndef USE_ROCM
+constexpr float kFP8E4M3Max = 448.0f;
+#else  // USE_ROCM
+#if HIP_FP8_TYPE_FNUZ
+constexpr float kFP8E4M3Max = 224.0f;
+#else   // HIP_FP8_TYPE_E4M3
+constexpr float kFP8E4M3Max = 448.0f;
+#endif  // HIP_FP8_TYPE_FNUZ
+#endif  // USE_ROCM

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

@@ -0,0 +1,369 @@
+/// \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>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/extra/c_env_api.h>
+
+#include <concepts>
+#include <cstddef>
+#include <type_traits>
+#ifndef USE_ROCM
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+#else
+#include <hip/hip_bf16.h>
+#include <hip/hip_fp16.h>
+#include <hip/hip_runtime.h>
+#ifndef __grid_constant__
+#define __grid_constant__
+#endif
+using cudaError_t = hipError_t;
+using cudaStream_t = hipStream_t;
+using cudaLaunchConfig_t = hipLaunchConfig_t;
+using cudaLaunchAttribute = hipLaunchAttribute;
+inline constexpr auto cudaSuccess = hipSuccess;
+#define cudaStreamPerThread hipStreamPerThread
+#define cudaGetErrorString hipGetErrorString
+#define cudaGetLastError hipGetLastError
+#define cudaLaunchKernel hipLaunchKernel
+#define cudaMemcpyAsync hipMemcpyAsync
+#define cudaMemcpyHostToDevice hipMemcpyHostToDevice
+#define cudaMemcpyDeviceToHost hipMemcpyDeviceToHost
+#endif
+
+#ifndef USE_ROCM
+using fp32_t = float;
+using fp16_t = __half;
+using bf16_t = __nv_bfloat16;
+using fp8_e4m3_t = __nv_fp8_e4m3;
+using fp8_e5m2_t = __nv_fp8_e5m2;
+
+using fp32x2_t = float2;
+using fp16x2_t = __half2;
+using bf16x2_t = __nv_bfloat162;
+using fp8x2_e4m3_t = __nv_fp8x2_e4m3;
+using fp8x2_e5m2_t = __nv_fp8x2_e5m2;
+
+using fp32x4_t = float4;
+#else
+using fp32_t = float;
+using fp16_t = __half;
+using bf16_t = __hip_bfloat16;
+using fp8_e4m3_t = uint8_t;
+using fp8_e5m2_t = uint8_t;
+using fp32x2_t = float2;
+using fp16x2_t = half2;
+using bf16x2_t = __hip_bfloat162;
+using fp8x2_e4m3_t = uint16_t;
+using fp8x2_e5m2_t = uint16_t;
+using fp32x4_t = float4;
+#endif
+
+/*
+ * LDG Support
+ */
+#ifndef USE_ROCM
+#define SGLANG_LDG(arg) __ldg(arg)
+#else
+#define SGLANG_LDG(arg) *(arg)
+#endif
+
+// DLPack device type for the current platform
+#ifndef USE_ROCM
+inline constexpr auto kDLGPU = kDLCUDA;
+inline constexpr auto kDLGPUHost = kDLCUDAHost;
+#else
+inline constexpr auto kDLGPU = kDLROCM;
+inline constexpr auto kDLGPUHost = kDLROCMHost;
+#endif
+
+namespace device {
+
+/// \brief Macro: forced-inline device function qualifier.
+#define SGL_DEVICE __forceinline__ __device__
+
+// Architecture detection: SGL_CUDA_ARCH is injected by load_jit() and is
+// available in both host and device compilation passes, whereas __CUDA_ARCH__
+// is only defined by nvcc during the device pass.
+#if !defined(USE_ROCM)
+#if !defined(SGL_CUDA_ARCH)
+#error "SGL_CUDA_ARCH is not defined. JIT compilation must inject -DSGL_CUDA_ARCH via load_jit()."
+#endif
+#if defined(__CUDA_ARCH__)
+static_assert(
+    __CUDA_ARCH__ == SGL_CUDA_ARCH, "SGL_CUDA_ARCH mismatch: injected arch flag does not match device target");
+#endif
+#define SGL_ARCH_HOPPER_OR_GREATER (SGL_CUDA_ARCH >= 900)
+#define SGL_ARCH_BLACKWELL_OR_GREATER ((SGL_CUDA_ARCH >= 1000) && (CUDA_VERSION >= 12090))
+#else  // USE_ROCM
+#define SGL_ARCH_HOPPER_OR_GREATER 0
+#define SGL_ARCH_BLACKWELL_OR_GREATER 0
+#endif
+
+// Maximum vector size in bytes supported by current architecture.
+// Pre-Blackwell / AMD: 128-bit (16 bytes)
+// Blackwell or greater: 256-bit (32 bytes)
+inline constexpr std::size_t kMaxVecBytes = SGL_ARCH_BLACKWELL_OR_GREATER ? 32 : 16;
+
+/// \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).
+#ifndef USE_ROCM
+inline constexpr auto kFullMask = 0xffffffffu;
+#else
+inline constexpr auto kFullMask = 0xffffffffffffffffULL;
+#endif
+
+/**
+ * \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 SGL_ARCH_HOPPER_OR_GREATER
+  if constexpr (kUsePDL) {
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+  }
+#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 SGL_ARCH_HOPPER_OR_GREATER
+  if constexpr (kUsePDL) {
+    asm volatile("griddepcontrol.launch_dependents;" :::);
+  }
+#endif
+}
+
+template <std::integral T, std::integral U>
+SGL_DEVICE constexpr auto div_ceil(T a, U b) {
+  return (a + b - 1) / b;
+}
+
+/**
+ * \brief Load data with the specified type and offset from a void pointer.
+ * \tparam T The type to load.
+ * \param ptr The base pointer.
+ * \param offset The offset in number of elements of type T.
+ */
+template <typename T>
+SGL_DEVICE T load_as(const void* ptr, int64_t offset = 0) {
+  return static_cast<const T*>(ptr)[offset];
+}
+
+/**
+ * \brief Store data with the specified type and offset to a void pointer.
+ * \tparam T The type to store.
+ * \param ptr The base pointer.
+ * \param val The value to store.
+ * \param offset The offset in number of elements of type T.
+ * \note we use type_identity_t to force the caller to explicitly specify
+ * the template parameter `T`, which can avoid accidentally using the wrong type.
+ */
+template <typename T>
+SGL_DEVICE void store_as(void* ptr, std::type_identity_t<T> val, int64_t offset = 0) {
+  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
+
+template <typename T = char, std::integral... U>
+SGL_DEVICE auto offset(void* ptr, U... offset) -> void* {
+  return static_cast<T*>(ptr) + (... + offset);
+}
+
+template <typename T = char, std::integral... U>
+SGL_DEVICE auto offset(const void* ptr, U... offset) -> const void* {
+  return static_cast<const T*>(ptr) + (... + offset);
+}
+
+}  // namespace pointer
+
+/// PTX pragma that lets the compiler spill registers into otherwise-unused
+/// shared memory instead of local memory. The radix kernels run at occupancy 2
+/// (32 regs/thread) and rely on this to avoid local-memory traffic.
+SGL_DEVICE void enable_smem_spilling() {
+#if defined(__CUDA_ARCH__) && CUDART_VERSION >= 13000
+  asm(".pragma \"enable_smem_spilling\";");
+#endif
+}
+
+}  // namespace device
+
+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]];
+    ::host::panic(location, "CUDA error: ", ::cudaGetErrorString(error));
+  }
+}
+
+/// \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, arg0, arg1);
+ *   host::LaunchKernel(grid, block, stream)
+ *       .config({.use_pdl = true, .cluster_dim = cluster_dim})(my_kernel, arg0);
+ * \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 {
+ private:
+  struct KernelConfig {
+    bool use_pdl = false;
+    std::optional<dim3> cluster_dim = std::nullopt;
+  };
+
+ public:
+  explicit LaunchKernel(
+      dim3 grid_dim,
+      dim3 block_dim,
+      DLDevice device,
+      std::size_t dynamic_shared_mem_bytes = 0,
+      DebugInfo location = {}) noexcept
+      : m_config(s_make_config(grid_dim, block_dim, resolve_device(device), dynamic_shared_mem_bytes)),
+        m_location(location) {}
+
+  explicit LaunchKernel(
+      dim3 grid_dim,
+      dim3 block_dim,
+      cudaStream_t stream,
+      std::size_t dynamic_shared_mem_bytes = 0,
+      DebugInfo location = {}) noexcept
+      : m_config(s_make_config(grid_dim, block_dim, stream, dynamic_shared_mem_bytes)), m_location(location) {}
+
+  LaunchKernel(const LaunchKernel&) = delete;
+  LaunchKernel& operator=(const LaunchKernel&) = delete;
+
+  static auto resolve_device(DLDevice device) -> cudaStream_t {
+    return static_cast<cudaStream_t>(::TVMFFIEnvGetStream(device.device_type, device.device_id));
+  }
+
+  auto enable_pdl(bool enabled = true) -> LaunchKernel& {
+#ifdef USE_ROCM
+    (void)enabled;
+    m_config.numAttrs = 0;
+#else
+    if (enabled) {
+      auto& attr = m_attrs[m_config.numAttrs++];
+      attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
+      attr.val.programmaticStreamSerializationAllowed = true;
+      m_config.attrs = m_attrs;
+    }
+#endif
+    return *this;
+  }
+
+  auto enable_cluster(dim3 cluster_dim) -> LaunchKernel& {
+#ifdef USE_ROCM
+    (void)cluster_dim;
+#else
+    auto& attr = m_attrs[m_config.numAttrs++];
+    attr.id = cudaLaunchAttributeClusterDimension;
+    attr.val.clusterDim = {cluster_dim.x, cluster_dim.y, cluster_dim.z};
+    m_config.attrs = m_attrs;
+#endif
+    return *this;
+  }
+
+  /**
+   * \brief Configure the kernel launch with the given options.
+   * \param config The kernel configuration options.
+   * \return A reference to this `LaunchKernel` for chaining.
+   * \note This is a convenience method that applies multiple configurations at once.
+   * We are in favor of this instead of `enable_pdl` and `enable_cluster`.
+   * We enforce use of designated initializers for better readability.
+   */
+  auto config(const KernelConfig& config) -> LaunchKernel& {
+    if (config.use_pdl) this->enable_pdl(true);
+    if (config.cluster_dim) this->enable_cluster(*config.cluster_dim);
+    return *this;
+  }
+
+  template <typename T, typename... Args>
+  auto operator()(T&& kernel, Args&&... args) const -> void {
+#ifdef USE_ROCM
+    hipLaunchKernelGGL(
+        std::forward<T>(kernel),
+        m_config.gridDim,
+        m_config.blockDim,
+        m_config.dynamicSmemBytes,
+        m_config.stream,
+        std::forward<Args>(args)...);
+    RuntimeDeviceCheck(m_location);
+#else
+    RuntimeDeviceCheck(::cudaLaunchKernelEx(&m_config, kernel, std::forward<Args>(args)...), m_location);
+#endif
+  }
+
+  template <typename T, typename... Args>
+  auto launch(T&& kernel, Args&&... args) const -> void {
+    return (*this)(std::forward<T>(kernel), std::forward<Args>(args)...);
+  }
+
+ private:
+  static auto s_make_config(  // Make a config for kernel launch
+      dim3 grid_dim,
+      dim3 block_dim,
+      cudaStream_t stream,
+      std::size_t smem) -> cudaLaunchConfig_t {
+    auto config = ::cudaLaunchConfig_t{};
+    config.gridDim = grid_dim;
+    config.blockDim = block_dim;
+    config.dynamicSmemBytes = smem;
+    config.stream = stream;
+    config.numAttrs = 0;
+    return config;
+  }
+
+  cudaLaunchConfig_t m_config;
+  const DebugInfo m_location;
+  cudaLaunchAttribute m_attrs[2];
+};
+
+}  // namespace host

python/sglang/jit_kernel/include/sgl_kernel/utils.h

@@ -0,0 +1,186 @@
+/// \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
+#ifdef __CUDACC__
+#include <cuda.h>
+#if CUDA_VERSION <= 12010
+
+#pragma push_macro("__cpp_consteval")
+#pragma push_macro("_NODISCARD")
+#pragma push_macro("__builtin_LINE")
+
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wbuiltin-macro-redefined"
+#define __cpp_consteval 201811L
+#pragma clang diagnostic pop
+
+#ifdef _NODISCARD
+#undef _NODISCARD
+#define _NODISCARD
+#endif
+
+#define consteval constexpr
+
+#include "source_location.h"
+
+#undef consteval
+#pragma pop_macro("__cpp_consteval")
+#pragma pop_macro("_NODISCARD")
+#else  // __CUDACC__ && CUDA_VERSION > 12010
+#include "source_location.h"
+#endif
+#else  // no __CUDACC__
+#include "source_location.h"
+#endif
+
+#include <dlpack/dlpack.h>
+
+#include <concepts>
+#include <cstddef>
+#include <ostream>
+#include <ranges>
+#include <sstream>
+#include <utility>
+
+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)) {}
+  auto root_cause() const -> std::string_view {
+    const auto str = std::string_view{m_message};
+    const auto pos = str.find(": ");
+    return pos == std::string_view::npos ? str : str.substr(pos + 2);
+  }
+
+ private:
+  std::string m_message;
+};
+
+/// \brief Unconditionally abort with a formatted error message.
+template <typename... Args>
+[[noreturn]]
+inline auto panic(DebugInfo location, Args&&... args) -> void {
+  std::ostringstream os;
+  os << "Runtime check failed at " << location.file_name() << ":" << location.line();
+  if constexpr (sizeof...(args) > 0) {
+    os << ": ";
+    (os << ... << std::forward<Args>(args));
+  } else {
+    os << " in " << location.function_name();
+  }
+  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>
+  explicit RuntimeCheck(Cond&& condition, Args&&... args, DebugInfo location = {}) {
+    if (condition) return;
+    [[unlikely]] ::host::panic(location, std::forward<Args>(args)...);
+  }
+  template <typename Cond>
+  explicit RuntimeCheck(DebugInfo location, Cond&& condition, Args&&... args) {
+    if (condition) return;
+    [[unlikely]] ::host::panic(location, std::forward<Args>(args)...);
+  }
+};
+
+template <typename... Args>
+struct Panic {
+  explicit Panic(Args&&... args, DebugInfo location = {}) {
+    ::host::panic(location, std::forward<Args>(args)...);
+  }
+  explicit Panic(DebugInfo location, Args&&... args) {
+    ::host::panic(location, std::forward<Args>(args)...);
+  }
+  [[noreturn]] ~Panic() {
+    std::terminate();
+  }
+};
+
+template <typename Cond, typename... Args>
+explicit RuntimeCheck(Cond&&, Args&&...) -> RuntimeCheck<Args...>;
+
+template <typename Cond, typename... Args>
+explicit RuntimeCheck(DebugInfo, Cond&&, Args&&...) -> RuntimeCheck<Args...>;
+
+template <typename... Args>
+explicit Panic(Args&&...) -> Panic<Args...>;
+
+template <typename... Args>
+explicit Panic(DebugInfo, Args&&...) -> Panic<Args...>;
+
+namespace pointer {
+
+// we only allow void * pointer arithmetic for safety
+
+template <typename T = char, std::integral... U>
+inline auto offset(void* ptr, U... offset) -> void* {
+  return static_cast<T*>(ptr) + (... + offset);
+}
+
+template <typename T = char, std::integral... U>
+inline auto offset(const void* ptr, U... offset) -> const void* {
+  return static_cast<const T*>(ptr) + (... + offset);
+}
+
+}  // 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);
+}
+
+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);
+}
+
+}  // namespace host

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

@@ -0,0 +1,118 @@
+/// \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>
+
+#include <cstddef>
+#include <cstdint>
+
+namespace device {
+
+namespace details {
+
+/// \brief Maps byte-width to the corresponding unsigned integer type.
+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;
+};
+
+/// \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:
+  static_assert(
+      (N > 0 && (N & (N - 1)) == 0) && sizeof(T) * N <= kMaxVecBytes,
+      "CUDA vector size exceeds arch limit: max 16 bytes on pre-Blackwell/AMD, "
+      "32 bytes on Blackwell or greater");
+  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`.
+  SGL_DEVICE void load(const void* ptr, int64_t offset = 0) {
+    m_storage = reinterpret_cast<const storage_t*>(ptr)[offset];
+  }
+  /// \brief Vectorized store to `ptr` at the given element `offset`.
+  SGL_DEVICE void store(void* ptr, int64_t offset = 0) const {
+    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
+    for (std::size_t i = 0; i < N; ++i) {
+      m_storage.data[i] = store_value;
+    }
+  }
+
+  SGL_DEVICE auto operator[](std::size_t idx) -> T& {
+    return reinterpret_cast<T*>(&m_storage)[idx];
+  }
+  SGL_DEVICE auto operator[](std::size_t idx) const -> T {
+    return reinterpret_cast<const T*>(&m_storage)[idx];
+  }
+  SGL_DEVICE auto data() -> T* {
+    return reinterpret_cast<T*>(&m_storage);
+  }
+  SGL_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

@@ -0,0 +1,56 @@
+/// \file warp.cuh
+/// \brief Warp-level reduction primitives.
+
+#pragma once
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/utils.cuh>
+
+namespace device::warp {
+
+/// \brief Full warp active mask.
+#ifndef USE_ROCM
+static constexpr uint32_t kFullMask = 0xffffffffu;
+using mask_t = uint32_t;
+#else
+static constexpr uint64_t kFullMask = 0xffffffffffffffffULL;
+using mask_t = uint64_t;
+#endif
+
+/**
+ * \brief Warp-level sum reduction.
+ *
+ * On CUDA: uses __shfl_xor_sync with width=32.
+ * On HIP: uses __shfl_xor with explicit width parameter (supports wave64 sub-groups).
+ */
+template <uint32_t kNumThreads = kWarpThreads, typename T>
+SGL_DEVICE T reduce_sum(T value, mask_t active_mask = kFullMask) {
+  static_assert(kNumThreads >= 1 && kNumThreads <= kWarpThreads);
+  static_assert(std::has_single_bit(kNumThreads), "must be pow of 2");
+#pragma unroll
+  for (int mask = kNumThreads / 2; mask > 0; mask >>= 1)
+#ifndef USE_ROCM
+    value = value + __shfl_xor_sync(active_mask, value, mask, 32);
+#else
+    value = value + __shfl_xor(value, mask, kNumThreads);
+#endif
+  return value;
+}
+
+/**
+ * \brief Warp-level max reduction.
+ */
+template <uint32_t kNumThreads = kWarpThreads, typename T>
+SGL_DEVICE T reduce_max(T value, mask_t active_mask = kFullMask) {
+  static_assert(kNumThreads >= 1 && kNumThreads <= kWarpThreads);
+  static_assert(std::has_single_bit(kNumThreads), "must be pow of 2");
+#pragma unroll
+  for (int mask = kNumThreads / 2; mask > 0; mask >>= 1)
+#ifndef USE_ROCM
+    value = math::max(value, __shfl_xor_sync(active_mask, value, mask, 32));
+#else
+    value = math::max(value, __shfl_xor(value, mask, kNumThreads));
+#endif
+  return value;
+}
+
+}  // namespace device::warp