[Refactor] Clean up JIT kernel utilites (#16884)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
2026-01-13 17:54:16 +08:00
parent 740d3c0b39
commit ba9f6d8f26
30 changed files with 928 additions and 513 deletions
--- a/python/sglang/jit_kernel/csrc/add_constant.cuh
+++ b/python/sglang/jit_kernel/csrc/add_constant.cuh
@@ -1,6 +1,7 @@
-#include <sgl_kernel/tensor.h>   // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For div_ceil, RuntimeCheck
+
 #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>
--- a/python/sglang/jit_kernel/csrc/cuda_wait_value.cuh
+++ b/python/sglang/jit_kernel/csrc/cuda_wait_value.cuh
@@ -1,9 +1,9 @@
 #include <sgl_kernel/tensor.h>
+
 #include <sgl_kernel/utils.cuh>

-#include <cuda_runtime_api.h>
-
 #include <cstdint>
+#include <cuda_runtime_api.h>

 namespace {

--- a/python/sglang/jit_kernel/csrc/elementwise/kvcache.cuh
+++ b/python/sglang/jit_kernel/csrc/elementwise/kvcache.cuh
@@ -1,8 +1,9 @@
 #include <sgl_kernel/tensor.h>
-#include <sgl_kernel/utils.cuh>
 #include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/tile.cuh>
+#include <sgl_kernel/utils.cuh>
 #include <sgl_kernel/vec.cuh>
-#include <sgl_kernel/warp.cuh>

 #include <dlpack/dlpack.h>
 #include <tvm/ffi/container/tensor.h>
@@ -27,8 +28,17 @@ struct StoreKVCacheParams {
 constexpr uint32_t kNumWarps = 4;
 constexpr uint32_t kThreadsPerBlock = kNumWarps * device::kWarpThreads;

+/**
+ * \brief Use a single warp to copy key and value data from source to destination.
+ * Each thread in the warp copies a portion of the data in a coalesced manner.
+ * \tparam kElementBytes The size of each key/value element in bytes.
+ * \param k_src Pointer to the source key data.
+ * \param v_src Pointer to the source value data.
+ * \param k_dst Pointer to the destination key data.
+ * \param v_dst Pointer to the destination value data.
+ */
 template <int64_t kElementBytes>
-__device__ void copy_impl(
+SGL_DEVICE void copy_kv_warp(
    const void* __restrict__ k_src,
    const void* __restrict__ v_src,
    void* __restrict__ k_dst,
@@ -42,31 +52,39 @@ __device__ void copy_impl(

  static_assert(kAlignment > 0, "Element size must be multiple of 4 bytes");

-  using vec_t = aligned_vector<uint32_t, kAlignment / 4>;
+  using vec_t = AlignedStorage<uint32_t, kAlignment / 4>;
  constexpr auto kLoopBytes = sizeof(vec_t) * kWarpThreads;
  constexpr auto kLoopCount = kElementBytes / kLoopBytes;

+  const auto gmem = tile::Memory<vec_t>::warp();
+
 #pragma unroll kLoopCount
  for (int64_t i = 0; i < kLoopCount; ++i) {
-    const auto k = warp::load<vec_t>(pointer::offset(k_src, i * kLoopBytes));
-    const auto v = warp::load<vec_t>(pointer::offset(v_src, i * kLoopBytes));
-    warp::store(pointer::offset(k_dst, i * kLoopBytes), k);
-    warp::store(pointer::offset(v_dst, i * kLoopBytes), v);
+    const auto k = gmem.load(k_src, i);
+    const auto v = gmem.load(v_src, i);
+    gmem.store(k_dst, k, i);
+    gmem.store(v_dst, v, i);
  }

  // handle the epilogue if any
  if constexpr (kLoopCount * kLoopBytes < kElementBytes) {
-    constexpr auto kOffset = kLoopCount * kLoopBytes;
-    if ((threadIdx.x % kWarpThreads) * sizeof(vec_t) < kElementBytes - kOffset) {
-      const auto k = warp::load<vec_t>(pointer::offset(k_src, kOffset));
-      const auto v = warp::load<vec_t>(pointer::offset(v_src, kOffset));
-      warp::store(pointer::offset(k_dst, kOffset), k);
-      warp::store(pointer::offset(v_dst, kOffset), v);
+    if (gmem.in_bound(kElementBytes / sizeof(vec_t), kLoopCount)) {
+      const auto k = gmem.load(k_src, kLoopCount);
+      const auto v = gmem.load(v_src, kLoopCount);
+      gmem.store(k_dst, k, kLoopCount);
+      gmem.store(v_dst, v, kLoopCount);
    }
  }
 }

-// Each warp handles one item
+/**
+ * \brief Kernel to store key-value pairs into the KV cache.
+ * Each element is split into multiple parts to allow parallel memory copy.
+ * \tparam kElementBytes The size of each key/value element in bytes.
+ * \tparam kSplit The number of warps that handle each element.
+ * \tparam kUsePDL Whether to use PDL feature.
+ * \tparam T The data type of the indices (`int32_t` or `int64_t`).
+ */
 template <int64_t kElementBytes, int kSplit, bool kUsePDL, typename T>
 __global__ void store_kvcache(const __grid_constant__ StoreKVCacheParams params) {
  using namespace device;
@@ -89,7 +107,7 @@ __global__ void store_kvcache(const __grid_constant__ StoreKVCacheParams params)
  const auto k_dst = pointer::offset(k_cache, index * stride_cache, split_id * kSplitSize);
  const auto v_dst = pointer::offset(v_cache, index * stride_cache, split_id * kSplitSize);

-  copy_impl<kSplitSize>(k_src, v_src, k_dst, v_dst);
+  copy_kv_warp<kSplitSize>(k_src, v_src, k_dst, v_dst);
  PDLTriggerSecondary<kUsePDL>();
 }

--- a/python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh
+++ b/python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh
@@ -1,42 +1,22 @@
-#include <sgl_kernel/runtime.cuh>
 #include <sgl_kernel/tensor.h>
-#include <sgl_kernel/utils.cuh>
 #include <sgl_kernel/utils.h>
-#include <sgl_kernel/vec.cuh>
-#include <sgl_kernel/warp.cuh>

-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/tile.cuh>
+#include <sgl_kernel/utils.cuh>
+
+#include <sgl_kernel/impl/norm.cuh>
+
 #include <dlpack/dlpack.h>
 #include <tvm/ffi/container/tensor.h>

-#include <cstddef>
 #include <cstdint>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
 #include <type_traits>

 namespace {

-[[maybe_unused]]
-__device__ auto to_float2(nv_bfloat162 x) -> float2 {
-  return __bfloat1622float2(x);
-}
-
-[[maybe_unused]]
-__device__ auto to_float2(half2 x) -> float2 {
-  return __half22float2(x);
-}
-
-template <typename T>
-__device__ auto from_float2(float2 x) -> T {
-  if constexpr (std::is_same_v<T, nv_bfloat162>) {
-    return __float22bfloat162_rn(x);
-  } else if constexpr (std::is_same_v<T, half2>) {
-    return __float22half2_rn(x);
-  } else {
-    static_assert(sizeof(T) == 0, "Unsupported type");
-  }
-}
-
 struct QKNormParams {
  void* __restrict__ q;
  void* __restrict__ k;  // k is offset by (-num_qo_heads * head_dim) elements
@@ -50,52 +30,15 @@ struct QKNormParams {
  uint32_t num_tokens;
 };

-template <int64_t kHeadDim, typename PackedFloat>
-__always_inline __device__ void apply_norm(void* __restrict__ input, const void* __restrict__ weight, float eps) {
-  using namespace device;
-
-  constexpr std::size_t kLoopCount = kHeadDim / (kWarpThreads * 2);
-  static_assert(kHeadDim % (kWarpThreads * 2) == 0);
-
-  float sum_of_squares = 0.0f;
-
-  using vec_t = aligned_vector<PackedFloat, kLoopCount>;
-  const auto input_vec = warp::load<vec_t>(input);
-
-#pragma unroll
-  for (auto i = 0u; i < kLoopCount; ++i) {
-    const auto fp16_input = input_vec[i];
-    const auto fp32_input = to_float2(fp16_input);
-    sum_of_squares += fp32_input.x * fp32_input.x;
-    sum_of_squares += fp32_input.y * fp32_input.y;
-  }
-
-  sum_of_squares = warp::reduce_sum(sum_of_squares);
-  const auto norm_factor = rsqrtf(sum_of_squares / kHeadDim + eps);
-  const auto weight_vec = warp::load<vec_t>(weight);
-  vec_t output_vec;
-
-#pragma unroll
-  for (auto i = 0u; i < kLoopCount; ++i) {
-    const auto fp32_input = to_float2(input_vec[i]);
-    const auto fp32_weight = to_float2(weight_vec[i]);
-    output_vec[i] = from_float2<PackedFloat>({
-        fp32_input.x * norm_factor * fp32_weight.x,
-        fp32_input.y * norm_factor * fp32_weight.y,
-    });
-  }
-
-  warp::store(input, output_vec);
-}
-
 constexpr uint32_t kWarpsPerBlock = 4;
 constexpr uint32_t kThreadsPerBlock = kWarpsPerBlock * device::kWarpThreads;

-template <int64_t kHeadDim, bool kUsePDL, typename PackedFloat, typename Float>
+template <int64_t kHeadDim, bool kUsePDL, typename Float>
 __global__ void fused_qknorm(const QKNormParams __grid_constant__ params) {
  using namespace device;
+  using Storage = norm::StorageType<Float, kHeadDim>;

-  static_assert(sizeof(Float) == 2 && sizeof(PackedFloat) == 4, "Only support FP16/BF16");
+  static_assert(sizeof(Float) == 2, "Only support FP16/BF16");
  const auto& [q, k, q_stride, k_stride, num_qo_heads, num_kv_heads, eps, q_weight, k_weight, num_tokens] = params;

  const auto num_blks = gridDim.x;
@@ -103,6 +46,7 @@ __global__ void fused_qknorm(const QKNormParams __grid_constant__ params) {
  const auto num_q_and_k_heads = num_qo_heads + num_kv_heads;
  const auto num_works = num_q_and_k_heads * num_tokens;
  const auto start_worker_id = blockIdx.x * kWarpsPerBlock + threadIdx.x / kWarpThreads;
+  const auto gmem = tile::Memory<Storage>::warp();

  PDLWaitPrimary<kUsePDL>();  // wait for primary kernel

@@ -113,7 +57,10 @@ __global__ void fused_qknorm(const QKNormParams __grid_constant__ params) {
    const auto input = load_q ? pointer::offset(q, 2 * (token_id * q_stride + head_id * kHeadDim))
                              : pointer::offset(k, 2 * (token_id * k_stride + head_id * kHeadDim));
    const auto weight = load_q ? q_weight : k_weight;
-    apply_norm<kHeadDim, PackedFloat>(input, weight, eps);
+    const auto input_vec = gmem.load(input);
+    const auto weight_vec = gmem.load(weight);
+    const auto output_vec = norm::apply_norm_warp<kHeadDim>(input_vec, weight_vec, eps);
+    gmem.store(input, output_vec);
  }

  PDLTriggerSecondary<kUsePDL>();  // launch secondary kernel
@@ -121,13 +68,9 @@ __global__ void fused_qknorm(const QKNormParams __grid_constant__ params) {

 template <int64_t kHeadDim, bool kUsePDL, typename DType>
 struct QKNormKernel {
-  static_assert(
-      std::is_same_v<DType, half> || std::is_same_v<DType, nv_bfloat16>,
-      "Unsupported DType: QKNormKernel only supports half and nv_bfloat16.");
-  using DType2 = host::PackedDType<DType, 2>::type;
-
-  // only initialize once (static variable) to avoid overhead
-  static constexpr auto kernel = fused_qknorm<kHeadDim, kUsePDL, DType2, DType>;
+  static_assert(std::is_same_v<DType, fp16_t> || std::is_same_v<DType, bf16_t>);
+  static_assert(!host::norm::should_use_cta<DType, kHeadDim>(), "Head dim too large for QKNorm");
+  static constexpr auto kernel = fused_qknorm<kHeadDim, kUsePDL, DType>;

  static void
  run(const tvm::ffi::TensorView q,
@@ -143,38 +86,29 @@ struct QKNormKernel {
    auto D = SymbolicSize{"head_dim"};
    auto Sq = SymbolicSize{"q_stride"};
    auto Sk = SymbolicSize{"k_stride"};
-    auto dtype = SymbolicDType{};
    auto device = SymbolicDevice{};
+    D.set_value(kHeadDim);
+    device.set_options<kDLCUDA>();

-    /*
-     * We need the .template disambiguator here because this call happens in a dependent context.
-     * After switching to with_dtype<DType>(...) (where DType is a template parameter), the chained expression becomes
-     * dependent. In C++, when calling a member function template via ./-> on a dependent expression, the compiler may
-     * otherwise parse <kDLCUDA> as the < operator instead of template arguments. Adding .template forces correct
-     * parsing and fixes compilation errors (often seen with NVCC/clang). Ref:
-     * https://en.cppreference.com/w/cpp/language/dependent_name
-     */
    TensorMatcher({N, Q, D})  // q input
        .with_strides({Sq, D, 1})
-        .with_dtype<DType>(dtype)
-        .template with_device<kDLCUDA>(device)
+        .with_dtype<DType>()
+        .with_device(device)
        .verify(q);
    TensorMatcher({N, K, D})  // k input
        .with_strides({Sk, D, 1})
-        .with_dtype<DType>(dtype)
-        .template with_device<kDLCUDA>(device)
+        .with_dtype<DType>()
+        .with_device(device)
        .verify(k);
    TensorMatcher({D})  // weight
-        .with_dtype<DType>(dtype)
-        .template with_device<kDLCUDA>(device)
+        .with_dtype<DType>()
+        .with_device(device)
        .verify(q_weight)
        .verify(k_weight);

    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
    const auto num_qo_heads = static_cast<uint32_t>(Q.unwrap());
    const auto num_kv_heads = static_cast<uint32_t>(K.unwrap());
-    const auto head_dim = D.unwrap();
-    RuntimeCheck(head_dim == kHeadDim, "Wrong head_dim: ", head_dim, ". Expected:", kHeadDim);

    // NOTE: we offset the k here to reduce computation cost in the kernel
    const auto params = QKNormParams{
--- a/python/sglang/jit_kernel/csrc/elementwise/rmsnorm.cuh
+++ b/python/sglang/jit_kernel/csrc/elementwise/rmsnorm.cuh
@@ -0,0 +1,109 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/tile.cuh>
+#include <sgl_kernel/utils.cuh>
+
+#include <sgl_kernel/impl/norm.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+namespace {
+
+struct RMSNormParams {
+  const void* input;
+  const void* __restrict__ weight;
+  void* output;
+  int64_t input_stride;
+  int64_t output_stride;
+  uint32_t num_tokens;
+  float eps;
+};
+
+template <int64_t kDim, bool kUsePDL, typename Float>
+__global__ void rmsnorm_cta(const RMSNormParams __grid_constant__ params) {
+  using namespace device;
+  using Storage = norm::StorageType<Float, kDim>;
+
+  constexpr auto kNumThreads = host::norm::get_cta_threads<Float, kDim>();
+  constexpr auto kNumWarps = kNumThreads / kWarpThreads;
+
+  const auto& [input, weight_ptr, output, input_stride, output_stride, num_tokens, eps] = params;
+  const auto gmem = tile::Memory<Storage>::cta(kNumThreads);
+  __shared__ float smem[norm::kSmemBufferSize];
+
+  PDLWaitPrimary<kUsePDL>();  // wait for primary kernel
+
+  void* output_ptr = nullptr;
+  Storage output_vec;
+  for (uint32_t i = blockIdx.x; i < num_tokens; i += gridDim.x) {
+    const auto input_ptr = pointer::offset<Float>(input, i * input_stride);
+    const auto input_vec = gmem.load(input_ptr);
+    const auto weight_vec = gmem.load(weight_ptr);
+    if (output_ptr != nullptr) {
+      gmem.store(output_ptr, output_vec);
+    }
+    output_ptr = pointer::offset<Float>(output, i * output_stride);
+    output_vec = norm::apply_norm_cta<kDim>(input_vec, weight_vec, eps, smem, kNumWarps);
+  }
+  gmem.store(output_ptr, output_vec);
+
+  PDLTriggerSecondary<kUsePDL>();  // launch secondary kernel
+}
+
+template <int64_t kDim, bool kUsePDL, typename DType>
+struct RMSNormKernel {
+  static_assert(host::norm::should_use_cta<DType, kDim>(), "Hidden size invalid for RMSNorm");
+  static constexpr auto kernel = rmsnorm_cta<kDim, kUsePDL, DType>;
+
+  static void
+  run(const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView weight,
+      const tvm::ffi::TensorView output,
+      float eps) {
+    using namespace host;
+    auto N = SymbolicSize{"num_tokens"};
+    auto D = SymbolicSize{"hidden_size"};
+    auto SI = SymbolicSize{"input_stride"};
+    auto SO = SymbolicSize{"output_stride"};
+    auto device = SymbolicDevice{};
+    D.set_value(kDim);
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, D})  // input
+        .with_strides({SI, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(input);
+    TensorMatcher({D})  // weight
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(weight);
+    TensorMatcher({N, D})  // output
+        .with_strides({SO, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(output);
+
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto params = RMSNormParams{
+        .input = input.data_ptr(),
+        .weight = weight.data_ptr(),
+        .output = output.data_ptr(),
+        .input_stride = SI.unwrap(),
+        .output_stride = SO.unwrap(),
+        .num_tokens = num_tokens,
+        .eps = eps,
+    };
+
+    static constexpr auto kNumThreads = norm::get_cta_threads<DType, kDim>();
+    static const uint32_t max_occupancy = runtime::get_blocks_per_sm(kernel, kNumThreads);
+    static const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id);
+    const auto num_blocks = std::min<uint32_t>(num_tokens, max_occupancy * kNumSM);
+    LaunchKernel(num_blocks, kNumThreads, device.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace
--- a/python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh
+++ b/python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh
@@ -1,54 +1,62 @@
-#include <sgl_kernel/fp8_utils.cuh>
 #include <sgl_kernel/tensor.h>
-#include <sgl_kernel/utils.cuh>
 #include <sgl_kernel/utils.h>

+#include <sgl_kernel/atomic.cuh>
+#include <sgl_kernel/cta.cuh>
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/tile.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
 #include <cstddef>
 #include <cstdint>
-#include <cub/block/block_reduce.cuh>
-#include <flashinfer/vec_dtypes.cuh>

 namespace {

-using device::atomicMaxFloat;
-using device::blockReduceMax;
-using device::FP8_E4M3_MAX;
+constexpr size_t kBlockSize = 256;

+// each warp will handle 512B data
 template <typename T>
 __global__ void
 per_tensor_absmax_kernel(const T* __restrict__ input, float* __restrict__ output_s, const int64_t num_elements) {
+  using namespace device;
+  constexpr uint32_t VEC_SIZE = 16 / sizeof(T);
+
+  const int64_t gid = blockIdx.x * gridDim.x + threadIdx.x;
+
  float max_value = 0.0f;
-  unsigned int tid = threadIdx.x;
-  unsigned int gid = blockIdx.x * blockDim.x + threadIdx.x;
-  const int grid_size = blockDim.x * gridDim.x;
-
-  constexpr uint32_t vec_size = 16 / sizeof(T);
-  using vec_t = flashinfer::vec_t<T, vec_size>;
-
-  const int32_t num_vec_elems = num_elements / vec_size;
-
-  for (int32_t i = gid; i < num_vec_elems; i += grid_size) {
-    vec_t input_vec;
-    input_vec.cast_load(input + i * vec_size);
-
+  if (gid * VEC_SIZE + VEC_SIZE <= num_elements) {
+    using vec_t = AlignedVector<T, VEC_SIZE>;
+    const auto gmem_in = tile::Memory<vec_t>::thread();
+    const auto input_vec = gmem_in.load(input, gid);
 #pragma unroll
-    for (uint32_t j = 0; j < vec_size; ++j) {
-      float val = static_cast<float>(input_vec[j]);
-      max_value = fmaxf(max_value, fabsf(val));
+    for (uint32_t i = 0; i < VEC_SIZE; ++i) {
+      const float value = static_cast<float>(input_vec[i]);
+      max_value = math::max(max_value, math::abs(value));
+    }
+  } else if (gid * VEC_SIZE < num_elements) {
+    [[unlikely]];  // poorly aligned case, do not optimize
+    const auto remainder = num_elements - gid * VEC_SIZE;
+    for (uint32_t i = 0; i < remainder; ++i) {
+      const float value = static_cast<float>(input[gid * VEC_SIZE + i]);
+      max_value = math::max(max_value, math::abs(value));
    }
  }

-  const int32_t remaining_start = num_vec_elems * vec_size;
-  for (int32_t idx = remaining_start + gid; idx < num_elements; idx += grid_size) {
-    float val = static_cast<float>(input[idx]);
-    max_value = fmaxf(max_value, fabsf(val));
+  // reduce within block and then atomic reduce between blocks
+  __shared__ float smem[kWarpThreads];
+  cta::reduce_max(max_value, smem);
+  if (threadIdx.x == 0) {
+    const auto max_value = smem[0];
+    atomic::max(output_s, max_value / math::FP8_E4M3_MAX);
  }
+}

-  max_value = blockReduceMax(max_value);
-
-  if (tid == 0) {
-    atomicMaxFloat(output_s, max_value / FP8_E4M3_MAX);
-  }
+[[maybe_unused]]
+SGL_DEVICE float fp8_e4m3_clip(float val) {
+  namespace math = device::math;
+  return math::max(math::min(val, math::FP8_E4M3_MAX), -math::FP8_E4M3_MAX);
 }

 template <typename T, typename DST_DTYPE>
@@ -57,123 +65,75 @@ __global__ void per_tensor_quant_fp8_kernel(
    DST_DTYPE* __restrict__ output,
    const float* __restrict__ scale,
    const int64_t num_elements) {
-  const int gid = blockIdx.x * blockDim.x + threadIdx.x;
-  const int grid_size = blockDim.x * gridDim.x;
+  using namespace device;
+  constexpr uint32_t VEC_SIZE = 16 / sizeof(T);
+
+  const int64_t gid = blockIdx.x * blockDim.x + threadIdx.x;
  const float scale_val = 1.0f / (*scale);

-  const uint32_t VEC_SIZE = 16;
-  using vec_t = flashinfer::vec_t<T, VEC_SIZE>;
-
-  const int32_t num_vec_elems = num_elements / VEC_SIZE;
-
-  for (int32_t i = gid; i < num_vec_elems; i += grid_size) {
-    vec_t input_vec;
-    input_vec.cast_load(input + i * VEC_SIZE);
-
-    DST_DTYPE output_arr[VEC_SIZE];
+  if (gid * VEC_SIZE + VEC_SIZE <= num_elements) {
+    using input_vec_t = AlignedVector<T, VEC_SIZE>;
+    using output_vec_t = AlignedVector<DST_DTYPE, VEC_SIZE>;
+    const auto gmem_in = tile::Memory<input_vec_t>::thread();
+    const auto gmem_out = tile::Memory<output_vec_t>::thread();
+    const auto input_vec = gmem_in.load(input, gid);
+    output_vec_t output_vec;
 #pragma unroll
-    for (uint32_t j = 0; j < VEC_SIZE; ++j) {
-      float val = fmax(fmin(static_cast<float>(input_vec[j]) * scale_val, FP8_E4M3_MAX), -FP8_E4M3_MAX);
-#if !defined(USE_ROCM) || defined(HIP_FP8_TYPE_E4M3)
-      output_arr[j] = static_cast<DST_DTYPE>(val);
-#else
-      output_arr[j] = c10::Float8_e4m3fnuz(
-          __hip_cvt_float_to_fp8(val, fp8::fp8_type::__default_saturation, fp8::fp8_type::__default_interpret),
-          c10::Float8_e4m3fnuz::from_bits());
-#endif
+    for (uint32_t i = 0; i < VEC_SIZE; ++i) {
+      const float value = fp8_e4m3_clip(static_cast<float>(input_vec[i]) * scale_val);
+      output_vec[i] = static_cast<DST_DTYPE>(value);
+    }
+    gmem_out.store(output, output_vec, gid);
+  } else if (gid * VEC_SIZE < num_elements) {
+    [[unlikely]];  // poorly aligned case, do not optimize
+    const auto remainder = num_elements - gid * VEC_SIZE;
+    for (uint32_t i = 0; i < remainder; ++i) {
+      const float value = fp8_e4m3_clip(static_cast<float>(input[gid * VEC_SIZE + i]) * scale_val);
+      output[gid * VEC_SIZE + i] = static_cast<DST_DTYPE>(value);
    }
-    *(uint4*)(output + i * VEC_SIZE) = *(uint4*)output_arr;
-  }
-
-  const int32_t remaining_start = num_vec_elems * VEC_SIZE;
-  for (int32_t idx = remaining_start + gid; idx < num_elements; idx += grid_size) {
-    float val = fmax(-FP8_E4M3_MAX, fmin(static_cast<float>(input[idx]) * scale_val, FP8_E4M3_MAX));
-#if !defined(USE_ROCM) || defined(HIP_FP8_TYPE_E4M3)
-    output[idx] = static_cast<DST_DTYPE>(val);
-#else
-    output[idx] = c10::Float8_e4m3fnuz(
-        __hip_cvt_float_to_fp8(val, fp8::fp8_type::__default_saturation, fp8::fp8_type::__default_interpret),
-        c10::Float8_e4m3fnuz::from_bits());
-#endif
  }
 }

-constexpr size_t kBlockSize = 256;
-
-template <bool kIsStatic>
+template <bool kIsStatic, typename DType>
 void per_tensor_quant_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, tvm::ffi::TensorView output_s) {
  using namespace host;

-  const DLDevice device = input.device();
-  RuntimeCheck(device.device_type == kDLCUDA, "input must be on CUDA");
-  RuntimeCheck(input.is_contiguous(), "input must be contiguous");
-
-  const int64_t ndim = input.dim();
-  RuntimeCheck(ndim >= 1, "input.ndim must be >= 1, but got ", ndim);
-
-  RuntimeCheck(output_q.device() == device, "output_q must be on the same device as input");
-  RuntimeCheck(output_q.is_contiguous(), "output_q must be contiguous");
-  RuntimeCheck(output_q.dim() == ndim, "output_q.ndim must match input.ndim");
-  for (int64_t i = 0; i < ndim; ++i) {
-    RuntimeCheck(
-        output_q.size(i) == input.size(i),
-        "output_q.shape mismatch at dim ",
-        i,
-        ": expected ",
-        input.size(i),
-        " but got ",
-        output_q.size(i));
-  }
+  auto device = SymbolicDevice{};
+  auto N = SymbolicSize{"num_elements"};
+  device.set_options<kDLCUDA>();

+  TensorMatcher({N})  //
+      .with_dtype<DType>()
+      .with_device(device)
+      .verify(input);
+  TensorMatcher({N})  //
+      .with_dtype<fp8_e4m3_t>()
+      .with_device(device)
+      .verify(output_q);
  TensorMatcher({1})  //
      .with_dtype<float>()
-      .with_device<kDLCUDA>()
+      .with_device(device)
      .verify(output_s);
-  RuntimeCheck(output_s.device() == device, "output_s must be on the same device as input");

-  const DLDataType in_dtype = input.dtype();
-  const bool in_ok = (in_dtype.code == kDLFloat && in_dtype.bits == 32) ||
-                     (in_dtype.code == kDLFloat && in_dtype.bits == 16) ||
-                     (in_dtype.code == kDLBfloat && in_dtype.bits == 16);
-  RuntimeCheck(in_ok, "input dtype must be fp32/fp16/bf16, but got ", in_dtype);
+  const auto num_elements = N.unwrap();

-  const DLDataType out_dtype = output_q.dtype();
-  RuntimeCheck(
-      out_dtype.code == kDLFloat8_e4m3fn && out_dtype.bits == 8,
-      "output_q dtype must be fp8_e4m3fn, but got ",
-      out_dtype);
+  constexpr size_t kElementsPerBlock = kBlockSize * (16 / sizeof(DType));
+  const uint32_t num_blocks = div_ceil(num_elements, kElementsPerBlock);

-  size_t total_elements = 1;
-  for (const auto s : input.shape()) {
-    RuntimeCheck(s > 0, "Input tensor must be non-empty");
-    total_elements *= static_cast<size_t>(s);
+  if constexpr (!kIsStatic) {
+    LaunchKernel(num_blocks, kBlockSize, device.unwrap())(
+        per_tensor_absmax_kernel<DType>,
+        static_cast<const DType*>(input.data_ptr()),
+        static_cast<float*>(output_s.data_ptr()),
+        static_cast<int64_t>(num_elements));
  }
-  const size_t num_blocks = std::min((total_elements + kBlockSize - 1) / kBlockSize, size_t(1024));

-  auto launch_kernels = [&]<typename T>() {
-    if constexpr (!kIsStatic) {
-      LaunchKernel(num_blocks, kBlockSize, device)(
-          per_tensor_absmax_kernel<T>,
-          static_cast<const T*>(input.data_ptr()),
-          static_cast<float*>(output_s.data_ptr()),
-          static_cast<int64_t>(total_elements));
-    }
-
-    LaunchKernel(num_blocks, kBlockSize, device)(
-        per_tensor_quant_fp8_kernel<T, __nv_fp8_e4m3>,
-        static_cast<const T*>(input.data_ptr()),
-        static_cast<__nv_fp8_e4m3*>(output_q.data_ptr()),
-        static_cast<const float*>(output_s.data_ptr()),
-        static_cast<int64_t>(total_elements));
-  };
-
-  if (in_dtype.code == kDLFloat && in_dtype.bits == 32) {
-    launch_kernels.template operator()<float>();
-  } else if (in_dtype.code == kDLBfloat && in_dtype.bits == 16) {
-    launch_kernels.template operator()<__nv_bfloat16>();
-  } else if (in_dtype.code == kDLFloat && in_dtype.bits == 16) {
-    launch_kernels.template operator()<__half>();
-  }
+  LaunchKernel(num_blocks, kBlockSize, device.unwrap())(
+      per_tensor_quant_fp8_kernel<DType, fp8_e4m3_t>,
+      static_cast<const DType*>(input.data_ptr()),
+      static_cast<fp8_e4m3_t*>(output_q.data_ptr()),
+      static_cast<const float*>(output_s.data_ptr()),
+      static_cast<int64_t>(num_elements));
 }

 }  // namespace
--- a/python/sglang/jit_kernel/csrc/hicache.cuh
+++ b/python/sglang/jit_kernel/csrc/hicache.cuh
@@ -1,7 +1,8 @@
 #include <sgl_kernel/tensor.h>
-#include <sgl_kernel/utils.cuh>
 #include <sgl_kernel/utils.h>

+#include <sgl_kernel/utils.cuh>
+
 #include <dlpack/dlpack.h>

 #include <algorithm>
--- a/python/sglang/jit_kernel/csrc/test_utils.h
+++ b/python/sglang/jit_kernel/csrc/test_utils.h
@@ -1,22 +0,0 @@
-#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