[Minor] Enhance JIT kernel and add dev docs (#14570)

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

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{};
 

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