minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Delete sgl-kernel AOT router GEMM and fused A GEMM (#30280)

Browse Source 
Co-authored-by: Brayden Zhong <brayden@radixark.ai>
Co-authored-by: root <root@sgl-b300-inference.datacrunch.io>
This commit is contained in:
Brayden Zhong
2026-07-22 08:44:59 +08:00
committed by GitHub
co-authored by Brayden Zhong root
parent 8ae0eb83fc
commit 03342e7732
19 changed files with 10 additions and 1834 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py
+4 -14
View File
@@ -1,5 +1,4 @@
"""Benchmark for DeepSeek V3 fused QKV-A GEMM: CuTe DSL vs CUDA JIT vs
sgl_kernel AOT vs torch.
"""Benchmark for DeepSeek V3 fused QKV-A GEMM: CuTe DSL vs CUDA JIT vs torch.
Run on SM90+ (Hopper or later):
python test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py
@@ -8,7 +7,6 @@ Run on SM90+ (Hopper or later):
import torch
import torch.nn.functional as F
import triton.testing
from sgl_kernel import dsv3_fused_a_gemm as sgl_kernel_dsv3_fused_a_gemm
from sglang.jit_kernel.benchmark import marker
from sglang.jit_kernel.cutedsl_dsv3_fused_a_gemm import (
@@ -16,7 +14,6 @@ from sglang.jit_kernel.cutedsl_dsv3_fused_a_gemm import (
)
from sglang.jit_kernel.dsv3_fused_a_gemm import dsv3_fused_a_gemm
from sglang.jit_kernel.utils import get_jit_cuda_arch, is_hip_runtime
from sglang.srt.utils.common import is_sm120_supported
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.utils import is_in_ci
@@ -31,14 +28,11 @@ DEVICE = "cuda"
HD_OUT = 2112
HD_IN_LIST = [6144, 7168]
AOT_HD_IN = 7168
HAS_AOT = not is_sm120_supported()
NUM_TOKENS_LIST = [1, 8, 16] if IS_CI else list(range(1, 17))
LINE_VALS = ["cutedsl", "jit", "sgl_kernel", "torch"]
LINE_NAMES = ["CuTe DSL", "CUDA JIT", "sgl_kernel AOT", "torch F.linear"]
STYLES = [("blue", "-"), ("orange", "--"), ("red", ":"), ("green", "-.")]
LINE_VALS = ["cutedsl", "jit", "torch"]
LINE_NAMES = ["CuTe DSL", "CUDA JIT", "torch F.linear"]
STYLES = [("blue", "-"), ("orange", "--"), ("green", "-.")]
def _median_us(fn, *args) -> float:
@@ -53,15 +47,11 @@ def _median_us(fn, *args) -> float:
def _bench(num_tokens, provider, hd_in):
if provider == "sgl_kernel" and not (HAS_AOT and hd_in == AOT_HD_IN):
return float("nan")
mat_a = torch.randn((num_tokens, hd_in), dtype=DTYPE, device=DEVICE)
mat_b = torch.randn((HD_OUT, hd_in), dtype=DTYPE, device=DEVICE).transpose(0, 1)
fn_map = {
"cutedsl": cutedsl_dsv3_fused_a_gemm,
"jit": dsv3_fused_a_gemm,
"sgl_kernel": sgl_kernel_dsv3_fused_a_gemm,
"torch": lambda a, b: F.linear(a, b.T),
}
return _median_us(fn_map[provider], mat_a, mat_b)
test/registered/jit/benchmark/bench_dsv3_router_gemm.py
+2 -16
View File
@@ -1,4 +1,4 @@
"""Benchmark for DeepSeek V3 router GEMM (JIT kernel vs sgl_kernel AOT vs torch).
"""Benchmark for DeepSeek V3 router GEMM (JIT kernel vs torch).
Run on a Hopper (SM90+) GPU:
python -m sglang.jit_kernel.benchmark.bench_dsv3_router_gemm
@@ -7,11 +7,6 @@ Run on a Hopper (SM90+) GPU:
import torch
import torch.nn.functional as F
try:
from sgl_kernel import dsv3_router_gemm as sgl_kernel_dsv3_router_gemm
except ImportError:
sgl_kernel_dsv3_router_gemm = None
from sglang.jit_kernel.benchmark import marker
from sglang.jit_kernel.benchmark.utils import create_random
from sglang.jit_kernel.dsv3_router_gemm import dsv3_router_gemm
@@ -23,9 +18,6 @@ register_cuda_ci(
)
register_amd_ci(est_time=5, stage="jit-kernel-benchmark", runner_config="amd")
# sgl_kernel AOT kernel is specialized for hidden_dim=7168 only.
SGL_KERNEL_HIDDEN_DIM = 7168
def _torch(mat_a, mat_b, out_dtype):
return F.linear(mat_a, mat_b).to(out_dtype)
@@ -33,7 +25,6 @@ def _torch(mat_a, mat_b, out_dtype):
FN_MAP = {
"jit": dsv3_router_gemm,
"sgl_kernel": sgl_kernel_dsv3_router_gemm,
"torch": _torch,
}
@@ -42,13 +33,8 @@ FN_MAP = {
@marker.parametrize("hidden_dim", [6144, 7168], [7168])
@marker.parametrize("num_tokens", list(range(1, 17)), [1, 8, 16])
@marker.parametrize("out_dtype", [torch.bfloat16, torch.float32])
@marker.benchmark("provider", ["jit", "sgl_kernel", "torch"])
@marker.benchmark("provider", ["jit", "torch"])
def benchmark(num_experts, hidden_dim, num_tokens, out_dtype, provider):
if provider == "sgl_kernel":
if sgl_kernel_dsv3_router_gemm is None:
marker.skip("sgl_kernel dsv3_router_gemm not available in this build")
if hidden_dim != SGL_KERNEL_HIDDEN_DIM:
marker.skip("sgl_kernel AOT only supports hidden_dim=7168")
mat_a = create_random(num_tokens, hidden_dim)
mat_b = create_random(num_experts, hidden_dim)
return marker.do_bench(