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

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)