[Feature] JIT rmsnorm update (with claude) (#21834)

2026-04-01 23:40:00 +08:00
parent 4f5b55e379
commit 20f4193589
8 changed files with 321 additions and 393 deletions
@@ -4,6 +4,7 @@
 #include <sgl_kernel/runtime.cuh>
 #include <sgl_kernel/tile.cuh>
 #include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>

 #include <sgl_kernel/impl/norm.cuh>

@@ -47,6 +48,130 @@ __global__ void rmsnorm_cta(const RMSNormParams __grid_constant__ params) {
  PDLTriggerSecondary<kUsePDL>();  // launch secondary kernel
 }

+// Pre-Blackwell: 16B vector, each thread loads/stores twice
+template <int64_t kDim, bool kUsePDL, typename Float>
+__global__ __launch_bounds__(kDim / 16) void rmsnorm_cta_double(const RMSNormParams __grid_constant__ params) {
+  using namespace device;
+  using Float2 = packed_t<Float>;
+  using Storage = AlignedVector<Float2, 4>;
+
+  constexpr auto kNumThreads = kDim / 16;
+  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[32];
+
+  PDLWaitPrimary<kUsePDL>();
+
+  const auto input_ptr = pointer::offset<Float>(input, blockIdx.x * input_stride);
+  const auto output_ptr = pointer::offset<Float>(output, blockIdx.x * output_stride);
+
+  const auto input_first = gmem.load(input_ptr, 0);
+  const auto input_second = gmem.load(input_ptr, 1);
+  const auto weight_first = gmem.load(weight_ptr, 0);
+  const auto weight_second = gmem.load(weight_ptr, 1);
+
+  float sum_of_squares = 0.0f;
+#pragma unroll
+  for (auto j = 0u; j < 4u; ++j) {
+    const auto [x, y] = cast<fp32x2_t>(input_first[j]);
+    sum_of_squares += x * x + y * y;
+  }
+#pragma unroll
+  for (auto j = 0u; j < 4u; ++j) {
+    const auto [x, y] = cast<fp32x2_t>(input_second[j]);
+    sum_of_squares += x * x + y * y;
+  }
+
+  sum_of_squares = warp::reduce_sum(sum_of_squares);
+  const auto warp_id = threadIdx.x / kWarpThreads;
+  smem[warp_id] = sum_of_squares;
+  __syncthreads();
+  if (warp_id == 0) {
+    const auto tx = threadIdx.x;
+    const auto local_sum = tx < kNumWarps ? smem[tx] : 0.0f;
+    sum_of_squares = warp::reduce_sum(local_sum);
+    smem[tx] = math::rsqrt(sum_of_squares / kDim + eps);
+  }
+  __syncthreads();
+  const float norm_factor = smem[warp_id];
+
+  Storage output_first, output_second;
+#pragma unroll
+  for (auto j = 0u; j < 4u; ++j) {
+    const auto [ix, iy] = cast<fp32x2_t>(input_first[j]);
+    const auto [wx, wy] = cast<fp32x2_t>(weight_first[j]);
+    output_first[j] = cast<Float2>(fp32x2_t{ix * norm_factor * wx, iy * norm_factor * wy});
+  }
+#pragma unroll
+  for (auto j = 0u; j < 4u; ++j) {
+    const auto [ix, iy] = cast<fp32x2_t>(input_second[j]);
+    const auto [wx, wy] = cast<fp32x2_t>(weight_second[j]);
+    output_second[j] = cast<Float2>(fp32x2_t{ix * norm_factor * wx, iy * norm_factor * wy});
+  }
+
+  gmem.store(output_ptr, output_first, 0);
+  gmem.store(output_ptr, output_second, 1);
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+// Blackwell: 32B vector, each thread loads/stores once
+template <int64_t kDim, bool kUsePDL, typename Float>
+__global__ __launch_bounds__(kDim / 16) void rmsnorm_cta_wide(const RMSNormParams __grid_constant__ params) {
+  using namespace device;
+  using Float2 = packed_t<Float>;
+  using Storage = AlignedVector<Float2, 8>;
+
+  constexpr auto kNumThreads = kDim / 16;
+  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[32];
+
+  PDLWaitPrimary<kUsePDL>();
+
+  const auto input_ptr = pointer::offset<Float>(input, blockIdx.x * input_stride);
+  const auto output_ptr = pointer::offset<Float>(output, blockIdx.x * output_stride);
+
+  const auto input_vec = gmem.load(input_ptr);
+  const auto weight_vec = gmem.load(weight_ptr);
+
+  float sum_of_squares = 0.0f;
+#pragma unroll
+  for (auto j = 0u; j < 8u; ++j) {
+    const auto [x, y] = cast<fp32x2_t>(input_vec[j]);
+    sum_of_squares += x * x + y * y;
+  }
+
+  sum_of_squares = warp::reduce_sum(sum_of_squares);
+  const auto warp_id = threadIdx.x / kWarpThreads;
+  smem[warp_id] = sum_of_squares;
+  __syncthreads();
+  if (warp_id == 0) {
+    const auto tx = threadIdx.x;
+    const auto local_sum = tx < kNumWarps ? smem[tx] : 0.0f;
+    sum_of_squares = warp::reduce_sum(local_sum);
+    smem[tx] = math::rsqrt(sum_of_squares / kDim + eps);
+  }
+  __syncthreads();
+  const float norm_factor = smem[warp_id];
+
+  Storage output_vec;
+#pragma unroll
+  for (auto j = 0u; j < 8u; ++j) {
+    const auto [ix, iy] = cast<fp32x2_t>(input_vec[j]);
+    const auto [wx, wy] = cast<fp32x2_t>(weight_vec[j]);
+    output_vec[j] = cast<Float2>(fp32x2_t{ix * norm_factor * wx, iy * norm_factor * wy});
+  }
+
+  gmem.store(output_ptr, output_vec);
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
 template <int64_t kDim, bool kUsePDL, typename Float>
 __global__ void rmsnorm_warp(const RMSNormParams __grid_constant__ params) {
  using namespace device;
@@ -178,4 +303,59 @@ struct RMSNormKernel {
  }
 };

+template <int64_t kDim, bool kUsePDL, typename DType>
+struct RMSNormHalfKernel {
+  static_assert(kDim % 512 == 0 && sizeof(DType) == 2);
+#if SGL_ARCH_BLACKWELL_OR_GREATER
+  static constexpr auto kernel = rmsnorm_cta_wide<kDim, kUsePDL, DType>;
+#else
+  static constexpr auto kernel = rmsnorm_cta_double<kDim, kUsePDL, DType>;
+#endif
+  static constexpr auto kBlockSize = static_cast<uint32_t>(kDim / 16);
+
+  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,
+    };
+
+    LaunchKernel(num_tokens, kBlockSize, device.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
 }  // namespace