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

[JIT kernel] hd=512,1024 in JIT QK norm (cta based) (#17515)

Browse Source 
Signed-off-by: vincentzed <207368749+vincentzed@users.noreply.github.com>
This commit is contained in:
Yi Zhong
2026-02-16 03:07:24 -05:00
committed by GitHub
parent 206accd15d
commit ed22720c07
3 changed files with 134 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -74,8 +74,6 @@ def torch_impl_qknorm(
k.copy_(k.float() * k_norm * k_weight.float())
HEAD_DIM = 128
BS_RANGE = get_benchmark_range(
full_range=[2**n for n in range(0, 14)],
ci_range=[16],
@@ -88,17 +86,21 @@ KV_HEAD_RANGE = get_benchmark_range(
full_range=[1, 2, 4, 8],
ci_range=[1],
)
HEAD_DIM_RANGE = get_benchmark_range(
full_range=[128, 256, 512, 1024],
ci_range=[128],
)
LINE_VALS = ["aot", "jit", "fi", "torch"]
LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "FlashInfer", "PyTorch"]
STYLES = [("orange", "-"), ("blue", "--"), ("green", "-."), ("red", ":")]
configs = list(itertools.product(GQA_RANGE, KV_HEAD_RANGE, BS_RANGE))
configs = list(itertools.product(HEAD_DIM_RANGE, GQA_RANGE, KV_HEAD_RANGE, BS_RANGE))
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["GQA", "num_kv_heads", "batch_size"],
x_names=["head_dim", "GQA", "num_kv_heads", "batch_size"],
x_vals=configs,
line_arg="provider",
line_vals=LINE_VALS,
@@ -109,16 +111,18 @@ 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):
def benchmark(
head_dim: int, GQA: int, num_kv_heads: int, batch_size: int, provider: str
):
num_qo_heads = GQA * num_kv_heads
q = torch.randn(
(batch_size, num_qo_heads, HEAD_DIM), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE
(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
(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)
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,