feat: integrate norm kernels into sgl-kernel (#3052)

sgl-kernel/src/sgl-kernel/__init__.py

@@ -1,6 +1,9 @@
 from sgl_kernel.ops import (
     custom_dispose,
     custom_reduce,
+    fused_add_rmsnorm,
+    gemma_fused_add_rmsnorm,
+    gemma_rmsnorm,
     get_graph_buffer_ipc_meta,
     init_custom_reduce,
     int8_scaled_mm,
@@ -12,14 +15,17 @@ from sgl_kernel.ops import (
 )
 
 __all__ = [
-    "moe_align_block_size",
-    "init_custom_reduce",
     "custom_dispose",
     "custom_reduce",
-    "int8_scaled_mm",
-    "sampling_scaling_penalties",
+    "fused_add_rmsnorm",
+    "gemma_fused_add_rmsnorm",
+    "gemma_rmsnorm",
     "get_graph_buffer_ipc_meta",
+    "init_custom_reduce",
+    "int8_scaled_mm",
+    "moe_align_block_size",
     "register_graph_buffers",
-    "rotary_embedding",
     "rmsnorm",
+    "rotary_embedding",
+    "sampling_scaling_penalties",
 ]

sgl-kernel/src/sgl-kernel/csrc/sgl_kernel_ops.cu

@@ -33,6 +33,16 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query, torch::Ten
 // rms norm
 void rmsnorm(at::Tensor& output, at::Tensor& input, at::Tensor& weight, double eps, int64_t cuda_stream);
 
+// fused rms norm
+void fused_add_rmsnorm(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps, int64_t cuda_stream);
+
+// gemma rms norm
+void gemma_rmsnorm(at::Tensor& output, at::Tensor& input, at::Tensor& weight, double eps, int64_t cuda_stream);
+
+// fused gemma rms norm
+void gemma_fused_add_rmsnorm(at::Tensor& input, at::Tensor& residual, at::Tensor& weight, double eps,
+                             int64_t cuda_stream);
+
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   // trt_reduce
   m.def("init_custom_ar", &init_custom_ar, "init custom allreduce meta (CUDA)");
@@ -50,4 +60,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("rotary_embedding", &rotary_embedding, "Rotary Embedding (CUDA)");
   // rms norm
   m.def("rmsnorm", &rmsnorm, "RMSNorm (CUDA)");
+  // fused rms norm
+  m.def("fused_add_rmsnorm", &fused_add_rmsnorm, "Fused Add RMSNorm (CUDA)");
+  // gemma rms norm
+  m.def("gemma_rmsnorm", &gemma_rmsnorm, "Gemma RMSNorm (CUDA)");
+  // fused gemma rms norm
+  m.def("gemma_fused_add_rmsnorm", &gemma_fused_add_rmsnorm, "Gemma Fused Add RMSNorm (CUDA)");
 }

sgl-kernel/src/sgl-kernel/ops/__init__.py

@@ -3,6 +3,9 @@ from typing import Optional
 import torch
 from sgl_kernel.ops._kernels import all_reduce as _all_reduce
 from sgl_kernel.ops._kernels import dispose as _dispose
+from sgl_kernel.ops._kernels import fused_add_rmsnorm as _fused_add_rmsnorm
+from sgl_kernel.ops._kernels import gemma_fused_add_rmsnorm as _gemma_fused_add_rmsnorm
+from sgl_kernel.ops._kernels import gemma_rmsnorm as _gemma_rmsnorm
 from sgl_kernel.ops._kernels import (
     get_graph_buffer_ipc_meta as _get_graph_buffer_ipc_meta,
 )
@@ -17,6 +20,10 @@ from sgl_kernel.ops._kernels import (
 )
 
 
+def get_cuda_stream(device: torch.device) -> int:
+    return torch.cuda.current_stream(device).cuda_stream
+
+
 def init_custom_reduce(
     rank_id, num_devices, rank_data, buffers, tmp_buffers, barrier_in, barrier_out
 ):
@@ -88,9 +95,35 @@ def rmsnorm(
     eps: float = 1e-6,
     out: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
-    if out is None:
-        out = torch.empty_like(input)
-    stream = torch.cuda.current_stream().cuda_stream
-    stream_int = int(stream)
-    _rmsnorm(out, input, weight, eps, stream_int)
-    return out
+    with input.device as device:
+        if out is None:
+            out = torch.empty_like(input)
+        _rmsnorm(out, input, weight, eps, get_cuda_stream(device))
+        return out
+
+
+def fused_add_rmsnorm(
+    input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6
+) -> None:
+    with input.device as device:
+        _fused_add_rmsnorm(input, residual, weight, eps, get_cuda_stream(device))
+
+
+def gemma_rmsnorm(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    with input.device as device:
+        if out is None:
+            out = torch.empty_like(input)
+        _gemma_rmsnorm(out, input, weight, eps, get_cuda_stream(device))
+        return out
+
+
+def gemma_fused_add_rmsnorm(
+    input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6
+) -> None:
+    with input.device as device:
+        _gemma_fused_add_rmsnorm(input, residual, weight, eps, get_cuda_stream(device))