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

Release sglang kernel 0.4.0 (#20440)

Browse Source 
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
This commit is contained in:
Xiaoyu Zhang
2026-03-16 20:34:58 +08:00
committed by GitHub
parent 3d58cd16d9
commit 15097c5c3b
36 changed files with 188 additions and 146 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/pyproject.toml
View File

@@ -60,7 +60,7 @@ dependencies = [
"sentencepiece",
"setproctitle",
"sgl-fa4==4.0.5",
"sgl-kernel==0.3.21",
"sglang-kernel==0.4.0",
"soundfile==0.13.1",
"tiktoken",
"timm==1.0.16",

2
python/sglang/check_env.py
View File

@@ -20,7 +20,7 @@ def is_cuda_v2():
# List of packages to check versions
PACKAGE_LIST = [
"sglang",
"sgl_kernel",
"sglang-kernel",
"flashinfer_python",
"flashinfer_cubin",
"flashinfer_jit_cache",

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

@@ -10,12 +10,9 @@ from sglang.jit_kernel.awq_marlin_repack import (
from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark
from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
try:
from sgl_kernel import awq_marlin_moe_repack as aot_awq_marlin_moe_repack
AOT_AVAILABLE = True
except ImportError:
AOT_AVAILABLE = False
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()
@@ -66,7 +63,9 @@ def check_correctness():
)
out_jit = jit_awq_marlin_moe_repack(b_q_weight, perm, size_k, SIZE_N, NUM_BITS)
out_aot = aot_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)")
@@ -111,7 +110,7 @@ def benchmark(num_experts, size_k, size_n, num_bits, provider):
b_q_weight, perm, size_k, size_n, num_bits
)
elif provider == "aot":
fn = lambda: aot_awq_marlin_moe_repack(
fn = lambda: torch.ops.sgl_kernel.awq_marlin_moe_repack.default(
b_q_weight, perm, size_k, size_n, num_bits
)
else:

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

@@ -10,12 +10,9 @@ from sglang.jit_kernel.awq_marlin_repack import (
from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark
from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
try:
from sgl_kernel import awq_marlin_repack as aot_awq_marlin_repack
AOT_AVAILABLE = True
except ImportError:
AOT_AVAILABLE = False
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()
@@ -50,7 +47,9 @@ def check_correctness():
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 = aot_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)")
@@ -96,7 +95,9 @@ def benchmark(size_k, size_n, num_bits, provider):
if provider == "jit":
fn = lambda: jit_awq_marlin_repack(q_w_awq, size_k, size_n, num_bits)
elif provider == "aot":
fn = lambda: aot_awq_marlin_repack(q_w_awq, size_k, size_n, num_bits)
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}")

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

@@ -8,12 +8,9 @@ from sglang.jit_kernel.gptq_marlin import gptq_marlin_gemm as jit_gptq_marlin_ge
from sglang.srt.layers.quantization.marlin_utils import marlin_make_workspace
from sglang.test.test_marlin_utils import marlin_quantize
try:
from sgl_kernel import gptq_marlin_gemm as aot_gptq_marlin_gemm
AOT_AVAILABLE = True
except ImportError:
AOT_AVAILABLE = False
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()
@@ -51,13 +48,35 @@ def _run_gemm(fn, a):
)
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_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)")
@@ -97,7 +116,7 @@ def benchmark(size_m, provider):
if provider == "jit":
fn = lambda: _run_gemm(jit_gptq_marlin_gemm, a)
elif provider == "aot":
fn = lambda: _run_gemm(aot_gptq_marlin_gemm, a)
fn = lambda: _run_gemm_aot(a)
else:
raise ValueError(f"Unknown provider: {provider}")

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

@@ -7,12 +7,9 @@ from sglang.jit_kernel.benchmark.utils import is_in_ci, 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
try:
from sgl_kernel import gptq_marlin_repack as aot_fn
AOT_AVAILABLE = True
except ImportError:
AOT_AVAILABLE = False
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()
@@ -44,7 +41,9 @@ def check_correctness():
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 = aot_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)")
@@ -83,7 +82,9 @@ def benchmark(size_k, provider):
if provider == "jit":
fn = lambda: jit_fn(q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS)
elif provider == "aot":
fn = lambda: aot_fn(q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS)
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}")

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

@@ -8,12 +8,9 @@ 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
try:
from sgl_kernel import moe_wna16_marlin_gemm as _aot_import # noqa: F401
AOT_AVAILABLE = True
except (ImportError, AttributeError):
AOT_AVAILABLE = False
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()

16
python/sglang/jit_kernel/tests/test_awq_marlin_moe_repack.py
View File

@@ -8,12 +8,14 @@ from sglang.jit_kernel.awq_marlin_repack import (
)
from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
try:
from sgl_kernel import awq_marlin_moe_repack as aot_awq_marlin_moe_repack
AOT_AVAILABLE = True
except ImportError:
AOT_AVAILABLE = False
def _has_aot_awq_marlin_moe_repack() -> bool:
return hasattr(torch.ops.sgl_kernel, "awq_marlin_moe_repack") and hasattr(
torch.ops.sgl_kernel.awq_marlin_moe_repack, "default"
)
AOT_AVAILABLE = _has_aot_awq_marlin_moe_repack()
def awq_pack(
@@ -68,7 +70,9 @@ def test_awq_marlin_moe_repack_jit_vs_aot(
perm = torch.empty((num_experts, 0), dtype=torch.int32, device="cuda")
out_jit = jit_awq_marlin_moe_repack(b_q_weight, perm, size_k, size_n, num_bits)
out_aot = aot_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()

16
python/sglang/jit_kernel/tests/test_awq_marlin_repack.py
View File

@@ -9,12 +9,14 @@ from sglang.jit_kernel.awq_marlin_repack import (
from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
from sglang.test.test_marlin_utils import get_weight_perm, marlin_weights
try:
from sgl_kernel import awq_marlin_repack as aot_awq_marlin_repack
AOT_AVAILABLE = True
except ImportError:
AOT_AVAILABLE = False
def _has_aot_awq_marlin_repack() -> bool:
return hasattr(torch.ops.sgl_kernel, "awq_marlin_repack") and hasattr(
torch.ops.sgl_kernel.awq_marlin_repack, "default"
)
AOT_AVAILABLE = _has_aot_awq_marlin_repack()
def awq_pack(
@@ -58,7 +60,9 @@ def test_awq_marlin_repack_jit_vs_aot(num_bits, k_tiles, n_tiles, group_size):
q_w_awq = awq_pack(q_w, num_bits, size_k, size_n)
out_jit = jit_awq_marlin_repack(q_w_awq, size_k, size_n, num_bits)
out_aot = aot_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()

15
python/sglang/jit_kernel/tests/test_moe_wna16_marlin.py
View File

@@ -2,7 +2,6 @@ import itertools
import pytest
import torch
from sgl_kernel import moe_wna16_marlin_gemm as aot_moe_wna16_marlin_gemm
from sgl_kernel.scalar_type import scalar_types
from sglang.jit_kernel.moe_wna16_marlin import moe_wna16_marlin_gemm
@@ -10,6 +9,15 @@ from sglang.srt.layers.moe.fused_moe_triton import moe_align_block_size
from sglang.test.test_marlin_utils import awq_marlin_quantize, marlin_quantize
def _has_aot_moe_wna16_marlin_gemm() -> bool:
return hasattr(torch.ops.sgl_kernel, "moe_wna16_marlin_gemm") and hasattr(
torch.ops.sgl_kernel.moe_wna16_marlin_gemm, "default"
)
AOT_AVAILABLE = _has_aot_moe_wna16_marlin_gemm()
def stack_and_dev(tensors: list[torch.Tensor]):
dev = tensors[0].device
return torch.stack(tensors, dim=0).to(dev)
@@ -143,7 +151,7 @@ def _run_single_gemm_aot(
is_k_full,
use_atomic_add,
):
return aot_moe_wna16_marlin_gemm(
return torch.ops.sgl_kernel.moe_wna16_marlin_gemm.default(
a,
c,
qweight,
@@ -224,6 +232,9 @@ TEST_CASES = generate_test_cases()
def test_moe_wna16_marlin_gemm(
m, n, k, e, topk, dtype, group_size, act_order, quant_type
):
if not AOT_AVAILABLE:
pytest.skip("sgl_kernel moe_wna16_marlin_gemm AOT op not available")
torch.manual_seed(0)
has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]

6
python/sglang/srt/entrypoints/engine.py
View File

@@ -1142,9 +1142,9 @@ def _set_envs_and_config(server_args: ServerArgs):
)
if _is_cuda:
assert_pkg_version(
"sgl-kernel",
"0.3.21",
"Please reinstall the latest version with `pip install sgl-kernel --force-reinstall`",
"sglang-kernel",
"0.4.0",
"Please reinstall the latest version with `pip install sglang-kernel --force-reinstall`",
)
# Signal handlers can only be registered from the main thread.