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)