Migrate norm kernels to FlashInfer JIT implementation (#18871)

python/sglang/jit_kernel/benchmark/bench_norm.py

@@ -0,0 +1,97 @@
+import itertools
+
+import torch
+import triton
+import triton.testing
+from flashinfer.norm import fused_add_rmsnorm as fi_fused_add_rmsnorm
+from flashinfer.norm import rmsnorm as fi_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
+from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm
+
+IS_CI = is_in_ci()
+
+DTYPE = torch.bfloat16
+DEVICE = "cuda"
+
+# JIT rmsnorm: hidden_size in {64,128,256} or (multiple of 256, <=8192)
+# JIT fused_add_rmsnorm: hidden_size % 8 == 0, <=8192
+# Use multiples of 256 <=8192 to satisfy both kernels
+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 = [("blue", "--"), ("green", "-.")]
+
+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="rmsnorm-performance",
+        args={},
+    )
+)
+def benchmark_rmsnorm(hidden_size: int, batch_size: int, provider: str):
+    input = torch.randn((batch_size, hidden_size), dtype=DTYPE, device=DEVICE)
+    weight = torch.randn(hidden_size, dtype=DTYPE, device=DEVICE)
+    FN_MAP = {
+        "jit": lambda: jit_rmsnorm(input.clone(), weight),
+        "fi": lambda: fi_rmsnorm(input.clone(), weight, out=input.clone()),
+    }
+    fn = FN_MAP[provider]
+    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
+
+
+@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_fused_add_rmsnorm(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": lambda: jit_fused_add_rmsnorm(
+            input.clone(), residual.clone(), weight, torch.finfo(DTYPE).eps
+        ),
+        "fi": lambda: fi_fused_add_rmsnorm(
+            input.clone(), residual.clone(), weight, eps=torch.finfo(DTYPE).eps
+        ),
+    }
+    fn = FN_MAP[provider]
+    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__":
+    print("Benchmarking rmsnorm...")
+    benchmark_rmsnorm.run(print_data=True)
+
+    print("Benchmarking fused_add_rmsnorm...")
+    benchmark_fused_add_rmsnorm.run(print_data=True)

python/sglang/jit_kernel/tests/test_norm_jit.py

@@ -0,0 +1,85 @@
+# Adapted from sgl-kernel/tests/test_norm.py
+
+import pytest
+import torch
+
+# JIT rmsnorm: fp16/bf16 only; hidden_size must be a multiple of 256, > 256, and <=8192
+RMSNORM_HIDDEN_SIZES = [512, 1024, 3072, 3584, 4096, 8192]
+
+# JIT fused_add_rmsnorm: fp16/bf16 only; hidden_size % 8 == 0, <=8192
+FUSED_ADD_RMSNORM_HIDDEN_SIZES = [1024, 3072, 3584, 4096, 8192]
+
+BS_LIST = [1, 19, 99, 989]
+
+
+def _jit_rmsnorm(input, weight, output, eps):
+    from sglang.jit_kernel.norm import rmsnorm
+
+    rmsnorm(input, weight, output=output, eps=eps)
+
+
+def _fi_rmsnorm(input, weight, out, eps):
+    from flashinfer.norm import rmsnorm
+
+    rmsnorm(input, weight, out=out, eps=eps)
+
+
+def _jit_fused_add_rmsnorm(input, residual, weight, eps):
+    from sglang.jit_kernel.norm import fused_add_rmsnorm
+
+    fused_add_rmsnorm(input, residual, weight, eps)
+
+
+def _fi_fused_add_rmsnorm(input, residual, weight, eps):
+    from flashinfer.norm import fused_add_rmsnorm
+
+    fused_add_rmsnorm(input, residual, weight, eps=eps)
+
+
+@pytest.mark.parametrize("batch_size", BS_LIST)
+@pytest.mark.parametrize("hidden_size", RMSNORM_HIDDEN_SIZES)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("specify_out", [True, False])
+def test_rmsnorm_jit(batch_size, hidden_size, dtype, specify_out):
+    eps = 1e-6
+    x = torch.randn(batch_size, hidden_size, device="cuda", dtype=dtype)
+    w = torch.randn(hidden_size, device="cuda", dtype=dtype)
+
+    # flashinfer reference
+    x_ref = x.clone()
+    _fi_rmsnorm(x_ref, w, out=x_ref, eps=eps)
+
+    if specify_out:
+        y = torch.empty_like(x)
+        _jit_rmsnorm(x, w, output=y, eps=eps)
+    else:
+        y = x.clone()
+        _jit_rmsnorm(y, w, output=y, eps=eps)
+
+    torch.testing.assert_close(y, x_ref, rtol=1e-2, atol=1e-2)
+
+
+@pytest.mark.parametrize("batch_size", BS_LIST)
+@pytest.mark.parametrize("hidden_size", FUSED_ADD_RMSNORM_HIDDEN_SIZES)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_fused_add_rmsnorm_jit(batch_size, hidden_size, dtype):
+    eps = 1e-6
+    x = torch.randn(batch_size, hidden_size, dtype=dtype, device="cuda")
+    residual = torch.randn_like(x)
+    weight = torch.randn(hidden_size, dtype=dtype, device="cuda")
+
+    # flashinfer reference
+    x_ref = x.clone()
+    r_ref = residual.clone()
+    _fi_fused_add_rmsnorm(x_ref, r_ref, weight, eps=eps)
+
+    x_jit = x.clone()
+    r_jit = residual.clone()
+    _jit_fused_add_rmsnorm(x_jit, r_jit, weight, eps)
+
+    torch.testing.assert_close(x_jit, x_ref, rtol=1e-2, atol=1e-2)
+    torch.testing.assert_close(r_jit, r_ref, rtol=1e-2, atol=1e-2)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

sgl-kernel/python/sgl_kernel/elementwise.py

@@ -1,9 +1,76 @@
 from dataclasses import dataclass
-from typing import List, Optional
+from typing import Optional
 
 import torch
 from sgl_kernel.utils import is_arch_support_pdl
 
+try:
+    import flashinfer.norm as _flashinfer_norm
+
+    _has_flashinfer = True
+except ImportError:
+    _has_flashinfer = False
+
+_FLASHINFER_NORM_SUPPORTED_DTYPES = {torch.float16, torch.bfloat16}
+
+
+def _rmsnorm_internal(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    out: Optional[torch.Tensor],
+    enable_pdl: Optional[bool],
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(input)
+    if enable_pdl is None:
+        enable_pdl = is_arch_support_pdl()
+    torch.ops.sgl_kernel.rmsnorm.default(out, input, weight, eps, enable_pdl)
+    return out
+
+
+def _fused_add_rmsnorm_internal(
+    input: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    enable_pdl: Optional[bool],
+) -> None:
+    if enable_pdl is None:
+        enable_pdl = is_arch_support_pdl()
+    torch.ops.sgl_kernel.fused_add_rmsnorm.default(
+        input, residual, weight, eps, enable_pdl
+    )
+
+
+def _gemma_rmsnorm_internal(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    out: Optional[torch.Tensor],
+    enable_pdl: Optional[bool],
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(input)
+    if enable_pdl is None:
+        enable_pdl = is_arch_support_pdl()
+    torch.ops.sgl_kernel.gemma_rmsnorm.default(out, input, weight, eps, enable_pdl)
+    return out
+
+
+def _gemma_fused_add_rmsnorm_internal(
+    input: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    enable_pdl: Optional[bool],
+) -> None:
+    if enable_pdl is None:
+        enable_pdl = is_arch_support_pdl()
+    torch.ops.sgl_kernel.gemma_fused_add_rmsnorm.default(
+        input, residual, weight, eps, enable_pdl
+    )
+
 
 # These implementations extensively draw from and build upon the FlashInfer project https://github.com/flashinfer-ai/flashinfer
 # Kudos to @yzh119
@@ -38,12 +105,23 @@ def rmsnorm(
     output: torch.Tensor
         Normalized tensor, shape (batch_size, hidden_size).
     """
-    if out is None:
+    # torch.compiler.is_dynamo_compiling(): FlashInfer norm paths are not safe under
-        out = torch.empty_like(input)
+    # torch.compile(..., fullgraph=True). Dynamo traces into FlashInfer's JIT module
-    if enable_pdl is None:
+    # loading path, which calls Path.exists() / os.stat() — both untraceable — causing
-        enable_pdl = is_arch_support_pdl()
+    # the entire compilation to fail. We fall back to the internal implementation while
-    torch.ops.sgl_kernel.rmsnorm.default(out, input, weight, eps, enable_pdl)
+    # tracing as a temporary workaround. Once the upstream fix is merged and we upgrade
-    return out
+    # FlashInfer, this check can be removed.
+    # See: https://github.com/flashinfer-ai/flashinfer/issues/2734
+    #      https://github.com/flashinfer-ai/flashinfer/pull/2733
+    if (
+        input.device.type == "musa"
+        or not _has_flashinfer
+        or input.dtype not in _FLASHINFER_NORM_SUPPORTED_DTYPES
+        or torch.compiler.is_dynamo_compiling()
+    ):
+        return _rmsnorm_internal(input, weight, eps, out, enable_pdl)
+    else:
+        return _flashinfer_norm.rmsnorm(input, weight, eps, out, enable_pdl)
 
 
 def fused_add_rmsnorm(
@@ -76,11 +154,16 @@ def fused_add_rmsnorm(
         <https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#programmatic-dependent-launch-and-synchronization>`_
         If None, will be automatically enabled on Hopper architecture.
     """
-    if enable_pdl is None:
+    # See is_dynamo_compiling() comment in rmsnorm() above.
-        enable_pdl = is_arch_support_pdl()
+    if (
-    torch.ops.sgl_kernel.fused_add_rmsnorm.default(
+        input.device.type == "musa"
-        input, residual, weight, eps, enable_pdl
+        or not _has_flashinfer
-    )
+        or input.dtype not in _FLASHINFER_NORM_SUPPORTED_DTYPES
+        or torch.compiler.is_dynamo_compiling()
+    ):
+        _fused_add_rmsnorm_internal(input, residual, weight, eps, enable_pdl)
+    else:
+        _flashinfer_norm.fused_add_rmsnorm(input, residual, weight, eps, enable_pdl)
 
 
 def gemma_rmsnorm(
@@ -114,12 +197,16 @@ def gemma_rmsnorm(
     output: torch.Tensor
         Gemma Normalized tensor, shape (batch_size, hidden_size).
     """
-    if out is None:
+    # See is_dynamo_compiling() comment in rmsnorm() above.
-        out = torch.empty_like(input)
+    if (
-    if enable_pdl is None:
+        input.device.type == "musa"
-        enable_pdl = is_arch_support_pdl()
+        or not _has_flashinfer
-    torch.ops.sgl_kernel.gemma_rmsnorm.default(out, input, weight, eps, enable_pdl)
+        or input.dtype not in _FLASHINFER_NORM_SUPPORTED_DTYPES
-    return out
+        or torch.compiler.is_dynamo_compiling()
+    ):
+        return _gemma_rmsnorm_internal(input, weight, eps, out, enable_pdl)
+    else:
+        return _flashinfer_norm.gemma_rmsnorm(input, weight, eps, out, enable_pdl)
 
 
 def gemma_fused_add_rmsnorm(
@@ -152,9 +239,16 @@ def gemma_fused_add_rmsnorm(
         <https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#programmatic-dependent-launch-and-synchronization>`_
         If None, will be automatically enabled on Hopper architecture.
     """
-    if enable_pdl is None:
+    # See is_dynamo_compiling() comment in rmsnorm() above.
-        enable_pdl = is_arch_support_pdl()
+    if (
-    torch.ops.sgl_kernel.gemma_fused_add_rmsnorm.default(
+        input.device.type == "musa"
+        or not _has_flashinfer
+        or input.dtype not in _FLASHINFER_NORM_SUPPORTED_DTYPES
+        or torch.compiler.is_dynamo_compiling()
+    ):
+        _gemma_fused_add_rmsnorm_internal(input, residual, weight, eps, enable_pdl)
+    else:
+        _flashinfer_norm.gemma_fused_add_rmsnorm(
             input, residual, weight, eps, enable_pdl
         )