diff --git a/python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py b/python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py index 09b6ccb3f..40c73f5df 100644 --- a/python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py +++ b/python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py @@ -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__": diff --git a/python/sglang/jit_kernel/benchmark/bench_awq_marlin_moe_repack.py b/python/sglang/jit_kernel/benchmark/bench_awq_marlin_moe_repack.py index 120c177d5..f6cd8561e 100644 --- a/python/sglang/jit_kernel/benchmark/bench_awq_marlin_moe_repack.py +++ b/python/sglang/jit_kernel/benchmark/bench_awq_marlin_moe_repack.py @@ -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__": diff --git a/python/sglang/jit_kernel/benchmark/bench_awq_marlin_repack.py b/python/sglang/jit_kernel/benchmark/bench_awq_marlin_repack.py index 51403363d..b2f1d9e76 100644 --- a/python/sglang/jit_kernel/benchmark/bench_awq_marlin_repack.py +++ b/python/sglang/jit_kernel/benchmark/bench_awq_marlin_repack.py @@ -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__": diff --git a/python/sglang/jit_kernel/benchmark/bench_concat_mla.py b/python/sglang/jit_kernel/benchmark/bench_concat_mla.py index 87c7ae56f..7e5c5c43a 100644 --- a/python/sglang/jit_kernel/benchmark/bench_concat_mla.py +++ b/python/sglang/jit_kernel/benchmark/bench_concat_mla.py @@ -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__": diff --git a/python/sglang/jit_kernel/benchmark/bench_fused_add_rmsnorm.py b/python/sglang/jit_kernel/benchmark/bench_fused_add_rmsnorm.py index e4c0a7e8e..ddce16b92 100644 --- a/python/sglang/jit_kernel/benchmark/bench_fused_add_rmsnorm.py +++ b/python/sglang/jit_kernel/benchmark/bench_fused_add_rmsnorm.py @@ -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__": diff --git a/python/sglang/jit_kernel/benchmark/bench_fused_norm_scale_shift.py b/python/sglang/jit_kernel/benchmark/bench_fused_norm_scale_shift.py index e1c7d72af..f6bf7f57f 100644 --- a/python/sglang/jit_kernel/benchmark/bench_fused_norm_scale_shift.py +++ b/python/sglang/jit_kernel/benchmark/bench_fused_norm_scale_shift.py @@ -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__": diff --git a/python/sglang/jit_kernel/benchmark/bench_gptq_marlin.py b/python/sglang/jit_kernel/benchmark/bench_gptq_marlin.py index a4c46ff03..77c1d61af 100644 --- a/python/sglang/jit_kernel/benchmark/bench_gptq_marlin.py +++ b/python/sglang/jit_kernel/benchmark/bench_gptq_marlin.py @@ -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__": diff --git a/python/sglang/jit_kernel/benchmark/bench_gptq_marlin_repack.py b/python/sglang/jit_kernel/benchmark/bench_gptq_marlin_repack.py index 5ec65efec..1253036f6 100644 --- a/python/sglang/jit_kernel/benchmark/bench_gptq_marlin_repack.py +++ b/python/sglang/jit_kernel/benchmark/bench_gptq_marlin_repack.py @@ -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__": diff --git a/python/sglang/jit_kernel/benchmark/bench_hicache.py b/python/sglang/jit_kernel/benchmark/bench_hicache.py index 8929ce706..c825a7f85 100644 --- a/python/sglang/jit_kernel/benchmark/bench_hicache.py +++ b/python/sglang/jit_kernel/benchmark/bench_hicache.py @@ -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 diff --git a/python/sglang/jit_kernel/benchmark/bench_moe_wna16_marlin.py b/python/sglang/jit_kernel/benchmark/bench_moe_wna16_marlin.py index ab7e1ec9c..5a0f480c8 100644 --- a/python/sglang/jit_kernel/benchmark/bench_moe_wna16_marlin.py +++ b/python/sglang/jit_kernel/benchmark/bench_moe_wna16_marlin.py @@ -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__": diff --git a/python/sglang/jit_kernel/benchmark/bench_norm.py b/python/sglang/jit_kernel/benchmark/bench_norm.py index 0e0e80a6f..2bccd672d 100644 --- a/python/sglang/jit_kernel/benchmark/bench_norm.py +++ b/python/sglang/jit_kernel/benchmark/bench_norm.py @@ -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__": diff --git a/python/sglang/jit_kernel/benchmark/bench_qknorm.py b/python/sglang/jit_kernel/benchmark/bench_qknorm.py index 76e217f09..d38fe0b53 100644 --- a/python/sglang/jit_kernel/benchmark/bench_qknorm.py +++ b/python/sglang/jit_kernel/benchmark/bench_qknorm.py @@ -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) diff --git a/python/sglang/jit_kernel/benchmark/bench_qknorm_across_heads.py b/python/sglang/jit_kernel/benchmark/bench_qknorm_across_heads.py index 64d6bd921..ab660f489 100644 --- a/python/sglang/jit_kernel/benchmark/bench_qknorm_across_heads.py +++ b/python/sglang/jit_kernel/benchmark/bench_qknorm_across_heads.py @@ -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__": diff --git a/python/sglang/jit_kernel/benchmark/bench_renorm.py b/python/sglang/jit_kernel/benchmark/bench_renorm.py index d8b67fb83..9c654ab7f 100644 --- a/python/sglang/jit_kernel/benchmark/bench_renorm.py +++ b/python/sglang/jit_kernel/benchmark/bench_renorm.py @@ -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__": diff --git a/python/sglang/jit_kernel/benchmark/bench_rmsnorm.py b/python/sglang/jit_kernel/benchmark/bench_rmsnorm.py index 74b929da6..b55a2c06b 100644 --- a/python/sglang/jit_kernel/benchmark/bench_rmsnorm.py +++ b/python/sglang/jit_kernel/benchmark/bench_rmsnorm.py @@ -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) diff --git a/python/sglang/jit_kernel/benchmark/bench_rope.py b/python/sglang/jit_kernel/benchmark/bench_rope.py index 8d5cdae27..1e8d324e3 100644 --- a/python/sglang/jit_kernel/benchmark/bench_rope.py +++ b/python/sglang/jit_kernel/benchmark/bench_rope.py @@ -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 diff --git a/python/sglang/jit_kernel/benchmark/utils.py b/python/sglang/jit_kernel/benchmark/utils.py index a4d3f2773..36af0e334 100644 --- a/python/sglang/jit_kernel/benchmark/utils.py +++ b/python/sglang/jit_kernel/benchmark/utils.py @@ -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