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

[CI] Remove obsolete AOT-only jit-kernel benchmarks after sgl-kernel 4.0 (#21075)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-03-21 13:40:42 +08:00
committed by GitHub
parent 5f3393c04c
commit c076968c52
5 changed files with 0 additions and 701 deletions
Download Patch File Download Diff File Expand all files Collapse all files

125
python/sglang/jit_kernel/benchmark/bench_awq_marlin_moe_repack.py
View File

@@ -1,125 +0,0 @@
import numpy as np
import torch
import triton
import triton.testing
from sgl_kernel.scalar_type import scalar_types
from sglang.jit_kernel.awq_marlin_repack import (
awq_marlin_moe_repack as jit_awq_marlin_moe_repack,
)
from sglang.jit_kernel.benchmark.utils import run_benchmark
from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
from sglang.utils import is_in_ci
AOT_AVAILABLE = hasattr(torch.ops.sgl_kernel, "awq_marlin_moe_repack") and hasattr(
torch.ops.sgl_kernel.awq_marlin_moe_repack, "default"
)
IS_CI = is_in_ci()
NUM_BITS = 4
GROUP_SIZE = 128
SIZE_N = 4096
def awq_pack(q_w, num_bits, size_k, size_n):
if num_bits == 4:
interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7])
elif num_bits == 8:
interleave = np.array([0, 2, 1, 3])
else:
raise Exception("num_bits must be 4 or 8, got {}".format(num_bits))
q_w = q_w.reshape((-1, len(interleave)))[:, interleave].ravel()
q_w = q_w.reshape((-1, size_n)).contiguous()
return pack_cols(q_w, num_bits, size_k, size_n)
def make_moe_weights(num_experts, size_k, size_n, num_bits, group_size):
pack_factor = 32 // num_bits
b_q_weight = torch.empty(
(num_experts, size_k, size_n // pack_factor),
dtype=torch.int32,
device="cuda",
)
for e in range(num_experts):
b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
w_ref, q_w, s, zp = quantize_weights(
b_weight, scalar_types.uint4, min(group_size, size_k), zero_points=True
)
b_q_weight[e] = awq_pack(q_w, num_bits, size_k, size_n)
perm = torch.empty((num_experts, 0), dtype=torch.int32, device="cuda")
return b_q_weight, perm
def check_correctness():
if not AOT_AVAILABLE:
print("sgl_kernel AOT not available, skipping correctness check")
return
num_experts = 4
size_k = 1024
b_q_weight, perm = make_moe_weights(
num_experts, size_k, SIZE_N, NUM_BITS, GROUP_SIZE
)
out_jit = jit_awq_marlin_moe_repack(b_q_weight, perm, size_k, SIZE_N, NUM_BITS)
out_aot = torch.ops.sgl_kernel.awq_marlin_moe_repack.default(
b_q_weight, perm, size_k, SIZE_N, NUM_BITS
)
torch.cuda.synchronize()
torch.testing.assert_close(out_jit, out_aot, rtol=0, atol=0)
print("Correctness check passed (JIT vs AOT)")
if IS_CI:
expert_range = [2, 4]
else:
expert_range = [2, 4, 8, 16]
if AOT_AVAILABLE:
line_vals = ["jit", "aot"]
line_names = ["JIT Kernel", "AOT Kernel"]
styles = [("blue", "-"), ("green", "-")]
else:
line_vals = ["jit"]
line_names = ["JIT Kernel"]
styles = [("blue", "-")]
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["num_experts"],
x_vals=expert_range,
line_arg="provider",
line_vals=line_vals,
line_names=line_names,
styles=styles,
ylabel="us",
plot_name="awq-marlin-moe-repack-performance",
args={"size_k": 4096, "size_n": SIZE_N, "num_bits": NUM_BITS},
)
)
def benchmark(num_experts, size_k, size_n, num_bits, provider):
group_size = min(GROUP_SIZE, size_k)
b_q_weight, perm = make_moe_weights(
num_experts, size_k, size_n, num_bits, group_size
)
if provider == "jit":
fn = lambda: jit_awq_marlin_moe_repack(
b_q_weight, perm, size_k, size_n, num_bits
)
elif provider == "aot":
fn = lambda: torch.ops.sgl_kernel.awq_marlin_moe_repack.default(
b_q_weight, perm, size_k, size_n, num_bits
)
else:
raise ValueError(f"Unknown provider: {provider}")
return run_benchmark(fn)
if __name__ == "__main__":
check_correctness()
benchmark.run(print_data=True)

110
python/sglang/jit_kernel/benchmark/bench_awq_marlin_repack.py
View File

@@ -1,110 +0,0 @@
import numpy as np
import torch
import triton
import triton.testing
from sgl_kernel.scalar_type import scalar_types
from sglang.jit_kernel.awq_marlin_repack import (
awq_marlin_repack as jit_awq_marlin_repack,
)
from sglang.jit_kernel.benchmark.utils import run_benchmark
from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
from sglang.utils import is_in_ci
AOT_AVAILABLE = hasattr(torch.ops.sgl_kernel, "awq_marlin_repack") and hasattr(
torch.ops.sgl_kernel.awq_marlin_repack, "default"
)
IS_CI = is_in_ci()
SIZE_K = 4096
SIZE_N = 4096
NUM_BITS = 4
GROUP_SIZE = 128
def awq_pack(q_w, num_bits, size_k, size_n):
if num_bits == 4:
interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7])
elif num_bits == 8:
interleave = np.array([0, 2, 1, 3])
else:
raise Exception("num_bits must be 4 or 8, got {}".format(num_bits))
q_w = q_w.reshape((-1, len(interleave)))[:, interleave].ravel()
q_w = q_w.reshape((-1, size_n)).contiguous()
return pack_cols(q_w, num_bits, size_k, size_n)
_b_weight = torch.randn((SIZE_K, SIZE_N), dtype=torch.float16, device="cuda")
_w_ref, _q_w, _s, _zp = quantize_weights(
_b_weight, scalar_types.uint4, GROUP_SIZE, zero_points=True
)
_q_w_awq = awq_pack(_q_w, NUM_BITS, SIZE_K, SIZE_N)
def check_correctness():
if not AOT_AVAILABLE:
print("sgl_kernel AOT not available, skipping correctness check")
return
out_jit = jit_awq_marlin_repack(_q_w_awq, SIZE_K, SIZE_N, NUM_BITS)
out_aot = torch.ops.sgl_kernel.awq_marlin_repack.default(
_q_w_awq, SIZE_K, SIZE_N, NUM_BITS
)
torch.cuda.synchronize()
torch.testing.assert_close(out_jit, out_aot, rtol=0, atol=0)
print("Correctness check passed (JIT vs AOT)")
if IS_CI:
k_range = [1024, 4096]
else:
k_range = [512, 1024, 2048, 4096, 8192]
if AOT_AVAILABLE:
line_vals = ["jit", "aot"]
line_names = ["JIT Kernel", "AOT Kernel"]
styles = [("blue", "-"), ("green", "-")]
else:
line_vals = ["jit"]
line_names = ["JIT Kernel"]
styles = [("blue", "-")]
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["size_k"],
x_vals=k_range,
line_arg="provider",
line_vals=line_vals,
line_names=line_names,
styles=styles,
ylabel="us",
plot_name="awq-marlin-repack-performance",
args={"size_n": SIZE_N, "num_bits": NUM_BITS},
)
)
def benchmark(size_k, size_n, num_bits, provider):
group_size = min(GROUP_SIZE, size_k)
b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
w_ref, q_w, s, zp = quantize_weights(
b_weight, scalar_types.uint4, group_size, zero_points=True
)
q_w_awq = awq_pack(q_w, num_bits, size_k, size_n)
if provider == "jit":
fn = lambda: jit_awq_marlin_repack(q_w_awq, size_k, size_n, num_bits)
elif provider == "aot":
fn = lambda: torch.ops.sgl_kernel.awq_marlin_repack.default(
q_w_awq, size_k, size_n, num_bits
)
else:
raise ValueError(f"Unknown provider: {provider}")
return run_benchmark(fn)
if __name__ == "__main__":
check_correctness()
benchmark.run(print_data=True)

129
python/sglang/jit_kernel/benchmark/bench_gptq_marlin.py
View File

@@ -1,129 +0,0 @@
import torch
import triton
import triton.testing
from sgl_kernel.scalar_type import scalar_types
from sglang.jit_kernel.benchmark.utils import run_benchmark
from sglang.jit_kernel.gptq_marlin import gptq_marlin_gemm as jit_gptq_marlin_gemm
from sglang.srt.layers.quantization.marlin_utils import marlin_make_workspace
from sglang.test.test_marlin_utils import marlin_quantize
from sglang.utils import is_in_ci
AOT_AVAILABLE = hasattr(torch.ops.sgl_kernel, "gptq_marlin_gemm") and hasattr(
torch.ops.sgl_kernel.gptq_marlin_gemm, "default"
)
IS_CI = is_in_ci()
SIZE_K = 4096
SIZE_N = 4096
GROUP_SIZE = 128
QUANT_TYPE = scalar_types.uint4b8
_b_weight = torch.randn((SIZE_K, SIZE_N), dtype=torch.float16, device="cuda")
_w_ref, _marlin_q_w, _marlin_s, _g_idx, _sort_indices, _ = marlin_quantize(
_b_weight, QUANT_TYPE, GROUP_SIZE, act_order=False
)
_workspace = marlin_make_workspace(_w_ref.device)
def _run_gemm(fn, a):
return fn(
a,
None,
_marlin_q_w,
_marlin_s,
None,
None,
_g_idx,
_sort_indices,
_workspace,
QUANT_TYPE,
a.shape[0],
SIZE_N,
SIZE_K,
is_k_full=True,
use_atomic_add=False,
use_fp32_reduce=False,
is_zp_float=False,
)
def _run_gemm_aot(a):
return torch.ops.sgl_kernel.gptq_marlin_gemm.default(
a,
None,
_marlin_q_w,
_marlin_s,
None,
None,
_g_idx,
_sort_indices,
_workspace,
QUANT_TYPE.id,
a.shape[0],
SIZE_N,
SIZE_K,
True,
False,
False,
False,
)
def check_correctness():
if not AOT_AVAILABLE:
print("sgl_kernel AOT not available, skipping correctness check")
return
a = torch.randn((16, SIZE_K), dtype=torch.float16, device="cuda")
out_jit = _run_gemm(jit_gptq_marlin_gemm, a)
out_aot = _run_gemm_aot(a)
torch.testing.assert_close(out_jit, out_aot, rtol=1e-3, atol=1e-3)
print("Correctness check passed (JIT vs AOT)")
if IS_CI:
m_range = [1, 16, 128]
else:
m_range = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]
if AOT_AVAILABLE:
line_vals = ["jit", "aot"]
line_names = ["JIT Kernel", "AOT Kernel"]
styles = [("blue", "-"), ("green", "-")]
else:
line_vals = ["jit"]
line_names = ["JIT Kernel"]
styles = [("blue", "-")]
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["size_m"],
x_vals=m_range,
line_arg="provider",
line_vals=line_vals,
line_names=line_names,
styles=styles,
ylabel="us",
plot_name="gptq-marlin-gemm-performance",
args={},
)
)
def benchmark(size_m, provider):
device = torch.device("cuda")
a = torch.randn((size_m, SIZE_K), dtype=torch.float16, device=device)
if provider == "jit":
fn = lambda: _run_gemm(jit_gptq_marlin_gemm, a)
elif provider == "aot":
fn = lambda: _run_gemm_aot(a)
else:
raise ValueError(f"Unknown provider: {provider}")
return run_benchmark(fn)
if __name__ == "__main__":
check_correctness()
benchmark.run(print_data=True)

97
python/sglang/jit_kernel/benchmark/bench_gptq_marlin_repack.py
View File

@@ -1,97 +0,0 @@
import torch
import triton
import triton.testing
from sgl_kernel.scalar_type import scalar_types
from sglang.jit_kernel.benchmark.utils import run_benchmark
from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack as jit_fn
from sglang.srt.layers.quantization.utils import gptq_quantize_weights, pack_rows
from sglang.utils import is_in_ci
AOT_AVAILABLE = hasattr(torch.ops.sgl_kernel, "gptq_marlin_repack") and hasattr(
torch.ops.sgl_kernel.gptq_marlin_repack, "default"
)
IS_CI = is_in_ci()
SIZE_N = 4096
NUM_BITS = 4
QUANT_TYPE = scalar_types.uint4b8
GROUP_SIZE = 128
_cache = {}
def _get_inputs(size_k):
if size_k not in _cache:
size_n = SIZE_N
b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
_, q_w, _, _, _ = gptq_quantize_weights(
b_weight, QUANT_TYPE, GROUP_SIZE, act_order=False
)
q_w_gptq = pack_rows(q_w, NUM_BITS, size_k, size_n)
sort_indices = torch.empty(0, dtype=torch.int, device="cuda")
_cache[size_k] = (q_w_gptq, sort_indices)
return _cache[size_k]
def check_correctness():
if not AOT_AVAILABLE:
print("sgl_kernel AOT not available, skipping correctness check")
return
size_k = 4096
q_w_gptq, sort_indices = _get_inputs(size_k)
out_jit = jit_fn(q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS)
out_aot = torch.ops.sgl_kernel.gptq_marlin_repack.default(
q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS
)
torch.testing.assert_close(out_jit, out_aot, rtol=0, atol=0)
print("Correctness check passed (JIT vs AOT)")
if IS_CI:
k_range = [128, 1024, 4096]
else:
k_range = [128, 256, 512, 1024, 2048, 4096, 8192]
if AOT_AVAILABLE:
line_vals = ["jit", "aot"]
line_names = ["JIT Kernel", "AOT Kernel"]
styles = [("blue", "-"), ("green", "-")]
else:
line_vals = ["jit"]
line_names = ["JIT Kernel"]
styles = [("blue", "-")]
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["size_k"],
x_vals=k_range,
line_arg="provider",
line_vals=line_vals,
line_names=line_names,
styles=styles,
ylabel="us",
plot_name="gptq-marlin-repack-performance",
args={},
)
)
def benchmark(size_k, provider):
q_w_gptq, sort_indices = _get_inputs(size_k)
if provider == "jit":
fn = lambda: jit_fn(q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS)
elif provider == "aot":
fn = lambda: torch.ops.sgl_kernel.gptq_marlin_repack.default(
q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS
)
else:
raise ValueError(f"Unknown provider: {provider}")
return run_benchmark(fn)
if __name__ == "__main__":
check_correctness()
benchmark.run(print_data=True)

240
python/sglang/jit_kernel/benchmark/bench_moe_wna16_marlin.py
View File

@@ -1,240 +0,0 @@
import torch
import triton
import triton.testing
from sgl_kernel.scalar_type import scalar_types
from sglang.jit_kernel.benchmark.utils import run_benchmark
from sglang.jit_kernel.moe_wna16_marlin import moe_wna16_marlin_gemm as jit_fn
from sglang.srt.layers.moe.fused_moe_triton import moe_align_block_size
from sglang.test.test_marlin_utils import marlin_quantize
from sglang.utils import is_in_ci
AOT_AVAILABLE = hasattr(torch.ops.sgl_kernel, "moe_wna16_marlin_gemm") and hasattr(
torch.ops.sgl_kernel.moe_wna16_marlin_gemm, "default"
)
IS_CI = is_in_ci()
def stack_and_dev(tensors):
dev = tensors[0].device
return torch.stack(tensors, dim=0).to(dev)
E = 8
SIZE_K = 4096
SIZE_N = 4096
GROUP_SIZE = 128
TOPK = 2
QUANT_TYPE = scalar_types.uint4b8
DTYPE = torch.float16
BLOCK_SIZE_M = 64
torch.manual_seed(0)
_qweight_l, _scales_l, _w_ref_l = [], [], []
for i in range(E):
_w = torch.randn((SIZE_N, SIZE_K), dtype=DTYPE, device="cuda") / 20
_perm = torch.randperm(SIZE_K)
_w_ref, _qw, _s, _, _, _ = marlin_quantize(_w, QUANT_TYPE, GROUP_SIZE, False, _perm)
_w_ref_l.append(_w_ref.T)
_qweight_l.append(_qw)
_scales_l.append(_s)
_qweight = stack_and_dev(_qweight_l).contiguous()
_scales = stack_and_dev(_scales_l)
_sms = torch.cuda.get_device_properties("cuda").multi_processor_count
def _make_inputs(size_m):
a = torch.randn((size_m, SIZE_K), dtype=DTYPE, device="cuda") / 10
score = torch.randn((size_m, E), dtype=DTYPE, device="cuda")
score_softmax = torch.softmax(score, dim=-1, dtype=torch.float32)
topk_weights, topk_ids = torch.topk(score_softmax, TOPK)
sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
topk_ids, BLOCK_SIZE_M, E
)
max_workspace_size = (SIZE_N // 64) * (sorted_token_ids.size(0) // BLOCK_SIZE_M)
max_workspace_size = min(max_workspace_size, _sms * 4)
workspace = torch.zeros(max_workspace_size, dtype=torch.int, device="cuda")
c = torch.empty((size_m * TOPK, SIZE_N), dtype=DTYPE, device="cuda")
return (
a,
c,
topk_weights,
topk_ids,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
workspace,
)
def _run_jit(
a,
c,
topk_weights,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
workspace,
size_m,
):
return jit_fn(
a,
c,
_qweight,
None,
_scales,
None,
None,
None,
None,
workspace,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
topk_weights,
moe_block_size=BLOCK_SIZE_M,
top_k=TOPK,
mul_topk_weights=False,
is_ep=False,
b_q_type=QUANT_TYPE,
size_m=size_m,
size_n=SIZE_N,
size_k=SIZE_K,
is_k_full=True,
use_atomic_add=True,
use_fp32_reduce=True,
is_zp_float=False,
)
def _run_aot(
a,
c,
topk_weights,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
workspace,
size_m,
):
return torch.ops.sgl_kernel.moe_wna16_marlin_gemm.default(
a,
c,
_qweight,
None,
_scales,
None,
None,
None,
None,
workspace,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
topk_weights,
moe_block_size=BLOCK_SIZE_M,
top_k=TOPK,
mul_topk_weights=False,
is_ep=False,
b_q_type_id=QUANT_TYPE.id,
size_m=size_m,
size_n=SIZE_N,
size_k=SIZE_K,
is_k_full=True,
use_atomic_add=True,
use_fp32_reduce=True,
is_zp_float=False,
)
def check_correctness():
if not AOT_AVAILABLE:
print("sgl_kernel AOT not available, skipping correctness check")
return
size_m = 16
a, c, topk_weights, topk_ids, sorted_token_ids, expert_ids, ntp, workspace = (
_make_inputs(size_m)
)
c_jit = c.clone()
c_aot = c.clone()
_run_jit(
a, c_jit, topk_weights, sorted_token_ids, expert_ids, ntp, workspace, size_m
)
_run_aot(
a, c_aot, topk_weights, sorted_token_ids, expert_ids, ntp, workspace, size_m
)
torch.testing.assert_close(c_jit, c_aot, rtol=1e-3, atol=1e-3)
print("Correctness check passed (JIT vs AOT)")
if IS_CI:
m_range = [1, 16, 128]
else:
m_range = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]
if AOT_AVAILABLE:
line_vals = ["jit", "aot"]
line_names = ["JIT Kernel", "AOT Kernel"]
styles = [("blue", "-"), ("green", "-")]
else:
line_vals = ["jit"]
line_names = ["JIT Kernel"]
styles = [("blue", "-")]
@triton.testing.perf_report(
triton.testing.Benchmark(
x_names=["size_m"],
x_vals=m_range,
line_arg="provider",
line_vals=line_vals,
line_names=line_names,
styles=styles,
ylabel="us",
plot_name="moe-wna16-marlin-gemm-performance",
args={},
)
)
def benchmark(size_m, provider):
a, c, topk_weights, topk_ids, sorted_token_ids, expert_ids, ntp, workspace = (
_make_inputs(size_m)
)
if provider == "jit":
fn = lambda: _run_jit(
a,
c.clone(),
topk_weights,
sorted_token_ids,
expert_ids,
ntp,
workspace,
size_m,
)
elif provider == "aot":
fn = lambda: _run_aot(
a,
c.clone(),
topk_weights,
sorted_token_ids,
expert_ids,
ntp,
workspace,
size_m,
)
else:
raise ValueError(f"Unknown provider: {provider}")
return run_benchmark(fn)
if __name__ == "__main__":
check_correctness()
benchmark.run(print_data=True)