add hicache jit test (#17847)

Signed-off-by: Xuchun Shang <xuchun.shang@linux.alibaba.com>

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)

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:
+# Benchmark configuration
-    element_range = [16384]
+element_range = get_benchmark_range(
-else:
+    full_range=[2**n for n in range(10, 20)],
-    element_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 run_benchmark(fn)
-
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
 if __name__ == "__main__":

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 = get_benchmark_range(
-    BS_RANGE = [16]
+    full_range=[2**n for n in range(0, 14)],
-    GQA_RANGE = [4]
+    ci_range=[16],
-    KV_HEAD_RANGE = [1]
+)
-else:
+GQA_RANGE = get_benchmark_range(
-    BS_RANGE = [2**n for n in range(0, 14)]
+    full_range=[4, 8],
-    GQA_RANGE = [4, 8]
+    ci_range=[4],
-    KV_HEAD_RANGE = [1, 2, 4, 8]
+)
+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(
+def benchmark(batch_size: int, GQA: int, num_kv_heads: int, provider: str):
-    batch_size: int, GQA: int, num_kv_heads: int, provider: str
-) -> Tuple[float, float, float]:
     num_qo_heads = GQA * num_kv_heads
-    q = torch.randn((batch_size, num_qo_heads, HEAD_DIM), dtype=DTYPE, device=DEVICE)
+    q = torch.randn(
-    k = torch.randn((batch_size, num_kv_heads, HEAD_DIM), dtype=DTYPE, device=DEVICE)
+        (batch_size, num_qo_heads, HEAD_DIM), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE
-    q_weight = torch.randn(HEAD_DIM, dtype=DTYPE, device=DEVICE)
+    )
-    k_weight = torch.randn(HEAD_DIM, dtype=DTYPE, device=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]
+    return run_benchmark(fn)
-    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)  # type: ignore
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
 if __name__ == "__main__":

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
+BS_LIST = get_benchmark_range(
-DEVICE = "cuda"
+    full_range=[2**n for n in range(0, 14)],
-
+    ci_range=[16],
-if IS_CI:
+)
-    BS_LIST = [16]
+HIDDEN_SIZE_LIST = get_benchmark_range(
-    HIDDEN_SIZE_LIST = [512, 2048]
+    full_range=[1536, 3072, 4096, 5120, 8192],
-else:
+    ci_range=[512, 2048],
-    BS_LIST = [2**n for n in range(0, 14)]
+)
-    HIDDEN_SIZE_LIST = [1536, 3072, 4096, 5120, 8192]
 
 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)
+    input = torch.randn(
-    weight = torch.randn(hidden_size, dtype=DTYPE, device=DEVICE)
+        (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]
+    return run_benchmark(fn)
-    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)  # type: ignore
-    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
 
 
 if __name__ == "__main__":

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 = get_benchmark_range(
-    BS_RANGE = [16]
+    full_range=[2**n for n in range(0, 15)],
-    ITEM_SIZE = [1024]
+    ci_range=[16],
-else:
+)
-    BS_RANGE = [2**n for n in range(0, 15)]
+ITEM_SIZE = get_benchmark_range(
-    ITEM_SIZE = [64, 128, 256, 512, 1024]
+    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)
+    k = torch.randn(
-    v = torch.randn((NUM_LAYERS, batch_size, item_size), dtype=DTYPE, device=DEVICE)
+        (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]
+    # Custom time calculation: divide by NUM_LAYERS
-    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)  # type: ignore
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+        fn, quantiles=DEFAULT_QUANTILES
+    )
     return (
         1000 * ms / NUM_LAYERS,
         1000 * max_ms / NUM_LAYERS,

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