[3/n jit_kernel restruct] Clean up benchmark naming and benchmarking helpers (#20250)

This commit is contained in:
Xiaoyu Zhang
2026-03-10 16:39:03 +08:00
committed by GitHub
parent c812504b92
commit 8517da5d08
17 changed files with 82 additions and 114 deletions
@@ -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)