[JIT] Refactor dtype traits into DTypeTrait and unify warp reductions (#30838)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: BBuf <xiaoyu.zhang@radixark.ai>
Co-authored-by: jessiewei7 <jessiewei747@gmail.com>
Co-authored-by: root <root@GPUC5A6.maas>

.claude/skills/add-jit-kernel/SKILL.md

@@ -35,7 +35,8 @@ Add a new operation that scales each element of a tensor by a scalar factor:
 #include <sgl_kernel/utils.h>
 ```
 
-- **`host::RuntimeCheck(cond, args...)`** — Assert a condition at runtime; throws `PanicError` with file/line info on failure. Prefer this over bare `assert`.
+- **`CHECK_HOST(cond) << "msg " << value`** — **Preferred** runtime check: stream-style, throws `PanicError` with file/line info on failure. Zero overhead on the true path — the message expressions are only evaluated when the check fails.
+- **`host::RuntimeCheck(cond, args...)`** — Function-style alternative to `CHECK_HOST`. Note its message args are always evaluated (even when the check passes), so prefer `CHECK_HOST` — especially on hot paths.
 - **`host::Panic(args...)`** — Unconditionally throw a `PanicError` with a descriptive message.
 - **`host::div_ceil(a, b)`** — Integer ceiling division `(a + b - 1) / b`.
 - **`host::irange(n)`** / **`host::irange(start, end)`** — Range views for cleaner loops.
@@ -57,6 +58,7 @@ Add a new operation that scales each element of a tensor by a scalar factor:
   - Checks the CUDA error with file/line info after launch via `operator()(kernel, args...)`.
   - Supports `.enable_pdl(bool)` for PDL (Programmatic Dependent Launch, SM90+).
 - **`host::RuntimeDeviceCheck(cudaError_t)`** — Check a CUDA error; throw on failure.
+- **`CHECK_CUDA(expr) << "context"`** — Stream-style CUDA error check; evaluates `expr` once and throws `PanicError` with `cudaGetErrorString` + file/line info if it is not `cudaSuccess`. Extra streamed context is optional.
 
 ### `tensor.h` — Tensor validation (`TensorMatcher`, Symbolic types)
 
@@ -90,18 +92,21 @@ const size_t n = N.unwrap();
 const DLDevice dev = device.unwrap();
 ```
 
-### `type.cuh` — `dtype_trait<T>` and `packed_t<T>`
+### `type.cuh` — `DTypeTrait<T>`, `packed_t<T>`, and reduction traits
 
 ```cpp
 #include <sgl_kernel/type.cuh>
 ```
 
-- **`dtype_trait<T>`** — Static trait struct for each scalar type. Provides:
+- **`DTypeTrait<T>`** — Static trait struct, specialized for integral types, `fp32_t`, `fp16_t`, `bf16_t`, `fp8_e4m3_t`, and their packed x2/x4 variants. Provides:
-  - `dtype_trait<T>::from(value)` — convert from another type (e.g. `fp32_t` → `fp16_t`)
+  - `DTypeTrait<T>::from(value)` — convert from another type via the right CUDA intrinsic (e.g. `fp32_t` → `fp16_t`)
-  - `dtype_trait<T>::abs/sqrt/rsqrt/exp/sin/cos(x)` — type-dispatched unary math (primarily for `fp32_t`)
+  - `DTypeTrait<T>::abs/max/min` — type-dispatched math (fp32, fp16/bf16 scalar and x2, integrals)
-  - `dtype_trait<T>::max/min(x, y)` — type-dispatched binary math (primarily for `fp32_t`)
+  - `DTypeTrait<T>::sqrt/rsqrt/exp/sin/cos(x)` — `fp32_t` only
+  - Metadata: `packed_t` / `unpacked_t` / `kVecSize` (packed layout), `kFloatMax` (dtype max as float, e.g. 448.0f for fp8-e4m3), `kZeroBits`
 - **`packed_t<T>`** — Two-element packed alias: `packed_t<fp16_t>` = `fp16x2_t`, `packed_t<bf16_t>` = `bf16x2_t`, `packed_t<fp32_t>` = `fp32x2_t`. Use for vectorized loads/stores.
-- **`device::cast<To, From>(value)`** — Type-safe cast using `dtype_trait`, e.g. `cast<fp32x2_t, fp16x2_t>(v)`.
+- **`device::cast<To, From>(value)`** — Type-safe cast using `DTypeTrait`, e.g. `cast<fp32x2_t, fp16x2_t>(v)`.
+- **`device::unpack(value)`** — View a packed value as an `unpacked_t[kVecSize]` array reference (e.g. `fp32x2_t` → `fp32_t[2]`); element writes propagate back to the packed value.
+- **`device::ReductionOp` (`SUM`/`MAX`/`MIN`) and `device::ReductionTrait<Op, T>::reduce(x, y)`** — One binary reduction step, dispatched through `DTypeTrait` (packed types reduce elementwise). This is the engine behind `warp::reduce`; use it directly when writing custom reductions.
 
 ### `vec.cuh` — Vectorized memory access (`AlignedVector`)
 
@@ -135,8 +140,8 @@ For a **2D tile**, either flatten `(row, col)` into a linear tile index first, o
 #include <sgl_kernel/math.cuh>
 ```
 
-- `device::math::max/min<T>(a, b)` — type-dispatched binary math via `dtype_trait`
+- `device::math::max/min<T>(a, b)` — type-dispatched binary math via `DTypeTrait`
-- `device::math::abs/sqrt/rsqrt/exp/sin/cos<T>(x)` — type-dispatched unary math via `dtype_trait`
+- `device::math::abs/sqrt/rsqrt/exp/sin/cos<T>(x)` — type-dispatched unary math via `DTypeTrait`
 
 ### `warp.cuh` — Warp-level primitives
 
@@ -144,8 +149,8 @@ For a **2D tile**, either flatten `(row, col)` into a linear tile index first, o
 #include <sgl_kernel/warp.cuh>
 ```
 
-- `device::warp::reduce_sum<T>(value)` — warp-level sum reduction via `__shfl_xor_sync`
+- `device::warp::reduce<Op, kNumThreads, kInner>(value, active_mask)` — generic warp reduction via `__shfl_xor_sync`. `Op` is a `device::ReductionOp` (`SUM`/`MAX`/`MIN`); `kNumThreads` is a power-of-two group size (default 32 = full warp); `kInner=true` (default) reduces within each `kNumThreads`-sized group, `kInner=false` reduces across groups (lanes at the same offset in different groups).
-- `device::warp::reduce_max<T>(value)` — warp-level max reduction
+- `device::warp::reduce_sum/reduce_max/reduce_min<kNumThreads, kInner>(value)` — convenience wrappers over `reduce`. Work for any type with a `ReductionTrait`: floats, integers, and packed x2 types.
 
 ### `cta.cuh` — CTA-level primitives
 
@@ -203,8 +208,8 @@ The implementation fully uses the project abstractions described above:
 // NOTE: Comments for headers are not common in practice.
 // It is only shown here for tutorial purposes to highlight the key abstractions.
 #include <sgl_kernel/tensor.h>   // For TensorMatcher, SymbolicSize, SymbolicDevice
-#include <sgl_kernel/type.cuh>   // For dtype_trait, fp16_t, bf16_t, fp32_t
+#include <sgl_kernel/type.cuh>   // For DTypeTrait, fp16_t, bf16_t, fp32_t
-#include <sgl_kernel/utils.h>    // For RuntimeCheck, div_ceil
+#include <sgl_kernel/utils.h>    // For CHECK_HOST, div_ceil
 #include <sgl_kernel/utils.cuh>  // For LaunchKernel, SGL_DEVICE
 #include <sgl_kernel/vec.cuh>    // For AlignedVector
 
@@ -280,7 +285,7 @@ void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, float factor) {
   const uint32_t n = static_cast<uint32_t>(N.unwrap());
   const DLDevice device = device_.unwrap();
 
-  RuntimeCheck(n > 0, "scale: num_elements must be > 0, got ", n);
+  CHECK_HOST(n > 0) << "scale: num_elements must be > 0, got " << n;
 
   // 2. Choose vector width for 128-bit loads (16 bytes)
   //    fp16/bf16: 8 elements x 2 bytes = 16 bytes
@@ -316,10 +321,10 @@ void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, float factor) {
 - Use `TensorMatcher` for all tensor validation; never manually check shape/dtype/device
 - Use `AlignedVector` for vectorised 128-bit loads/stores — significant bandwidth win
 - Use `LaunchKernel` — it resolves the stream and checks errors automatically
-- Use `RuntimeCheck` for runtime assertions with useful error messages
+- Use `CHECK_HOST(cond) << ...` for runtime assertions with useful error messages (zero overhead when the check passes)
 - Prefer passing runtime scalars like `factor` directly unless compile-time specialisation is genuinely required
 - `fp16_t` / `bf16_t` / `fp32_t` are the project's type aliases (from `utils.cuh`)
-- `device::cast<To, From>` or `dtype_trait<T>::from(val)` for cross-type conversions
+- `device::cast<To, From>` or `DTypeTrait<T>::from(val)` for cross-type conversions
 - `device::math::` functions for device math instead of bare `__` intrinsics if possible.
 - Try to use `PDL` feature. In some cases, this will benefit the performance.
 
@@ -626,9 +631,9 @@ cd test && python3 run_suite.py --hw cuda --suite base-b-kernel-benchmark-test-1
 - `python/sglang/jit_kernel/include/sgl_kernel/utils.cuh` — type aliases, `LaunchKernel`, `SGL_DEVICE`
 - `python/sglang/jit_kernel/include/sgl_kernel/vec.cuh` — `AlignedVector`
 - `python/sglang/jit_kernel/include/sgl_kernel/tile.cuh` — `tile::Memory`
-- `python/sglang/jit_kernel/include/sgl_kernel/type.cuh` — `dtype_trait`, `packed_t`, `device::cast`
+- `python/sglang/jit_kernel/include/sgl_kernel/type.cuh` — `DTypeTrait`, `packed_t`, `device::cast`, `device::unpack`, `ReductionTrait`
 - `python/sglang/jit_kernel/include/sgl_kernel/math.cuh` — `device::math::`
-- `python/sglang/jit_kernel/include/sgl_kernel/warp.cuh` — `warp::reduce_sum/max`
+- `python/sglang/jit_kernel/include/sgl_kernel/warp.cuh` — `warp::reduce<Op>` and `reduce_sum/max/min` wrappers
 - `python/sglang/jit_kernel/include/sgl_kernel/cta.cuh` — `cta::reduce_max`
 - `python/sglang/jit_kernel/include/sgl_kernel/atomic.cuh` — `atomic::max`
 - `python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh` — occupancy / SM count helpers

docs_new/docs/developer_guide/development_jit_kernel_guide.mdx

@@ -61,19 +61,24 @@ void test() {
 
 #### Runtime Checking
 
-`RuntimeCheck` validates conditions at runtime. It accepts optional arguments for error reporting.
+`CHECK_HOST` is the preferred runtime check: stream-style, and zero overhead when the
-If the check fails, these arguments are output to aid debugging.
+check passes — the message expressions are only evaluated on failure.
-`RuntimeDeviceCheck` verifies the status of the last kernel launch.
+`RuntimeCheck` is the function-style alternative; note its message arguments are always
+evaluated, even when the check passes.
+`RuntimeDeviceCheck` verifies the status of the last kernel launch, and `CHECK_CUDA`
+is its stream-style equivalent for checking a `cudaError_t` with extra context.
 
 ```C++ Example
 #include <sgl_kernel/utils.h>
 #include <sgl_kernel/utils.cuh>
 
 void test() {
+  CHECK_HOST(1 + 1 == 2) << 1 + 1 << " != " << 2;  // preferred
   host::RuntimeCheck(1 + 1 == 2, 1 + 1, " != ", 2);
   host::RuntimeDeviceCheck();
   // check the provided `cudaError_t`
   host::RuntimeDeviceCheck(cudaGetLastError());
+  CHECK_CUDA(cudaGetLastError()) << "after my_kernel launch";
 }
 
 ```
@@ -161,7 +166,7 @@ Write your CUDA kernel in [jit_kernel/csrc/add_constant.cuh](https://github.com/
 ```cpp Example
 #include <sgl_kernel/tensor.h>   // For TensorMatcher, SymbolicSize, SymbolicDevice
 #include <sgl_kernel/utils.cuh>  // For LaunchKernel
-#include <sgl_kernel/utils.h>    // For div_ceil, RuntimeCheck
+#include <sgl_kernel/utils.h>    // For div_ceil, CHECK_HOST
 
 #include <dlpack/dlpack.h>
 #include <tvm/ffi/container/tensor.h>
@@ -199,8 +204,8 @@ void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
   const size_t num_elements = N.unwrap();
   const size_t grid_size = div_ceil(num_elements, kBlockSize);
   const DLDevice device = device_.unwrap();
-  // some extra runtime checks using host::RuntimeCheck
+  // some extra runtime checks using CHECK_HOST
-  RuntimeCheck(num_elements > 0, "We only support non-empty tensors, got num_elements = ", num_elements);
+  CHECK_HOST(num_elements > 0) << "We only support non-empty tensors, got num_elements = " << num_elements;
 
   // 3. Launch the kernel. Error code will be automatically checked.
   LaunchKernel(grid_size, kBlockSize, device /*, dynamic_smem*/)(
@@ -289,12 +294,12 @@ and its key APIs.
     <tr>
       <td><code>utils.h</code></td>
       <td><code>host</code></td>
-      <td>Host-side essentials: <code>RuntimeCheck</code>, <code>Panic</code>, <code>div_ceil</code>, <code>irange</code></td>
+      <td>Host-side essentials: <code>RuntimeCheck</code>, <code>CHECK_HOST(cond) &lt;&lt; ...</code>, <code>Panic</code>, <code>div_ceil</code>, <code>irange</code></td>
     </tr>
     <tr>
       <td><code>utils.cuh</code></td>
       <td><code>device</code> / <code>host</code></td>
-      <td>Type aliases (<code>fp16_t</code>, <code>bf16_t</code>, ...), <code>SGL_DEVICE</code> macro, PDL helpers, <code>LaunchKernel</code>, <code>RuntimeDeviceCheck</code></td>
+      <td>Type aliases (<code>fp16_t</code>, <code>bf16_t</code>, ...), <code>SGL_DEVICE</code> macro, PDL helpers, <code>LaunchKernel</code>, <code>RuntimeDeviceCheck</code>, <code>CHECK_CUDA(expr) &lt;&lt; ...</code></td>
     </tr>
     <tr>
       <td><code>source_location.h</code></td>
@@ -347,7 +352,7 @@ and its key APIs.
     <tr>
       <td><code>type.cuh</code></td>
       <td>(global) / <code>device</code></td>
-      <td><code>dtype_trait&lt;T&gt;</code>, <code>packed_t&lt;T&gt;</code>, <code>device::cast&lt;To&gt;(from)</code></td>
+      <td><code>DTypeTrait&lt;T&gt;</code>, <code>packed_t&lt;T&gt;</code>, <code>device::cast&lt;To&gt;(from)</code></td>
     </tr>
   </tbody>
 </table>
@@ -390,7 +395,7 @@ and its key APIs.
     <tr>
       <td><code>warp.cuh</code></td>
       <td><code>device::warp</code></td>
-      <td><code>reduce_sum</code>, <code>reduce_max</code> via <code>__shfl_xor_sync</code></td>
+      <td><code>reduce&lt;Op, kNumThreads, kInner&gt;</code> (SUM/MAX/MIN, grouped or inter-group) and <code>reduce_sum</code> / <code>reduce_max</code> / <code>reduce_min</code> wrappers via <code>__shfl_xor_sync</code></td>
     </tr>
     <tr>
       <td><code>cta.cuh</code></td>

python/sglang/jit_kernel/__main__.py

@@ -1,16 +1,30 @@
 import argparse
 import logging
 import os
+import re
+import shutil
+import subprocess
 
 from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path
 
-from sglang.jit_kernel.utils import (
+from sglang.jit_kernel.utils import get_jit_cuda_arch, override_jit_cuda_arch
-    _REGISTERED_DEPENDENCIES,
+from sglang.jit_kernel.utils.arch import get_default_target_flags
-    DEFAULT_INCLUDE,
+from sglang.jit_kernel.utils.compile import DEFAULT_INCLUDE
-    _get_default_target_flags,
+from sglang.jit_kernel.utils.deps import REGISTERED_DEPENDENCIES
-    get_jit_cuda_arch,
+
-    override_jit_cuda_arch,
+
+def _clangd_major_version() -> int | None:
+    clangd = shutil.which("clangd")
+    if clangd is None:
+        return None
+    try:
+        result = subprocess.run(
+            [clangd, "--version"], capture_output=True, text=True, check=True
         )
+    except (OSError, subprocess.CalledProcessError):
+        return None
+    match = re.search(r"clangd version (\d+)", result.stdout)
+    return int(match.group(1)) if match else None
 
 
 def generate_clangd():
@@ -28,7 +42,7 @@ def generate_clangd():
         "--dep",
         nargs="*",
         default=[],
-        choices=_REGISTERED_DEPENDENCIES.keys(),
+        choices=REGISTERED_DEPENDENCIES.keys(),
         help="Extra dependency libraries to include.",
     )
     parser.add_argument(
@@ -42,9 +56,9 @@ def generate_clangd():
 
     dep_include_paths = []
     for dep in args.dependencies:
-        if dep not in _REGISTERED_DEPENDENCIES:
+        if dep not in REGISTERED_DEPENDENCIES:
             raise ValueError(f"Dependency {dep} is not registered.")
-        dep_include_paths += _REGISTERED_DEPENDENCIES[dep]()
+        dep_include_paths += REGISTERED_DEPENDENCIES[dep]()
 
     include_paths = [
         *DEFAULT_INCLUDE,
@@ -70,9 +84,15 @@ def generate_clangd():
         f"--cuda-gpu-arch=sm_{major}{minor}",
         "-Wall",
         "-Wextra",
-        *_get_default_target_flags(),
+        *get_default_target_flags(),
         *[f"-isystem{path}" for path in include_paths],
     ]
+
+    # NOTE: for local clangd (fix the missing cluster related macros)
+    if major >= 9:
+        compile_flags.append("-D_CG_LIMIT_INCLUDED_DEPENDENCIES=1")
+        compile_flags.append("-D_CG_HAS_CLUSTER_GROUP=1")
+
     # NOTE: skip these flags because clangd don't recognize them
     UNSUPPORTED_FLAGS = {"--expt-relaxed-constexpr"}
     compile_flags = [flag for flag in compile_flags if flag not in UNSUPPORTED_FLAGS]
@@ -83,6 +103,10 @@ CompileFlags:
     {compile_flags_str}
   ]
 """
+    # Documentation.CommentFormat lands in clangd 21.
+    clangd_major = _clangd_major_version()
+    if clangd_major is not None and clangd_major >= 21:
+        clangd_content += "Documentation:\n  CommentFormat: Doxygen\n"
     if os.path.exists(".clangd") and not args.overwrite:
         logger.warning(".clangd file already exists, nothing done.")
         logger.warning("Use --overwrite to force overwrite the existing .clangd file.")

python/sglang/jit_kernel/benchmark/marker.py

@@ -279,6 +279,9 @@ class Benchmark(Generic[F]):
                     if not DISABLE_LOG_BANDWIDTH:
                         bandwidths.append(float("nan"))
                     continue
+                except BaseException:
+                    print(f"Benchmark failed at {system}, kwargs =", kwargs)
+                    raise
                 latencies.append(result.times[0] / self._unit_scale)
                 if not DISABLE_LOG_BANDWIDTH and result.memory_footprint is not None:
                     should_log_bandwidth = True

python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh

@@ -114,32 +114,6 @@ SGL_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t val) {
   return val;
 }
 
-/// Warp-wide max/min for integer types. `device::warp::reduce_max` routes through
-/// `dtype_trait<T>::max` which is only specialized for FP types.
-SGL_DEVICE uint32_t warp_reduce_max_u32(uint32_t val) {
-#pragma unroll
-  for (uint32_t mask = 16; mask > 0; mask >>= 1) {
-#ifndef USE_ROCM
-    val = max(val, __shfl_xor_sync(device::kFullMask, val, mask, 32));
-#else
-    val = max(val, __shfl_xor(val, mask, 32));
-#endif
-  }
-  return val;
-}
-
-SGL_DEVICE uint32_t warp_reduce_min_u32(uint32_t val) {
-#pragma unroll
-  for (uint32_t mask = 16; mask > 0; mask >>= 1) {
-#ifndef USE_ROCM
-    val = min(val, __shfl_xor_sync(device::kFullMask, val, mask, 32));
-#else
-    val = min(val, __shfl_xor(val, mask, 32));
-#endif
-  }
-  return val;
-}
-
 __global__ __launch_bounds__(1024, 1)  //
     void plan_compress_prefill_kernel0(const Prefill0Params params) {
   using namespace device;
@@ -185,12 +159,12 @@ __global__ __launch_bounds__(1024, 1)  //
   // For min, treat threads outside `batch_size` as +inf so they don't pull the min down.
   const uint32_t e_for_max = static_cast<uint32_t>(extend_len);
   const uint32_t e_for_min = (tx < params.batch_size) ? e_for_max : 0xFFFFFFFFu;
-  warp_max[warp_id] = warp_reduce_max_u32(e_for_max);
+  warp_max[warp_id] = warp::reduce_max(e_for_max);
-  warp_min[warp_id] = warp_reduce_min_u32(e_for_min);
+  warp_min[warp_id] = warp::reduce_min(e_for_min);
   __syncthreads();
   if (warp_id == 0) {
-    s_max_extend = warp_reduce_max_u32(warp_max[lane_id]);
+    s_max_extend = warp::reduce_max(warp_max[lane_id]);
-    s_min_extend = warp_reduce_min_u32(warp_min[lane_id]);
+    s_min_extend = warp::reduce_min(warp_min[lane_id]);
   }
   __syncthreads();
 

python/sglang/jit_kernel/csrc/diffusion/residual_gate_add.cuh

@@ -15,7 +15,7 @@
 #include <sgl_kernel/tensor.h>  // For host dtype helpers and TensorView metadata
 #include <sgl_kernel/utils.h>   // For RuntimeCheck and div_ceil
 
-#include <sgl_kernel/type.cuh>   // For dtype_trait conversions
+#include <sgl_kernel/type.cuh>   // For DTypeTrait conversions
 #include <sgl_kernel/utils.cuh>  // For LaunchKernel and CUDA dtype aliases
 #include <sgl_kernel/vec.cuh>    // For device::AlignedVector
 
@@ -99,8 +99,8 @@ __device__ __forceinline__ float to_float<bf16_t>(bf16_t v) {
 
 template <typename T>
 __device__ __forceinline__ T residual_gate_value(T residual, T update, T gate) {
-  const T product = dtype_trait<T>::from(to_float(update) * to_float(gate));
+  const T product = DTypeTrait<T>::from(to_float(update) * to_float(gate));
-  return dtype_trait<T>::from(to_float(residual) + to_float(product));
+  return DTypeTrait<T>::from(to_float(residual) + to_float(product));
 }
 
 template <typename T, int kVec>

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

@@ -2,14 +2,12 @@
 /// \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
+/// dispatching through `DTypeTrait<T>`. All functions are forced-inline
 /// device functions.
 
 #pragma once
 #include <sgl_kernel/type.cuh>
 
-#include <cmath>
-
 namespace device::math {
 
 /// \brief Constant: log2(e)
@@ -27,49 +25,49 @@ 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);
+  return DTypeTrait<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);
+  return DTypeTrait<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);
+  return DTypeTrait<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);
+  return DTypeTrait<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);
+  return DTypeTrait<T>::rsqrt(a);
 }
 
 /// \brief Returns e^a.
 template <typename T>
 SGL_DEVICE T exp(T a) {
-  return dtype_trait<T>::exp(a);
+  return DTypeTrait<T>::exp(a);
 }
 
 /// \brief Returns sin(a).
 template <typename T>
 SGL_DEVICE T sin(T a) {
-  return dtype_trait<T>::sin(a);
+  return DTypeTrait<T>::sin(a);
 }
 
 /// \brief Returns cos(a).
 template <typename T>
 SGL_DEVICE T cos(T a) {
-  return dtype_trait<T>::cos(a);
+  return DTypeTrait<T>::cos(a);
 }
 
 }  // namespace device::math

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

@@ -52,10 +52,10 @@ struct DTypeRef;
 struct DeviceRef;
 
 template <typename T>
-struct _dtype_trait {};
+struct DLDataTypeTrait {};
 
 template <std::integral T>
-struct _dtype_trait<T> {
+struct DLDataTypeTrait<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),
@@ -63,45 +63,45 @@ struct _dtype_trait<T> {
 };
 
 template <std::floating_point T>
-struct _dtype_trait<T> {
+struct DLDataTypeTrait<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> {
+struct DLDataTypeTrait<fp16_t> {
   inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLFloat, .bits = 16, .lanes = 1};
 };
 template <>
-struct _dtype_trait<bf16_t> {
+struct DLDataTypeTrait<bf16_t> {
   inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLBfloat, .bits = 16, .lanes = 1};
 };
 template <>
-struct _dtype_trait<fp8_e4m3_t> {
+struct DLDataTypeTrait<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> {
+struct DLDataTypeTrait<fp16_t> {
   inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLFloat, .bits = 16, .lanes = 1};
 };
 template <>
-struct _dtype_trait<bf16_t> {
+struct DLDataTypeTrait<bf16_t> {
   inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLBfloat, .bits = 16, .lanes = 1};
 };
 #endif
 
 template <DLDeviceType Code>
-struct _device_trait {
+struct DLDeviceTrait {
   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...};
+inline constexpr auto kDTypeList = std::array<DLDataType, sizeof...(Ts)>{DLDataTypeTrait<Ts>::value...};
 
 template <DLDeviceType... Codes>
-inline constexpr auto kDeviceList = std::array<DLDevice, sizeof...(Codes)>{_device_trait<Codes>::value...};
+inline constexpr auto kDeviceList = std::array<DLDevice, sizeof...(Codes)>{DLDeviceTrait<Codes>::value...};
 
 template <typename T>
 struct PrintAbleSpan {
@@ -176,7 +176,7 @@ inline auto& operator<<(std::ostream& os, PrintAbleSpan<T> span) {
 /// \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;
+  return dtype == details::DLDataTypeTrait<T>::value;
 }
 
 /**

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

@@ -1,37 +1,35 @@
 /// \file type.cuh
 /// \brief Dtype trait system for CUDA scalar/packed types.
 ///
-/// `dtype_trait<T>` provides per-type metadata: packed type alias,
+/// `DTypeTrait<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>
+#include <concepts>
-struct dtype_trait {};
+#include <cstddef>
+#include <limits>
+#include <type_traits>
 
-#define SGL_REGISTER_DTYPE_TRAIT(TYPE, PACK2, ...)  \
+template <typename T>
-  template <>                                       \
+struct DTypeTrait {};
-  struct dtype_trait<TYPE> {                        \
+
-    using self_t = TYPE;                            \
+#define SGL_REGISTER_PACKED(SELF, PACKED) \
-    using packed_t = PACK2;                         \
+  using self_t = SELF;                    \
+  using packed_t = PACKED
+
+#define SGL_REGISTER_UNPACK(UNPACK, N) \
+  using unpacked_t = UNPACK;           \
+  static constexpr size_t kVecSize = N
+
+#define SGL_REGISTER_FROM_DEFAULT()               \
   template <typename S>                           \
   SGL_DEVICE static self_t from(const S& value) { \
-      return static_cast<TYPE>(value);              \
+    return static_cast<self_t>(value);            \
   }                                               \
-    __VA_ARGS__                                     \
+  static_assert(true)
-  }
-
-#define SGL_REGISTER_TYPE_END static_assert(true)
 
 #define SGL_REGISTER_FROM_FUNCTION(FROM, FN)     \
   SGL_DEVICE static self_t from(const FROM& x) { \
@@ -45,14 +43,33 @@ struct dtype_trait {};
   }                                                \
   static_assert(true)
 
+// Also emits a `kHas_<NAME>` flag so reduction dispatch can detect the op via
+// plain member SFINAE (see details::HasMax below) - hipcc mis-evaluates
+// requires-expressions that probe device functions, so detection must only
+// ever look at data members.
 #define SGL_REGISTER_BINARY_FUNCTION(NAME, FN)                      \
+  static constexpr bool kHas_##NAME = true;                         \
   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(
+template <std::integral T>
-    fp32_t, fp32x2_t, SGL_REGISTER_TYPE_END;  //
+struct DTypeTrait<T> {
+  SGL_REGISTER_PACKED(T, void);
+  SGL_REGISTER_UNPACK(T, 1);
+  SGL_REGISTER_FROM_DEFAULT();
+  SGL_REGISTER_UNARY_FUNCTION(abs, ::abs);
+  SGL_REGISTER_BINARY_FUNCTION(max, ::max);
+  SGL_REGISTER_BINARY_FUNCTION(min, ::min);
+  static constexpr T kZeroBits = 0;
+};
+
+template <>
+struct DTypeTrait<fp32_t> {
+  SGL_REGISTER_PACKED(fp32_t, fp32x2_t);
+  SGL_REGISTER_UNPACK(fp32_t, 1);
+  SGL_REGISTER_FROM_DEFAULT();
   SGL_REGISTER_FROM_FUNCTION(fp16_t, __half2float);
   SGL_REGISTER_FROM_FUNCTION(bf16_t, __bfloat162float);
   SGL_REGISTER_UNARY_FUNCTION(abs, fabsf);
@@ -62,59 +79,276 @@ SGL_REGISTER_DTYPE_TRAIT(
   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_BINARY_FUNCTION(min, fminf);
-SGL_REGISTER_DTYPE_TRAIT(fp16_t, fp16x2_t);
+  static constexpr float kFloatMax = std::numeric_limits<float>::max();
-SGL_REGISTER_DTYPE_TRAIT(bf16_t, bf16x2_t);
+  static constexpr uint32_t kZeroBits = 0x00000000;
+};
 
-/// TODO: Add ROCM implementation
+template <>
-SGL_REGISTER_DTYPE_TRAIT(
+struct DTypeTrait<fp32x2_t> {
-    fp32x2_t, fp32x4_t, SGL_REGISTER_TYPE_END; SGL_REGISTER_FROM_FUNCTION(fp16x2_t, __half22float2);
+  SGL_REGISTER_PACKED(fp32x2_t, fp32x4_t);
-    SGL_REGISTER_FROM_FUNCTION(bf16x2_t, __bfloat1622float2););
+  SGL_REGISTER_UNPACK(fp32_t, 2);
+  SGL_REGISTER_FROM_DEFAULT();
+  SGL_REGISTER_FROM_FUNCTION(fp16x2_t, __half22float2);
+  SGL_REGISTER_FROM_FUNCTION(bf16x2_t, __bfloat1622float2);
+};
 
-SGL_REGISTER_DTYPE_TRAIT(
+template <>
-    fp16x2_t, void, SGL_REGISTER_TYPE_END; SGL_REGISTER_FROM_FUNCTION(fp32x2_t, __float22half2_rn););
+struct DTypeTrait<fp32x4_t> {
+  SGL_REGISTER_PACKED(fp32x4_t, void);
+  SGL_REGISTER_UNPACK(fp32_t, 4);
+  SGL_REGISTER_FROM_DEFAULT();
+};
 
-SGL_REGISTER_DTYPE_TRAIT(
+template <>
-    bf16x2_t, void, SGL_REGISTER_TYPE_END; SGL_REGISTER_FROM_FUNCTION(fp32x2_t, __float22bfloat162_rn););
+struct DTypeTrait<fp16_t> {
+  SGL_REGISTER_PACKED(fp16_t, fp16x2_t);
+  SGL_REGISTER_UNPACK(fp16_t, 1);
+  SGL_REGISTER_FROM_DEFAULT();
+  SGL_REGISTER_FROM_FUNCTION(fp32_t, __float2half_rn);
+  SGL_REGISTER_UNARY_FUNCTION(abs, __habs);
+  SGL_REGISTER_BINARY_FUNCTION(max, __hmax);
+  SGL_REGISTER_BINARY_FUNCTION(min, __hmin);
+  // CUDA fp16 max clamp value
+  static constexpr float kFloatMax = 65504.0f;
+  static constexpr uint16_t kZeroBits = 0x0000;
+};
+
+template <>
+struct DTypeTrait<fp16x2_t> {
+  SGL_REGISTER_PACKED(fp16x2_t, void);
+  SGL_REGISTER_UNPACK(fp16_t, 2);
+  SGL_REGISTER_FROM_DEFAULT();
+  SGL_REGISTER_FROM_FUNCTION(fp32x2_t, __float22half2_rn);
+  SGL_REGISTER_UNARY_FUNCTION(abs, __habs2);
+#ifndef USE_ROCM
+  SGL_REGISTER_BINARY_FUNCTION(add, __hadd2);
+  SGL_REGISTER_BINARY_FUNCTION(max, __hmax2);
+  SGL_REGISTER_BINARY_FUNCTION(min, __hmin2);
+#else
+  // HIP only provides __hmax2/__hmin2 for __hip_bfloat162, not __half2.
+  // No `add` registered on HIP (packed SUM falls back to lane-wise scalar).
+  static constexpr bool kHas_max = true;
+  static constexpr bool kHas_min = true;
+  SGL_DEVICE static self_t max(const self_t& x, const self_t& y) {
+    return self_t{__hmax(x.x, y.x), __hmax(x.y, y.y)};
+  }
+  SGL_DEVICE static self_t min(const self_t& x, const self_t& y) {
+    return self_t{__hmin(x.x, y.x), __hmin(x.y, y.y)};
+  }
+#endif
+};
+
+template <>
+struct DTypeTrait<bf16_t> {
+  SGL_REGISTER_PACKED(bf16_t, bf16x2_t);
+  SGL_REGISTER_UNPACK(bf16_t, 1);
+  SGL_REGISTER_FROM_DEFAULT();
+#ifndef USE_ROCM
+  SGL_REGISTER_FROM_FUNCTION(fp32_t, __float2bfloat16_rn);
+#else
+  // HIP has no _rn-suffixed variant; __float2bfloat16 rounds to nearest.
+  SGL_REGISTER_FROM_FUNCTION(fp32_t, __float2bfloat16);
+#endif
+  SGL_REGISTER_UNARY_FUNCTION(abs, __habs);
+  SGL_REGISTER_BINARY_FUNCTION(max, __hmax);
+  SGL_REGISTER_BINARY_FUNCTION(min, __hmin);
+  // CUDA bf16 max clamp value
+  static constexpr float kFloatMax = 3.38953139e38f;
+  static constexpr uint16_t kZeroBits = 0x0000;
+};
+
+template <>
+struct DTypeTrait<bf16x2_t> {
+  SGL_REGISTER_PACKED(bf16x2_t, void);
+  SGL_REGISTER_UNPACK(bf16_t, 2);
+  SGL_REGISTER_FROM_DEFAULT();
+  SGL_REGISTER_FROM_FUNCTION(fp32x2_t, __float22bfloat162_rn);
+  SGL_REGISTER_UNARY_FUNCTION(abs, __habs2);
+#ifndef USE_ROCM
+  // No `add` on HIP: bf162 __hadd2 is unverified there (packed SUM falls
+  // back to lane-wise scalar).
+  SGL_REGISTER_BINARY_FUNCTION(add, __hadd2);
+#endif
+  SGL_REGISTER_BINARY_FUNCTION(max, __hmax2);
+  SGL_REGISTER_BINARY_FUNCTION(min, __hmin2);
+};
 
 #ifndef USE_ROCM
-SGL_REGISTER_DTYPE_TRAIT(fp8_e4m3_t, fp8x2_e4m3_t);
+template <>
+struct DTypeTrait<fp8_e4m3_t> {
+  SGL_REGISTER_PACKED(fp8_e4m3_t, fp8x2_e4m3_t);
+  SGL_REGISTER_UNPACK(fp8_e4m3_t, 1);
+  SGL_REGISTER_FROM_DEFAULT();
+  // NOTE: CUDA fp8 support explicit cast (i.e. use default from is ok)
+
+  static constexpr float kFloatMax = 448.0f;  // CUDA fp8 max clamp value
+  static constexpr uint8_t kZeroBits = 0x00;
+};
+
+template <>
+struct DTypeTrait<fp8x2_e4m3_t> {
+  SGL_REGISTER_PACKED(fp8x2_e4m3_t, fp8x4_e4m3_t);
+  SGL_REGISTER_UNPACK(fp8_e4m3_t, 2);
+  SGL_REGISTER_FROM_DEFAULT();
+  // NOTE: CUDA fp8 support explicit cast (i.e. use default from is ok)
+};
+
+template <>
+struct DTypeTrait<fp8x4_e4m3_t> {
+  SGL_REGISTER_PACKED(fp8x4_e4m3_t, void);
+  SGL_REGISTER_UNPACK(fp8_e4m3_t, 4);
+  SGL_REGISTER_FROM_DEFAULT();
+  // NOTE: CUDA fp8 support explicit cast (i.e. use default from is ok)
+};
 #endif
 
-#undef SGL_REGISTER_DTYPE_TRAIT
+#undef SGL_REGISTER_PACKED
+#undef SGL_REGISTER_UNPACK
+#undef SGL_REGISTER_FROM_DEFAULT
 #undef SGL_REGISTER_FROM_FUNCTION
+#undef SGL_REGISTER_UNARY_FUNCTION
+#undef SGL_REGISTER_BINARY_FUNCTION
 
 /// \brief Alias: the packed (x2) type for `T`.
 template <typename T>
-using packed_t = typename dtype_trait<T>::packed_t;
+using packed_t = typename DTypeTrait<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
+ * Dispatches through `DTypeTrait<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);
+  return DTypeTrait<To>::from(value);
 }
 
+/**
+ * \brief View a packed value as an array of its `unpacked_t` elements.
+ *
+ * Returns a reference to `value` reinterpreted as `unpacked_t[kVecSize]`,
+ * so element writes propagate back to the original packed value.
+ * Constness of `value` is preserved.
+ */
+template <typename T>
+SGL_DEVICE auto& unpack(T& value) {
+  using Trait = DTypeTrait<std::remove_const_t<T>>;
+  using U = typename Trait::unpacked_t;
+  constexpr size_t kVecSize = Trait::kVecSize;
+  static_assert(sizeof(T) == sizeof(U) * kVecSize, "packed type must be layout-compatible");
+  using A = std::conditional_t<std::is_const_v<T>, const U, U>;
+  return reinterpret_cast<A(&)[kVecSize]>(value);
+}
+
+enum class ReductionOp : uint8_t { SUM, MAX, MIN };
+
+template <ReductionOp Op, typename T>
+struct ReductionTrait {};
+
+namespace details {
+
+// Op detection via the `kHas_*` data members emitted by
+// SGL_REGISTER_BINARY_FUNCTION. Deliberately classic void_t member SFINAE:
+// hipcc mis-evaluates requires-expressions in device instantiation contexts
+// (observed: even `requires { a + b; }` on float came out false), so detection
+// must never probe function-call expressions.
+template <typename T, typename = void>
+struct HasAdd : std::false_type {};
+template <typename T>
+struct HasAdd<T, std::void_t<decltype(DTypeTrait<T>::kHas_add)>> : std::true_type {};
+
+template <typename T, typename = void>
+struct HasMax : std::false_type {};
+template <typename T>
+struct HasMax<T, std::void_t<decltype(DTypeTrait<T>::kHas_max)>> : std::true_type {};
+
+template <typename T, typename = void>
+struct HasMin : std::false_type {};
+template <typename T>
+struct HasMin<T, std::void_t<decltype(DTypeTrait<T>::kHas_min)>> : std::true_type {};
+
+template <ReductionOp Op, typename T>
+SGL_DEVICE T reduce_recursive(const T& x, const T& y) {
+  using U = typename DTypeTrait<T>::unpacked_t;
+  constexpr size_t kVecSize = DTypeTrait<T>::kVecSize;
+  static_assert(kVecSize > 1, "unsupported scalar type for reduction");
+  using Trait = ReductionTrait<Op, U>;
+  auto& x_unpacked = ::device::unpack(x);
+  auto& y_unpacked = ::device::unpack(y);
+  T result{};
+  auto& z_unpacked = ::device::unpack(result);
+#pragma unroll
+  for (size_t i = 0; i < kVecSize; ++i) {
+    z_unpacked[i] = Trait::reduce(x_unpacked[i], y_unpacked[i]);
+  }
+  return result;
+}
+
+}  // namespace details
+
+// Dispatch rules, chosen so correctness never depends on detection:
+// scalars (kVecSize == 1) call the trait member / operator directly - a
+// missing op is a clear compile error at the call line; packed types use the
+// native op when the trait registered one and fall back to lane-wise
+// recursion otherwise (worst case for a mis-detecting compiler is a slightly
+// slower but still correct lane-wise path).
+template <typename T>
+struct ReductionTrait<ReductionOp::SUM, T> {
+  SGL_DEVICE static T reduce(const T& x, const T& y) {
+    if constexpr (details::HasAdd<T>::value) {
+      return DTypeTrait<T>::add(x, y);
+    } else if constexpr (DTypeTrait<T>::kVecSize == 1) {
+      return static_cast<T>(x + y);
+    } else {
+      return details::reduce_recursive<ReductionOp::SUM>(x, y);
+    }
+  }
+};
+
+template <typename T>
+struct ReductionTrait<ReductionOp::MAX, T> {
+  SGL_DEVICE static T reduce(const T& x, const T& y) {
+    if constexpr (DTypeTrait<T>::kVecSize == 1) {
+      return DTypeTrait<T>::max(x, y);
+    } else if constexpr (details::HasMax<T>::value) {
+      return DTypeTrait<T>::max(x, y);
+    } else {
+      return details::reduce_recursive<ReductionOp::MAX>(x, y);
+    }
+  }
+};
+
+template <typename T>
+struct ReductionTrait<ReductionOp::MIN, T> {
+  SGL_DEVICE static T reduce(const T& x, const T& y) {
+    if constexpr (DTypeTrait<T>::kVecSize == 1) {
+      return DTypeTrait<T>::min(x, y);
+    } else if constexpr (details::HasMin<T>::value) {
+      return DTypeTrait<T>::min(x, y);
+    } else {
+      return details::reduce_recursive<ReductionOp::MIN>(x, y);
+    }
+  }
+};
+
 }  // namespace device
 
 // ---------------------------------------------------------------------------
-// FP8 max clamp value — platform-dependent
+// 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;
+inline constexpr float kFP8E4M3Max = 448.0f;
 #else  // USE_ROCM
 #if HIP_FP8_TYPE_FNUZ
-constexpr float kFP8E4M3Max = 224.0f;
+inline constexpr float kFP8E4M3Max = 224.0f;
 #else   // HIP_FP8_TYPE_E4M3
-constexpr float kFP8E4M3Max = 448.0f;
+inline constexpr float kFP8E4M3Max = 448.0f;
 #endif  // HIP_FP8_TYPE_FNUZ
 #endif  // USE_ROCM

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

@@ -65,6 +65,8 @@ using fp16x2_t = __half2;
 using bf16x2_t = __nv_bfloat162;
 using fp8x2_e4m3_t = __nv_fp8x2_e4m3;
 using fp8x2_e5m2_t = __nv_fp8x2_e5m2;
+using fp8x4_e4m3_t = __nv_fp8x4_e4m3;
+using fp8x4_e5m2_t = __nv_fp8x4_e5m2;
 
 using fp32x4_t = float4;
 #else
@@ -78,6 +80,8 @@ using fp16x2_t = half2;
 using bf16x2_t = __hip_bfloat162;
 using fp8x2_e4m3_t = uint16_t;
 using fp8x2_e5m2_t = uint16_t;
+using fp8x4_e4m3_t = uint32_t;
+using fp8x4_e5m2_t = uint32_t;
 using fp32x4_t = float4;
 #endif
 
@@ -214,9 +218,7 @@ SGL_DEVICE auto offset(const void* ptr, U... offset) -> const void* {
 
 }  // namespace pointer
 
-/// PTX pragma that lets the compiler spill registers into otherwise-unused
+/// PTX pragma that lets the compiler spill registers into shared memory
-/// 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\";");
@@ -377,4 +379,11 @@ struct LaunchKernel {
   cudaLaunchAttribute m_attrs[2];
 };
 
+// The empty-true-branch if/else form keeps a trailing `else` in user code
+// bound to the user's `if`, not to the macro's.
+#define CHECK_CUDA(COND)                                              \
+  if (const auto error = (COND); error == ::cudaSuccess) [[likely]] { \
+  } else                                                              \
+    ::host::Error() << "CUDA error: " << ::cudaGetErrorString(error) << ". "
+
 }  // namespace host

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

@@ -1,14 +1,5 @@
 /// \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
 
@@ -83,7 +74,7 @@ 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();
+  os << "Failed at " << location.file_name() << ":" << location.line();
   if constexpr (sizeof...(args) > 0) {
     os << ": ";
     (os << ... << std::forward<Args>(args));
@@ -183,4 +174,38 @@ inline auto irange(T start, T end) {
   return stdv::iota(start, end);
 }
 
+/** \brief Error class for stream-style error logging. */
+struct Error {
+  Error(DebugInfo location = {}) {
+    m_oss << "Failed at " << location.file_name() << ":" << location.line() << ": ";
+  }
+
+  template <typename T>
+  Error& operator<<(T&& arg) {
+    m_oss << std::forward<T>(arg);
+    return *this;
+  }
+
+  [[noreturn]]
+  ~Error() noexcept(false) {
+    throw PanicError(std::move(m_oss).str());
+  }
+
+ private:
+  std::ostringstream m_oss;
+};
+
+/**
+ * \brief 0-overhead CHECK macro for host code. This can avoid unnecessary
+ * instantiation of error messages when the condition is true.
+ *
+ * Usage: CHECK_HOST(ptr != nullptr) << "Pointer must not be null";
+ */
+// The empty-true-branch if/else form keeps a trailing `else` in user code
+// bound to the user's `if`, not to the macro's.
+#define CHECK_HOST(COND) \
+  if (COND) [[likely]] { \
+  } else                 \
+    ::host::Error()
+
 }  // namespace host

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

@@ -5,6 +5,9 @@
 #include <sgl_kernel/math.cuh>
 #include <sgl_kernel/utils.cuh>
 
+#include <cstdint>
+#include <type_traits>
+
 namespace device::warp {
 
 /// \brief Full warp active mask.
@@ -17,40 +20,103 @@ using mask_t = uint64_t;
 #endif
 
 /**
- * \brief Warp-level sum reduction.
+ * \brief Warp-level reduction.
  *
- * On CUDA: uses __shfl_xor_sync with width=32.
+ * On CUDA: uses __shfl_xor_sync with width=32. Full-warp reductions
+ * use a single `redux.sync` instruction when the target supports it.
  * On HIP: uses __shfl_xor with explicit width parameter (supports wave64 sub-groups).
+ * \tparam OP Reduction operation to perform (SUM, MAX, MIN).
+ * \tparam kNumThreads Number of threads as a group.
+ * \tparam kInner Whether to perform within a group or not.
+ * \tparam T Type of the value to reduce.
+ *
+ * \param value The value to reduce.
+ * \param active_mask The active mask of threads participating in the reduction.
+ *
+ * \note We will divide into groups of `kNumThreads`.
+ * e.g. kNumThreads = 8, we have 0..7, 8..15, 16..23, 24..31 as groups.
+ * By reduction is performed within a group. Inter-group reduction will reduce
+ * over the same offset in different groups. e.g. {0, 8, 16, 24} in the above example.
  */
-template <uint32_t kNumThreads = kWarpThreads, typename T>
+template <ReductionOp OP, uint32_t kNumThreads = kWarpThreads, bool kInner = true, typename T>
-SGL_DEVICE T reduce_sum(T value, mask_t active_mask = kFullMask) {
+SGL_DEVICE T reduce(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
+  using Trait = ReductionTrait<OP, T>;
-  for (int mask = kNumThreads / 2; mask > 0; mask >>= 1)
+
-#ifndef USE_ROCM
+#ifdef SGL_CUDA_ARCH
-    value = value + __shfl_xor_sync(active_mask, value, mask, 32);
+  // CUDA target only
-#else
+  constexpr bool kFullReduction = (kNumThreads == kWarpThreads && kInner) || (kNumThreads == 1 && !kInner);
-    value = value + __shfl_xor(value, mask, kNumThreads);
+  if constexpr (kFullReduction) {
+#if SGL_CUDA_ARCH >= 800
+    // 32 bit integer reduction
+    if constexpr (std::is_integral_v<T> && sizeof(T) <= 4) {
+      if constexpr (OP == ReductionOp::SUM) {
+        return __reduce_add_sync(active_mask, value);
+      } else if constexpr (OP == ReductionOp::MAX) {
+        return __reduce_max_sync(active_mask, value);
+      } else if constexpr (OP == ReductionOp::MIN) {
+        return __reduce_min_sync(active_mask, value);
+      }
+    }
 #endif
+#if SGL_CUDA_ARCH >= 1000 && SGL_CUDA_ARCH < 1100
+    // 32-bit float reduction
+    if constexpr (std::is_same_v<T, float>) {
+      if constexpr (OP == ReductionOp::MAX) {
+        float result;
+        asm("redux.sync.max.f32 %0, %1, %2;" : "=f"(result) : "f"(value), "r"(active_mask));
+        return result;
+      } else if constexpr (OP == ReductionOp::MIN) {
+        float result;
+        asm("redux.sync.min.f32 %0, %1, %2;" : "=f"(result) : "f"(value), "r"(active_mask));
+        return result;
+      }
+    }
+#endif
+  }
+#endif  // redux.sync for CUDA only
+
+  if constexpr (kInner) {
+#pragma unroll
+    for (uint32_t mask = kNumThreads / 2; mask >= 1; mask >>= 1) {
+#ifndef USE_ROCM
+      value = Trait::reduce(value, __shfl_xor_sync(active_mask, value, mask, 32));
+#else
+      value = Trait::reduce(value, __shfl_xor(value, mask, kNumThreads));
+#endif
+    }
+  } else {
+#pragma unroll
+    for (uint32_t mask = kNumThreads; mask <= kWarpThreads / 2; mask <<= 1) {
+#ifndef USE_ROCM
+      value = Trait::reduce(value, __shfl_xor_sync(active_mask, value, mask, 32));
+#else
+      // Inter-group shuffle crosses kNumThreads-sized sub-groups, so the
+      // shuffle width must span the whole warp.
+      value = Trait::reduce(value, __shfl_xor(value, mask, kWarpThreads));
+#endif
+    }
+  }
   return value;
 }
 
-/**
+/** \brief Warp-level sum reduction. */
- * \brief Warp-level max reduction.
+template <uint32_t kNumThreads = kWarpThreads, bool kInner = true, typename T>
- */
+SGL_DEVICE T reduce_sum(T value, mask_t active_mask = kFullMask) {
-template <uint32_t kNumThreads = kWarpThreads, typename T>
+  return reduce<ReductionOp::SUM, kNumThreads, kInner>(value, active_mask);
+}
+
+/** \brief Warp-level max reduction. */
+template <uint32_t kNumThreads = kWarpThreads, bool kInner = true, typename T>
 SGL_DEVICE T reduce_max(T value, mask_t active_mask = kFullMask) {
-  static_assert(kNumThreads >= 1 && kNumThreads <= kWarpThreads);
+  return reduce<ReductionOp::MAX, kNumThreads, kInner>(value, active_mask);
-  static_assert(std::has_single_bit(kNumThreads), "must be pow of 2");
+}
-#pragma unroll
+
-  for (int mask = kNumThreads / 2; mask > 0; mask >>= 1)
+/** \brief Warp-level min reduction. */
-#ifndef USE_ROCM
+template <uint32_t kNumThreads = kWarpThreads, bool kInner = true, typename T>
-    value = math::max(value, __shfl_xor_sync(active_mask, value, mask, 32));
+SGL_DEVICE T reduce_min(T value, mask_t active_mask = kFullMask) {
-#else
+  return reduce<ReductionOp::MIN, kNumThreads, kInner>(value, active_mask);
-    value = math::max(value, __shfl_xor(value, mask, kNumThreads));
-#endif
-  return value;
 }
 
 }  // namespace device::warp

python/sglang/jit_kernel/per_token_group_quant_8bit.py

@@ -16,27 +16,26 @@ from sglang.srt.utils.custom_op import register_custom_op
 if TYPE_CHECKING:
     from tvm_ffi.module import Module
 
-from sglang.jit_kernel.utils import CPP_DTYPE_MAP as OUTPUT_DTYPE_MAP
-
 
 @cache_once
 def _jit_per_token_group_quant_8bit_module(
     dtype: torch.dtype, output_type: torch.dtype, group_size: int
 ) -> Module:
     dtype_arg = make_cpp_args(dtype)
+    out_arg = make_cpp_args(output_type)
     gs_arg = make_cpp_args(group_size)
     pdl_arg = make_cpp_args(is_arch_support_pdl())
-    out_cpp = OUTPUT_DTYPE_MAP[output_type]
     return load_jit(
         "per_token_group_quant_8bit",
         *dtype_arg,
+        *out_arg,
         *gs_arg,
         *pdl_arg,
         cuda_files=["gemm/per_token_group_quant_8bit.cuh"],
         cuda_wrappers=[
             (
                 "per_token_group_quant_8bit",
-                f"per_token_group_quant_8bit<{dtype_arg}, {out_cpp}, {gs_arg}, {pdl_arg}>",
+                f"per_token_group_quant_8bit<{dtype_arg}, {out_arg}, {gs_arg}, {pdl_arg}>",
             )
         ],
     )

python/sglang/jit_kernel/sparse_mla_q8kv8_prefill_sm90.py

@@ -40,7 +40,7 @@ def _q8kv8_cuda_flags() -> list[str]:
     #     torch.utils.cpp_extension's AOT path does (COMMON_NVCC_FLAGS). The JIT
     #     toolchain never defines them, so undefining is a no-op.
     #   * --expt-relaxed-constexpr and -O3: already supplied by the JIT default
-    #     target flags (see utils._get_default_target_flags).
+    #     target flags (see utils.arch.get_default_target_flags).
     #   * --expt-extended-lambda, -lineinfo, -D_USE_MATH_DEFINES: not required
     #     by this single-translation-unit kernel.
     return [

python/sglang/jit_kernel/utils/__init__.py

@@ -0,0 +1,31 @@
+"""Public interface of sglang.jit_kernel.utils."""
+
+from sglang.jit_kernel.utils.arch import (
+    get_jit_cuda_arch,
+    is_arch_support_pdl,
+    override_jit_cuda_arch,
+)
+from sglang.jit_kernel.utils.common import (
+    cache_once,
+    get_ci_test_range,
+    is_hip_runtime,
+    is_musa_runtime,
+    lazy_register_class,
+    should_run_full_tests,
+)
+from sglang.jit_kernel.utils.compile import KERNEL_PATH, load_jit, make_cpp_args
+
+__all__ = [
+    "should_run_full_tests",
+    "get_ci_test_range",
+    "cache_once",
+    "lazy_register_class",
+    "is_hip_runtime",
+    "is_musa_runtime",
+    "make_cpp_args",
+    "load_jit",
+    "override_jit_cuda_arch",
+    "get_jit_cuda_arch",
+    "is_arch_support_pdl",
+    "KERNEL_PATH",
+]

python/sglang/jit_kernel/utils/arch.py

@@ -0,0 +1,119 @@
+"""CUDA/ROCm architecture detection and default compile target flags."""
+
+from __future__ import annotations
+
+import logging
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import List
+
+import torch
+
+from sglang.jit_kernel.utils.common import (
+    cache_once,
+    is_hip_runtime,
+    is_musa_runtime,
+)
+from sglang.srt.utils.common import get_cuda_version
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ArchInfo:
+    major: int
+    minor: int
+    suffix: str
+
+    @property
+    def target_name(self) -> str:
+        return f"{self.major}.{self.minor}{self.suffix}"
+
+    @property
+    def jit_flag(self) -> str:
+        return f"-DSGL_CUDA_ARCH={self.major * 100 + self.minor * 10}"
+
+
+def _cuda_arch_suffix(major: int, minor: int) -> str:
+    """Mirror FlashInfer's `_normalize_cuda_arch`: 9.x/10.x+ -> "a"; 12.0 -> "f"
+    and 12.x (x>0) -> "a" (SM120/SM121 need separate cubins to avoid
+    cudaErrorIllegalInstruction, requires CUDA >= 12.9); below 9.0 -> plain.
+    Unlike FlashInfer, pre-12.9 CUDA falls back to plain instead of raising.
+    """
+    if major == 9:
+        return "a"
+    if major == 12:
+        if get_cuda_version() < (12, 9):
+            return ""
+        return "f" if minor == 0 else "a"
+    if major >= 10:
+        return "a"
+    return ""
+
+
+@cache_once
+def _init_jit_cuda_arch_once():
+    global _CUDA_ARCH
+    try:
+        device = torch.cuda.current_device()
+        major, minor = torch.cuda.get_device_capability(device)
+    except Exception:
+        logger.warning("Cannot detect CUDA architecture.")
+        major, minor = 0, 0  # invalid value to trigger compile error if used
+    # JIT builds target the exact local GPU, so the arch-specific target is
+    # always correct on Hopper+ and unlocks arch-only instructions (redux.f32).
+    # HIP/MUSA capability numbers aren't CUDA SM versions and stay unsuffixed.
+    suffix = (
+        ""
+        if (is_hip_runtime() or is_musa_runtime())
+        else _cuda_arch_suffix(major, minor)
+    )
+    _CUDA_ARCH = ArchInfo(major, minor, suffix)
+
+
+def get_default_target_flags() -> List[str]:
+    if is_hip_runtime():
+        flags = ["-DUSE_ROCM", "-std=c++20", "-O3"]
+        # Detect FP8 type based on GPU architecture
+        try:
+            device = torch.cuda.current_device()
+            gcn_arch = torch.cuda.get_device_properties(device).gcnArchName
+            if "gfx942" in gcn_arch:
+                flags.append("-DHIP_FP8_TYPE_FNUZ=1")
+            else:
+                flags.append("-DHIP_FP8_TYPE_E4M3=1")
+        except Exception:
+            flags.append("-DHIP_FP8_TYPE_E4M3=1")
+        return flags
+    else:
+        return [
+            get_jit_cuda_arch().jit_flag,
+            "-std=c++20",
+            "-O3",
+            "--expt-relaxed-constexpr",
+        ]
+
+
+@contextmanager
+def override_jit_cuda_arch(major: int, minor: int, suffix: str = ""):
+    """A context manager to temporarily override CUDA architecture."""
+    global _CUDA_ARCH
+    old_value = get_jit_cuda_arch()
+    _CUDA_ARCH = ArchInfo(major, minor, suffix)
+    try:
+        yield
+    finally:
+        _CUDA_ARCH = old_value
+
+
+def get_jit_cuda_arch() -> ArchInfo:
+    """Get the current CUDA architecture info."""
+    _init_jit_cuda_arch_once()
+    return _CUDA_ARCH
+
+
+@cache_once
+def is_arch_support_pdl() -> bool:
+    if is_hip_runtime() or is_musa_runtime():
+        return False
+    return get_jit_cuda_arch().major >= 9

python/sglang/jit_kernel/utils/common.py

@@ -0,0 +1,79 @@
+"""Shared helpers: caching decorator, CI test gating, and runtime detection."""
+
+from __future__ import annotations
+
+import functools
+from typing import Any, Callable, Dict, List, TypeVar
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.utils import is_in_ci
+
+F = TypeVar("F", bound=Callable[..., Any])
+T = TypeVar("T")
+
+
+def should_run_full_tests() -> bool:
+    return envs.SGLANG_JIT_KERNEL_RUN_FULL_TESTS.get()
+
+
+def get_ci_test_range(full_range: List[Any], ci_range: List[Any]) -> List[Any]:
+    if should_run_full_tests():
+        return full_range
+    return ci_range if is_in_ci() else full_range
+
+
+def cache_once(fn: F) -> F:
+    """
+    NOTE: `functools.lru_cache` is not compatible with `torch.compile`
+    So we manually implement a simple cache_once decorator to replace it.
+    """
+    result_map = {}
+
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        key = (args, tuple(sorted(kwargs.items())))
+        if key not in result_map:
+            result_map[key] = fn(*args, **kwargs)
+        return result_map[key]
+
+    return wrapper  # type: ignore
+
+
+@cache_once
+def is_hip_runtime() -> bool:
+    return bool(torch.version.hip)
+
+
+@cache_once
+def is_musa_runtime() -> bool:
+    return hasattr(torch.version, "musa") and torch.version.musa is not None
+
+
+_REGISTERED_CLASSES: Dict[type, type] = {}
+
+
+def lazy_register_class(name: str, init_fn: Callable[[], None]) -> Callable[[T], T]:
+    """A decorator to lazily register a tvm-ffi object class on first use.
+
+    `init_fn` runs once (typically JIT-compiling and registering the C++
+    reflection) right before the class is registered under the FFI type key
+    `name`; afterwards instantiation proceeds normally.
+    """
+
+    def decorator(cls: T) -> T:
+        def __new__(cls, *args, **kwargs):
+            import tvm_ffi
+
+            if cls not in _REGISTERED_CLASSES:
+                init_fn()  # lazy initialization before registration once
+                _REGISTERED_CLASSES[cls] = tvm_ffi.register_object(name)(cls)
+            cls = _REGISTERED_CLASSES[cls]
+            return original_new(cls, *args, **kwargs)
+
+        original_new = cls.__new__
+        cls.__new__ = __new__
+        return cls
+
+    return decorator

python/sglang/jit_kernel/utils/compile.py

@@ -1,6 +1,7 @@
+"""JIT compilation: load_jit, the build cache, and C++ template arguments."""
+
 from __future__ import annotations
 
-import functools
 import hashlib
 import importlib.util
 import logging
@@ -8,89 +9,20 @@ import os
 import pathlib
 import re
 from contextlib import contextmanager
-from dataclasses import dataclass
+from typing import TYPE_CHECKING, List, Tuple, TypeAlias, Union
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Dict,
-    List,
-    Optional,
-    Tuple,
-    TypeAlias,
-    TypeVar,
-    Union,
-)
 
 import torch
 
-from sglang.srt.environ import envs
+from sglang.jit_kernel.utils.arch import get_default_target_flags, get_jit_cuda_arch
-from sglang.utils import is_in_ci
+from sglang.jit_kernel.utils.common import cache_once, is_hip_runtime
+from sglang.jit_kernel.utils.deps import REGISTERED_DEPENDENCIES
 
 if TYPE_CHECKING:
     from tvm_ffi import Module
 
-F = TypeVar("F", bound=Callable[..., Any])
-
 logger = logging.getLogger(__name__)
 
 
-def should_run_full_tests() -> bool:
-    return envs.SGLANG_JIT_KERNEL_RUN_FULL_TESTS.get()
-
-
-def get_ci_test_range(full_range: List[Any], ci_range: List[Any]) -> List[Any]:
-    if should_run_full_tests():
-        return full_range
-    return ci_range if is_in_ci() else full_range
-
-
-def cache_once(fn: F) -> F:
-    """
-    NOTE: `functools.lru_cache` is not compatible with `torch.compile`
-    So we manually implement a simple cache_once decorator to replace it.
-    """
-    result_map = {}
-
-    @functools.wraps(fn)
-    def wrapper(*args, **kwargs):
-        key = (args, tuple(sorted(kwargs.items())))
-        if key not in result_map:
-            result_map[key] = fn(*args, **kwargs)
-        return result_map[key]
-
-    return wrapper  # type: ignore
-
-
-_REGISTERED_CLASSES: Dict[type, type] = {}
-T = TypeVar("T")
-
-
-def lazy_register_class(name: str, init_fn: Callable[[], None]) -> Callable[[T], T]:
-    """A decorator to lazily register a tvm-ffi object class on first use.
-
-    `init_fn` runs once (typically JIT-compiling and registering the C++
-    reflection) right before the class is registered under the FFI type key
-    `name`; afterwards instantiation proceeds normally.
-    """
-
-    def decorator(cls: T) -> T:
-        def __new__(cls, *args, **kwargs):
-            import tvm_ffi
-
-            if cls not in _REGISTERED_CLASSES:
-                init_fn()  # lazy initialization before registration once
-                _REGISTERED_CLASSES[cls] = tvm_ffi.register_object(name)(cls)
-            cls = _REGISTERED_CLASSES[cls]
-            return original_new(cls, *args, **kwargs)
-
-        original_new = cls.__new__
-        cls.__new__ = __new__
-        return cls
-
-    return decorator
-
-
 def _make_wrapper(tup: Tuple[str, str]) -> str:
     export_name, kernel_name = tup
     return f"TVM_FFI_DLL_EXPORT_TYPED_FUNC({export_name}, ({kernel_name}));"
@@ -140,7 +72,10 @@ def _local_jit_source_hash(source_files: List[str]) -> str:
 
 @cache_once
 def _resolve_kernel_path() -> pathlib.Path:
-    cur_dir = pathlib.Path(__file__).parent.resolve()
+    # Resolve via the package spec so the lookup is location-independent.
+    spec = importlib.util.find_spec("sglang.jit_kernel")
+    assert spec is not None and spec.origin is not None
+    cur_dir = pathlib.Path(spec.origin).parent.resolve()
 
     # first, try this directory structure
     def _environment_install():
@@ -172,28 +107,28 @@ class CPPArgList(list[str]):
 
 
 CPP_DTYPE_MAP = {
-    torch.float: "fp32_t",
+    torch.float64: "double",
+    torch.float32: "fp32_t",
     torch.float16: "fp16_t",
-    torch.float8_e4m3fn: "fp8_e4m3_t",
     torch.bfloat16: "bf16_t",
+    # The fnuz variants are the ROCm-side torch dtypes; fp8_*_t resolves to
+    # the matching HIP type there (see HIP_FP8_TYPE_* in utils.cuh).
+    torch.float8_e4m3fn: "fp8_e4m3_t",
+    torch.float8_e4m3fnuz: "fp8_e4m3_t",
+    torch.float8_e5m2: "fp8_e5m2_t",
+    torch.float8_e5m2fnuz: "fp8_e5m2_t",
     torch.int8: "int8_t",
+    torch.int16: "int16_t",
     torch.int32: "int32_t",
     torch.int64: "int64_t",
+    torch.uint8: "uint8_t",
+    torch.uint16: "uint16_t",
+    torch.uint32: "uint32_t",
+    torch.uint64: "uint64_t",
+    torch.bool: "bool",
 }
 
 
-# AMD/ROCm note:
-@cache_once
-def is_hip_runtime() -> bool:
-    return bool(torch.version.hip)
-
-
-# MThreads/MUSA note:
-@cache_once
-def is_musa_runtime() -> bool:
-    return hasattr(torch.version, "musa") and torch.version.musa is not None
-
-
 def make_cpp_args(*args: CPP_TEMPLATE_TYPE) -> CPPArgList:
     def _convert(arg: CPP_TEMPLATE_TYPE) -> str:
         if isinstance(arg, bool):
@@ -294,9 +229,9 @@ def load_jit(
     cuda_files = [str((KERNEL_PATH / "csrc" / f).resolve()) for f in cuda_files]
 
     for dep in set(extra_dependencies or []):
-        if dep not in _REGISTERED_DEPENDENCIES:
+        if dep not in REGISTERED_DEPENDENCIES:
             raise ValueError(f"Dependency {dep} is not registered.")
-        extra_include_paths += _REGISTERED_DEPENDENCIES[dep]()
+        extra_include_paths += REGISTERED_DEPENDENCIES[dep]()
 
     module_name = "sgl_kernel_jit_" + "_".join(str(arg) for arg in args)
     if cpp_files or cuda_files:
@@ -338,7 +273,7 @@ def load_jit(
                 cpp_sources=cpp_sources,
                 cuda_sources=cuda_sources,
                 extra_cflags=DEFAULT_CFLAGS + extra_cflags,
-                extra_cuda_cflags=_get_default_target_flags() + extra_cuda_cflags,
+                extra_cuda_cflags=get_default_target_flags() + extra_cuda_cflags,
                 extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags,
                 extra_include_paths=DEFAULT_INCLUDE + extra_include_paths,
                 build_directory=build_directory,
@@ -351,40 +286,13 @@ def load_jit(
                 cpp_files=cpp_files,
                 cuda_files=cuda_files,
                 extra_cflags=DEFAULT_CFLAGS + extra_cflags,
-                extra_cuda_cflags=_get_default_target_flags() + extra_cuda_cflags,
+                extra_cuda_cflags=get_default_target_flags() + extra_cuda_cflags,
                 extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags,
                 extra_include_paths=DEFAULT_INCLUDE + extra_include_paths,
                 build_directory=build_directory,
             )
 
 
-@dataclass
-class ArchInfo:
-    major: int
-    minor: int
-    suffix: str
-
-    @property
-    def target_name(self) -> str:
-        return f"{self.major}.{self.minor}{self.suffix}"
-
-    @property
-    def jit_flag(self) -> str:
-        return f"-DSGL_CUDA_ARCH={self.major * 100 + self.minor * 10}"
-
-
-@cache_once
-def _init_jit_cuda_arch_once():
-    global _CUDA_ARCH
-    try:
-        device = torch.cuda.current_device()
-        major, minor = torch.cuda.get_device_capability(device)
-    except Exception:
-        logger.warning("Cannot detect CUDA architecture.")
-        major, minor = 0, 0  # invalid value to trigger compile error if used
-    _CUDA_ARCH = ArchInfo(major, minor, "")
-
-
 @contextmanager
 def _jit_compile_context():
     if is_hip_runtime():
@@ -400,200 +308,3 @@ def _jit_compile_context():
             os.environ.pop(env_key, None)
         else:
             os.environ[env_key] = old_value
-
-
-# NOTE: this might also be used in __main__.py for compile flags export
-def _get_default_target_flags() -> List[str]:
-    if is_hip_runtime():
-        flags = ["-DUSE_ROCM", "-std=c++20", "-O3"]
-        # Detect FP8 type based on GPU architecture
-        try:
-            device = torch.cuda.current_device()
-            gcn_arch = torch.cuda.get_device_properties(device).gcnArchName
-            if "gfx942" in gcn_arch:
-                flags.append("-DHIP_FP8_TYPE_FNUZ=1")
-            else:
-                flags.append("-DHIP_FP8_TYPE_E4M3=1")
-        except Exception:
-            flags.append("-DHIP_FP8_TYPE_E4M3=1")
-        return flags
-    else:
-        return [
-            get_jit_cuda_arch().jit_flag,
-            "-std=c++20",
-            "-O3",
-            "--expt-relaxed-constexpr",
-        ]
-
-
-@contextmanager
-def override_jit_cuda_arch(major: int, minor: int, suffix: str = ""):
-    """A context manager to temporarily override CUDA architecture."""
-    global _CUDA_ARCH
-    old_value = get_jit_cuda_arch()
-    _CUDA_ARCH = ArchInfo(major, minor, suffix)
-    try:
-        yield
-    finally:
-        _CUDA_ARCH = old_value
-
-
-def get_jit_cuda_arch() -> ArchInfo:
-    """Get the current CUDA architecture info."""
-    _init_jit_cuda_arch_once()
-    return _CUDA_ARCH
-
-
-@cache_once
-def is_arch_support_pdl() -> bool:
-    if is_hip_runtime() or is_musa_runtime():
-        return False
-    return get_jit_cuda_arch().major >= 9
-
-
-def _find_package_root(package: str) -> Optional[pathlib.Path]:
-    spec = importlib.util.find_spec(package)
-    if spec is None or spec.origin is None:
-        return None
-    return pathlib.Path(spec.origin).resolve().parent
-
-
-# NOTE: this might also be used in __main__.py for compile flags export
-_REGISTERED_DEPENDENCIES: Dict[str, Callable[[], List[str]]] = {}
-
-
-def register_dependency(name: str):
-    def decorator(f: Callable[[], List[str]]) -> Callable[[], List[str]]:
-        if name in _REGISTERED_DEPENDENCIES:
-            raise ValueError(f"Dependency {name} already registered")
-        _REGISTERED_DEPENDENCIES[name] = f
-        return f
-
-    return decorator
-
-
-@register_dependency("flashinfer")
-def get_flashinfer_include_paths() -> List[str]:
-    include_paths: List[str] = []
-    flashinfer_root = _find_package_root("flashinfer")
-    if flashinfer_root is None:
-        raise RuntimeError(
-            "Cannot find flashinfer package. Please install flashinfer to get"
-            "the required headers for JIT compilation."
-        )
-
-    flashinfer_data = flashinfer_root / "data"
-    candidates = [
-        flashinfer_data / "include",
-        flashinfer_data / "csrc",
-        flashinfer_data / "cutlass" / "include",
-        flashinfer_data / "cutlass" / "tools" / "util" / "include",
-        flashinfer_data / "spdlog" / "include",
-    ]
-
-    for path in candidates:
-        if not path.exists():
-            raise RuntimeError(
-                f"Required header path {path} for flashinfer dependency not found."
-                " Please check your flashinfer installation."
-            )
-        include_paths.append(str(path))
-    return include_paths
-
-
-def get_mathdx_root() -> Optional[pathlib.Path]:
-    """Locate the NVIDIA Math-DX install (cuBLASDx headers).
-
-    Searches in order:
-      1. ``$MATHDX_HOME`` env var (extracted Math-DX archive root).
-      2. The ``nvidia-mathdx`` PyPI package, if installed.
-    """
-    env_home = os.environ.get("MATHDX_HOME")
-    if env_home:
-        candidate = pathlib.Path(env_home).expanduser().resolve()
-        if (candidate / "include").exists():
-            return candidate
-
-    # The ``nvidia-mathdx`` wheel installs as the namespace package
-    # ``nvidia.mathdx`` (no __init__, so spec.origin is None); resolve it via
-    # submodule_search_locations rather than _find_package_root, which only
-    # handles regular packages.
-    spec = importlib.util.find_spec("nvidia.mathdx")
-    if spec is not None:
-        roots = list(spec.submodule_search_locations or [])
-        if spec.origin is not None:
-            roots.append(str(pathlib.Path(spec.origin).parent))
-        for root in roots:
-            candidate = pathlib.Path(root).resolve()
-            if (candidate / "include").exists():
-                return candidate
-
-    return None
-
-
-@register_dependency("mathdx")
-def get_mathdx_include_paths() -> List[str]:
-    root = get_mathdx_root()
-    if root is None:
-        raise RuntimeError(
-            "Cannot find NVIDIA Math-DX (cuBLASDx) headers. "
-            "Install the `nvidia-mathdx` package "
-            "(`pip install nvidia-mathdx`) or set MATHDX_HOME to an "
-            "extracted Math-DX archive root."
-        )
-    candidates = [root / "include"]
-    cutlass = root / "external" / "cutlass" / "include"
-    if cutlass.exists():
-        candidates.append(cutlass)
-    return [str(p) for p in candidates]
-
-
-@register_dependency("cutlass")
-def get_cutlass_include_paths() -> List[str]:
-    include_paths: List[str] = []
-
-    flashinfer_root = _find_package_root("flashinfer")
-    if flashinfer_root is not None:
-        candidates = [
-            flashinfer_root / "data" / "cutlass" / "include",
-            flashinfer_root / "data" / "cutlass" / "tools" / "util" / "include",
-        ]
-        for path in candidates:
-            if path.exists():
-                include_paths.append(str(path))
-
-    deep_gemm_root = _find_package_root("deep_gemm")
-    if deep_gemm_root is not None:
-        candidate = deep_gemm_root / "include"
-        if candidate.exists():
-            include_paths.append(str(candidate))
-
-    # De-duplicate while preserving order.
-    unique_paths = []
-    seen = set()
-    for path in include_paths:
-        if path in seen:
-            continue
-        seen.add(path)
-        unique_paths.append(path)
-
-    if not unique_paths:
-        raise RuntimeError(
-            "Cannot find CUTLASS headers required for JIT compilation. "
-            "Please install flashinfer or deep_gemm with CUTLASS headers."
-        )
-    return unique_paths
-
-
-__all__ = [
-    "should_run_full_tests",
-    "get_ci_test_range",
-    "cache_once",
-    "is_hip_runtime",
-    "make_cpp_args",
-    "load_jit",
-    "override_jit_cuda_arch",
-    "get_jit_cuda_arch",
-    "is_arch_support_pdl",
-    "register_dependency",
-]

python/sglang/jit_kernel/utils/deps.py

@@ -0,0 +1,141 @@
+"""Header-only dependency registration (flashinfer, cutlass, mathdx, ...)."""
+
+from __future__ import annotations
+
+import importlib.util
+import os
+import pathlib
+from typing import Callable, Dict, List, Optional
+
+
+def _find_package_root(package: str) -> Optional[pathlib.Path]:
+    spec = importlib.util.find_spec(package)
+    if spec is None or spec.origin is None:
+        return None
+    return pathlib.Path(spec.origin).resolve().parent
+
+
+REGISTERED_DEPENDENCIES: Dict[str, Callable[[], List[str]]] = {}
+
+
+def register_dependency(name: str):
+    def decorator(f: Callable[[], List[str]]) -> Callable[[], List[str]]:
+        if name in REGISTERED_DEPENDENCIES:
+            raise ValueError(f"Dependency {name} already registered")
+        REGISTERED_DEPENDENCIES[name] = f
+        return f
+
+    return decorator
+
+
+@register_dependency("flashinfer")
+def get_flashinfer_include_paths() -> List[str]:
+    include_paths: List[str] = []
+    flashinfer_root = _find_package_root("flashinfer")
+    if flashinfer_root is None:
+        raise RuntimeError(
+            "Cannot find flashinfer package. Please install flashinfer to get"
+            "the required headers for JIT compilation."
+        )
+
+    flashinfer_data = flashinfer_root / "data"
+    candidates = [
+        flashinfer_data / "include",
+        flashinfer_data / "csrc",
+        flashinfer_data / "cutlass" / "include",
+        flashinfer_data / "cutlass" / "tools" / "util" / "include",
+        flashinfer_data / "spdlog" / "include",
+    ]
+
+    for path in candidates:
+        if not path.exists():
+            raise RuntimeError(
+                f"Required header path {path} for flashinfer dependency not found."
+                " Please check your flashinfer installation."
+            )
+        include_paths.append(str(path))
+    return include_paths
+
+
+def get_mathdx_root() -> Optional[pathlib.Path]:
+    """Locate the NVIDIA Math-DX install (cuBLASDx headers).
+
+    Searches in order:
+      1. ``$MATHDX_HOME`` env var (extracted Math-DX archive root).
+      2. The ``nvidia-mathdx`` PyPI package, if installed.
+    """
+    env_home = os.environ.get("MATHDX_HOME")
+    if env_home:
+        candidate = pathlib.Path(env_home).expanduser().resolve()
+        if (candidate / "include").exists():
+            return candidate
+
+    # The ``nvidia-mathdx`` wheel installs as the namespace package
+    # ``nvidia.mathdx`` (no __init__, so spec.origin is None); resolve it via
+    # submodule_search_locations rather than _find_package_root, which only
+    # handles regular packages.
+    spec = importlib.util.find_spec("nvidia.mathdx")
+    if spec is not None:
+        roots = list(spec.submodule_search_locations or [])
+        if spec.origin is not None:
+            roots.append(str(pathlib.Path(spec.origin).parent))
+        for root in roots:
+            candidate = pathlib.Path(root).resolve()
+            if (candidate / "include").exists():
+                return candidate
+
+    return None
+
+
+@register_dependency("mathdx")
+def get_mathdx_include_paths() -> List[str]:
+    root = get_mathdx_root()
+    if root is None:
+        raise RuntimeError(
+            "Cannot find NVIDIA Math-DX (cuBLASDx) headers. "
+            "Install the `nvidia-mathdx` package "
+            "(`pip install nvidia-mathdx`) or set MATHDX_HOME to an "
+            "extracted Math-DX archive root."
+        )
+    candidates = [root / "include"]
+    cutlass = root / "external" / "cutlass" / "include"
+    if cutlass.exists():
+        candidates.append(cutlass)
+    return [str(p) for p in candidates]
+
+
+@register_dependency("cutlass")
+def get_cutlass_include_paths() -> List[str]:
+    include_paths: List[str] = []
+
+    flashinfer_root = _find_package_root("flashinfer")
+    if flashinfer_root is not None:
+        candidates = [
+            flashinfer_root / "data" / "cutlass" / "include",
+            flashinfer_root / "data" / "cutlass" / "tools" / "util" / "include",
+        ]
+        for path in candidates:
+            if path.exists():
+                include_paths.append(str(path))
+
+    deep_gemm_root = _find_package_root("deep_gemm")
+    if deep_gemm_root is not None:
+        candidate = deep_gemm_root / "include"
+        if candidate.exists():
+            include_paths.append(str(candidate))
+
+    # De-duplicate while preserving order.
+    unique_paths = []
+    seen = set()
+    for path in include_paths:
+        if path in seen:
+            continue
+        seen.add(path)
+        unique_paths.append(path)
+
+    if not unique_paths:
+        raise RuntimeError(
+            "Cannot find CUTLASS headers required for JIT compilation. "
+            "Please install flashinfer or deep_gemm with CUTLASS headers."
+        )
+    return unique_paths