+2
3f5e2c7688
3f5e2c7688
Co-authored-by: wunhuang <wunhuang@amd.com> Co-authored-by: Thomas Wang <1am9trash@gmail.com> Co-authored-by: Xinyi Song <86638975+RolaoDenthu@users.noreply.github.com> Co-authored-by: HaiShaw <hixiao@gmail.com> Co-authored-by: amd-danli103 <danli103@amd.com> Co-authored-by: Lin, Soga <soga.lin@amd.com> Co-authored-by: Raiden-Makoto <Raiden-Makoto@users.noreply.github.com> Co-authored-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com> Co-authored-by: yichiche@amd.com <jacky.cheng> Co-authored-by: yctseng0211 <yctseng@amd.com> Co-authored-by: Bingxu Chen <bingxche@amd.com>
76 lines
2.1 KiB
Python
76 lines
2.1 KiB
Python
"""Benchmark for DeepSeek-V4 fused norm + RoPE kernels."""
|
|
|
|
import itertools
|
|
|
|
import sgl_kernel
|
|
import torch
|
|
import triton
|
|
import triton.testing
|
|
|
|
try:
|
|
from sglang.utils import is_in_ci
|
|
|
|
IS_CI = is_in_ci()
|
|
except ImportError:
|
|
IS_CI = False
|
|
|
|
batch_sizes = [1] if IS_CI else [1, 4, 16, 64, 256]
|
|
num_heads_list = [8] if IS_CI else [8, 16, 64]
|
|
head_dims = [192] if IS_CI else [128, 192]
|
|
|
|
configs = list(itertools.product(batch_sizes, num_heads_list, head_dims))
|
|
|
|
|
|
def torch_rmsnorm_rope(
|
|
q: torch.Tensor, freqs_cis: torch.Tensor, positions: torch.Tensor, eps: float
|
|
) -> torch.Tensor:
|
|
"""Naive PyTorch reference: RMSNorm + RoPE."""
|
|
rms = torch.sqrt(q.float().pow(2).mean(dim=-1, keepdim=True) + eps)
|
|
q_normed = (q.float() / rms).to(q.dtype)
|
|
return q_normed
|
|
|
|
|
|
@triton.testing.perf_report(
|
|
triton.testing.Benchmark(
|
|
x_names=["batch_size", "num_heads", "head_dim"],
|
|
x_vals=configs,
|
|
line_arg="provider",
|
|
line_vals=["sglang", "torch"],
|
|
line_names=["SGL Kernel", "PyTorch"],
|
|
styles=[("green", "-"), ("red", "--")],
|
|
ylabel="µs (median)",
|
|
plot_name="dsv4-q-norm-rope-performance",
|
|
args={},
|
|
)
|
|
)
|
|
def benchmark_q_norm_rope(batch_size, num_heads, head_dim, provider):
|
|
torch.manual_seed(42)
|
|
eps = 1e-6
|
|
max_pos = 8192
|
|
rope_dim = 64
|
|
|
|
q_input = torch.randn(
|
|
batch_size, num_heads, head_dim, dtype=torch.bfloat16, device="cuda"
|
|
)
|
|
q_output = torch.empty_like(q_input)
|
|
freqs_cis = torch.randn(max_pos, rope_dim, dtype=torch.float32, device="cuda")
|
|
positions = torch.randint(
|
|
0, max_pos, (batch_size,), dtype=torch.int32, device="cuda"
|
|
)
|
|
|
|
if provider == "sglang":
|
|
fn = lambda: sgl_kernel.dsv4_fused_q_norm_rope(
|
|
q_input, freqs_cis, positions, eps, q_output
|
|
)
|
|
else:
|
|
fn = lambda: torch_rmsnorm_rope(q_input, freqs_cis, positions, eps)
|
|
|
|
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
|
fn, quantiles=[0.5, 0.2, 0.8]
|
|
)
|
|
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
|
|
|
|
|
|
if __name__ == "__main__":
|
|
benchmark_q_norm_rope.run(print_data=True)
|