From 3d68bd9d9b6ed994a2de1f7d9d2bd2956f8f1ac9 Mon Sep 17 00:00:00 2001 From: Xuchun Shang Date: Fri, 6 Feb 2026 16:54:33 +0800 Subject: [PATCH] add hicache jit test (#17847) Signed-off-by: Xuchun Shang --- .../jit_kernel/benchmark/bench_hicache.py | 406 ++++++++++++++++++ .../benchmark/bench_per_tensor_quant_fp8.py | 20 +- .../jit_kernel/benchmark/bench_qknorm.py | 52 +-- .../jit_kernel/benchmark/bench_rmsnorm.py | 36 +- .../jit_kernel/benchmark/bench_store_cache.py | 45 +- python/sglang/jit_kernel/benchmark/utils.py | 36 +- 6 files changed, 522 insertions(+), 73 deletions(-) create mode 100644 python/sglang/jit_kernel/benchmark/bench_hicache.py diff --git a/python/sglang/jit_kernel/benchmark/bench_hicache.py b/python/sglang/jit_kernel/benchmark/bench_hicache.py new file mode 100644 index 000000000..0218b562a --- /dev/null +++ b/python/sglang/jit_kernel/benchmark/bench_hicache.py @@ -0,0 +1,406 @@ +"""Benchmark for HiCache JIT kernel performance. + +This benchmark tests the performance of KV cache transfer operations +between GPU and CPU (host pinned memory), comparing: +- SGL AOT Kernel: Pre-compiled transfer_kv kernels from sgl_kernel +- SGL JIT Kernel: JIT-compiled hicache kernels +- PyTorch Indexing: Plain PyTorch index copy +- PyTorch 2 Stream: PyTorch implementation using 2 CUDA streams + +Tests cover: +- One Layer: CPU->GPU +- All Layer: GPU->CPU + +Note: Uses do_bench instead of do_bench_cudagraph since CUDA graph +capture doesn't support CPU-GPU memory transfers. +""" + +import itertools +from typing import Tuple + +import torch +import triton +import triton.testing +from sgl_kernel import transfer_kv_all_layer, transfer_kv_per_layer + +from sglang.jit_kernel.benchmark.utils import ( + DEFAULT_DTYPE, + DEFAULT_QUANTILES, + get_benchmark_range, +) +from sglang.jit_kernel.hicache import ( + can_use_hicache_jit_kernel, + transfer_hicache_all_layer, + transfer_hicache_one_layer, +) + + +def sglang_aot_transfer_one( + k_cache_dst: torch.Tensor, + v_cache_dst: torch.Tensor, + indices_dst: torch.Tensor, + k_cache_src: torch.Tensor, + v_cache_src: torch.Tensor, + indices_src: torch.Tensor, + item_size: int, +) -> None: + """SGL AOT Kernel for single layer transfer.""" + transfer_kv_per_layer( + k_cache_src, + k_cache_dst, + v_cache_src, + v_cache_dst, + indices_src, + indices_dst, + item_size, + ) + + +def sglang_jit_transfer_one( + k_cache_dst: torch.Tensor, + v_cache_dst: torch.Tensor, + indices_dst: torch.Tensor, + k_cache_src: torch.Tensor, + v_cache_src: torch.Tensor, + indices_src: torch.Tensor, + element_dim: int, +) -> None: + """SGL JIT Kernel for single layer transfer.""" + transfer_hicache_one_layer( + k_cache_dst, + v_cache_dst, + indices_dst, + k_cache_src, + v_cache_src, + indices_src, + element_dim=element_dim, + ) + + +def sglang_aot_transfer_all( + k_ptrs_dst: torch.Tensor, + v_ptrs_dst: torch.Tensor, + indices_dst: torch.Tensor, + k_ptrs_src: torch.Tensor, + v_ptrs_src: torch.Tensor, + indices_src: torch.Tensor, + item_size: int, + num_layers: int, +) -> None: + """SGL AOT Kernel for all layer transfer.""" + transfer_kv_all_layer( + k_ptrs_src, + k_ptrs_dst, + v_ptrs_src, + v_ptrs_dst, + indices_src, + indices_dst, + item_size, + num_layers, + ) + + +def sglang_jit_transfer_all( + k_ptrs_dst: torch.Tensor, + v_ptrs_dst: torch.Tensor, + indices_dst: torch.Tensor, + k_ptrs_src: torch.Tensor, + v_ptrs_src: torch.Tensor, + indices_src: torch.Tensor, + stride_bytes: int, + element_size: int, +) -> None: + """SGL JIT Kernel for all layer transfer.""" + transfer_hicache_all_layer( + k_ptrs_dst, + v_ptrs_dst, + indices_dst, + k_ptrs_src, + v_ptrs_src, + indices_src, + kv_cache_src_stride_bytes=stride_bytes, + kv_cache_dst_stride_bytes=stride_bytes, + element_size=element_size, + ) + + +def pytorch_transfer( + k_cache_dst: torch.Tensor, + v_cache_dst: torch.Tensor, + indices_dst_on_dst: torch.Tensor, + k_cache_src: torch.Tensor, + v_cache_src: torch.Tensor, + indices_src_on_src: torch.Tensor, +) -> None: + """PyTorch indexing baseline.""" + dst_device = k_cache_dst.device + k_cache_dst[indices_dst_on_dst] = k_cache_src[indices_src_on_src].to(dst_device) + v_cache_dst[indices_dst_on_dst] = v_cache_src[indices_src_on_src].to(dst_device) + + +alt_stream = torch.cuda.Stream() + + +def torch_streams_transfer( + k_cache_dst: torch.Tensor, + v_cache_dst: torch.Tensor, + indices_dst_on_dst: torch.Tensor, + k_cache_src: torch.Tensor, + v_cache_src: torch.Tensor, + indices_src_on_src: torch.Tensor, +) -> None: + """PyTorch 2 Stream baseline.""" + dst_device = k_cache_dst.device + current_stream = torch.cuda.current_stream() + alt_stream.wait_stream(current_stream) + k_cache_dst[indices_dst_on_dst] = k_cache_src[indices_src_on_src].to(dst_device) + with torch.cuda.stream(alt_stream): + v_cache_dst[indices_dst_on_dst] = v_cache_src[indices_src_on_src].to(dst_device) + current_stream.wait_stream(alt_stream) + + +# Benchmark configuration +GPU_CACHE_SIZE = 32 * 1024 # 32K tokens on GPU +HOST_CACHE_SIZE = 128 * 1024 # 128K tokens on CPU +NUM_LAYERS = 8 + +BS_RANGE = get_benchmark_range( + full_range=[2**n for n in range(0, 15)], + ci_range=[16], +) +ELEMENT_SIZE_RANGE = get_benchmark_range( + full_range=[64, 128, 256, 512, 1024], + ci_range=[1024], +) + +LINE_VALS = ["aot", "jit", "pytorch", "torch_streams"] +LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "PyTorch", "PyTorch 2 Stream"] +STYLES = [("orange", "-"), ("blue", "--"), ("red", ":"), ("green", "-.")] + +CONFIGS = list(itertools.product(ELEMENT_SIZE_RANGE, BS_RANGE)) + + +# ============================================================================= +# One Layer Benchmarks +# ============================================================================= + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["element_size", "batch_size"], + x_vals=CONFIGS, + line_arg="provider", + line_vals=LINE_VALS, + line_names=LINE_NAMES, + styles=STYLES, + ylabel="us", + plot_name="hicache-one-layer-h2d", + args={}, + ) +) +def benchmark_one_layer_h2d( + element_size: int, batch_size: int, provider: str +) -> Tuple[float, float, float]: + """One Layer: Host (CPU) -> Device (GPU).""" + k_cache_src = torch.randn( + (HOST_CACHE_SIZE, element_size), + dtype=DEFAULT_DTYPE, + device="cpu", + pin_memory=True, + ) + v_cache_src = torch.randn( + (HOST_CACHE_SIZE, element_size), + dtype=DEFAULT_DTYPE, + device="cpu", + pin_memory=True, + ) + k_cache_dst = torch.randn( + (GPU_CACHE_SIZE, element_size), dtype=DEFAULT_DTYPE, device="cuda" + ) + v_cache_dst = torch.randn( + (GPU_CACHE_SIZE, element_size), dtype=DEFAULT_DTYPE, device="cuda" + ) + + indices_src_gpu = torch.randperm(HOST_CACHE_SIZE, device="cuda")[:batch_size] + indices_dst_gpu = torch.randperm(GPU_CACHE_SIZE, device="cuda")[:batch_size] + indices_src_cpu = indices_src_gpu.cpu() + torch.cuda.synchronize() + + element_bytes = element_size * k_cache_src.element_size() + + FN_MAP = { + "aot": lambda: sglang_aot_transfer_one( + k_cache_dst, + v_cache_dst, + indices_dst_gpu, + k_cache_src, + v_cache_src, + indices_src_gpu, + element_bytes, + ), + "jit": lambda: sglang_jit_transfer_one( + k_cache_dst, + v_cache_dst, + indices_dst_gpu, + k_cache_src, + v_cache_src, + indices_src_gpu, + element_size, + ), + "pytorch": lambda: pytorch_transfer( + k_cache_dst, + v_cache_dst, + indices_dst_gpu, + k_cache_src, + v_cache_src, + indices_src_cpu, + ), + "torch_streams": lambda: torch_streams_transfer( + k_cache_dst, + v_cache_dst, + indices_dst_gpu, + k_cache_src, + v_cache_src, + indices_src_cpu, + ), + } + + if provider == "jit" and not can_use_hicache_jit_kernel(element_size=element_bytes): + return (float("nan"), float("nan"), float("nan")) + + ms, min_ms, max_ms = triton.testing.do_bench( + FN_MAP[provider], quantiles=DEFAULT_QUANTILES + ) + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + +# ============================================================================= +# All Layer Benchmarks +# ============================================================================= + + +def _create_ptr_tensor(tensors, device="cuda"): + """Create a tensor of data pointers.""" + return torch.tensor( + [t.data_ptr() for t in tensors], + dtype=torch.uint64, + device=device, + ) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["element_size", "batch_size"], + x_vals=CONFIGS, + line_arg="provider", + line_vals=LINE_VALS, + line_names=LINE_NAMES, + styles=STYLES, + ylabel="us", + plot_name="hicache-all-layer-d2h", + args={}, + ) +) +def benchmark_all_layer_d2h( + element_size: int, batch_size: int, provider: str +) -> Tuple[float, float, float]: + """All Layer: Device (GPU) -> Host (CPU).""" + k_caches_src = torch.randn( + (NUM_LAYERS, GPU_CACHE_SIZE, element_size), dtype=DEFAULT_DTYPE, device="cuda" + ) + v_caches_src = torch.randn( + (NUM_LAYERS, GPU_CACHE_SIZE, element_size), dtype=DEFAULT_DTYPE, device="cuda" + ) + k_caches_dst = torch.randn( + (NUM_LAYERS, HOST_CACHE_SIZE, element_size), + dtype=DEFAULT_DTYPE, + device="cpu", + pin_memory=True, + ) + v_caches_dst = torch.randn( + (NUM_LAYERS, HOST_CACHE_SIZE, element_size), + dtype=DEFAULT_DTYPE, + device="cpu", + pin_memory=True, + ) + + indices_src_gpu = torch.randperm(GPU_CACHE_SIZE, device="cuda")[:batch_size] + indices_dst_gpu = torch.randperm(HOST_CACHE_SIZE, device="cuda")[:batch_size] + indices_dst_cpu = indices_dst_gpu.cpu() + torch.cuda.synchronize() + + element_bytes = element_size * k_caches_src.element_size() + + k_ptrs_src = _create_ptr_tensor([k_caches_src[i] for i in range(NUM_LAYERS)]) + v_ptrs_src = _create_ptr_tensor([v_caches_src[i] for i in range(NUM_LAYERS)]) + k_ptrs_dst = _create_ptr_tensor([k_caches_dst[i] for i in range(NUM_LAYERS)]) + v_ptrs_dst = _create_ptr_tensor([v_caches_dst[i] for i in range(NUM_LAYERS)]) + + FN_MAP = { + "aot": lambda: sglang_aot_transfer_all( + k_ptrs_dst, + v_ptrs_dst, + indices_dst_gpu, + k_ptrs_src, + v_ptrs_src, + indices_src_gpu, + element_bytes, + NUM_LAYERS, + ), + "jit": lambda: sglang_jit_transfer_all( + k_ptrs_dst, + v_ptrs_dst, + indices_dst_gpu, + k_ptrs_src, + v_ptrs_src, + indices_src_gpu, + element_bytes, + element_bytes, + ), + "pytorch": lambda: [ + pytorch_transfer( + k_caches_dst[i], + v_caches_dst[i], + indices_dst_cpu, + k_caches_src[i], + v_caches_src[i], + indices_src_gpu, + ) + for i in range(NUM_LAYERS) + ], + "torch_streams": lambda: [ + torch_streams_transfer( + k_caches_dst[i], + v_caches_dst[i], + indices_dst_cpu, + k_caches_src[i], + v_caches_src[i], + indices_src_gpu, + ) + for i in range(NUM_LAYERS) + ], + } + + if provider == "jit" and not can_use_hicache_jit_kernel(element_size=element_bytes): + return (float("nan"), float("nan"), float("nan")) + + ms, min_ms, max_ms = triton.testing.do_bench( + FN_MAP[provider], quantiles=DEFAULT_QUANTILES + ) + return ( + 1000 * ms / NUM_LAYERS, + 1000 * max_ms / NUM_LAYERS, + 1000 * min_ms / NUM_LAYERS, + ) + + +if __name__ == "__main__": + print("=" * 60) + print("One Layer: Host -> Device (CPU -> GPU)") + print("=" * 60) + benchmark_one_layer_h2d.run(print_data=True) + + print("\n" + "=" * 60) + print("All Layer: Device -> Host (GPU -> CPU) [per-layer avg]") + print("=" * 60) + benchmark_all_layer_d2h.run(print_data=True) diff --git a/python/sglang/jit_kernel/benchmark/bench_per_tensor_quant_fp8.py b/python/sglang/jit_kernel/benchmark/bench_per_tensor_quant_fp8.py index 1fb0e45cb..9549526d5 100644 --- a/python/sglang/jit_kernel/benchmark/bench_per_tensor_quant_fp8.py +++ b/python/sglang/jit_kernel/benchmark/bench_per_tensor_quant_fp8.py @@ -4,7 +4,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 get_benchmark_range, run_benchmark from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8 try: @@ -22,8 +22,6 @@ try: except ImportError: _is_hip = False -IS_CI = is_in_ci() - fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn @@ -69,11 +67,11 @@ def calculate_diff(batch_size: int, seq_len: int): triton.testing.assert_close(vllm_scale, sglang_scale, rtol=1e-3, atol=1e-3) -if IS_CI: - element_range = [16384] -else: - element_range = [2**n for n in range(10, 20)] - +# Benchmark configuration +element_range = get_benchmark_range( + full_range=[2**n for n in range(10, 20)], + ci_range=[16384], +) if VLLM_AVAILABLE: line_vals = ["vllm", "sglang"] @@ -104,8 +102,6 @@ def benchmark(element_count, provider): x = torch.randn(element_count, 4096, device=device, dtype=dtype) - quantiles = [0.5, 0.2, 0.8] - if provider == "vllm": fn = lambda: vllm_scaled_fp8_quant(x.clone()) elif provider == "sglang": @@ -113,9 +109,7 @@ def benchmark(element_count, 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_qknorm.py b/python/sglang/jit_kernel/benchmark/bench_qknorm.py index 7f5821e92..70f61a4ca 100644 --- a/python/sglang/jit_kernel/benchmark/bench_qknorm.py +++ b/python/sglang/jit_kernel/benchmark/bench_qknorm.py @@ -1,17 +1,19 @@ import itertools -from typing import Tuple import torch 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 ( + DEFAULT_DEVICE, + DEFAULT_DTYPE, + get_benchmark_range, + run_benchmark, +) from sglang.jit_kernel.norm import fused_inplace_qknorm from sglang.srt.utils import get_current_device_stream_fast -IS_CI = is_in_ci() - alt_stream = torch.cuda.Stream() @@ -73,17 +75,19 @@ def torch_impl_qknorm( HEAD_DIM = 128 -DTYPE = torch.bfloat16 -DEVICE = "cuda" -if IS_CI: - BS_RANGE = [16] - GQA_RANGE = [4] - KV_HEAD_RANGE = [1] -else: - BS_RANGE = [2**n for n in range(0, 14)] - GQA_RANGE = [4, 8] - KV_HEAD_RANGE = [1, 2, 4, 8] +BS_RANGE = get_benchmark_range( + full_range=[2**n for n in range(0, 14)], + ci_range=[16], +) +GQA_RANGE = get_benchmark_range( + full_range=[4, 8], + ci_range=[4], +) +KV_HEAD_RANGE = get_benchmark_range( + full_range=[1, 2, 4, 8], + ci_range=[1], +) LINE_VALS = ["aot", "jit", "fi", "torch"] LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "FlashInfer", "PyTorch"] @@ -105,14 +109,16 @@ configs = list(itertools.product(GQA_RANGE, KV_HEAD_RANGE, BS_RANGE)) args={}, ) ) -def benchmark( - batch_size: int, GQA: int, num_kv_heads: int, provider: str -) -> Tuple[float, float, float]: +def benchmark(batch_size: int, GQA: int, num_kv_heads: int, provider: str): num_qo_heads = GQA * num_kv_heads - q = torch.randn((batch_size, num_qo_heads, HEAD_DIM), dtype=DTYPE, device=DEVICE) - k = torch.randn((batch_size, num_kv_heads, HEAD_DIM), dtype=DTYPE, device=DEVICE) - q_weight = torch.randn(HEAD_DIM, dtype=DTYPE, device=DEVICE) - k_weight = torch.randn(HEAD_DIM, dtype=DTYPE, device=DEVICE) + q = torch.randn( + (batch_size, num_qo_heads, HEAD_DIM), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE + ) + k = torch.randn( + (batch_size, num_kv_heads, HEAD_DIM), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE + ) + q_weight = torch.randn(HEAD_DIM, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE) + k_weight = torch.randn(HEAD_DIM, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE) FN_MAP = { "aot": sglang_aot_qknorm, "jit": sglang_jit_qknorm, @@ -120,9 +126,7 @@ def benchmark( "torch": torch_impl_qknorm, } 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_rmsnorm.py b/python/sglang/jit_kernel/benchmark/bench_rmsnorm.py index 1863c189f..74b929da6 100644 --- a/python/sglang/jit_kernel/benchmark/bench_rmsnorm.py +++ b/python/sglang/jit_kernel/benchmark/bench_rmsnorm.py @@ -6,11 +6,14 @@ import triton.testing from flashinfer import rmsnorm as fi_rmsnorm from sgl_kernel import rmsnorm -from sglang.jit_kernel.benchmark.utils import is_in_ci +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 -IS_CI = is_in_ci() - def sglang_aot_rmsnorm( input: torch.Tensor, @@ -44,15 +47,14 @@ def torch_impl_rmsnorm( input.copy_(input.float() * norm * weight.float()) -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] +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"] @@ -75,8 +77,10 @@ configs = list(itertools.product(HIDDEN_SIZE_LIST, BS_LIST)) ) ) def benchmark(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) + 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, @@ -84,9 +88,7 @@ def benchmark(hidden_size: int, batch_size: int, provider: str): "torch": torch_impl_rmsnorm, } fn = lambda: FN_MAP[provider](input.clone(), 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_store_cache.py b/python/sglang/jit_kernel/benchmark/bench_store_cache.py index d7d014f2e..2aeb7708d 100644 --- a/python/sglang/jit_kernel/benchmark/bench_store_cache.py +++ b/python/sglang/jit_kernel/benchmark/bench_store_cache.py @@ -6,11 +6,14 @@ import triton import triton.testing from sgl_kernel import set_kv_buffer_kernel -from sglang.jit_kernel.benchmark.utils import is_in_ci +from sglang.jit_kernel.benchmark.utils import ( + DEFAULT_DEVICE, + DEFAULT_DTYPE, + DEFAULT_QUANTILES, + get_benchmark_range, +) from sglang.jit_kernel.kvcache import store_cache -IS_CI = is_in_ci() - def sglang_aot_store_cache( k: torch.Tensor, @@ -62,17 +65,17 @@ def torch_streams_store_cache( current_stream.wait_stream(alt_stream) -DTYPE = torch.bfloat16 -DEVICE = "cuda" NUM_LAYERS = 8 CACHE_SIZE = 2 * 1024 * 1024 // NUM_LAYERS -if IS_CI: - BS_RANGE = [16] - ITEM_SIZE = [1024] -else: - BS_RANGE = [2**n for n in range(0, 15)] - ITEM_SIZE = [64, 128, 256, 512, 1024] +BS_RANGE = get_benchmark_range( + full_range=[2**n for n in range(0, 15)], + ci_range=[16], +) +ITEM_SIZE = get_benchmark_range( + full_range=[64, 128, 256, 512, 1024], + ci_range=[1024], +) LINE_VALS = ["aot", "jit", "torch_compile", "torch_streams"] LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "PyTorch Compile", "PyTorch 2 Stream"] @@ -97,15 +100,19 @@ CONFIGS = list(itertools.product(ITEM_SIZE, BS_RANGE)) def benchmark( batch_size: int, item_size: int, provider: str ) -> Tuple[float, float, float]: - k = torch.randn((NUM_LAYERS, batch_size, item_size), dtype=DTYPE, device=DEVICE) - v = torch.randn((NUM_LAYERS, batch_size, item_size), dtype=DTYPE, device=DEVICE) + k = torch.randn( + (NUM_LAYERS, batch_size, item_size), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE + ) + v = torch.randn( + (NUM_LAYERS, batch_size, item_size), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE + ) k_cache = torch.randn( - (NUM_LAYERS, CACHE_SIZE, item_size), dtype=DTYPE, device=DEVICE + (NUM_LAYERS, CACHE_SIZE, item_size), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE ) v_cache = torch.randn( - (NUM_LAYERS, CACHE_SIZE, item_size), dtype=DTYPE, device=DEVICE + (NUM_LAYERS, CACHE_SIZE, item_size), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE ) - indices = torch.randperm(CACHE_SIZE, device=DEVICE)[:batch_size] + indices = torch.randperm(CACHE_SIZE, device=DEFAULT_DEVICE)[:batch_size] torch.cuda.synchronize() FN_MAP = { @@ -120,8 +127,10 @@ def benchmark( for i in range(NUM_LAYERS): impl(k[i], v[i], k_cache[i], v_cache[i], indices) - quantiles = [0.5, 0.2, 0.8] - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles) # type: ignore + # Custom time calculation: divide by NUM_LAYERS + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + fn, quantiles=DEFAULT_QUANTILES + ) return ( 1000 * ms / NUM_LAYERS, 1000 * max_ms / NUM_LAYERS, diff --git a/python/sglang/jit_kernel/benchmark/utils.py b/python/sglang/jit_kernel/benchmark/utils.py index 5055c700f..a4d3f2773 100644 --- a/python/sglang/jit_kernel/benchmark/utils.py +++ b/python/sglang/jit_kernel/benchmark/utils.py @@ -1,8 +1,42 @@ +"""Common utilities for jit_kernel benchmark files.""" + import os +from typing import Callable, List, Tuple + +import torch +import triton.testing + +# Common constants +DEFAULT_DTYPE = torch.bfloat16 +DEFAULT_DEVICE = "cuda" +DEFAULT_QUANTILES = [0.5, 0.2, 0.8] -def is_in_ci(): +def is_in_ci() -> bool: + """Check if running in CI environment.""" return ( os.getenv("CI", "false").lower() == "true" or os.getenv("GITHUB_ACTIONS", "false").lower() == "true" ) + + +def get_benchmark_range(full_range: List, ci_range: List) -> List: + """Return appropriate benchmark range based on CI environment.""" + return ci_range if is_in_ci() else full_range + + +def run_benchmark( + fn: Callable, quantiles: List[float] = None +) -> 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] + + Returns: + Tuple of (median_us, max_us, min_us) + """ + 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