minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[Feature] JIT rmsnorm update (with claude) (#21834)

Browse Source
This commit is contained in:
DarkSharpness
2026-04-01 23:40:00 +08:00
committed by GitHub
parent 4f5b55e379
commit 20f4193589
8 changed files with 321 additions and 393 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/benchmark/bench_fused_add_rmsnorm.py
-75
View File
@@ -1,75 +0,0 @@
import itertools
import torch
import triton
import triton.testing
from flashinfer import fused_add_rmsnorm as fi_fused_add_rmsnorm
from sglang.jit_kernel.benchmark.utils import run_benchmark
from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.utils import is_in_ci
register_cuda_ci(est_time=6, suite="stage-b-kernel-benchmark-1-gpu-large")
IS_CI = is_in_ci()
def sglang_jit_fused_add_rmsnorm(
input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
) -> None:
jit_fused_add_rmsnorm(input, residual, weight, eps)
def flashinfer_fused_add_rmsnorm(
input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
) -> None:
fi_fused_add_rmsnorm(input, residual, weight, eps=eps)
DTYPE = torch.bfloat16
DEVICE = "cuda"
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", "flashinfer"]
LINE_NAMES = ["SGL JIT Kernel", "FlashInfer"]
STYLES = [("orange", "-"), ("blue", "--"), ("green", "-."), ("red", ":")]
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="fused-add-rmsnorm-performance",
args={},
)
)
def benchmark(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": sglang_jit_fused_add_rmsnorm,
"flashinfer": flashinfer_fused_add_rmsnorm,
}
fn = lambda: FN_MAP[provider](
input.clone(), residual.clone(), weight, torch.finfo(torch.bfloat16).eps
)
return run_benchmark(fn)
if __name__ == "__main__":
benchmark.run(print_data=True)
python/sglang/jit_kernel/benchmark/bench_norm.py
+42 -39
View File
@@ -6,40 +6,39 @@ 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 run_benchmark
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm
from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.utils import is_in_ci
register_cuda_ci(est_time=5, suite="stage-b-kernel-benchmark-1-gpu-large")
register_cuda_ci(est_time=30, suite="stage-b-kernel-benchmark-1-gpu-large")
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]
BS_LIST = get_benchmark_range(
full_range=[2**n for n in range(0, 14)],
ci_range=[16, 32],
)
HIDDEN_SIZE_LIST = get_benchmark_range(
full_range=sorted([1536, *range(1024, 8192 + 1, 1024)]),
ci_range=[512, 2048],
)
LINE_VALS = ["jit", "flashinfer"]
LINE_NAMES = ["SGL JIT Kernel", "FlashInfer"]
LINE_VALS = ["flashinfer", "jit"]
LINE_NAMES = ["FlashInfer", "SGL JIT Kernel"]
STYLES = [("blue", "--"), ("green", "-.")]
NUM_LAYERS = 4 # avoid L2 effect
configs = list(itertools.product(HIDDEN_SIZE_LIST, BS_LIST))
configs_0 = list(itertools.product(HIDDEN_SIZE_LIST + [16384], BS_LIST))
configs_1 = list(itertools.product(HIDDEN_SIZE_LIST, BS_LIST))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["hidden_size", "batch_size"],
x_vals=configs,
x_vals=configs_0,
line_arg="provider",
line_vals=LINE_VALS,
line_names=LINE_NAMES,
@@ -50,20 +49,24 @@ configs = list(itertools.product(HIDDEN_SIZE_LIST, BS_LIST))
)
)
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),
"flashinfer": lambda: fi_rmsnorm(input.clone(), weight, out=input.clone()),
}
fn = FN_MAP[provider]
return run_benchmark(fn)
input = torch.randn(
(NUM_LAYERS, batch_size, hidden_size), dtype=DTYPE, device=DEVICE
)
weight = torch.randn((NUM_LAYERS, hidden_size), dtype=DTYPE, device=DEVICE)
FN_MAP = {"jit": jit_rmsnorm, "flashinfer": fi_rmsnorm}
def f():
fn = FN_MAP[provider]
for i in range(NUM_LAYERS):
fn(input[i], weight[i], out=input[i])
return run_benchmark(f, scale=NUM_LAYERS)
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["hidden_size", "batch_size"],
x_vals=configs,
x_vals=configs_1,
line_arg="provider",
line_vals=LINE_VALS,
line_names=LINE_NAMES,
@@ -74,19 +77,19 @@ def benchmark_rmsnorm(hidden_size: int, batch_size: int, provider: str):
)
)
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
),
"flashinfer": lambda: fi_fused_add_rmsnorm(
input.clone(), residual.clone(), weight, eps=torch.finfo(DTYPE).eps
),
}
fn = FN_MAP[provider]
return run_benchmark(fn)
input = torch.randn(
(NUM_LAYERS, batch_size, hidden_size), dtype=DTYPE, device=DEVICE
)
residual = torch.randn_like(input)
weight = torch.randn((NUM_LAYERS, hidden_size), dtype=DTYPE, device=DEVICE)
FN_MAP = {"jit": jit_fused_add_rmsnorm, "flashinfer": fi_fused_add_rmsnorm}
def f():
fn = FN_MAP[provider]
for i in range(NUM_LAYERS):
fn(input[i], residual[i], weight[i])
return run_benchmark(f, scale=NUM_LAYERS)
if __name__ == "__main__":
python/sglang/jit_kernel/benchmark/bench_rmsnorm.py
-98
View File
@@ -1,98 +0,0 @@
import itertools
import torch
import triton
import triton.testing
from flashinfer import rmsnorm as fi_rmsnorm
from sgl_kernel import rmsnorm
from sglang.jit_kernel.benchmark.utils import (
DEFAULT_DEVICE,
DEFAULT_DTYPE,
get_benchmark_range,
run_benchmark,
)
from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=21, suite="stage-b-kernel-benchmark-1-gpu-large")
def sglang_aot_rmsnorm(
input: torch.Tensor,
weight: torch.Tensor,
) -> None:
rmsnorm(input, weight, out=input)
def sglang_jit_rmsnorm(
input: torch.Tensor,
weight: torch.Tensor,
) -> None:
jit_rmsnorm(input, weight, output=input)
def flashinfer_rmsnorm(
input: torch.Tensor,
weight: torch.Tensor,
) -> None:
fi_rmsnorm(input, weight, out=input)
@torch.compile()
def torch_impl_rmsnorm(
input: torch.Tensor,
weight: torch.Tensor,
eps: float = 1e-6,
) -> None:
mean = input.float().pow(2).mean(dim=-1, keepdim=True)
norm = (mean + eps).rsqrt()
input.copy_(input.float() * norm * weight.float())
BS_LIST = get_benchmark_range(
full_range=[2**n for n in range(0, 14)],
ci_range=[16],
)
HIDDEN_SIZE_LIST = get_benchmark_range(
full_range=[1536, 3072, 4096, 5120, 8192],
ci_range=[512, 2048],
)
LINE_VALS = ["aot", "jit", "flashinfer", "torch"]
LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "FlashInfer", "PyTorch"]
STYLES = [("orange", "-"), ("blue", "--"), ("green", "-."), ("red", ":")]
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(hidden_size: int, batch_size: int, provider: str):
input = torch.randn(
(batch_size, hidden_size), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE
)
weight = torch.randn(hidden_size, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE)
FN_MAP = {
"aot": sglang_aot_rmsnorm,
"jit": sglang_jit_rmsnorm,
"flashinfer": flashinfer_rmsnorm,
"torch": torch_impl_rmsnorm,
}
fn = lambda: FN_MAP[provider](input.clone(), weight)
return run_benchmark(fn)
if __name__ == "__main__":
benchmark.run(print_data=True)
python/sglang/jit_kernel/benchmark/utils.py
+12 -7
View File
@@ -1,6 +1,6 @@
"""Common utilities for jit_kernel benchmark files."""
from typing import Callable, List, Tuple
from typing import Callable, List, Sequence, Tuple
import torch
import triton.testing
@@ -19,25 +19,30 @@ def get_benchmark_range(full_range: List, ci_range: List) -> List:
def run_benchmark(
fn: Callable, quantiles: List[float] = None
fn: Callable,
quantiles: Sequence[float] = (),
scale: float = 1.0,
) -> Tuple[float, float, float]:
"""Execute benchmark using CUDA graph and return times in microseconds.
Args:
fn: Function to benchmark
quantiles: Quantiles for timing measurements [median, min, max]
scale: Scale the result down (usually num_layers).
Returns:
Tuple of (median_us, max_us, min_us)
"""
quantiles = quantiles or DEFAULT_QUANTILES
quantiles = list(quantiles or DEFAULT_QUANTILES)
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
return 1000 * ms / scale, 1000 * max_ms / scale, 1000 * min_ms / scale
def run_benchmark_no_cudagraph(
fn: Callable, quantiles: List[float] = None
fn: Callable,
quantiles: Sequence[float] = (),
scale: float = 1.0,
) -> Tuple[float, float, float]:
quantiles = quantiles or DEFAULT_QUANTILES
quantiles = list(quantiles or DEFAULT_QUANTILES)
ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=quantiles)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
return 1000 * ms / scale, 1000 * max_ms / scale, 1000 * min_ms / scale