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

[4/n jit_kernel restruct] speed up CI tests and add benchmark workflow (#20268)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-03-10 21:37:41 +08:00
committed by GitHub
parent 5a7c1b8ec6
commit 60cc06297e
20 changed files with 208 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py
View File

@@ -1,12 +1,11 @@
import itertools
import os
import torch
import triton
import triton.testing
from sglang.jit_kernel.awq_dequantize import awq_dequantize as jit_awq_dequantize
from sglang.jit_kernel.benchmark.utils import run_benchmark
from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark
try:
from sgl_kernel import awq_dequantize as aot_awq_dequantize
@@ -15,10 +14,7 @@ try:
except ImportError:
AOT_AVAILABLE = False
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
if IS_CI:
qweight_row_range = [128]

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

@@ -1,5 +1,3 @@
import os
import numpy as np
import torch
import triton
@@ -9,7 +7,7 @@ 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.jit_kernel.benchmark.utils import is_in_ci, run_benchmark
from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
try:
@@ -19,10 +17,7 @@ try:
except ImportError:
AOT_AVAILABLE = False
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
NUM_BITS = 4
GROUP_SIZE = 128

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

@@ -1,5 +1,3 @@
import os
import numpy as np
import torch
import triton
@@ -9,7 +7,7 @@ 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.jit_kernel.benchmark.utils import is_in_ci, run_benchmark
from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
try:
@@ -19,10 +17,7 @@ try:
except ImportError:
AOT_AVAILABLE = False
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
SIZE_K = 4096
SIZE_N = 4096

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

@@ -1,11 +1,9 @@
import os
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.benchmark.utils import is_in_ci, 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
@@ -17,10 +15,7 @@ try:
except ImportError:
AOT_AVAILABLE = False
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
SIZE_K = 4096
SIZE_N = 4096

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

@@ -1,11 +1,9 @@
import os
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.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
@@ -16,10 +14,7 @@ try:
except ImportError:
AOT_AVAILABLE = False
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
SIZE_N = 4096
NUM_BITS = 4

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

@@ -1,11 +1,9 @@
import os
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.benchmark.utils import is_in_ci, 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
@@ -17,10 +15,7 @@ try:
except (ImportError, AttributeError):
AOT_AVAILABLE = False
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
def stack_and_dev(tensors):

8
python/sglang/jit_kernel/benchmark/bench_nvfp4_blockwise_moe.py
View File

@@ -1,5 +1,6 @@
from __future__ import annotations
import sys
from typing import Any
import torch
@@ -11,9 +12,11 @@ from sglang.jit_kernel.nvfp4 import (
scaled_fp4_experts_quant,
scaled_fp4_quant,
)
from sglang.srt.utils import is_sm100_supported
FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
_NVFP4_SUPPORTED = is_sm100_supported()
def _round_up(x: int, y: int) -> int:
@@ -168,6 +171,8 @@ def _aot_cutlass_fp4_group_mm(case: dict[str, Any]) -> torch.Tensor:
def _probe_legacy_aot_group_mm() -> tuple[bool, str]:
if not torch.cuda.is_available():
return False, "CUDA is not available."
if not _NVFP4_SUPPORTED:
return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+."
try:
import sgl_kernel # noqa: F401
except Exception as e:
@@ -243,6 +248,9 @@ def benchmark(total_tokens, n, k, num_experts, provider):
if __name__ == "__main__":
if not _NVFP4_SUPPORTED:
print("[skip] NVFP4 blockwise MoE benchmark requires sm100+ with CUDA 12.8+.")
sys.exit(0)
if not _AOT_GROUP_MM_AVAILABLE:
print(
f"[info] legacy AOT grouped_mm baseline unavailable: {_AOT_GROUP_MM_REASON}"

11
python/sglang/jit_kernel/benchmark/bench_nvfp4_quant.py
View File

@@ -1,14 +1,18 @@
from __future__ import annotations
import sys
import torch
import triton
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
from sglang.jit_kernel.nvfp4 import scaled_fp4_quant
from sglang.srt.utils import is_sm100_supported
FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
BLOCK_SIZE = 16
_NVFP4_SUPPORTED = is_sm100_supported()
try:
from flashinfer import fp4_quantize as flashinfer_fp4_quantize
@@ -59,6 +63,8 @@ def _aot_scaled_fp4_quant(input: torch.Tensor, input_global_scale: torch.Tensor)
def _probe_legacy_aot_quant() -> tuple[bool, str]:
if not torch.cuda.is_available():
return False, "CUDA is not available."
if not _NVFP4_SUPPORTED:
return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+."
try:
import sgl_kernel # noqa: F401
except Exception as e:
@@ -88,6 +94,8 @@ def _probe_flashinfer_quant() -> tuple[bool, str]:
return False, "import flashinfer.fp4_quantize failed."
if not torch.cuda.is_available():
return False, "CUDA is not available."
if not _NVFP4_SUPPORTED:
return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+."
try:
x = torch.randn((16, 64), dtype=torch.bfloat16, device="cuda")
global_scale = (
@@ -172,6 +180,9 @@ def benchmark(m, n, provider):
if __name__ == "__main__":
if not _NVFP4_SUPPORTED:
print("[skip] NVFP4 quant benchmark requires sm100+ with CUDA 12.8+.")
sys.exit(0)
if not _FLASHINFER_QUANT_AVAILABLE:
print(
f"[info] flashinfer quant baseline unavailable: {_FLASHINFER_QUANT_REASON}"

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

@@ -1,14 +1,18 @@
from __future__ import annotations
import sys
import torch
import triton
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant
from sglang.srt.utils import is_sm100_supported
FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
BLOCK_SIZE = 16
_NVFP4_SUPPORTED = is_sm100_supported()
K_E2M1_TO_FLOAT = [
0.0,
@@ -72,6 +76,8 @@ def _aot_cutlass_scaled_fp4_mm(
def _probe_legacy_aot_scaled_mm() -> tuple[bool, str]:
if not torch.cuda.is_available():
return False, "CUDA is not available."
if not _NVFP4_SUPPORTED:
return False, "NVFP4 benchmarks require sm100+ with CUDA 12.8+."
try:
import sgl_kernel # noqa: F401
except Exception as e:
@@ -168,6 +174,9 @@ def benchmark(m, n, k, provider):
if __name__ == "__main__":
if not _NVFP4_SUPPORTED:
print("[skip] NVFP4 scaled_mm benchmark requires sm100+ with CUDA 12.8+.")
sys.exit(0)
if not _AOT_SCALED_MM_AVAILABLE:
print(
f"[info] legacy AOT scaled_mm baseline unavailable: {_AOT_SCALED_MM_REASON}"

10
python/sglang/jit_kernel/benchmark/bench_per_token_group_quant_8bit.py
View File

@@ -1,5 +1,4 @@
import itertools
import os
from typing import Any, Dict, List
import torch
@@ -8,6 +7,7 @@ from sgl_kernel.test_utils import create_per_token_group_quant_test_data
from sglang.jit_kernel.benchmark.utils import (
get_benchmark_range,
is_in_ci,
)
from sglang.jit_kernel.per_token_group_quant_8bit import (
per_token_group_quant_8bit as sglang_per_token_group_quant_8bit,
@@ -21,16 +21,12 @@ from sglang.srt.layers.quantization.fp8_kernel import (
from sglang.srt.utils import is_hip
from sglang.srt.utils.bench_utils import bench_kineto
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
_is_hip = is_hip()
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
NUM_TESTS = 300 if IS_CI else 30
NUM_TESTS = 30 if IS_CI else 300
GROUP_SIZE_RANGE = [128]
DST_DTYPE_RANGE = [fp8_type_]

8
python/sglang/jit_kernel/benchmark/utils.py
View File

@@ -1,11 +1,12 @@
"""Common utilities for jit_kernel benchmark files."""
import os
from typing import Callable, List, Tuple
import torch
import triton.testing
from sglang.jit_kernel.utils import is_in_ci as jit_kernel_is_in_ci
# Common constants
DEFAULT_DTYPE = torch.bfloat16
DEFAULT_DEVICE = "cuda"
@@ -14,10 +15,7 @@ DEFAULT_QUANTILES = [0.5, 0.2, 0.8]
def is_in_ci() -> bool:
"""Check if running in CI environment."""
return (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
return jit_kernel_is_in_ci()
def get_benchmark_range(full_range: List, ci_range: List) -> List: