diff --git a/.gitignore b/.gitignore
index 3326fca2c..03e86f9e0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -252,3 +252,9 @@ outputs/
 
 # Eval Cache
 .longbench_cache/
+
+# CUDA kernel develop, profile and debug
+.clangd
+*.nsys-rep
+*.ncu-rep
+*.nvcudmp
diff --git a/docs/developer_guide/JIT_kernels.md b/docs/developer_guide/JIT_kernels.md
new file mode 100644
index 000000000..44f298b9c
--- /dev/null
+++ b/docs/developer_guide/JIT_kernels.md
@@ -0,0 +1,258 @@
+# Development Guide for JIT Kernels
+
+## Environment Setup
+
+We strongly recommend using `clangd` as the language server for JIT kernel development.
+For Ubuntu/Debian, you can download clangd from [apt.llvm.org](https://apt.llvm.org/).
+If you are using VS Code, we recommend installing the `clangd` extension for better IDE integration.
+
+All JIT-related files are located in `python/sglang/jit_kernel`.
+Unlike `sgl-kernel`, which compiles CUDA/C++ binaries ahead of time (AOT), just-in-time (JIT) kernels are compiled at runtime.
+Consequently, a static `compile_commands.json` cannot be generated.
+To enable code completion with `clangd`, run `python -m sglang.jit_kernel` to generate a `.clangd` configuration file in your current directory.
+After generating the file, restart the clangd language server. It should now recognize all JIT kernel files.
+
+## Code Structure
+
+### C++ Implementation
+
+C++ source code is located in `python/sglang/jit_kernel/csrc`.
+Reusable functions should be placed in `python/sglang/jit_kernel/include`.
+
+We use [tvm-ffi](https://github.com/apache/tvm-ffi) for efficient foreign language bindings.
+Refer to the [documentation](https://tvm.apache.org/ffi/) for advanced usage, such as exporting C++ objects.
+Typically, `tvm::ffi::TensorView` is sufficient for passing PyTorch Tensors from Python.
+
+### Python Interface
+
+Python interfaces are defined in `python/sglang/jit_kernel`.
+The `load_jit` utility function in `python/sglang/jit_kernel/utils.py` loads and returns the compiled module.
+To export a C++ function (e.g., `cpp_func`), pass `cuda_wrappers=[("func", "cpp_func")]` to `load_jit`.
+The function can then be called in Python as `module.func`.
+
+### C++ Utilities
+
+The following C++ utilities are available:
+
+#### Integer Range
+
+Similar to PyTorch, we provide an `irange` function to represent an integer range.
+
+```C++
+#include <sgl_kernel/utils.h>
+
+void test() {
+  for (auto i : host::irange(100)) { // [0, 100)
+    // do something
+  }
+  for (auto i : host::irange(0, 100)) { // [0, 100)
+    // do something
+  }
+}
+
+```
+
+#### Runtime Checking
+
+`RuntimeCheck` validates conditions at runtime. It accepts optional arguments for error reporting.
+If the check fails, these arguments are output to aid debugging.
+`RuntimeDeviceCheck` verifies the status of the last kernel launch.
+
+```C++
+#include <sgl_kernel/utils.h>
+#include <sgl_kernel/utils.cuh>
+
+void test() {
+  host::RuntimeCheck(1 + 1 == 2, 1 + 1, " != ", 2);
+  host::RuntimeDeviceCheck();
+  // check the provided `cudaError_t`
+  host::RuntimeDeviceCheck(cudaGetLastError());
+}
+
+```
+
+#### Tensor Checking
+
+`TensorMatcher` provides a readable way to validate and extract tensor shape information.
+
+```cpp
+#include <sgl_kernel/tensor.h>
+
+void test(const tvm::ffi::TensorView k_cache, const tvm::ffi::TensorView v_cache) {
+  using namespace host;
+
+  auto D = SymbolicSize{"D"};  // cache dimension
+  auto N = SymbolicSize{"N"};  // kvcache stride
+  auto dtype = SymbolicDType{};
+  auto device = SymbolicDevice{};
+
+  TensorMatcher({-1, D})  //
+      .with_strides({N, 1})
+      .with_dtype<int32_t, int64_t>(dtype)
+      .with_device<kDLCUDA, kDLCPU>(device)
+      .verify(k_cache)
+      .verify(v_cache);
+}
+```
+
+Configure the `TensorMatcher` with expected stride, dtype, and device properties before verification.
+- If `with_strides` is omitted, the tensor is expected to be contiguous.
+- Template arguments in `with_dtype` restrict the allowed data types.
+- Template arguments in `with_device` restrict the allowed devices.
+- Values passed to `with_xxx` methods enforce equality checks.
+- Passing `-1` for size or stride allows matching any value.
+
+A `Symbolic` variable must resolve to the same value across all verifications.
+Use `.unwrap()` to retrieve the matched value after verification.
+
+> Note: `TensorMatcher` is a temporary expression and should not be stored in a variable.
+
+> Tip: Add `//` at the end of the `TensorMatcher` chain to enforce proper indentation.
+
+#### Kernel Launching
+
+`LaunchKernel::resolve_device` retrieves the current `cudaStream` from PyTorch.
+Kernels can also be launched directly using `LaunchKernel`.
+
+```cpp
+#include <sgl_kernel/utils.cuh>
+
+#include <dlpack/dlpack.h>
+
+__global__ void kernel() {}
+
+void test() {
+  const auto num_blocks = 1;
+  const auto num_threads = 32;
+  const auto dynamic_smem = 0;
+
+  DLDevice dev;  // suppose this is initialized properly
+  host::LaunchKernel(num_blocks, num_threads, dev)(kernel);
+
+  cudaStream_t stream = host::LaunchKernel::resolve_device(dev);
+  host::LaunchKernel(num_blocks, num_threads, stream, dynamic_smem)(kernel);
+}
+
+```
+
+## Add new kernels
+
+This section walks through a complete, end-to-end example of adding a new JIT kernel to the system.
+We use a simple add_constant kernel as a running example, which adds a constant integer value to every element of an input tensor.
+
+Conceptually, the Python interface looks like this:
+
+```python
+def add_constant(src: torch.Tensor, c: int):
+    return src + c
+```
+
+### STEP 1: Write the C++ kernel
+
+Write your CUDA kernel in [jit_kernel/csrc/add_constant.cuh](../../python/sglang/jit_kernel/csrc/add_constant.cuh). For demonstration purposes, we pass the constant value as a template parameter.
+
+```cpp
+#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 <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstddef>
+#include <cstdint>
+
+namespace {
+
+template <int32_t kConstant>
+__global__ void add_constant_kernel(int32_t* dst, const int32_t* src, size_t length) {
+  size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx < length) {
+    dst[idx] = src[idx] + kConstant;
+  }
+}
+
+constexpr size_t kBlockSize = 256;
+
+// You can also use struct with static method as an alternative
+template <int32_t kConstant>
+void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
+  using namespace host;
+
+  // 1. Validate input tensors
+  SymbolicSize N = {"num_elements"};
+  SymbolicDevice device_;
+  TensorMatcher({N})                  // 1D tensor, must be contiguous
+      .with_dtype<int32_t>()          // must be int32
+      .with_device<kDLCUDA>(device_)  // must be on CUDA device
+      .verify(dst)                    // check tensor dst
+      .verify(src);                   // check tensor src
+
+  // 2. Extract required parameters, prepare for kernel launch
+  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
+  RuntimeCheck(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*/)(
+      // kernel function
+      add_constant_kernel<kConstant>,
+      // kernel arguments
+      static_cast<int32_t*>(dst.data_ptr()),
+      static_cast<int32_t*>(src.data_ptr()),
+      num_elements);
+}
+
+}  // namespace
+
+```
+
+### STEP 2: Create Python Interfaces
+
+Next, expose the kernel through a Python wrapper.
+Create a new file at [jit_kernel/add_constant.py](../../python/sglang/jit_kernel/add_constant.py) and expose the needed interfaces.
+
+```python
+from __future__ import annotations
+
+import functools
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import load_jit, make_cpp_args
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@functools.cache
+def _jit_add_constant_module(constant: int) -> Module:
+    args = make_cpp_args(constant)  # pass all the template argument
+    return load_jit(
+        "add_constant",
+        *args,
+        cuda_files=["add_constant.cuh"],
+        cuda_wrappers=[("add_constant", f"add_constant<{args}>")],
+    )
+
+
+def add_constant(src: torch.Tensor, constant: int) -> torch.Tensor:
+    dst = torch.empty_like(src)
+    module = _jit_add_constant_module(constant)
+    module.add_constant(dst, src)
+    return dst
+
+```
+
+### STEP 3: Use your kernel
+
+Finally, import and use the kernel like a regular Python function:
+
+```python
+from sglang.jit_kernel.add_constant import add_constant
+```
+
+For a complete, runnable example, refer to [test_add_constant.py](../../python/sglang/jit_kernel/test_add_constant.py).
diff --git a/python/sglang/jit_kernel/__main__.py b/python/sglang/jit_kernel/__main__.py
new file mode 100644
index 000000000..bacf4f84e
--- /dev/null
+++ b/python/sglang/jit_kernel/__main__.py
@@ -0,0 +1,48 @@
+assert __name__ == "__main__"
+
+
+def generate_clangd():
+    import logging
+    import os
+    import subprocess
+
+    from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path
+
+    from sglang.jit_kernel.utils import DEFAULT_INCLUDE
+
+    logger = logging.getLogger()
+    logger.info("Generating .clangd file...")
+    include_paths = [find_include_path(), find_dlpack_include_path()] + DEFAULT_INCLUDE
+    status = subprocess.run(
+        args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"],
+        capture_output=True,
+        check=True,
+    )
+    compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0]
+    major, minor = compute_cap.split(".")
+    compile_flags = ",\n    ".join(
+        [
+            "-xcuda",
+            f"--cuda-gpu-arch=sm_{major}{minor}",
+            "-std=c++20",
+            "-Wall",
+            "-Wextra",
+        ]
+        + [f"-isystem{path}" for path in include_paths]
+    )
+    clangd_content = f"""
+CompileFlags:
+  Add: [
+    {compile_flags}
+  ]
+"""
+    if os.path.exists(".clangd"):
+        logger.warning(".clangd file already exists, nothing done.")
+        logger.warning(f"suggested content: {clangd_content}")
+    else:
+        with open(".clangd", "w") as f:
+            f.write(clangd_content)
+        logger.info(".clangd file generated.")
+
+
+generate_clangd()
diff --git a/python/sglang/jit_kernel/add_constant.py b/python/sglang/jit_kernel/add_constant.py
new file mode 100644
index 000000000..ac37eac5b
--- /dev/null
+++ b/python/sglang/jit_kernel/add_constant.py
@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+import functools
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import load_jit, make_cpp_args
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@functools.cache
+def _jit_add_constant_module(constant: int) -> Module:
+    args = make_cpp_args(constant)  # pass all the template argument
+    return load_jit(
+        "add_constant",
+        *args,
+        cuda_files=["add_constant.cuh"],
+        cuda_wrappers=[("add_constant", f"add_constant<{args}>")],
+    )
+
+
+def add_constant(src: torch.Tensor, constant: int) -> torch.Tensor:
+    dst = torch.empty_like(src)
+    module = _jit_add_constant_module(constant)
+    module.add_constant(dst, src)
+    return dst
diff --git a/python/sglang/jit_kernel/csrc/add_constant.cuh b/python/sglang/jit_kernel/csrc/add_constant.cuh
new file mode 100644
index 000000000..33f37d119
--- /dev/null
+++ b/python/sglang/jit_kernel/csrc/add_constant.cuh
@@ -0,0 +1,60 @@
+#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 <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstddef>
+#include <cstdint>
+
+namespace {
+
+template <int32_t kConstant>
+__global__ void add_constant_kernel(int32_t* dst, const int32_t* src, size_t length) {
+  size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx < length) {
+    dst[idx] = src[idx] + kConstant;
+  }
+}
+
+constexpr size_t kBlockSize = 256;
+
+// You can also use struct with static method as an alternative
+template <int32_t kConstant>
+void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
+  using namespace host;
+
+  // 1. Validate input tensors
+  SymbolicSize N = {"num_elements"};
+  SymbolicDevice device_;
+  TensorMatcher({N})                  // 1D tensor, must be contiguous
+      .with_dtype<int32_t>()          // must be int32
+      .with_device<kDLCUDA>(device_)  // must be on CUDA device
+      .verify(dst)                    // check tensor dst
+      .verify(src);                   // check tensor src
+
+  // 2. Extract required parameters, prepare for kernel launch
+  const size_t num_elements = N.unwrap();
+  const size_t grid_size = div_ceil(num_elements, kBlockSize);
+  const DLDevice device = device_.unwrap();
+  [[maybe_unused]]  // optional, can be omitted
+  const size_t dynamic_smem = 0;
+  [[maybe_unused]]  // optional, LaunchKernel can auto determine stream from device
+  const cudaStream_t stream = LaunchKernel::resolve_device(device);
+  // some extra runtime checks using host::RuntimeCheck
+  RuntimeCheck(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*/)(
+      // kernel function
+      add_constant_kernel<kConstant>,
+      // kernel arguments
+      static_cast<int32_t*>(dst.data_ptr()),
+      static_cast<int32_t*>(src.data_ptr()),
+      num_elements);
+  // You can also manually check the last CUDA error code via:
+  // RuntimeDeviceCheck();
+}
+
+}  // namespace
diff --git a/python/sglang/jit_kernel/csrc/hicache.cuh b/python/sglang/jit_kernel/csrc/hicache.cuh
index e52ecbd3a..e86fe9973 100644
--- a/python/sglang/jit_kernel/csrc/hicache.cuh
+++ b/python/sglang/jit_kernel/csrc/hicache.cuh
@@ -1,7 +1,6 @@
 #include <sgl_kernel/tensor.h>
 #include <sgl_kernel/utils.cuh>
 #include <sgl_kernel/utils.h>
-#include <sgl_kernel/warp.cuh>
 
 #include <dlpack/dlpack.h>
 
@@ -11,6 +10,105 @@
 #include <cstdint>
 #include <type_traits>
 
+namespace device::warp {
+
+namespace details {
+
+template <std::size_t kUnit>
+inline constexpr auto get_mem_package() {
+  if constexpr (kUnit == 16) {
+    return uint4{};
+  } else if constexpr (kUnit == 8) {
+    return uint2{};
+  } else if constexpr (kUnit == 4) {
+    return uint1{};
+  } else {
+    static_assert(kUnit == 16 || kUnit == 8 || kUnit == 4, "Unsupported memory package size");
+  }
+}
+
+template <std::size_t kBytes, std::size_t kUnit>
+using mem_package_t = decltype(get_mem_package<kUnit>());
+
+__always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 {
+  uint32_t tmp;
+  asm volatile("ld.global.cs.b32 %0,[%1];" : "=r"(tmp) : "l"(src));
+  return uint1{tmp};
+}
+
+__always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 {
+  uint32_t tmp0, tmp1;
+  asm volatile("ld.global.cs.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src));
+  return uint2{tmp0, tmp1};
+}
+
+__always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 {
+  uint32_t tmp0, tmp1, tmp2, tmp3;
+  asm volatile("ld.global.cs.v4.b32 {%0,%1,%2,%3},[%4];" : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3) : "l"(src));
+  return uint4{tmp0, tmp1, tmp2, tmp3};
+}
+
+__always_inline __device__ void store_nc(uint1* __restrict__ dst, const uint1& value) {
+  uint32_t tmp = value.x;
+  asm volatile("st.global.cs.b32 [%0],%1;" ::"l"(dst), "r"(tmp));
+}
+
+__always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& value) {
+  uint32_t tmp0 = value.x;
+  uint32_t tmp1 = value.y;
+  asm volatile("st.global.cs.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1));
+}
+
+__always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& value) {
+  uint32_t tmp0 = value.x;
+  uint32_t tmp1 = value.y;
+  uint32_t tmp2 = value.z;
+  uint32_t tmp3 = value.w;
+  asm volatile("st.global.cs.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3));
+}
+
+}  // namespace details
+
+template <std::size_t kBytes, std::size_t kUnit, std::size_t kThreads>
+__always_inline __device__ auto load_vec(const void* __restrict__ src) {
+  using Package = details::mem_package_t<kBytes, kUnit>;
+  constexpr auto kBytesPerLoop = sizeof(Package) * kThreads;
+  constexpr auto kLoopCount = kBytes / kBytesPerLoop;
+  static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes");
+
+  const auto src_packed = static_cast<const Package*>(src);
+  const auto lane_id = threadIdx.x % kThreads;
+  device_vec<Package, kLoopCount> vec;
+
+#pragma unroll kLoopCount
+  for (std::size_t i = 0; i < kLoopCount; ++i) {
+    const auto j = i * kThreads + lane_id;
+    vec.data[i] = details::load_nc(src_packed + j);
+  }
+
+  return vec;
+}
+
+template <std::size_t kBytes, std::size_t kUnit, std::size_t kThreads, typename Tp>
+__always_inline __device__ void store_vec(void* __restrict__ dst, const Tp& vec) {
+  using Package = details::mem_package_t<kBytes, kUnit>;
+  constexpr auto kBytesPerLoop = sizeof(Package) * kThreads;
+  constexpr auto kLoopCount = kBytes / kBytesPerLoop;
+  static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes");
+  static_assert(std::is_same_v<Tp, device_vec<Package, kLoopCount>>);
+
+  const auto dst_packed = static_cast<Package*>(dst);
+  const auto lane_id = threadIdx.x % kThreads;
+
+#pragma unroll kLoopCount
+  for (std::size_t i = 0; i < kLoopCount; ++i) {
+    const auto j = i * kThreads + lane_id;
+    details::store_nc(dst_packed + j, vec.data[i]);
+  }
+}
+
+}  // namespace device::warp
+
 namespace {
 
 struct HicacheKernelParams {
@@ -142,10 +240,10 @@ struct HiCacheKernel {
       const tvm::ffi::TensorView indices_src) {
     using namespace host;
 
-    auto D = SymbolicSize{"D"};  // cache dimension
-    auto N = SymbolicSize{"N"};  // src kv stride
-    auto M = SymbolicSize{"M"};  // dst kv stride
-    auto L = SymbolicSize{"L"};  // indices length
+    auto D = SymbolicSize{"head dimension"};
+    auto N = SymbolicSize{"src kv stride"};
+    auto M = SymbolicSize{"dst kv stride"};
+    auto L = SymbolicSize{"indices length"};
     auto cache_dtype = SymbolicDType{};
     auto indices_dtype = SymbolicDType{};
     auto indices_device = SymbolicDevice{};
@@ -213,8 +311,8 @@ struct HiCacheKernel {
       const std::size_t kv_dst_stride) {
     using namespace host;
 
-    auto N = SymbolicSize{"N"};  // num layers
-    auto L = SymbolicSize{"L"};  // indices length
+    auto N = SymbolicSize{"num_layers"};
+    auto L = SymbolicSize{"indices length"};
     auto dtype_ = SymbolicDType{};
     auto device_ = SymbolicDevice{};
 
diff --git a/python/sglang/jit_kernel/csrc/test_utils.h b/python/sglang/jit_kernel/csrc/test_utils.h
new file mode 100644
index 000000000..acaf824b2
--- /dev/null
+++ b/python/sglang/jit_kernel/csrc/test_utils.h
@@ -0,0 +1,22 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+namespace {
+
+[[maybe_unused]]
+void assert_same_shape(tvm::ffi::TensorView a, tvm::ffi::TensorView b) {
+  using namespace host;
+  auto N = SymbolicSize{"N"};
+  auto D = SymbolicSize{"D"};
+  TensorMatcher({N, D})  //
+      .with_dtype<float>()
+      .with_device<kDLCUDA>()
+      .verify(a)
+      .verify(b);
+  RuntimeCheck(N.unwrap() > 0 && D.unwrap() > 0);
+}
+
+}  // namespace
diff --git a/python/sglang/jit_kernel/include/sgl_kernel/tensor.h b/python/sglang/jit_kernel/include/sgl_kernel/tensor.h
index 8208149eb..9eb7bc6d5 100644
--- a/python/sglang/jit_kernel/include/sgl_kernel/tensor.h
+++ b/python/sglang/jit_kernel/include/sgl_kernel/tensor.h
@@ -13,7 +13,6 @@
 #include <initializer_list>
 #include <optional>
 #include <ranges>
-#include <source_location>
 #include <span>
 #include <sstream>
 #include <string>
@@ -21,13 +20,21 @@
 #include <type_traits>
 #include <utility>
 
+#ifdef __CUDACC__
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#endif
+
 namespace host {
 
-namespace stdr = std::ranges;
-namespace stdv = std::views;
-
 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;
@@ -37,7 +44,7 @@ struct dtype_trait {};
 
 template <std::integral T>
 struct dtype_trait<T> {
-  inline static constexpr auto value = DLDataType{
+  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};
@@ -45,22 +52,31 @@ struct dtype_trait<T> {
 
 template <std::floating_point T>
 struct dtype_trait<T> {
-  inline static constexpr auto value =
-      DLDataType{.code = DLDataTypeCode::kDLFloat, .bits = static_cast<std::uint8_t>(sizeof(T) * 8), .lanes = 1};
+  inline static constexpr DLDataType value = {
+      .code = DLDataTypeCode::kDLFloat, .bits = static_cast<std::uint8_t>(sizeof(T) * 8), .lanes = 1};
 };
 
-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);
+#ifdef __CUDACC__
+template <>
+struct dtype_trait<__half> {
+  inline static constexpr DLDataType value = {.code = DLDataTypeCode::kDLFloat, .bits = 16, .lanes = 1};
+};
+template <>
+struct dtype_trait<__nv_bfloat16> {
+  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)>{
-    DLDevice{.device_type = static_cast<DLDeviceType>(Codes), .device_id = kAnyDeviceID}...};
+inline constexpr auto kDeviceList = std::array<DLDevice, sizeof...(Codes)>{device_trait<Codes>::value...};
 
 template <typename T>
 struct PrintAbleSpan {
@@ -103,11 +119,13 @@ struct PrintableDevice {
 inline auto& operator<<(std::ostream& os, DLDevice device) {
   const auto& mapping = kDeviceStringMap;
   const auto entry = static_cast<std::size_t>(device.device_type);
-  host::RuntimeCheck(entry < mapping.size());
+  RuntimeCheck(entry < mapping.size());
   const auto name = mapping[entry];
-  host::RuntimeCheck(!name.empty(), "Unknown device: ", int(device.device_type));
+  RuntimeCheck(!name.empty(), "Unknown device: ", int(device.device_type));
   os << name;
-  if (device.device_id != kAnyDeviceID) os << "[" << device.device_id << "]";
+  if (device.device_id != kAnyDeviceID && device.device_type != DLDeviceType::kDLCPU) {
+    os << ":" << device.device_id;
+  }
   return os;
 }
 
@@ -118,7 +136,7 @@ inline auto& operator<<(std::ostream& os, PrintableDevice pd) {
 template <typename T>
 inline auto& operator<<(std::ostream& os, PrintAbleSpan<T> span) {
   os << "[";
-  for (const auto i : stdv::iota(std::size_t{0}, span.data.size())) {
+  for (const auto i : irange(span.data.size())) {
     if (i > 0) {
       os << ", ";
     }
@@ -133,37 +151,58 @@ inline auto& operator<<(std::ostream& os, PrintAbleSpan<T> span) {
 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 {
-    host::RuntimeCheck(!this->has_value(), "Size value already set");
+    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() const -> int64_t {
-    host::RuntimeCheck(this->has_value(), "Size value is not set");
+
+  auto unwrap(DebugInfo info = {}) const -> int64_t {
+    RuntimeCheck(info, this->has_value(), "Size value is not set");
     return m_value;
   }
 
-  SymbolicSize(const SymbolicSize&) = delete;
-  SymbolicSize& operator=(const SymbolicSize&) = delete;
-
-  auto verify(int64_t dim) -> void {
+  auto verify(int64_t value, const char* prefix, int64_t dim) -> void {
     if (this->has_value()) {
-      host::RuntimeCheck(m_value == dim, "Size mismatch: expected ", m_value, " but got ", dim);
+      if (m_value != value) {
+        [[unlikely]];
+        Panic("Size mismatch for ", m_name_str(prefix, dim), ": expected ", m_value, " but got ", value);
+      }
     } else {
-      this->set_value(dim);
+      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;
 };
@@ -175,27 +214,33 @@ inline auto operator==(DLDevice lhs, DLDevice rhs) -> bool {
 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 {
-    host::RuntimeCheck(!this->has_value(), "Dtype value already set");
-    host::RuntimeCheck(
+    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() const -> DLDataType {
-    host::RuntimeCheck(this->has_value(), "Dtype value is not set");
+
+  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...>;
@@ -203,7 +248,7 @@ struct SymbolicDType {
 
   auto verify(DLDataType dtype) -> void {
     if (this->has_value()) {
-      host::RuntimeCheck(m_value == dtype, "DType mismatch: expected ", m_value, " but got ", dtype);
+      RuntimeCheck(m_value == dtype, "DType mismatch: expected ", m_value, " but got ", dtype);
     } else {
       this->set_value(dtype);
     }
@@ -221,10 +266,12 @@ struct SymbolicDType {
 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 {
-    host::RuntimeCheck(!this->has_value(), "Device value already set");
-    host::RuntimeCheck(
+    RuntimeCheck(!this->has_value(), "Device value already set");
+    RuntimeCheck(
         m_check(value),
         "Device value [",
         details::PrintableDevice{value},
@@ -232,20 +279,24 @@ struct SymbolicDevice {
         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() const -> DLDevice {
-    host::RuntimeCheck(this->has_value(), "Device value is not set");
+
+  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...>;
@@ -253,7 +304,7 @@ struct SymbolicDevice {
 
   auto verify(DLDevice device) -> void {
     if (this->has_value()) {
-      host::RuntimeCheck(
+      RuntimeCheck(
           m_value == device,
           "Device mismatch: expected ",
           details::PrintableDevice{m_value},
@@ -313,19 +364,6 @@ struct SizeRef : BaseRef<SymbolicSize> {
       // otherwise, we can match any size
     }
   }
-
-  auto value_or_name(std::size_t dim) const -> std::string {
-    if (const auto value = (**this).get_value()) {
-      return std::to_string(*value);
-    } else {
-      const auto annotation = (**this).get_name();
-      if (annotation.empty()) {
-        return "dim#" + std::to_string(dim);
-      } else {
-        return static_cast<std::string>(annotation);
-      }
-    }
-  }
 };
 
 struct DTypeRef : BaseRef<SymbolicDType> {
@@ -361,7 +399,6 @@ struct TensorMatcher {
   using SizeRef = details::SizeRef;
   using DTypeRef = details::DTypeRef;
   using DeviceRef = details::DeviceRef;
-  using Loc_t = std::source_location;
 
  public:
   TensorMatcher(const TensorMatcher&) = delete;
@@ -371,8 +408,8 @@ struct TensorMatcher {
 
   auto with_strides(std::initializer_list<SizeRef> strides) && -> TensorMatcher&& {
     // no partial update allowed
-    host::RuntimeCheck(m_strides.size() == 0, "Strides already specified");
-    host::RuntimeCheck(m_shape.size() == strides.size(), "Strides size must match shape size");
+    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);
   }
@@ -381,6 +418,7 @@ struct TensorMatcher {
   auto with_dtype(DTypeRef&& dtype) && -> TensorMatcher&& {
     m_init_dtype();
     m_dtype.rebind(*dtype);
+    m_dtype->set_options<Ts...>();
     return std::move(*this);
   }
 
@@ -396,6 +434,7 @@ struct TensorMatcher {
   auto with_device(DeviceRef&& device) && -> TensorMatcher&& {
     m_init_device();
     m_device.rebind(*device);
+    m_device->set_options<Codes...>();
     return std::move(*this);
   }
 
@@ -408,70 +447,70 @@ struct TensorMatcher {
   }
 
   // once we start verification, we cannot modify anymore
-  auto verify(tvm::ffi::TensorView view, Loc_t loc = Loc_t::current()) const&& -> const TensorMatcher&& {
+  auto verify(tvm::ffi::TensorView view, DebugInfo info = {}) const&& -> const TensorMatcher&& {
     try {
-      this->m_verify_impl(view);
+      m_verify_impl(view);
     } catch (PanicError& e) {
       auto oss = std::ostringstream{};
-      oss << "Tensor match failed for " << this->debug_str() << " at " << loc.file_name() << ":" << loc.line()
-          << "\n- Root cause:  " << e.detail();
+      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);
   }
 
-  auto debug_str() const -> std::string {
-    auto oss = std::ostringstream{};
+ private:
+  static auto s_print_tensor(std::ostringstream& oss, tvm::ffi::TensorView view) -> void {
     oss << "Tensor<";
-    std::size_t dim = 0;
-    for (const auto& size_ref : m_shape) {
-      if (dim > 0) {
+    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 << size_ref.value_or_name(dim++);
+      oss << stride;
     }
-    oss << ">";
-    if (m_strides.size() > 0) {
-      oss << " [strides=<";
-      dim = 0;
-      for (const auto& stride_ref : m_strides) {
-        if (dim > 0) {
-          oss << ", ";
-        }
-        oss << stride_ref.value_or_name(dim++);
-      }
-      oss << ">]";
-    }
-    return std::move(oss).str();
+    oss << ">, dtype=" << view.dtype();
+    oss << ", device=" << details::PrintableDevice{view.device()} << "]";
   }
 
- private:
   auto m_verify_impl(tvm::ffi::TensorView view) const -> void {
     const auto dim = static_cast<std::size_t>(view.dim());
-    host::RuntimeCheck(dim == m_shape.size(), "Tensor dimension mismatch: expected ", m_shape.size(), " but got ", dim);
-    for (const auto i : stdv::iota(std::size_t{0}, dim)) {
-      m_shape[i]->verify(view.size(i));
+    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 (this->m_has_strides()) {
-      for (const auto i : stdv::iota(std::size_t{0}, dim)) {
-        m_strides[i]->verify(view.stride(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 {
-      host::RuntimeCheck(view.is_contiguous(), "Tensor is not contiguous as expected");
+      RuntimeCheck(view.is_contiguous(), "Tensor is not contiguous as expected");
     }
-    // since we may use the same matcher to verify again, we will force to check
+    // 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 {
-    host::RuntimeCheck(!m_has_dtype, "DType already specified");
+    RuntimeCheck(!m_has_dtype, "DType already specified");
     m_has_dtype = true;
   }
+
   auto m_init_device() -> void {
-    host::RuntimeCheck(!m_has_device, "Device already specified");
+    RuntimeCheck(!m_has_device, "Device already specified");
     m_has_device = true;
   }
+
   auto m_has_strides() const -> bool {
     return !m_strides.empty();
   }
diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh
index cf03d8c07..2f1051288 100644
--- a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh
+++ b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh
@@ -7,7 +7,6 @@
 
 #include <concepts>
 #include <cstddef>
-#include <source_location>
 #include <type_traits>
 
 namespace device {
@@ -32,60 +31,63 @@ __always_inline __device__ auto offset(const T* ptr, U... offset) -> const void*
 
 }  // namespace pointer
 
+template <typename T, std::size_t N>
+struct device_vec {
+  T data[N];
+};
+
 }  // namespace device
 
 namespace host {
 
-inline auto
-RuntimeDeviceCheck(::cudaError_t error, std::source_location location = std::source_location::current()) -> void {
+inline void RuntimeDeviceCheck(::cudaError_t error, DebugInfo location = {}) {
   if (error != ::cudaSuccess) {
     [[unlikely]];
     ::host::panic(location, "CUDA error: ", ::cudaGetErrorString(error));
   }
 }
 
-inline auto RuntimeCudaCheck(std::source_location location = std::source_location::current()) -> void {
+inline void RuntimeDeviceCheck(DebugInfo location = {}) {
   return RuntimeDeviceCheck(::cudaGetLastError(), location);
 }
 
-template <auto F>
-inline void set_smem_once(std::size_t smem_size) {
-  static const auto last_smem_size = [&] {
-    RuntimeDeviceCheck(::cudaFuncSetAttribute(F, ::cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
-    return smem_size;
-  }();
-  RuntimeCheck(
-      smem_size <= last_smem_size,
-      "Dynamic shared memory size exceeds the previously set maximum size: ",
-      last_smem_size,
-      " bytes");
-}
-
 struct LaunchKernel {
  public:
   explicit LaunchKernel(
-      dim3 grid_dim, dim3 block_dim, DLDevice device, std::size_t dynamic_shared_mem_bytes = 0) noexcept
-      : m_config(s_make_config(grid_dim, block_dim, resolve_device(device), dynamic_shared_mem_bytes)) {}
+      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) noexcept
-      : m_config(s_make_config(grid_dim, block_dim, stream, dynamic_shared_mem_bytes)) {}
+      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));
   }
 
-  LaunchKernel(const LaunchKernel&) = delete;
-  LaunchKernel& operator=(const LaunchKernel&) = delete;
-
   template <typename T, typename... Args>
   auto operator()(T&& kernel, Args&&... args) const -> void {
-    host::RuntimeDeviceCheck(::cudaLaunchKernelEx(&m_config, kernel, std::forward<Args>(args)...));
+    RuntimeDeviceCheck(::cudaLaunchKernelEx(&m_config, kernel, std::forward<Args>(args)...), m_location);
   }
 
  private:
-  static auto
-  s_make_config(dim3 grid_dim, dim3 block_dim, cudaStream_t stream, std::size_t smem) -> cudaLaunchConfig_t {
+  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;
@@ -94,8 +96,10 @@ struct LaunchKernel {
     config.numAttrs = 0;
     return config;
   }
+
   cudaLaunchConfig_t m_config;
-  /// TODO: We can add a queue to store the attributes if needed in the future.
+  const DebugInfo m_location;
+  /// TODO: We can add a queue to store the attributes (e.g. for PDL) if needed in the future.
 };
 
 }  // namespace host
diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.h b/python/sglang/jit_kernel/include/sgl_kernel/utils.h
index fd9723df6..f76c35bfd 100644
--- a/python/sglang/jit_kernel/include/sgl_kernel/utils.h
+++ b/python/sglang/jit_kernel/include/sgl_kernel/utils.h
@@ -1,23 +1,55 @@
 #pragma once
 
+// ref: https://forums.developer.nvidia.com/t/c-20s-source-location-compilation-error-when-using-nvcc-12-1/258026/3
+#ifdef __CUDACC__
+#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>
+
+#undef consteval
+#pragma pop_macro("__cpp_consteval")
+#pragma pop_macro("_NODISCARD")
+#else
+#include <source_location>
+#endif
+
 #include <dlpack/dlpack.h>
 
 #include <concepts>
+#include <cstddef>
 #include <ostream>
+#include <ranges>
 #include <source_location>
 #include <sstream>
 #include <utility>
 
 namespace host {
 
+struct DebugInfo : public std::source_location {
+  DebugInfo(std::source_location loc = std::source_location::current()) : std::source_location(loc) {}
+};
+
 struct PanicError : public std::runtime_error {
  public:
-  // copy and move constructors
   explicit PanicError(std::string msg) : runtime_error(msg), m_message(std::move(msg)) {}
-  auto detail() const -> std::string_view {
-    const auto sv = std::string_view{m_message};
-    const auto pos = sv.find(": ");
-    return pos == std::string_view::npos ? sv : sv.substr(pos + 2);
+  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:
@@ -26,7 +58,7 @@ struct PanicError : public std::runtime_error {
 
 template <typename... Args>
 [[noreturn]]
-inline auto panic(std::source_location location, Args&&... args) -> void {
+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) {
@@ -40,32 +72,42 @@ inline auto panic(std::source_location location, Args&&... args) -> void {
 
 template <typename... Args>
 struct RuntimeCheck {
-  using Loc_t = std::source_location;
   template <typename Cond>
-  explicit RuntimeCheck(Cond&& condition, Args&&... args, Loc_t location = Loc_t::current()) {
-    if (!condition) {
-      [[unlikely]];
-      ::host::panic(location, std::forward<Args>(args)...);
-    }
+  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 <std::signed_integral T, std::signed_integral U>
-inline constexpr auto div_ceil(T a, U b) {
-  return (a + b - 1) / b;
-}
+template <typename Cond, typename... Args>
+explicit RuntimeCheck(DebugInfo, Cond&&, Args&&...) -> RuntimeCheck<Args...>;
 
-template <std::unsigned_integral T, std::unsigned_integral U>
-inline constexpr auto div_ceil(T a, U b) {
-  return (a + b - 1) / b;
-}
+template <typename... Args>
+explicit Panic(Args&&...) -> Panic<Args...>;
 
-inline auto dtype_bytes(DLDataType dtype) -> std::size_t {
-  return static_cast<std::size_t>(dtype.bits / 8);
-}
+template <typename... Args>
+explicit Panic(DebugInfo, Args&&...) -> Panic<Args...>;
 
 namespace pointer {
 
@@ -85,4 +127,26 @@ inline auto offset(const T* ptr, U... offset) -> const void* {
 
 }  // namespace pointer
 
+template <std::integral T, std::integral U>
+inline constexpr auto div_ceil(T a, U b) {
+  return (a + b - 1) / b;
+}
+
+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;
+
+template <std::integral T>
+inline auto irange(T end) {
+  return stdv::iota(static_cast<T>(0), end);
+}
+
+template <std::integral T>
+inline auto irange(T start, T end) {
+  return stdv::iota(start, end);
+}
+
 }  // namespace host
diff --git a/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh b/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh
deleted file mode 100644
index 904531f30..000000000
--- a/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh
+++ /dev/null
@@ -1,145 +0,0 @@
-#pragma once
-#include <sgl_kernel/utils.cuh>
-
-#include <cstddef>
-#include <cstdint>
-#include <type_traits>
-
-namespace device::warp {
-
-namespace details {
-
-template <std::size_t kUnit>
-inline constexpr auto get_mem_package() {
-  if constexpr (kUnit == 16) {
-    return uint4{};
-  } else if constexpr (kUnit == 8) {
-    return uint2{};
-  } else if constexpr (kUnit == 4) {
-    return uint1{};
-  } else {
-    static_assert(kUnit == 16 || kUnit == 8 || kUnit == 4, "Unsupported memory package size");
-  }
-}
-
-inline constexpr auto default_unit_size(std::size_t x) -> std::size_t {
-  if (x % (16 * kWarpThreads) == 0) return 16;
-  if (x % (8 * kWarpThreads) == 0) return 8;
-  if (x % (4 * kWarpThreads) == 0) return 4;
-  return 0;  // trigger static assert in _get_mem_package
-}
-
-template <std::size_t kBytes, std::size_t kUnit>
-using mem_package_t = decltype(get_mem_package<kUnit>());
-
-template <typename T, std::size_t N>
-struct storage_vec {
-  T data[N];
-};
-
-__always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 {
-  uint32_t tmp;
-  asm volatile("ld.global.cs.b32 %0,[%1];" : "=r"(tmp) : "l"(src));
-  return uint1{tmp};
-}
-
-__always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 {
-  uint32_t tmp0, tmp1;
-  asm volatile("ld.global.cs.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src));
-  return uint2{tmp0, tmp1};
-}
-
-__always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 {
-  uint32_t tmp0, tmp1, tmp2, tmp3;
-  asm volatile("ld.global.cs.v4.b32 {%0,%1,%2,%3},[%4];" : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3) : "l"(src));
-  return uint4{tmp0, tmp1, tmp2, tmp3};
-}
-
-__always_inline __device__ void store_nc(uint1* __restrict__ dst, const uint1& value) {
-  uint32_t tmp = value.x;
-  asm volatile("st.global.cs.b32 [%0],%1;" ::"l"(dst), "r"(tmp));
-}
-
-__always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& value) {
-  uint32_t tmp0 = value.x;
-  uint32_t tmp1 = value.y;
-  asm volatile("st.global.cs.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1));
-}
-
-__always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& value) {
-  uint32_t tmp0 = value.x;
-  uint32_t tmp1 = value.y;
-  uint32_t tmp2 = value.z;
-  uint32_t tmp3 = value.w;
-  asm volatile("st.global.cs.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3));
-}
-
-}  // namespace details
-
-template <
-    std::size_t kBytes,
-    std::size_t kUnit = details::default_unit_size(kBytes),
-    std::size_t kThreads = ::device::kWarpThreads>
-__always_inline __device__ void copy(void* __restrict__ dst, const void* __restrict__ src) {
-  using Package = details::mem_package_t<kBytes, kUnit>;
-  constexpr auto kBytesPerLoop = sizeof(Package) * kThreads;
-  constexpr auto kLoopCount = kBytes / kBytesPerLoop;
-  static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes");
-
-  const auto dst_packed = static_cast<Package*>(dst);
-  const auto src_packed = static_cast<const Package*>(src);
-  const auto lane_id = threadIdx.x % kThreads;
-
-#pragma unroll kLoopCount
-  for (std::size_t i = 0; i < kLoopCount; ++i) {
-    const auto j = i * kThreads + lane_id;
-    dst_packed[j] = src_packed[j];
-  }
-}
-
-template <
-    std::size_t kBytes,
-    std::size_t kUnit = details::default_unit_size(kBytes),
-    std::size_t kThreads = ::device::kWarpThreads>
-__always_inline __device__ auto load_vec(const void* __restrict__ src) {
-  using Package = details::mem_package_t<kBytes, kUnit>;
-  constexpr auto kBytesPerLoop = sizeof(Package) * kThreads;
-  constexpr auto kLoopCount = kBytes / kBytesPerLoop;
-  static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes");
-
-  const auto src_packed = static_cast<const Package*>(src);
-  const auto lane_id = threadIdx.x % kThreads;
-  details::storage_vec<Package, kLoopCount> vec;
-
-#pragma unroll kLoopCount
-  for (std::size_t i = 0; i < kLoopCount; ++i) {
-    const auto j = i * kThreads + lane_id;
-    vec.data[i] = details::load_nc(src_packed + j);
-  }
-
-  return vec;
-}
-
-template <
-    std::size_t kBytes,
-    std::size_t kUnit = details::default_unit_size(kBytes),
-    std::size_t kThreads = ::device::kWarpThreads,
-    typename Tp>
-__always_inline __device__ void store_vec(void* __restrict__ dst, const Tp& vec) {
-  using Package = details::mem_package_t<kBytes, kUnit>;
-  constexpr auto kBytesPerLoop = sizeof(Package) * kThreads;
-  constexpr auto kLoopCount = kBytes / kBytesPerLoop;
-  static_assert(kBytes % kBytesPerLoop == 0, "kBytes must be multiple of 128 bytes");
-  static_assert(std::is_same_v<Tp, details::storage_vec<Package, kLoopCount>>);
-
-  const auto dst_packed = static_cast<Package*>(dst);
-  const auto lane_id = threadIdx.x % kThreads;
-
-#pragma unroll kLoopCount
-  for (std::size_t i = 0; i < kLoopCount; ++i) {
-    const auto j = i * kThreads + lane_id;
-    details::store_nc(dst_packed + j, vec.data[i]);
-  }
-}
-
-}  // namespace device::warp
diff --git a/python/sglang/jit_kernel/test_add_constant.py b/python/sglang/jit_kernel/test_add_constant.py
new file mode 100644
index 000000000..d588fc518
--- /dev/null
+++ b/python/sglang/jit_kernel/test_add_constant.py
@@ -0,0 +1,14 @@
+import torch
+
+from sglang.jit_kernel.add_constant import add_constant
+
+
+def main():
+    c = 1024
+    src = torch.arange(0, 1024 + 1, dtype=torch.int32).cuda()
+    dst = add_constant(src, c)
+    assert torch.all(dst == src + c)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/python/sglang/jit_kernel/utils.py b/python/sglang/jit_kernel/utils.py
index 6462cf41c..46390c161 100644
--- a/python/sglang/jit_kernel/utils.py
+++ b/python/sglang/jit_kernel/utils.py
@@ -70,6 +70,36 @@ def load_jit(
     extra_include_paths: List[str] | None = None,
     build_directory: str | None = None,
 ) -> Module:
+    """
+    Loading a JIT module from C++/CUDA source files.
+    We define a wrapper as a tuple of (export_name, kernel_name),
+    where `export_name` is the name used to called from Python,
+    and `kernel_name` is the name of the kernel class in C++/CUDA source.
+
+    :param args: Unique marker of the JIT module. Must be distinct for different kernels.
+    :type args: str
+    :param cpp_files: A list of C++ source files.
+    :type cpp_files: List[str] | None
+    :param cuda_files: A list of CUDA source files.
+    :type cuda_files: List[str] | None
+    :param cpp_wrappers: A list of C++ wrappers, defining the export name and kernel name.
+    :type cpp_wrappers: List[Tuple[str, str]] | None
+    :param cuda_wrappers: A list of CUDA wrappers, defining the export name and kernel name.
+    :type cuda_wrappers: List[Tuple[str, str]] | None
+    :param extra_cflags: Extra C++ compiler flags.
+    :type extra_cflags: List[str] | None
+    :param extra_cuda_cflags: Extra CUDA compiler flags.
+    :type extra_cuda_cflags: List[str] | None
+    :param extra_ldflags: Extra linker flags.
+    :type extra_ldflags: List[str] | None
+    :param extra_include_paths: Extra include paths.
+    :type extra_include_paths: List[str] | None
+    :param build_directory: The build directory for JIT compilation.
+    :type build_directory: str | None
+    :return: A just-in-time(JIT) compiled module.
+    :rtype: Module
+    """
+
     from tvm_ffi.cpp import load_inline
 
     cpp_files = cpp_files or []