diff --git a/python/sglang/jit_kernel/benchmark/bench_norm.py b/python/sglang/jit_kernel/benchmark/bench_norm.py new file mode 100644 index 000000000..0e0e80a6f --- /dev/null +++ b/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) diff --git a/python/sglang/jit_kernel/tests/test_norm_jit.py b/python/sglang/jit_kernel/tests/test_norm_jit.py new file mode 100644 index 000000000..271c5c947 --- /dev/null +++ b/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__]) diff --git a/sgl-kernel/python/sgl_kernel/elementwise.py b/sgl-kernel/python/sgl_kernel/elementwise.py index c7a6a0ed1..704cbf492 100644 --- a/sgl-kernel/python/sgl_kernel/elementwise.py +++ b/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: - 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 + # torch.compiler.is_dynamo_compiling(): FlashInfer norm paths are not safe under + # torch.compile(..., fullgraph=True). Dynamo traces into FlashInfer's JIT module + # loading path, which calls Path.exists() / os.stat() — both untraceable — causing + # the entire compilation to fail. We fall back to the internal implementation while + # tracing as a temporary workaround. Once the upstream fix is merged and we upgrade + # 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( `_ If None, will be automatically enabled on Hopper architecture. """ - 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 - ) + # See is_dynamo_compiling() comment in rmsnorm() above. + if ( + input.device.type == "musa" + 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: - 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 + # See is_dynamo_compiling() comment in rmsnorm() above. + 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 _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,11 +239,18 @@ def gemma_fused_add_rmsnorm( `_ If None, will be automatically enabled on Hopper architecture. """ - 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 - ) + # See is_dynamo_compiling() comment in rmsnorm() above. + if ( + 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 + ) def _check_shape(input: torch.Tensor, output: torch.Tensor) -> None: