96 lines
2.3 KiB
Python
96 lines
2.3 KiB
Python
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
|
|
|
|
|
|
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", "fi", "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,
|
|
"fi": 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)
|