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

add hicache jit test (#17847)

Browse Source 
Signed-off-by: Xuchun Shang <xuchun.shang@linux.alibaba.com>
This commit is contained in:
Xuchun Shang
2026-02-06 16:54:33 +08:00
committed by GitHub
parent f798ab9775
commit 3d68bd9d9b
6 changed files with 522 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files

406
python/sglang/jit_kernel/benchmark/bench_hicache.py Normal file
View File

@@ -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)

20
python/sglang/jit_kernel/benchmark/bench_per_tensor_quant_fp8.py
View File

@@ -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__":

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

@@ -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__":

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

@@ -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__":

45
python/sglang/jit_kernel/benchmark/bench_store_cache.py
View File

@@ -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,

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

@@ -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