107 lines
3.8 KiB
Python
107 lines
3.8 KiB
Python
import torch
|
|
import triton
|
|
import triton.testing
|
|
|
|
from sglang.jit_kernel.benchmark.utils import (
|
|
DEFAULT_DEVICE,
|
|
get_benchmark_range,
|
|
run_benchmark,
|
|
)
|
|
from sglang.jit_kernel.cast import downcast_fp8 as downcast_fp8_jit
|
|
from sglang.test.ci.ci_register import register_cuda_ci
|
|
|
|
register_cuda_ci(est_time=10, suite="stage-b-kernel-benchmark-1-gpu-large")
|
|
|
|
DEVICE = DEFAULT_DEVICE
|
|
DTYPE = torch.bfloat16
|
|
|
|
|
|
# ── Config ranges ──────────────────────────────────────────────────────────────
|
|
|
|
SL_LIST = get_benchmark_range(
|
|
full_range=[4, 16, 64, 256, 512, 1024, 2048],
|
|
ci_range=[4, 64],
|
|
)
|
|
|
|
HEAD_DIM_LIST = get_benchmark_range(
|
|
full_range=[(8, 128), (32, 128), (8, 256), (32, 256)],
|
|
ci_range=[(8, 128)],
|
|
)
|
|
|
|
CONFIGS = [(sl, h, d, sl * 2) for sl in SL_LIST for h, d in HEAD_DIM_LIST]
|
|
|
|
LINE_VALS = ["jit"]
|
|
LINE_NAMES = ["JIT (cast.cuh, 256 threads, 2D grid)"]
|
|
STYLES = [("orange", "-")]
|
|
|
|
|
|
# ── Perf report ────────────────────────────────────────────────────────────────
|
|
|
|
|
|
@triton.testing.perf_report(
|
|
triton.testing.Benchmark(
|
|
x_names=["input_sl", "head", "dim", "out_sl"],
|
|
x_vals=CONFIGS,
|
|
line_arg="provider",
|
|
line_vals=LINE_VALS,
|
|
line_names=LINE_NAMES,
|
|
styles=STYLES,
|
|
ylabel="us",
|
|
plot_name="downcast-fp8-jit",
|
|
args={},
|
|
)
|
|
)
|
|
def benchmark(input_sl, head, dim, out_sl, provider):
|
|
k = torch.randn(input_sl, head, dim, dtype=DTYPE, device=DEVICE)
|
|
v = torch.randn(input_sl, head, dim, dtype=DTYPE, device=DEVICE)
|
|
k_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device=DEVICE)
|
|
v_out = torch.zeros(out_sl, head, dim, dtype=torch.uint8, device=DEVICE)
|
|
k_scale = torch.tensor([1.0], dtype=torch.float32, device=DEVICE)
|
|
v_scale = torch.tensor([1.0], dtype=torch.float32, device=DEVICE)
|
|
loc = torch.arange(input_sl, dtype=torch.int64, device=DEVICE)
|
|
|
|
fn = lambda: downcast_fp8_jit(k, v, k_out, v_out, k_scale, v_scale, loc)
|
|
|
|
return run_benchmark(fn)
|
|
|
|
|
|
# ── Bandwidth analysis ─────────────────────────────────────────────────────────
|
|
|
|
|
|
def _report_bandwidth(input_sl, head, dim, dtype):
|
|
elem_bytes = torch.finfo(dtype).bits // 8
|
|
total_bytes = input_sl * head * dim * (2 * elem_bytes + 2)
|
|
|
|
k = torch.randn(input_sl, head, dim, dtype=dtype, device=DEVICE)
|
|
v = torch.randn(input_sl, head, dim, dtype=dtype, device=DEVICE)
|
|
k_out = torch.zeros(input_sl * 2, head, dim, dtype=torch.uint8, device=DEVICE)
|
|
v_out = torch.zeros(input_sl * 2, head, dim, dtype=torch.uint8, device=DEVICE)
|
|
k_scale = torch.tensor([1.0], dtype=torch.float32, device=DEVICE)
|
|
v_scale = torch.tensor([1.0], dtype=torch.float32, device=DEVICE)
|
|
loc = torch.arange(input_sl, dtype=torch.int64, device=DEVICE)
|
|
|
|
jit_fn = lambda: downcast_fp8_jit(k, v, k_out, v_out, k_scale, v_scale, loc)
|
|
|
|
jit_ms, _, _ = triton.testing.do_bench(jit_fn, quantiles=[0.5, 0.2, 0.8])
|
|
|
|
def fmt(ms):
|
|
return f"{ms*1000:6.2f}us {total_bytes/(ms*1e-3)/1e9:6.0f}GB/s"
|
|
|
|
print(f" sl={input_sl:5d} h={head:2d} d={dim:4d}" f" | jit {fmt(jit_ms)}")
|
|
|
|
|
|
def report_bandwidth():
|
|
print(f"\n{'='*95}")
|
|
print(" JIT (cast.cuh, 256 threads, 2D grid)")
|
|
print(f" dtype={DTYPE}, device={DEVICE}")
|
|
print(f"{'='*95}")
|
|
for sl in [64, 256, 1024, 2048]:
|
|
for h, d in [(8, 128), (32, 128), (8, 256), (32, 256)]:
|
|
_report_bandwidth(sl, h, d, DTYPE)
|
|
print()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
benchmark.run(print_data=True)
|
|
report_bandwidth()
|