[JIT Kernel]Support fused_add_rmsnorm in JIT Kernel (#17677)

python/sglang/jit_kernel/benchmark/bench_fused_add_rmsnorm.py

@@ -0,0 +1,73 @@
+import itertools
+
+import torch
+import triton
+import triton.testing
+from flashinfer import fused_add_rmsnorm as fi_fused_add_rmsnorm
+
+from sglang.jit_kernel.benchmark.utils import is_in_ci
+from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm
+
+IS_CI = is_in_ci()
+
+
+def sglang_jit_fused_add_rmsnorm(
+    input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
+) -> None:
+    jit_fused_add_rmsnorm(input, residual, weight, eps)
+
+
+def flashinfer_fused_add_rmsnorm(
+    input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
+) -> None:
+    fi_fused_add_rmsnorm(input, residual, weight, eps=eps)
+
+
+DTYPE = torch.bfloat16
+DEVICE = "cuda"
+
+if IS_CI:
+    BS_LIST = [16]
+    HIDDEN_SIZE_LIST = [512, 2048]
+else:
+    BS_LIST = [2**n for n in range(0, 14)]
+    HIDDEN_SIZE_LIST = [1536, 3072, 4096, 5120, 8192]
+
+LINE_VALS = ["jit", "fi"]
+LINE_NAMES = ["SGL JIT Kernel", "FlashInfer"]
+STYLES = [("orange", "-"), ("blue", "--"), ("green", "-."), ("red", ":")]
+
+configs = list(itertools.product(HIDDEN_SIZE_LIST, BS_LIST))
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["hidden_size", "batch_size"],
+        x_vals=configs,
+        line_arg="provider",
+        line_vals=LINE_VALS,
+        line_names=LINE_NAMES,
+        styles=STYLES,
+        ylabel="us",
+        plot_name="fused-add-rmsnorm-performance",
+        args={},
+    )
+)
+def benchmark(hidden_size: int, batch_size: int, provider: str):
+    input = torch.randn((batch_size, hidden_size), dtype=DTYPE, device=DEVICE)
+    residual = torch.randn((batch_size, hidden_size), dtype=DTYPE, device=DEVICE)
+    weight = torch.randn(hidden_size, dtype=DTYPE, device=DEVICE)
+    FN_MAP = {
+        "jit": sglang_jit_fused_add_rmsnorm,
+        "fi": flashinfer_fused_add_rmsnorm,
+    }
+    fn = lambda: FN_MAP[provider](
+        input.clone(), residual.clone(), weight, torch.finfo(torch.bfloat16).eps
+    )
+    quantiles = [0.5, 0.2, 0.8]
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)  # type: ignore
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    benchmark.run(print_data=True)

python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm.cuh

@@ -0,0 +1,186 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/tile.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <cooperative_groups/reduce.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cooperative_groups.h>
+#include <type_traits>
+
+namespace {
+
+template <typename T, int VEC_SIZE_IN_BYTE>
+struct VecTypeTrait;
+
+template <>
+struct VecTypeTrait<bf16_t, 16> {
+  using packed_t = packed_t<bf16_t>;
+  using vec_t = device::AlignedVector<packed_t, 4>;
+};
+
+template <>
+struct VecTypeTrait<fp16_t, 16> {
+  using packed_t = packed_t<fp16_t>;
+  using vec_t = device::AlignedVector<packed_t, 4>;
+};
+
+template <>
+struct VecTypeTrait<bf16_t, 32> {
+  using packed_t = packed_t<bf16_t>;
+  using vec_t = device::AlignedVector<packed_t, 8>;
+};
+
+template <>
+struct VecTypeTrait<fp16_t, 32> {
+  using packed_t = packed_t<fp16_t>;
+  using vec_t = device::AlignedVector<packed_t, 8>;
+};
+
+template <typename packed_t>
+SGL_DEVICE packed_t rms(packed_t& val, packed_t& weight, float rsqrt_square_sum) {
+  float2 valf = device::cast<fp32x2_t, packed_t>(val);
+  float2 weightf = device::cast<fp32x2_t, packed_t>(weight);
+  return device::cast<packed_t, fp32x2_t>(
+      make_float2(valf.x * weightf.x * rsqrt_square_sum, valf.y * weightf.y * rsqrt_square_sum));
+}
+
+template <typename T, int VEC_SIZE_IN_BYTE>
+__global__ void fused_add_rmsnorm_reg_kernel(
+    T* __restrict__ input, T* __restrict__ residual, const T* __restrict__ weight, int vec_hidden_size, float eps) {
+  constexpr int inner_loop = VEC_SIZE_IN_BYTE == 16 ? 4 : 8;
+
+  __shared__ float shared_memory[32];  // Used for CTA reduce
+
+  using vec_t = typename VecTypeTrait<T, VEC_SIZE_IN_BYTE>::vec_t;
+  using packed_t = typename VecTypeTrait<T, VEC_SIZE_IN_BYTE>::packed_t;
+  vec_t v;         // Save input
+  vec_t v_res;     // Save residual
+  vec_t v_weight;  // Save weight
+  vec_t v_out;     // Save output
+
+  auto token_id = blockIdx.x;
+  float2 acc_square = make_float2(0.0f, 0.0f);  // Sum of squares for each thread
+
+  if (threadIdx.x < vec_hidden_size) {
+    // Compute address
+    vec_t* p = reinterpret_cast<vec_t*>(input) + token_id * vec_hidden_size;
+    vec_t* p_res = reinterpret_cast<vec_t*>(residual) + token_id * vec_hidden_size;
+    const vec_t* p_weight = reinterpret_cast<const vec_t*>(weight);
+
+    // Load data
+    v = p[threadIdx.x];
+    v_res = p_res[threadIdx.x];
+    v_weight = p_weight[threadIdx.x];
+
+    for (int i = 0; i < inner_loop; i++) {
+      float2 val = device::cast<fp32x2_t, packed_t>(v[i]);
+      float2 res = device::cast<fp32x2_t, packed_t>(v_res[i]);
+      float2 inp_res = make_float2(val.x + res.x, val.y + res.y);
+      acc_square.x += inp_res.x * inp_res.x;
+      acc_square.y += inp_res.y * inp_res.y;
+      v[i] = device::cast<packed_t, fp32x2_t>(inp_res);
+    }
+
+    // Store inp+res to residual
+    p_res[threadIdx.x] = v;
+  }
+
+  // CTA Reduce
+  // Step 0: Warp Reduce
+  auto cg_warp = cooperative_groups::tiled_partition<32>(cooperative_groups::this_thread_block());
+  float warp_sum = cooperative_groups::reduce(cg_warp, acc_square.x + acc_square.y, cooperative_groups::plus<float>());
+
+  float* buffer = shared_memory;
+  if (threadIdx.x % 32 == 0) {
+    buffer[threadIdx.x / 32] = warp_sum;  // Write warp_sum to buffer
+  }
+
+  // Step 1: CTA Reduce
+  __syncthreads();
+  if (threadIdx.x < 32) {
+    float cta_sum = cooperative_groups::reduce(
+        cg_warp, (threadIdx.x < blockDim.x / 32) ? buffer[threadIdx.x] : 0.0f, cooperative_groups::plus<float>());
+    buffer[threadIdx.x] =
+        rsqrtf(eps + cta_sum * (1.0f / static_cast<float>(vec_hidden_size * (VEC_SIZE_IN_BYTE / sizeof(T)))));
+  }
+  __syncthreads();
+
+  // Compute RMSNorm
+  if (threadIdx.x < vec_hidden_size) {
+    float rsqrt_square_sum = buffer[threadIdx.x / 32];  // Read rsqrt from Shared Memory(Broadcast)
+    for (int i = 0; i < inner_loop; i++) {
+      v_out[i] = rms(v[i], v_weight[i], rsqrt_square_sum);
+    }
+    vec_t* p_out = reinterpret_cast<vec_t*>(input) + token_id * vec_hidden_size;
+    p_out[threadIdx.x] = v_out;
+  }
+}
+
+template <typename DType>
+struct FusedAddRMSNormKernel {
+  static void
+  run(const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView residual,
+      const tvm::ffi::TensorView weight,
+      float eps) {
+    using namespace host;
+    auto N = SymbolicSize{"num_tokens"};
+    auto D = SymbolicSize{"hidden_size"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, D})  // input
+        .with_strides({D, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(input);
+    TensorMatcher({D})  // weight
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(weight);
+    TensorMatcher({N, D})  // residual
+        .with_strides({D, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(residual);
+
+    auto cc_major = host::runtime::get_cc_major(device.unwrap().device_id);
+    int hidden_size = static_cast<int>(D.unwrap());
+    if ((cc_major <= 9 && hidden_size <= 8192) || (cc_major >= 10 && hidden_size <= 12288)) {
+      int max_vec_size_byte = cc_major >= 10 ? 32 : 16;
+      int elements_in_vec = max_vec_size_byte / sizeof(DType);
+      int vec_hidden_size = hidden_size / elements_in_vec;
+      uint threads = (vec_hidden_size + 31) / 32 * 32;
+
+      // Runtime check
+      host::RuntimeCheck(
+          hidden_size % elements_in_vec == 0,
+          "hidden_size",
+          hidden_size,
+          " can not align to elements_in_vec ",
+          elements_in_vec);
+
+      // Launch kernel
+      auto kernel =
+          max_vec_size_byte == 32 ? fused_add_rmsnorm_reg_kernel<DType, 32> : fused_add_rmsnorm_reg_kernel<DType, 16>;
+      LaunchKernel(static_cast<uint>(N.unwrap()), threads, device.unwrap())
+          .enable_pdl(false)(
+              kernel,
+              reinterpret_cast<DType*>(input.data_ptr()),
+              reinterpret_cast<DType*>(residual.data_ptr()),
+              reinterpret_cast<DType*>(weight.data_ptr()),
+              vec_hidden_size,
+              eps);
+    } else {
+      host::RuntimeCheck(false, "Large hidden_sizes are not supported for now.");
+    }
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/include/sgl_kernel/vec.cuh

@@ -45,7 +45,7 @@ template <typename T, std::size_t N>
 struct AlignedVector {
  private:
   /// NOTE: 1. must be pow of two 2. 16 * 8 = 128 byte, which is the max vector size supported by most devices
-  static_assert((N > 0 && (N & (N - 1)) == 0) && sizeof(T) * N <= 16, "CUDA only support at most 128B vector op");
+  static_assert((N > 0 && (N & (N - 1)) == 0) && sizeof(T) * N <= 32, "CUDA only support at most 256B vector op");
   using element_t = typename details::sized_int<T>;
   using storage_t = AlignedStorage<element_t, N>;
 

python/sglang/jit_kernel/norm.py

@@ -38,6 +38,17 @@ def _jit_rmsnorm_module(hidden_size: int, dtype: torch.dtype) -> Module:
     )
 
 
+@cache_once
+def _jit_fused_add_rmsnorm_module(dtype: torch.dtype) -> Module:
+    args = make_cpp_args(dtype)
+    return load_jit(
+        "fused_add_rmsnorm",
+        *args,
+        cuda_files=["elementwise/fused_add_rmsnorm.cuh"],
+        cuda_wrappers=[("fused_add_rmsnorm", f"FusedAddRMSNormKernel<{args}>::run")],
+    )
+
+
 @cache_once
 def can_use_fused_inplace_qknorm(head_dim: int, dtype: torch.dtype) -> bool:
     logger = logging.getLogger(__name__)
@@ -76,3 +87,13 @@ def rmsnorm(
     hidden_size = input.size(-1)
     module = _jit_rmsnorm_module(hidden_size, input.dtype)
     module.rmsnorm(input, weight, output, eps)
+
+
+def fused_add_rmsnorm(
+    input: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    module = _jit_fused_add_rmsnorm_module(input.dtype)
+    module.fused_add_rmsnorm(input, residual, weight, eps)

python/sglang/jit_kernel/tests/test_fused_add_rmsnorm.py

@@ -0,0 +1,55 @@
+import itertools
+
+import pytest
+import torch
+
+
+def sglang_jit_fused_add_rmsnorm(
+    input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
+) -> None:
+    from sglang.jit_kernel.norm import fused_add_rmsnorm
+
+    fused_add_rmsnorm(input, residual, weight, eps)
+
+
+def flashinfer_fused_add_rmsnorm(
+    input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
+) -> None:
+    from flashinfer.norm import fused_add_rmsnorm
+
+    fused_add_rmsnorm(input, residual, weight, eps=eps)
+
+
+BS_LIST = [2**n for n in range(0, 14)]
+BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)]
+HIDDEN_SIZE_LIST = [512, 1024, 1536, 2048, 3072, 4096, 5120, 6144, 7168, 8192]
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+
+
+@pytest.mark.parametrize(
+    "batch_size,hidden_size", list(itertools.product(BS_LIST, HIDDEN_SIZE_LIST))
+)
+def test_fused_add_rmsnorm(batch_size: int, hidden_size: int) -> None:
+    input = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=DTYPE)
+    residual = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=DTYPE)
+    weight = torch.randn(hidden_size, device=DEVICE, dtype=DTYPE)
+
+    input_sglang = input.clone()
+    residual_sglang = residual.clone()
+    input_flashinfer = input.clone()
+    residual_flashinfer = residual.clone()
+    sglang_jit_fused_add_rmsnorm(
+        input_sglang, residual_sglang, weight, torch.finfo(torch.bfloat16).eps
+    )
+    flashinfer_fused_add_rmsnorm(
+        input_flashinfer, residual_flashinfer, weight, torch.finfo(torch.bfloat16).eps
+    )
+    torch.testing.assert_close(input_sglang, input_flashinfer, atol=1e-2, rtol=1e-2)
+    torch.testing.assert_close(
+        residual_sglang, residual_flashinfer, atol=1e-2, rtol=1e-2
+    )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])