rms_norm.cu

#include <algorithm>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <cuda_fp8.h>
#include <cuda_runtime.h>
#include <float.h>
#include <stdio.h>
#include <stdlib.h>
#include <torch/extension.h>
#include <torch/types.h>
#include <vector>

#define WARP_SIZE 32
#define INT4(value) (reinterpret_cast<int4 *>(&(value))[0])
#define FLOAT4(value) (reinterpret_cast<float4 *>(&(value))[0])
#define HALF2(value) (reinterpret_cast<half2 *>(&(value))[0])
#define BFLOAT2(value) (reinterpret_cast<__nv_bfloat162 *>(&(value))[0])
#define LDST128BITS(value) (reinterpret_cast<float4 *>(&(value))[0])

// -------------------------------------- FP32
// -------------------------------------- Warp Reduce Sum
template <const int kWarpSize = WARP_SIZE>
__device__ __forceinline__ float warp_reduce_sum_f32(float val) {
#pragma unroll
  for (int mask = kWarpSize >> 1; mask >= 1; mask >>= 1) {
    val += __shfl_xor_sync(0xffffffff, val, mask);
  }
  return val;
}

// Block reduce sum/max/min device helper for Layer/RMS Norm/Softmax etc.
// grid 1D block 1D, grid(N/256), block(256)
template <const int NUM_THREADS = 256>
__device__ __forceinline__ float block_reduce_sum_f32(float val) {
  // always <= 32 warps per block (limited by 1024 threads per block)
  constexpr int NUM_WARPS = (NUM_THREADS + WARP_SIZE - 1) / WARP_SIZE;
  int warp = threadIdx.x / WARP_SIZE;
  int lane = threadIdx.x % WARP_SIZE;
  static __shared__ float shared[NUM_WARPS];

  val = warp_reduce_sum_f32<WARP_SIZE>(val);
  if (lane == 0)
    shared[warp] = val;
  __syncthreads();
  val = (lane < NUM_WARPS) ? shared[lane] : 0.0f;
  val = warp_reduce_sum_f32<NUM_WARPS>(val);
  return val;
}

// RMS Norm: x: NxK(K=256<1024), y': NxK, y'=x/rms(x) each row
// 1/rms(x) = rsqrtf( sum(x^2)/K ) each row
// grid(N*K/K), block(K<1024) N=batch_size*seq_len, K=hidden_size
// y=y'*g (g: scale)
template <const int NUM_THREADS = 256>
__global__ void rms_norm_f32_kernel(float *x, float *y, float g, int N, int K) {
  int tid = threadIdx.x; // 0..K-1
  int bid = blockIdx.x;  // 0..N-1
  int idx = bid * blockDim.x + threadIdx.x;
  const float epsilon = 1e-5f;

  __shared__ float s_variance;                 // shared within block
  float value = (idx < N * K) ? x[idx] : 0.0f; // load once only
  float variance = value * value;
  variance = block_reduce_sum_f32<NUM_THREADS>(variance);
  if (tid == 0)
    s_variance = rsqrtf(variance / (float)K + epsilon);
  // wait for s_variance in shared memory to be ready for all threads
  __syncthreads();
  if (idx < N * K)
    y[idx] = (value * s_variance) * g;
}

// RMS Norm Vec4: x: NxK(K=256<1024), y': NxK, y'=x/rms(x) each row
// 1/rms(x) = rsqrtf( sum(x^2)/K ) each row
// grid(N*K/K), block(K/4<1024) N=batch_size*seq_len, K=hidden_size
// y=y'*g (g: scale)
template <const int NUM_THREADS = 256 / 4>
__global__ void rms_norm_f32x4_kernel(float *x, float *y, float g, int N,
                                      int K) {
  int tid = threadIdx.x; // 0..K-1
  int bid = blockIdx.x;  // 0..N-1
  int idx = (bid * blockDim.x + threadIdx.x) * 4;
  const float epsilon = 1e-5f;

  __shared__ float s_variance; // shared within block
  float4 reg_x = FLOAT4(x[idx]);
  float variance = (idx < N * K) ? (reg_x.x * reg_x.x + reg_x.y * reg_x.y +
                                    reg_x.z * reg_x.z + reg_x.w * reg_x.w)
                                 : 0.0f;
  variance = block_reduce_sum_f32<NUM_THREADS>(variance);
  if (tid == 0)
    s_variance = rsqrtf(variance / (float)K + epsilon);
  // wait for s_variance in shared memory to be ready for all threads
  __syncthreads();
  float4 reg_y;
  reg_y.x = reg_x.x * s_variance * g;
  reg_y.y = reg_x.y * s_variance * g;
  reg_y.z = reg_x.z * s_variance * g;
  reg_y.w = reg_x.w * s_variance * g;
  if (idx < N * K)
    FLOAT4(y[idx]) = reg_y;
}

// -------------------------------------- FP16
// -------------------------------------- Warp Reduce Sum: Half
template <const int kWarpSize = WARP_SIZE>
__device__ __forceinline__ half warp_reduce_sum_f16_f16(half val) {
#pragma unroll
  for (int mask = kWarpSize >> 1; mask >= 1; mask >>= 1) {
    // val = __hadd(val, __shfl_xor_sync(0xffffffff, val, mask));
    val += __shfl_xor_sync(0xffffffff, val, mask);
  }
  return val;
}

template <const int kWarpSize = WARP_SIZE>
__device__ __forceinline__ float warp_reduce_sum_f16_f32(half val) {
  float val_f32 = __half2float(val);
#pragma unroll
  for (int mask = kWarpSize >> 1; mask >= 1; mask >>= 1) {
    val_f32 += __shfl_xor_sync(0xffffffff, val_f32, mask);
  }
  return val_f32;
}

template <const int NUM_THREADS = 256>
__device__ half block_reduce_sum_f16_f16(half val) {
  // always <= 32 warps per block (limited by 1024 threads per block)
  constexpr int NUM_WARPS = (NUM_THREADS + WARP_SIZE - 1) / WARP_SIZE;
  int warp = threadIdx.x / WARP_SIZE;
  int lane = threadIdx.x % WARP_SIZE;
  static __shared__ half shared[NUM_WARPS];
  // reduce using half dtype within warps
  val = warp_reduce_sum_f16_f16<WARP_SIZE>(val);
  if (lane == 0)
    shared[warp] = val;
  __syncthreads();
  val = (lane < NUM_WARPS) ? shared[lane] : __float2half(0.0f);
  val = warp_reduce_sum_f16_f16<NUM_WARPS>(val);
  return val; // half
}

template <const int NUM_THREADS = 256>
__device__ float block_reduce_sum_f16_f32(half val) {
  // always <= 32 warps per block (limited by 1024 threads per block)
  constexpr int NUM_WARPS = (NUM_THREADS + WARP_SIZE - 1) / WARP_SIZE;
  int warp = threadIdx.x / WARP_SIZE;
  int lane = threadIdx.x % WARP_SIZE;
  static __shared__ float shared[NUM_WARPS];
  // reduce using float dtype within warps
  float val_f32 = warp_reduce_sum_f16_f32<WARP_SIZE>(val);
  if (lane == 0)
    shared[warp] = val_f32;
  __syncthreads();
  val_f32 = (lane < NUM_WARPS) ? shared[lane] : 0.0f;
  val_f32 = warp_reduce_sum_f32<NUM_WARPS>(val_f32);
  return val_f32; // float
}

template <const int NUM_THREADS = 256>
__global__ void rms_norm_f16_f16_kernel(half *x, half *y, float g, int N,
                                        int K) {
  int tid = threadIdx.x; // 0..K-1
  int bid = blockIdx.x;  // 0..N-1
  int idx = bid * blockDim.x + threadIdx.x;
  const half epsilon = __float2half(1e-5f);
  const half g_ = __float2half(g);
  const half K_ = __int2half_rn(K);

  __shared__ half s_variance; // shared within block
  half value = (idx < N * K) ? x[idx] : __float2half(0.0f); // load once only
  half variance = value * value;
  variance = block_reduce_sum_f16_f16<NUM_THREADS>(variance);
  if (tid == 0)
    s_variance = hrsqrt(variance / (K_ + epsilon));
  // wait for s_variance in shared memory to be ready for all threads
  __syncthreads();
  if (idx < N * K)
    y[idx] = (value * s_variance) * g_;
}

template <const int NUM_THREADS = 256>
__global__ void rms_norm_f16x2_f16_kernel(half *x, half *y, float g, int N,
                                          int K) {
  int tid = threadIdx.x; // 0..K-1
  int bid = blockIdx.x;  // 0..N-1
  int idx = (bid * blockDim.x + threadIdx.x) * 2;
  const half epsilon = __float2half(1e-5f);
  const half g_ = __float2half(g);
  const half K_ = __int2half_rn(K);

  __shared__ half s_variance; // shared within block
  half2 reg_x = HALF2(x[idx]);
  half variance = (idx < N * K) ? (reg_x.x * reg_x.x + reg_x.y * reg_x.y)
                                : __float2half(0.0f);
  variance = block_reduce_sum_f16_f16<NUM_THREADS>(variance);
  if (tid == 0)
    s_variance = hrsqrt(variance / (K_ + epsilon));
  // wait for s_variance in shared memory to be ready for all threads
  __syncthreads();
  half2 reg_y;
  reg_y.x = reg_x.x * s_variance * g_;
  reg_y.y = reg_x.y * s_variance * g_;
  if (idx < N * K)
    HALF2(y[idx]) = reg_y;
}

#define HALF2_VARIANCE(reg, i)                                                 \
  (((idx + (i)) < N * K) ? ((reg).x * (reg).x + (reg).y * (reg).y)             \
                         : __float2half(0.0f))

#define FLOAT2_VARIANCE(reg, i)                                                \
  (((idx + (i)) < N * K) ? ((reg).x * (reg).x + (reg).y * (reg).y) : 0.0f)

#define HALF2_RMS_NORM(reg_y, reg_x, g)                                        \
  (reg_y).x = (reg_x).x * s_variance * (g);                                    \
  (reg_y).y = (reg_x).y * s_variance * (g);

#define FLOAT2_RMS_NORM(reg_y, reg_x, g)                                       \
  (reg_y).x = (reg_x).x * s_variance * (g);                                    \
  (reg_y).y = (reg_x).y * s_variance * (g);

template <const int NUM_THREADS = 256>
__global__ void rms_norm_f16x8_f16_kernel(half *x, half *y, float g, int N,
                                          int K) {
  int tid = threadIdx.x; // 0..K-1
  int bid = blockIdx.x;  // 0..N-1
  int idx = (bid * blockDim.x + threadIdx.x) * 8;
  const half epsilon = __float2half(1e-5f);
  const half g_ = __float2half(g);
  const half K_ = __int2half_rn(K);

  __shared__ half s_variance; // shared within block
  half2 reg_x_0 = HALF2(x[idx + 0]);
  half2 reg_x_1 = HALF2(x[idx + 2]);
  half2 reg_x_2 = HALF2(x[idx + 4]);
  half2 reg_x_3 = HALF2(x[idx + 6]);
  half variance = HALF2_VARIANCE(reg_x_0, 0);
  variance += HALF2_VARIANCE(reg_x_1, 2);
  variance += HALF2_VARIANCE(reg_x_2, 4);
  variance += HALF2_VARIANCE(reg_x_3, 6);
  variance = block_reduce_sum_f16_f16<NUM_THREADS>(variance);
  if (tid == 0)
    s_variance = hrsqrt(variance / (K_ + epsilon));
  // wait for s_variance in shared memory to be ready for all threads
  __syncthreads();
  // manual unroll
  half2 reg_y_0, reg_y_1, reg_y_2, reg_y_3;
  HALF2_RMS_NORM(reg_y_0, reg_x_0, g_);
  HALF2_RMS_NORM(reg_y_1, reg_x_1, g_);
  HALF2_RMS_NORM(reg_y_2, reg_x_2, g_);
  HALF2_RMS_NORM(reg_y_3, reg_x_3, g_);
  if ((idx + 0) < N * K) {
    HALF2(y[idx + 0]) = reg_y_0;
  }
  if ((idx + 2) < N * K) {
    HALF2(y[idx + 2]) = reg_y_1;
  }
  if ((idx + 4) < N * K) {
    HALF2(y[idx + 4]) = reg_y_2;
  }
  if ((idx + 6) < N * K) {
    HALF2(y[idx + 6]) = reg_y_3;
  }
}

template <const int NUM_THREADS = 256>
__global__ void rms_norm_f16x8_f32_kernel(half *x, half *y, float g, int N,
                                          int K) {
  int tid = threadIdx.x; // 0..K-1
  int bid = blockIdx.x;  // 0..N-1
  int idx = (bid * blockDim.x + threadIdx.x) * 8;
  const float epsilon = 1e-5f;

  __shared__ float s_variance; // shared within block
  // manual unroll and improve L2 cache hit rate.
  // Only   L2 cache: load 32  bytes in 1 memory issue (default)
  // Enable L1 cache: load 128 bytes in 1 memory issue (-Xptxas -dlcm=ca)
  // why try fp16x8 within 1 threads? ref:
  // https://zhuanlan.zhihu.com/p/641639133 0. first, tid_0 load 32 bytes in 1
  // memory issue and cache data into L2 cache.
  // 1. then, tid_1,...,tid_3 hit L2 cache and load data from L2 cache directly.
  float2 reg_x_0 = __half22float2(HALF2(x[idx + 0]));
  float2 reg_x_1 = __half22float2(HALF2(x[idx + 2]));
  float2 reg_x_2 = __half22float2(HALF2(x[idx + 4]));
  float2 reg_x_3 = __half22float2(HALF2(x[idx + 6]));

  float variance = FLOAT2_VARIANCE(reg_x_0, 0);
  variance += FLOAT2_VARIANCE(reg_x_1, 2);
  variance += FLOAT2_VARIANCE(reg_x_2, 4);
  variance += FLOAT2_VARIANCE(reg_x_3, 6);

  variance = block_reduce_sum_f32<NUM_THREADS>(variance);
  if (tid == 0)
    s_variance = rsqrtf(variance / ((float)K + epsilon));
  // wait for s_variance in shared memory to be ready for all threads
  __syncthreads();
  // manual unroll
  float2 reg_y_0, reg_y_1, reg_y_2, reg_y_3;
  FLOAT2_RMS_NORM(reg_y_0, reg_x_0, g);
  FLOAT2_RMS_NORM(reg_y_1, reg_x_1, g);
  FLOAT2_RMS_NORM(reg_y_2, reg_x_2, g);
  FLOAT2_RMS_NORM(reg_y_3, reg_x_3, g);
  if ((idx + 0) < N * K) {
    HALF2(y[idx + 0]) = __float22half2_rn(reg_y_0);
  }
  if ((idx + 2) < N * K) {
    HALF2(y[idx + 2]) = __float22half2_rn(reg_y_1);
  }
  if ((idx + 4) < N * K) {
    HALF2(y[idx + 4]) = __float22half2_rn(reg_y_2);
  }
  if ((idx + 6) < N * K) {
    HALF2(y[idx + 6]) = __float22half2_rn(reg_y_3);
  }
}

template <const int NUM_THREADS = 256>
__global__ void rms_norm_f16_f32_kernel(half *x, half *y, float g, int N,
                                        int K) {
  int tid = threadIdx.x; // 0..K-1
  int bid = blockIdx.x;  // 0..N-1
  int idx = bid * blockDim.x + threadIdx.x;
  const float epsilon = 1e-5f;

  __shared__ float s_variance; // shared within block
  float value = (idx < N * K) ? __half2float(x[idx]) : 0.0f; // load once only
  float variance = value * value;
  variance = block_reduce_sum_f32<NUM_THREADS>(variance);
  if (tid == 0)
    s_variance = rsqrtf(variance / ((float)K + epsilon));
  // wait for s_variance in shared memory to be ready for all threads
  __syncthreads();
  if (idx < N * K) {
    y[idx] = __float2half((value * s_variance) * g);
  }
}

template <const int NUM_THREADS = 256>
__global__ void rms_norm_f16x8_pack_f16_kernel(half *x, half *y, float g, int N,
                                               int K) {
  int tid = threadIdx.x; // 0..K-1
  int bid = blockIdx.x;  // 0..N-1
  int idx = (bid * blockDim.x + threadIdx.x) * 8;
  const half epsilon = __float2half(1e-5f);
  const half g_ = __float2half(g);
  const half K_ = __int2half_rn(K);
  const half z_ = __float2half(0.0f);

  __shared__ half s_variance; // shared within block
  // temporary register(memory), .local space in ptx, addressable
  half pack_x[8], pack_y[8]; // 8x16 bits=128 bits.
  // reinterpret as float4 and load 128 bits in 1 memory issue.
  LDST128BITS(pack_x[0]) = LDST128BITS(x[idx]); // load 128 bits

  half variance = z_;
#pragma unroll
  for (int i = 0; i < 8; ++i) {
    variance += ((idx + i) < N * K ? pack_x[i] * pack_x[i] : z_);
  }
  variance = block_reduce_sum_f16_f16<NUM_THREADS>(variance);
  if (tid == 0)
    s_variance = hrsqrt(variance / (K_ + epsilon));
  // wait for s_variance in shared memory to be ready for all threads
  __syncthreads();

#pragma unroll
  for (int i = 0; i < 8; ++i) {
    pack_y[i] = pack_x[i] * s_variance * g_;
  }
  // reinterpret as float4 and store 128 bits in 1 memory issue.
  if ((idx + 7) < N * K) {
    LDST128BITS(y[idx]) = LDST128BITS(pack_y[0]);
  }
  // TODO: support non 8-multiple K here
}

template <const int NUM_THREADS = 256>
__global__ void rms_norm_f16x8_pack_f32_kernel(half *x, half *y, float g, int N,
                                               int K) {
  int tid = threadIdx.x; // 0..K-1
  int bid = blockIdx.x;  // 0..N-1
  int idx = (bid * blockDim.x + threadIdx.x) * 8;
  const float epsilon = 1e-5f;
  __shared__ float s_variance; // shared within block
  // temporary register(memory), .local space in ptx, addressable
  half pack_x[8], pack_y[8]; // 8x16 bits=128 bits.
  // reinterpret as float4 and load 128 bits in 1 memory issue.
  LDST128BITS(pack_x[0]) = LDST128BITS(x[idx]); // load 128 bits

  float variance = 0.0f;
#pragma unroll
  for (int i = 0; i < 8; ++i) {
    float v = __half2float(pack_x[i]);
    variance += ((idx + i) < N * K ? v * v : 0.0f);
  }
  variance = block_reduce_sum_f32<NUM_THREADS>(variance);
  if (tid == 0)
    s_variance = rsqrtf(variance / ((float)K + epsilon));
  // wait for s_variance in shared memory to be ready for all threads
  __syncthreads();

#pragma unroll
  for (int i = 0; i < 8; i += 2) {
    float2 v2 = __half22float2(HALF2(pack_x[i]));
    float2 y2 = {v2.x * s_variance * g, v2.y * s_variance * g};
    HALF2(pack_y[i]) = __float22half2_rn(y2);
  }
  // reinterpret as float4 and store 128 bits in 1 memory issue.
  if ((idx + 7) < N * K) {
    LDST128BITS(y[idx]) = LDST128BITS(pack_y[0]);
  }
  // TODO: support non 8-multiple K here
}

// --------------------- PyTorch bindings for custom kernel
// -----------------------
#define STRINGFY(str) #str
#define TORCH_BINDING_COMMON_EXTENSION(func)                                   \
  m.def(STRINGFY(func), &func, STRINGFY(func));

#define CHECK_TORCH_TENSOR_DTYPE(T, th_type)                                   \
  if (((T).options().dtype() != (th_type))) {                                  \
    std::cout << "Tensor Info:" << (T).options() << std::endl;                 \
    throw std::runtime_error("values must be " #th_type);                      \
  }

#define CHECK_TORCH_TENSOR_SHAPE(T1, T2)                                       \
  assert((T1).dim() == (T2).dim());                                            \
  for (int i = 0; i < (T1).dim(); ++i) {                                       \
    if ((T2).size(i) != (T1).size(i)) {                                        \
      throw std::runtime_error("Tensor size mismatch!");                       \
    }                                                                          \
  }

#define LANUCH_RMS_NORM_F32_KERNEL(K)                                          \
  rms_norm_f32_kernel<(K)>                                                     \
      <<<grid, block>>>(reinterpret_cast<float *>(x.data_ptr()),               \
                        reinterpret_cast<float *>(y.data_ptr()), g, N, (K));

#define DISPATCH_RMS_NORM_F32_KERNEL(N, K)                                     \
  dim3 block((K));                                                             \
  dim3 grid((N));                                                              \
                                                                               \
  switch ((K)) {                                                               \
  case 64:                                                                     \
    LANUCH_RMS_NORM_F32_KERNEL(64)                                             \
    break;                                                                     \
  case 128:                                                                    \
    LANUCH_RMS_NORM_F32_KERNEL(128)                                            \
    break;                                                                     \
  case 256:                                                                    \
    LANUCH_RMS_NORM_F32_KERNEL(256)                                            \
    break;                                                                     \
  case 512:                                                                    \
    LANUCH_RMS_NORM_F32_KERNEL(512)                                            \
    break;                                                                     \
  case 1024:                                                                   \
    LANUCH_RMS_NORM_F32_KERNEL(1024)                                           \
    break;                                                                     \
  default:                                                                     \
    throw std::runtime_error("only support K: 64/128/256/512/1024");           \
    break;                                                                     \
  }

#define LANUCH_RMS_NORM_F32x4_KERNEL(K)                                        \
  rms_norm_f32x4_kernel<(K) / 4>                                               \
      <<<grid, block>>>(reinterpret_cast<float *>(x.data_ptr()),               \
                        reinterpret_cast<float *>(y.data_ptr()), g, N, (K));

#define DISPATCH_RMS_NORM_F32x4_KERNEL(N, K)                                   \
  dim3 block((K) / 4);                                                         \
  dim3 grid((N));                                                              \
                                                                               \
  switch ((K)) {                                                               \
  case 64:                                                                     \
    LANUCH_RMS_NORM_F32x4_KERNEL(64) break;                                    \
  case 128:                                                                    \
    LANUCH_RMS_NORM_F32x4_KERNEL(128) break;                                   \
  case 256:                                                                    \
    LANUCH_RMS_NORM_F32x4_KERNEL(256) break;                                   \
  case 512:                                                                    \
    LANUCH_RMS_NORM_F32x4_KERNEL(512) break;                                   \
  case 1024:                                                                   \
    LANUCH_RMS_NORM_F32x4_KERNEL(1024) break;                                  \
  case 2048:                                                                   \
    LANUCH_RMS_NORM_F32x4_KERNEL(2048) break;                                  \
  case 4096:                                                                   \
    LANUCH_RMS_NORM_F32x4_KERNEL(4096) break;                                  \
  default:                                                                     \
    throw std::runtime_error("only support K: 64/.../512/1024*4");             \
    break;                                                                     \
  }

void rms_norm_f32(torch::Tensor x, torch::Tensor y, float g) {
  CHECK_TORCH_TENSOR_DTYPE(x, torch::kFloat32)
  CHECK_TORCH_TENSOR_DTYPE(y, torch::kFloat32)
  CHECK_TORCH_TENSOR_SHAPE(x, y)
  const int N = x.size(0);
  const int K = x.size(1);
  DISPATCH_RMS_NORM_F32_KERNEL(N, K)
}

void rms_norm_f32x4(torch::Tensor x, torch::Tensor y, float g) {
  CHECK_TORCH_TENSOR_DTYPE(x, torch::kFloat32)
  CHECK_TORCH_TENSOR_DTYPE(y, torch::kFloat32)
  CHECK_TORCH_TENSOR_SHAPE(x, y)
  const int N = x.size(0);
  const int K = x.size(1);
  DISPATCH_RMS_NORM_F32x4_KERNEL(N, K)
}

// fp16
#define LANUCH_RMS_NORM_F16F16_KERNEL(K)                                       \
  rms_norm_f16_f16_kernel<(K)>                                                 \
      <<<grid, block>>>(reinterpret_cast<half *>(x.data_ptr()),                \
                        reinterpret_cast<half *>(y.data_ptr()), g, N, (K));

#define DISPATCH_RMS_NORM_F16F16_KERNEL(N, K)                                  \
  dim3 block((K));                                                             \
  dim3 grid((N));                                                              \
  switch ((K)) {                                                               \
  case 64:                                                                     \
    LANUCH_RMS_NORM_F16F16_KERNEL(64)                                          \
    break;                                                                     \
  case 128:                                                                    \
    LANUCH_RMS_NORM_F16F16_KERNEL(128)                                         \
    break;                                                                     \
  case 256:                                                                    \
    LANUCH_RMS_NORM_F16F16_KERNEL(256)                                         \
    break;                                                                     \
  case 512:                                                                    \
    LANUCH_RMS_NORM_F16F16_KERNEL(512)                                         \
    break;                                                                     \
  case 1024:                                                                   \
    LANUCH_RMS_NORM_F16F16_KERNEL(1024)                                        \
    break;                                                                     \
  default:                                                                     \
    throw std::runtime_error("only support K: 64/128/256/512/1024");           \
    break;                                                                     \
  }

#define LANUCH_RMS_NORM_F16F32_KERNEL(K)                                       \
  rms_norm_f16_f32_kernel<(K)>                                                 \
      <<<grid, block>>>(reinterpret_cast<half *>(x.data_ptr()),                \
                        reinterpret_cast<half *>(y.data_ptr()), g, N, (K));

#define DISPATCH_RMS_NORM_F16F32_KERNEL(N, K)                                  \
  dim3 block((K));                                                             \
  dim3 grid((N));                                                              \
  switch ((K)) {                                                               \
  case 64:                                                                     \
    LANUCH_RMS_NORM_F16F32_KERNEL(64)                                          \
    break;                                                                     \
  case 128:                                                                    \
    LANUCH_RMS_NORM_F16F32_KERNEL(128)                                         \
    break;                                                                     \
  case 256:                                                                    \
    LANUCH_RMS_NORM_F16F32_KERNEL(256)                                         \
    break;                                                                     \
  case 512:                                                                    \
    LANUCH_RMS_NORM_F16F32_KERNEL(512)                                         \
    break;                                                                     \
  case 1024:                                                                   \
    LANUCH_RMS_NORM_F16F32_KERNEL(1024)                                        \
    break;                                                                     \
  default:                                                                     \
    throw std::runtime_error("only support K: 64/128/256/512/1024");           \
    break;                                                                     \
  }

#define LANUCH_RMS_NORM_F16x2F16_KERNEL(K)                                     \
  rms_norm_f16x2_f16_kernel<(K) / 2>                                           \
      <<<grid, block>>>(reinterpret_cast<half *>(x.data_ptr()),                \
                        reinterpret_cast<half *>(y.data_ptr()), g, N, (K));

#define DISPATCH_RMS_NORM_F16x2F16_KERNEL(N, K)                                \
  dim3 block((K) / 2);                                                         \
  dim3 grid((N));                                                              \
  switch ((K)) {                                                               \
  case 64:                                                                     \
    LANUCH_RMS_NORM_F16x2F16_KERNEL(64) break;                                 \
  case 128:                                                                    \
    LANUCH_RMS_NORM_F16x2F16_KERNEL(128) break;                                \
  case 256:                                                                    \
    LANUCH_RMS_NORM_F16x2F16_KERNEL(256) break;                                \
  case 512:                                                                    \
    LANUCH_RMS_NORM_F16x2F16_KERNEL(512) break;                                \
  case 1024:                                                                   \
    LANUCH_RMS_NORM_F16x2F16_KERNEL(1024) break;                               \
  case 2048:                                                                   \
    LANUCH_RMS_NORM_F16x2F16_KERNEL(2048) break;                               \
  default:                                                                     \
    throw std::runtime_error("only support K: 64/128/.../1024*2");             \
    break;                                                                     \
  }

#define LANUCH_RMS_NORM_F16x8F16_KERNEL(K)                                     \
  rms_norm_f16x8_f16_kernel<(K) / 8>                                           \
      <<<grid, block>>>(reinterpret_cast<half *>(x.data_ptr()),                \
                        reinterpret_cast<half *>(y.data_ptr()), g, N, (K));

#define DISPATCH_RMS_NORM_F16x8F16_KERNEL(N, K)                                \
  dim3 block((K) / 8);                                                         \
  dim3 grid((N));                                                              \
  switch ((K)) {                                                               \
  case 64:                                                                     \
    LANUCH_RMS_NORM_F16x8F16_KERNEL(64) break;                                 \
  case 128:                                                                    \
    LANUCH_RMS_NORM_F16x8F16_KERNEL(128) break;                                \
  case 256:                                                                    \
    LANUCH_RMS_NORM_F16x8F16_KERNEL(256) break;                                \
  case 512:                                                                    \
    LANUCH_RMS_NORM_F16x8F16_KERNEL(512) break;                                \
  case 1024:                                                                   \
    LANUCH_RMS_NORM_F16x8F16_KERNEL(1024) break;                               \
  case 2048:                                                                   \
    LANUCH_RMS_NORM_F16x8F16_KERNEL(2048) break;                               \
  case 4096:                                                                   \
    LANUCH_RMS_NORM_F16x8F16_KERNEL(4096) break;                               \
  case 8192:                                                                   \
    LANUCH_RMS_NORM_F16x8F16_KERNEL(8192) break;                               \
  default:                                                                     \
    throw std::runtime_error("only support K: 64/128/.../1024*8");             \
    break;                                                                     \
  }

#define LANUCH_RMS_NORM_F16x8F32_KERNEL(K)                                     \
  rms_norm_f16x8_f16_kernel<(K) / 8>                                           \
      <<<grid, block>>>(reinterpret_cast<half *>(x.data_ptr()),                \
                        reinterpret_cast<half *>(y.data_ptr()), g, N, (K));

#define DISPATCH_RMS_NORM_F16x8F32_KERNEL(N, K)                                \
  dim3 block((K) / 8);                                                         \
  dim3 grid((N));                                                              \
  switch ((K)) {                                                               \
  case 64:                                                                     \
    LANUCH_RMS_NORM_F16x8F32_KERNEL(64) break;                                 \
  case 128:                                                                    \
    LANUCH_RMS_NORM_F16x8F32_KERNEL(128) break;                                \
  case 256:                                                                    \
    LANUCH_RMS_NORM_F16x8F32_KERNEL(256) break;                                \
  case 512:                                                                    \
    LANUCH_RMS_NORM_F16x8F32_KERNEL(512) break;                                \
  case 1024:                                                                   \
    LANUCH_RMS_NORM_F16x8F32_KERNEL(1024) break;                               \
  case 2048:                                                                   \
    LANUCH_RMS_NORM_F16x8F32_KERNEL(2048) break;                               \
  case 4096:                                                                   \
    LANUCH_RMS_NORM_F16x8F32_KERNEL(4096) break;                               \
  case 8192:                                                                   \
    LANUCH_RMS_NORM_F16x8F32_KERNEL(8192) break;                               \
  default:                                                                     \
    throw std::runtime_error("only support K: 64/128/.../1024*8");             \
    break;                                                                     \
  }

#define LANUCH_RMS_NORM_F16x8_PACK_F16_KERNEL(K)                               \
  rms_norm_f16x8_pack_f16_kernel<(K) / 8>                                      \
      <<<grid, block>>>(reinterpret_cast<half *>(x.data_ptr()),                \
                        reinterpret_cast<half *>(y.data_ptr()), g, N, (K));

#define DISPATCH_RMS_NORM_F16x8_PACK_F16_KERNEL(N, K)                          \
  dim3 block((K) / 8);                                                         \
  dim3 grid((N));                                                              \
  switch ((K)) {                                                               \
  case 64:                                                                     \
    LANUCH_RMS_NORM_F16x8_PACK_F16_KERNEL(64) break;                           \
  case 128:                                                                    \
    LANUCH_RMS_NORM_F16x8_PACK_F16_KERNEL(128) break;                          \
  case 256:                                                                    \
    LANUCH_RMS_NORM_F16x8_PACK_F16_KERNEL(256) break;                          \
  case 512:                                                                    \
    LANUCH_RMS_NORM_F16x8_PACK_F16_KERNEL(512) break;                          \
  case 1024:                                                                   \
    LANUCH_RMS_NORM_F16x8_PACK_F16_KERNEL(1024) break;                         \
  case 2048:                                                                   \
    LANUCH_RMS_NORM_F16x8_PACK_F16_KERNEL(2048) break;                         \
  case 4096:                                                                   \
    LANUCH_RMS_NORM_F16x8_PACK_F16_KERNEL(4096) break;                         \
  case 8192:                                                                   \
    LANUCH_RMS_NORM_F16x8_PACK_F16_KERNEL(8192) break;                         \
  default:                                                                     \
    throw std::runtime_error("only support K: 64/128/.../1024*8");             \
    break;                                                                     \
  }

#define LANUCH_RMS_NORM_F16x8_PACK_F32_KERNEL(K)                               \
  rms_norm_f16x8_pack_f32_kernel<(K) / 8>                                      \
      <<<grid, block>>>(reinterpret_cast<half *>(x.data_ptr()),                \
                        reinterpret_cast<half *>(y.data_ptr()), g, N, (K));

#define DISPATCH_RMS_NORM_F16x8_PACK_F32_KERNEL(N, K)                          \
  dim3 block((K) / 8);                                                         \
  dim3 grid((N));                                                              \
  switch ((K)) {                                                               \
  case 64:                                                                     \
    LANUCH_RMS_NORM_F16x8_PACK_F32_KERNEL(64) break;                           \
  case 128:                                                                    \
    LANUCH_RMS_NORM_F16x8_PACK_F32_KERNEL(128) break;                          \
  case 256:                                                                    \
    LANUCH_RMS_NORM_F16x8_PACK_F32_KERNEL(256) break;                          \
  case 512:                                                                    \
    LANUCH_RMS_NORM_F16x8_PACK_F32_KERNEL(512) break;                          \
  case 1024:                                                                   \
    LANUCH_RMS_NORM_F16x8_PACK_F32_KERNEL(1024) break;                         \
  case 2048:                                                                   \
    LANUCH_RMS_NORM_F16x8_PACK_F32_KERNEL(2048) break;                         \
  case 4096:                                                                   \
    LANUCH_RMS_NORM_F16x8_PACK_F32_KERNEL(4096) break;                         \
  case 8192:                                                                   \
    LANUCH_RMS_NORM_F16x8_PACK_F32_KERNEL(8192) break;                         \
  default:                                                                     \
    throw std::runtime_error("only support K: 64/128/.../1024*8");             \
    break;                                                                     \
  }

void rms_norm_f16_f16(torch::Tensor x, torch::Tensor y, float g) {
  CHECK_TORCH_TENSOR_DTYPE(x, torch::kHalf)
  CHECK_TORCH_TENSOR_DTYPE(y, torch::kHalf)
  CHECK_TORCH_TENSOR_SHAPE(x, y)
  const int N = x.size(0);
  const int K = x.size(1);
  DISPATCH_RMS_NORM_F16F16_KERNEL(N, K)
}

void rms_norm_f16x2_f16(torch::Tensor x, torch::Tensor y, float g) {
  CHECK_TORCH_TENSOR_DTYPE(x, torch::kHalf)
  CHECK_TORCH_TENSOR_DTYPE(y, torch::kHalf)
  CHECK_TORCH_TENSOR_SHAPE(x, y)
  const int N = x.size(0);
  const int K = x.size(1);
  DISPATCH_RMS_NORM_F16x2F16_KERNEL(N, K)
}

void rms_norm_f16x8_f16(torch::Tensor x, torch::Tensor y, float g) {
  CHECK_TORCH_TENSOR_DTYPE(x, torch::kHalf)
  CHECK_TORCH_TENSOR_DTYPE(y, torch::kHalf)
  CHECK_TORCH_TENSOR_SHAPE(x, y)
  const int N = x.size(0);
  const int K = x.size(1);
  DISPATCH_RMS_NORM_F16x8F16_KERNEL(N, K)
}

void rms_norm_f16x8_f32(torch::Tensor x, torch::Tensor y, float g) {
  CHECK_TORCH_TENSOR_DTYPE(x, torch::kHalf)
  CHECK_TORCH_TENSOR_DTYPE(y, torch::kHalf)
  CHECK_TORCH_TENSOR_SHAPE(x, y)
  const int N = x.size(0);
  const int K = x.size(1);
  DISPATCH_RMS_NORM_F16x8F32_KERNEL(N, K)
}

void rms_norm_f16_f32(torch::Tensor x, torch::Tensor y, float g) {
  CHECK_TORCH_TENSOR_DTYPE(x, torch::kHalf)
  CHECK_TORCH_TENSOR_DTYPE(y, torch::kHalf)
  CHECK_TORCH_TENSOR_SHAPE(x, y)
  const int N = x.size(0);
  const int K = x.size(1);
  DISPATCH_RMS_NORM_F16F32_KERNEL(N, K)
}

// pack
void rms_norm_f16x8_pack_f16(torch::Tensor x, torch::Tensor y, float g) {
  CHECK_TORCH_TENSOR_DTYPE(x, torch::kHalf)
  CHECK_TORCH_TENSOR_DTYPE(y, torch::kHalf)
  CHECK_TORCH_TENSOR_SHAPE(x, y)
  const int N = x.size(0);
  const int K = x.size(1);
  DISPATCH_RMS_NORM_F16x8_PACK_F16_KERNEL(N, K)
}

void rms_norm_f16x8_pack_f32(torch::Tensor x, torch::Tensor y, float g) {
  CHECK_TORCH_TENSOR_DTYPE(x, torch::kHalf)
  CHECK_TORCH_TENSOR_DTYPE(y, torch::kHalf)
  CHECK_TORCH_TENSOR_SHAPE(x, y)
  const int N = x.size(0);
  const int K = x.size(1);
  DISPATCH_RMS_NORM_F16x8_PACK_F32_KERNEL(N, K)
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  TORCH_BINDING_COMMON_EXTENSION(rms_norm_f32)
  TORCH_BINDING_COMMON_EXTENSION(rms_norm_f32x4)
  TORCH_BINDING_COMMON_EXTENSION(rms_norm_f16_f16)
  TORCH_BINDING_COMMON_EXTENSION(rms_norm_f16x2_f16)
  TORCH_BINDING_COMMON_EXTENSION(rms_norm_f16x8_f16)
  TORCH_BINDING_COMMON_EXTENSION(rms_norm_f16x8_pack_f16)
  TORCH_BINDING_COMMON_EXTENSION(rms_norm_f16x8_f32)
  TORCH_BINDING_COMMON_EXTENSION(rms_norm_f16x8_pack_f32)
  TORCH_BINDING_COMMON_EXTENSION(rms_norm_f16_f32)
}