chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[3/n jit_kernel restruct] Clean up benchmark naming and benchmarking helpers (#20250)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-03-10 16:39:03 +08:00
committed by GitHub
parent c812504b92
commit 8517da5d08
17 changed files with 82 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py
View File

@@ -6,6 +6,7 @@ import triton
import triton.testing
from sglang.jit_kernel.awq_dequantize import awq_dequantize as jit_awq_dequantize
from sglang.jit_kernel.benchmark.utils import run_benchmark
try:
from sgl_kernel import awq_dequantize as aot_awq_dequantize
@@ -19,7 +20,6 @@ IS_CI = (
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
# CI environment uses simplified parameters
if IS_CI:
qweight_row_range = [128]
qweight_cols_range = [16]
@@ -107,8 +107,6 @@ def benchmark(qweight_row, qweight_col, provider):
device=device,
)
quantiles = [0.5, 0.2, 0.8]
if provider == "jit":
fn = lambda: jit_awq_dequantize(qweight, scales, qzeros)
elif provider == "aot":
@@ -116,8 +114,7 @@ def benchmark(qweight_row, qweight_col, provider):
else:
raise ValueError(f"Unknown provider: {provider}")
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
return run_benchmark(fn)
if __name__ == "__main__":

7
python/sglang/jit_kernel/benchmark/bench_awq_marlin_moe_repack.py
View File

@@ -9,6 +9,7 @@ from sgl_kernel.scalar_type import scalar_types
from sglang.jit_kernel.awq_marlin_repack import (
awq_marlin_moe_repack as jit_awq_marlin_moe_repack,
)
from sglang.jit_kernel.benchmark.utils import run_benchmark
from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
try:
@@ -23,7 +24,6 @@ IS_CI = (
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
# Fixed parameters
NUM_BITS = 4
GROUP_SIZE = 128
SIZE_N = 4096
@@ -111,8 +111,6 @@ def benchmark(num_experts, size_k, size_n, num_bits, provider):
num_experts, size_k, size_n, num_bits, group_size
)
quantiles = [0.5, 0.2, 0.8]
if provider == "jit":
fn = lambda: jit_awq_marlin_moe_repack(
b_q_weight, perm, size_k, size_n, num_bits
@@ -124,8 +122,7 @@ def benchmark(num_experts, size_k, size_n, num_bits, provider):
else:
raise ValueError(f"Unknown provider: {provider}")
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
return run_benchmark(fn)
if __name__ == "__main__":

8
python/sglang/jit_kernel/benchmark/bench_awq_marlin_repack.py
View File

@@ -9,6 +9,7 @@ from sgl_kernel.scalar_type import scalar_types
from sglang.jit_kernel.awq_marlin_repack import (
awq_marlin_repack as jit_awq_marlin_repack,
)
from sglang.jit_kernel.benchmark.utils import run_benchmark
from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
try:
@@ -23,7 +24,6 @@ IS_CI = (
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
# Fixed problem dimensions
SIZE_K = 4096
SIZE_N = 4096
NUM_BITS = 4
@@ -43,7 +43,6 @@ def awq_pack(q_w, num_bits, size_k, size_n):
return pack_cols(q_w, num_bits, size_k, size_n)
# Quantize weights once
_b_weight = torch.randn((SIZE_K, SIZE_N), dtype=torch.float16, device="cuda")
_w_ref, _q_w, _s, _zp = quantize_weights(
_b_weight, scalar_types.uint4, GROUP_SIZE, zero_points=True
@@ -99,8 +98,6 @@ def benchmark(size_k, size_n, num_bits, provider):
)
q_w_awq = awq_pack(q_w, num_bits, size_k, size_n)
quantiles = [0.5, 0.2, 0.8]
if provider == "jit":
fn = lambda: jit_awq_marlin_repack(q_w_awq, size_k, size_n, num_bits)
elif provider == "aot":
@@ -108,8 +105,7 @@ def benchmark(size_k, size_n, num_bits, provider):
else:
raise ValueError(f"Unknown provider: {provider}")
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
return run_benchmark(fn)
if __name__ == "__main__":

11
python/sglang/jit_kernel/benchmark/bench_concat_mla.py
View File

@@ -6,13 +6,12 @@ import triton.testing
from sgl_kernel import concat_mla_absorb_q as aot_absorb_q
from sgl_kernel import concat_mla_k as aot_k
from sglang.jit_kernel.benchmark.utils import is_in_ci
from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark
from sglang.jit_kernel.concat_mla import concat_mla_absorb_q as jit_absorb_q
from sglang.jit_kernel.concat_mla import concat_mla_k as jit_k
IS_CI = is_in_ci()
# Constants
NUM_LOCAL_HEADS = 128
QK_NOPE_HEAD_DIM = 128
QK_ROPE_HEAD_DIM = 64
@@ -109,9 +108,7 @@ def bench_concat_mla_k(num_tokens: int, provider: str):
"torch": torch_concat_mla_k,
}
fn = lambda: FN_MAP[provider](k, k_nope, k_rope)
quantiles = [0.5, 0.2, 0.8]
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
return run_benchmark(fn)
if IS_CI:
@@ -153,9 +150,7 @@ def bench_concat_mla_absorb_q(dim_0: int, dim_1: int, provider: str):
}
fn = lambda: FN_MAP[provider](a, b)
quantiles = [0.5, 0.2, 0.8]
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
return run_benchmark(fn)
if __name__ == "__main__":

10
python/sglang/jit_kernel/benchmark/bench_fused_add_rmsnorm.py
View File

@@ -5,7 +5,7 @@ import triton
import triton.testing
from flashinfer import fused_add_rmsnorm as fi_fused_add_rmsnorm
from sglang.jit_kernel.benchmark.utils import is_in_ci
from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark
from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm
IS_CI = is_in_ci()
@@ -33,7 +33,7 @@ else:
BS_LIST = [2**n for n in range(0, 14)]
HIDDEN_SIZE_LIST = [1536, 3072, 4096, 5120, 8192]
LINE_VALS = ["jit", "fi"]
LINE_VALS = ["jit", "flashinfer"]
LINE_NAMES = ["SGL JIT Kernel", "FlashInfer"]
STYLES = [("orange", "-"), ("blue", "--"), ("green", "-."), ("red", ":")]
@@ -59,14 +59,12 @@ def benchmark(hidden_size: int, batch_size: int, provider: str):
weight = torch.randn(hidden_size, dtype=DTYPE, device=DEVICE)
FN_MAP = {
"jit": sglang_jit_fused_add_rmsnorm,
"fi": flashinfer_fused_add_rmsnorm,
"flashinfer": flashinfer_fused_add_rmsnorm,
}
fn = lambda: FN_MAP[provider](
input.clone(), residual.clone(), weight, torch.finfo(torch.bfloat16).eps
)
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
return run_benchmark(fn)
if __name__ == "__main__":

10
python/sglang/jit_kernel/benchmark/bench_fused_norm_scale_shift.py
View File

@@ -8,7 +8,7 @@ import torch
import triton
import triton.testing
from sglang.jit_kernel.benchmark.utils import is_in_ci
from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark_no_cudagraph
from sglang.multimodal_gen.runtime.layers.layernorm import (
LayerNormScaleShift,
RMSNormScaleShift,
@@ -78,9 +78,7 @@ def bench_fused_norm_scale_shift(
else:
fn = lambda: layer.forward_cuda(x, shift, scale)
quantiles = [0.5, 0.2, 0.8]
ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=quantiles)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms # convert to us
return run_benchmark_no_cudagraph(fn)
# ============================================================================
@@ -117,9 +115,7 @@ def bench_fused_scale_residual_norm_scale_shift(
else:
fn = lambda: layer.forward_cuda(residual, x, gate, shift, scale)
quantiles = [0.5, 0.2, 0.8]
ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=quantiles)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms # convert to us
return run_benchmark_no_cudagraph(fn)
if __name__ == "__main__":

8
python/sglang/jit_kernel/benchmark/bench_gptq_marlin.py
View File

@@ -5,6 +5,7 @@ import triton
import triton.testing
from sgl_kernel.scalar_type import scalar_types
from sglang.jit_kernel.benchmark.utils import run_benchmark
from sglang.jit_kernel.gptq_marlin import gptq_marlin_gemm as jit_gptq_marlin_gemm
from sglang.srt.layers.quantization.marlin_utils import marlin_make_workspace
from sglang.test.test_marlin_utils import marlin_quantize
@@ -21,13 +22,11 @@ IS_CI = (
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
# Fixed problem dimensions
SIZE_K = 4096
SIZE_N = 4096
GROUP_SIZE = 128
QUANT_TYPE = scalar_types.uint4b8
# Quantize weights once
_b_weight = torch.randn((SIZE_K, SIZE_N), dtype=torch.float16, device="cuda")
_w_ref, _marlin_q_w, _marlin_s, _g_idx, _sort_indices, _ = marlin_quantize(
_b_weight, QUANT_TYPE, GROUP_SIZE, act_order=False
@@ -100,8 +99,6 @@ def benchmark(size_m, provider):
device = torch.device("cuda")
a = torch.randn((size_m, SIZE_K), dtype=torch.float16, device=device)
quantiles = [0.5, 0.2, 0.8]
if provider == "jit":
fn = lambda: _run_gemm(jit_gptq_marlin_gemm, a)
elif provider == "aot":
@@ -109,8 +106,7 @@ def benchmark(size_m, provider):
else:
raise ValueError(f"Unknown provider: {provider}")
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
return run_benchmark(fn)
if __name__ == "__main__":

8
python/sglang/jit_kernel/benchmark/bench_gptq_marlin_repack.py
View File

@@ -5,6 +5,7 @@ import triton
import triton.testing
from sgl_kernel.scalar_type import scalar_types
from sglang.jit_kernel.benchmark.utils import run_benchmark
from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack as jit_fn
from sglang.srt.layers.quantization.utils import gptq_quantize_weights, pack_rows
@@ -20,13 +21,11 @@ IS_CI = (
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
# Fixed problem dimensions
SIZE_N = 4096
NUM_BITS = 4
QUANT_TYPE = scalar_types.uint4b8
GROUP_SIZE = 128
# Pre-compute quantized weight for each size_k in the sweep
_cache = {}
@@ -86,8 +85,6 @@ else:
def benchmark(size_k, provider):
q_w_gptq, sort_indices = _get_inputs(size_k)
quantiles = [0.5, 0.2, 0.8]
if provider == "jit":
fn = lambda: jit_fn(q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS)
elif provider == "aot":
@@ -95,8 +92,7 @@ def benchmark(size_k, provider):
else:
raise ValueError(f"Unknown provider: {provider}")
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
return run_benchmark(fn)
if __name__ == "__main__":

17
python/sglang/jit_kernel/benchmark/bench_hicache.py
View File

@@ -32,9 +32,6 @@ from sglang.jit_kernel.hicache import (
transfer_hicache_one_layer,
)
# NOTE: Adjustable hyperparameters for better benchmark stability
# NOTE: torch impl is too slow in benchmark
DISABLE_TORCH = os.environ.get("DISABLE_TORCH", "0") == "1"
PAGE_SIZE = 1
ENABLE_SORT = True
@@ -193,7 +190,7 @@ ELEMENT_SIZE_RANGE = get_benchmark_range(
ci_range=[1024],
)
LINE_VALS = ["aot", "jit", "pytorch"]
LINE_VALS = ["aot", "jit", "torch"]
LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "PyTorch"]
STYLES = [("orange", "-"), ("blue", "--"), ("red", ":")]
@@ -228,12 +225,10 @@ def benchmark_one_layer_h2d(
v_cache_src = cache_local.v_cache_host
k_cache_dst = cache_local.k_cache_cuda
v_cache_dst = cache_local.v_cache_cuda
# to avoid fluctutation, we set the seed as const
torch.manual_seed(batch_size * 65536 + element_size)
indices_src_gpu = gen_indices(batch_size, HOST_CACHE_SIZE)
indices_dst_gpu = gen_indices(batch_size, GPU_CACHE_SIZE)
# sort by host indices to improve host access performance
if ENABLE_SORT:
indices_src_gpu, mapping = indices_src_gpu.sort()
indices_dst_gpu = indices_dst_gpu[mapping]
@@ -267,7 +262,7 @@ def benchmark_one_layer_h2d(
)
for i in range(NUM_LAYERS)
],
"pytorch": lambda: [
"torch": lambda: [
pytorch_transfer(
k_cache_dst[i],
v_cache_dst[i],
@@ -283,7 +278,7 @@ def benchmark_one_layer_h2d(
if provider == "jit" and not can_use_hicache_jit_kernel(element_size=element_bytes):
return (float("nan"), float("nan"), float("nan"))
if DISABLE_TORCH and provider in ["pytorch"]:
if DISABLE_TORCH and provider in ["torch"]:
return (float("nan"), float("nan"), float("nan"))
ms, min_ms, max_ms = triton.testing.do_bench( # type: ignore
@@ -333,12 +328,10 @@ def benchmark_all_layer_d2h(
v_caches_src = cache_local.v_cache_cuda
k_caches_dst = cache_local.k_cache_host
v_caches_dst = cache_local.v_cache_host
# to avoid fluctutation, we set the seed as const
torch.manual_seed(batch_size * 65536 + element_size)
indices_src_gpu = gen_indices(batch_size, GPU_CACHE_SIZE)
indices_dst_gpu = gen_indices(batch_size, HOST_CACHE_SIZE)
# sort by host indices to improve host access performance
if ENABLE_SORT:
indices_dst_gpu, mapping = indices_dst_gpu.sort()
indices_src_gpu = indices_src_gpu[mapping]
@@ -373,7 +366,7 @@ def benchmark_all_layer_d2h(
element_bytes,
element_bytes,
),
"pytorch": lambda: [
"torch": lambda: [
pytorch_transfer(
k_caches_dst[i],
v_caches_dst[i],
@@ -389,7 +382,7 @@ def benchmark_all_layer_d2h(
if provider == "jit" and not can_use_hicache_jit_kernel(element_size=element_bytes):
return (float("nan"), float("nan"), float("nan"))
if DISABLE_TORCH and provider in ["pytorch"]:
if DISABLE_TORCH and provider in ["torch"]:
return (float("nan"), float("nan"), float("nan"))
ms, min_ms, max_ms = triton.testing.do_bench( # type: ignore

8
python/sglang/jit_kernel/benchmark/bench_moe_wna16_marlin.py
View File

@@ -5,6 +5,7 @@ import triton
import triton.testing
from sgl_kernel.scalar_type import scalar_types
from sglang.jit_kernel.benchmark.utils import run_benchmark
from sglang.jit_kernel.moe_wna16_marlin import moe_wna16_marlin_gemm as jit_fn
from sglang.srt.layers.moe.fused_moe_triton import moe_align_block_size
from sglang.test.test_marlin_utils import marlin_quantize
@@ -27,7 +28,6 @@ def stack_and_dev(tensors):
return torch.stack(tensors, dim=0).to(dev)
# Fixed problem dimensions
E = 8
SIZE_K = 4096
SIZE_N = 4096
@@ -37,7 +37,6 @@ QUANT_TYPE = scalar_types.uint4b8
DTYPE = torch.float16
BLOCK_SIZE_M = 64
# Quantize weights once (per-expert)
torch.manual_seed(0)
_qweight_l, _scales_l, _w_ref_l = [], [], []
for i in range(E):
@@ -215,8 +214,6 @@ def benchmark(size_m, provider):
_make_inputs(size_m)
)
quantiles = [0.5, 0.2, 0.8]
if provider == "jit":
fn = lambda: _run_jit(
a,
@@ -242,8 +239,7 @@ def benchmark(size_m, provider):
else:
raise ValueError(f"Unknown provider: {provider}")
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
return run_benchmark(fn)
if __name__ == "__main__":

16
python/sglang/jit_kernel/benchmark/bench_norm.py
View File

@@ -6,7 +6,7 @@ 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.benchmark.utils import is_in_ci, 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
@@ -25,7 +25,7 @@ else:
BS_LIST = [2**n for n in range(0, 14)]
HIDDEN_SIZE_LIST = [1536, 3072, 4096, 5120, 8192]
LINE_VALS = ["jit", "fi"]
LINE_VALS = ["jit", "flashinfer"]
LINE_NAMES = ["SGL JIT Kernel", "FlashInfer"]
STYLES = [("blue", "--"), ("green", "-.")]
@@ -50,12 +50,10 @@ def benchmark_rmsnorm(hidden_size: int, batch_size: int, provider: str):
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()),
"flashinfer": 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
return run_benchmark(fn)
@triton.testing.perf_report(
@@ -79,14 +77,12 @@ def benchmark_fused_add_rmsnorm(hidden_size: int, batch_size: int, provider: str
"jit": lambda: jit_fused_add_rmsnorm(
input.clone(), residual.clone(), weight, torch.finfo(DTYPE).eps
),
"fi": lambda: fi_fused_add_rmsnorm(
"flashinfer": 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
return run_benchmark(fn)
if __name__ == "__main__":

4
python/sglang/jit_kernel/benchmark/bench_qknorm.py
View File

@@ -91,7 +91,7 @@ HEAD_DIM_RANGE = get_benchmark_range(
ci_range=[128],
)
LINE_VALS = ["aot", "jit", "fi", "torch"]
LINE_VALS = ["aot", "jit", "flashinfer", "torch"]
LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "FlashInfer", "PyTorch"]
STYLES = [("orange", "-"), ("blue", "--"), ("green", "-."), ("red", ":")]
@@ -126,7 +126,7 @@ def benchmark(
FN_MAP = {
"aot": sglang_aot_qknorm,
"jit": sglang_jit_qknorm,
"fi": flashinfer_qknorm,
"flashinfer": flashinfer_qknorm,
"torch": torch_impl_qknorm,
}
fn = lambda: FN_MAP[provider](q, k, q_weight, k_weight)

10
python/sglang/jit_kernel/benchmark/bench_qknorm_across_heads.py
View File

@@ -6,7 +6,7 @@ import triton
import triton.testing
from sgl_kernel import rmsnorm
from sglang.jit_kernel.benchmark.utils import is_in_ci
from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark
from sglang.jit_kernel.norm import fused_inplace_qknorm_across_heads
from sglang.srt.utils import get_current_device_stream_fast
@@ -78,7 +78,7 @@ else:
BS_RANGE = [2**n for n in range(0, 14)]
HIDDEN_DIM_RANGE = [512, 1024, 2048, 4096, 8192]
LINE_VALS = ["jit", "aot", "fi", "torch"]
LINE_VALS = ["jit", "aot", "flashinfer", "torch"]
LINE_NAMES = ["SGL JIT Kernel", "SGL AOT Kernel", "FlashInfer", "PyTorch"]
STYLES = [("blue", "-"), ("orange", "--"), ("green", "-."), ("red", ":")]
@@ -108,13 +108,11 @@ def benchmark(
FN_MAP = {
"jit": sglang_jit_qknorm_across_heads,
"aot": sglang_aot_qknorm_across_heads,
"fi": flashinfer_qknorm_across_heads,
"flashinfer": flashinfer_qknorm_across_heads,
"torch": torch_impl_qknorm_across_heads,
}
fn = lambda: FN_MAP[provider](q, k, q_weight, k_weight)
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
return run_benchmark(fn)
if __name__ == "__main__":

17
python/sglang/jit_kernel/benchmark/bench_renorm.py
View File

@@ -5,7 +5,7 @@ import torch
import triton
import triton.testing
from sglang.jit_kernel.benchmark.utils import is_in_ci
from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark_no_cudagraph
def torch_top_k_renorm_probs(probs, top_k):
@@ -171,7 +171,7 @@ configs_p = list(itertools.product(batch_size_range, vocab_size_range, p_range))
x_vals=configs_k,
line_arg="provider",
line_vals=["torch", "sglang"],
line_names=["Torch Reference", "SGL Kernel (FlashInfer)"],
line_names=["Torch Reference", "SGL Kernel"],
styles=[("red", "-"), ("green", "-")],
ylabel="us",
plot_name="top-k-renorm-probs-performance",
@@ -195,8 +195,7 @@ def benchmark_top_k_renorm(batch_size, vocab_size, k, provider):
elif provider == "sglang":
fn = lambda: sgl_kernel.top_k_renorm_prob(probs.clone(), top_k_tensor)
ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=[0.5, 0.2, 0.8])
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
return run_benchmark_no_cudagraph(fn)
@triton.testing.perf_report(
@@ -205,7 +204,7 @@ def benchmark_top_k_renorm(batch_size, vocab_size, k, provider):
x_vals=configs_p,
line_arg="provider",
line_vals=["torch", "sglang"],
line_names=["Torch Reference", "SGL Kernel (FlashInfer)"],
line_names=["Torch Reference", "SGL Kernel"],
styles=[("red", "-"), ("blue", "-")],
ylabel="us",
plot_name="top-p-renorm-probs-performance",
@@ -225,8 +224,7 @@ def benchmark_top_p_renorm(batch_size, vocab_size, p, provider):
elif provider == "sglang":
fn = lambda: sgl_kernel.top_p_renorm_prob(probs.clone(), top_p_tensor)
ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=[0.5, 0.2, 0.8])
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
return run_benchmark_no_cudagraph(fn)
@triton.testing.perf_report(
@@ -235,7 +233,7 @@ def benchmark_top_p_renorm(batch_size, vocab_size, p, provider):
x_vals=configs_k,
line_arg="provider",
line_vals=["torch", "sglang"],
line_names=["Torch Reference", "SGL Kernel (FlashInfer)"],
line_names=["Torch Reference", "SGL Kernel"],
styles=[("red", "-"), ("orange", "-")],
ylabel="us",
plot_name="top-k-mask-logits-performance",
@@ -258,8 +256,7 @@ def benchmark_top_k_mask(batch_size, vocab_size, k, provider):
elif provider == "sglang":
fn = lambda: sgl_kernel.top_k_mask_logits(logits.clone(), top_k_tensor)
ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=[0.5, 0.2, 0.8])
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
return run_benchmark_no_cudagraph(fn)
if __name__ == "__main__":

4
python/sglang/jit_kernel/benchmark/bench_rmsnorm.py
View File

@@ -56,7 +56,7 @@ HIDDEN_SIZE_LIST = get_benchmark_range(
ci_range=[512, 2048],
)
LINE_VALS = ["aot", "jit", "fi", "torch"]
LINE_VALS = ["aot", "jit", "flashinfer", "torch"]
LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "FlashInfer", "PyTorch"]
STYLES = [("orange", "-"), ("blue", "--"), ("green", "-."), ("red", ":")]
@@ -84,7 +84,7 @@ def benchmark(hidden_size: int, batch_size: int, provider: str):
FN_MAP = {
"aot": sglang_aot_rmsnorm,
"jit": sglang_jit_rmsnorm,
"fi": flashinfer_rmsnorm,
"flashinfer": flashinfer_rmsnorm,
"torch": torch_impl_rmsnorm,
}
fn = lambda: FN_MAP[provider](input.clone(), weight)

43
python/sglang/jit_kernel/benchmark/bench_rope.py
View File

@@ -64,7 +64,7 @@ def flashinfer_rope(
)
def sglang_rope_v0(
def sglang_pos_enc_rope(
q: torch.Tensor,
k: torch.Tensor,
positions: torch.Tensor,
@@ -83,7 +83,7 @@ def sglang_rope_v0(
)
def sglang_rope_v1(
def sgl_kernel_rope(
q: torch.Tensor,
k: torch.Tensor,
positions: torch.Tensor,
@@ -102,7 +102,7 @@ def sglang_rope_v1(
)
def sglang_rope_v2(
def sglang_fused_rope(
q: torch.Tensor,
k: torch.Tensor,
positions: torch.Tensor,
@@ -118,7 +118,7 @@ def sglang_rope_v2(
# ---------------------------------------------------------------------------
def rope_v0_store(
def jit_rope_then_store(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
@@ -150,7 +150,7 @@ def rope_v0_store(
)
def rope_v1_store(
def sgl_kernel_fused_rope_store(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
@@ -181,7 +181,7 @@ def rope_v1_store(
)
def rope_v2_store(
def jit_fused_rope_store(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
@@ -221,8 +221,13 @@ IS_NEOX_RANGE = get_benchmark_range(
# Benchmark 1: RoPE only
# ---------------------------------------------------------------------------
ROPE_LINE_VALS = ["fi", "rope_v0", "rope_v1", "rope_v2"]
ROPE_LINE_NAMES = ["FlashInfer", "SGL RoPE v0", "SGL RoPE v1", "SGL RoPE v2"]
ROPE_LINE_VALS = ["flashinfer", "jit_pos_enc", "sgl_kernel", "jit_fused_rope"]
ROPE_LINE_NAMES = [
"FlashInfer",
"SGL JIT PosEnc",
"sgl-kernel",
"SGL JIT Fused RoPE",
]
ROPE_STYLES = [("green", "-."), ("red", "-"), ("orange", "-"), ("blue", "--")]
rope_configs = list(itertools.product(QK_HEAD_RANGE, IS_NEOX_RANGE, BS_RANGE))
@@ -263,10 +268,10 @@ def benchmark(batch_size: int, num_q_k_heads: str, is_neox: bool, provider: str)
torch.cuda.synchronize()
FN_MAP = {
"fi": flashinfer_rope,
"rope_v0": sglang_rope_v0,
"rope_v1": sglang_rope_v1,
"rope_v2": sglang_rope_v2,
"flashinfer": flashinfer_rope,
"jit_pos_enc": sglang_pos_enc_rope,
"sgl_kernel": sgl_kernel_rope,
"jit_fused_rope": sglang_fused_rope,
}
fn = lambda: FN_MAP[provider](q, k, positions, is_neox)
return run_benchmark(fn)
@@ -276,8 +281,12 @@ def benchmark(batch_size: int, num_q_k_heads: str, is_neox: bool, provider: str)
# Benchmark 2: RoPE + KV cache store
# ---------------------------------------------------------------------------
STORE_LINE_VALS = ["rope_v0_store", "rope_v1_store", "rope_v2_store"]
STORE_LINE_NAMES = ["SGL RoPE v0 + Store", "SGL RoPE v1 + Store", "SGL RoPE v2 + Store"]
STORE_LINE_VALS = ["jit_rope_then_store", "sgl_kernel_fused_store", "jit_fused_store"]
STORE_LINE_NAMES = [
"SGL JIT RoPE + Store",
"sgl-kernel Fused RoPE + Store",
"SGL JIT Fused RoPE + Store",
]
STORE_STYLES = [("red", "-"), ("orange", "-"), ("blue", "--")]
store_configs = list(itertools.product(QK_HEAD_RANGE, IS_NEOX_RANGE, BS_RANGE))
@@ -333,9 +342,9 @@ def benchmark_store(batch_size: int, num_q_k_heads: str, is_neox: bool, provider
torch.cuda.synchronize()
FN_MAP = {
"rope_v0_store": rope_v0_store,
"rope_v1_store": rope_v1_store,
"rope_v2_store": rope_v2_store,
"jit_rope_then_store": jit_rope_then_store,
"sgl_kernel_fused_store": sgl_kernel_fused_rope_store,
"jit_fused_store": jit_fused_rope_store,
}
fn = lambda: FN_MAP[provider](
q, k, v, k_cache, v_cache, positions, out_loc, is_neox

8
python/sglang/jit_kernel/benchmark/utils.py
View File

@@ -40,3 +40,11 @@ def run_benchmark(
quantiles = 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
def run_benchmark_no_cudagraph(
fn: Callable, quantiles: List[float] = None
) -> Tuple[float, float, float]:
quantiles = 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