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

refactor: consolidate is_in_ci (jit_kernel, sgl-kernel benchmarks, tests) (#21009)

Browse Source
This commit is contained in:
Lianmin Zheng
2026-03-20 05:55:36 -07:00
committed by GitHub
parent 9fbe6800aa
commit 104b10f70a
45 changed files with 104 additions and 168 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
sgl-kernel/benchmark/bench_activation.py
View File

@@ -13,6 +13,8 @@ import triton
import triton.testing
from sgl_kernel import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul
from sglang.utils import is_in_ci
# Optional vLLM import
try:
from vllm import _custom_ops as vllm_ops
@@ -22,11 +24,7 @@ except ImportError:
vllm_ops = None
VLLM_AVAILABLE = False
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
# gelu_quick is only available on HIP/ROCm platforms
try:

8
sgl-kernel/benchmark/bench_awq_dequant.py
View File

@@ -7,6 +7,8 @@ import triton
import triton.testing
from sgl_kernel import awq_dequantize
from sglang.utils import is_in_ci
# Optional vLLM import
try:
from vllm import _custom_ops as ops
@@ -16,11 +18,7 @@ except ImportError:
ops = None
VLLM_AVAILABLE = False
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
def vllm_awq_dequantize(

7
sgl-kernel/benchmark/bench_cutlass_mla.py
View File

@@ -8,12 +8,9 @@ import triton
from sgl_kernel import cutlass_mla_decode, cutlass_mla_get_workspace_size
from sglang.srt.utils import get_device_capability
from sglang.utils import is_in_ci
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
# CI environment uses simplified parameters
if IS_CI:

11
sgl-kernel/benchmark/bench_dsv3_fused_a_gemm.py
View File

@@ -1,11 +1,4 @@
import argparse
import os
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
import torch
import torch.nn.functional as F
@@ -13,6 +6,10 @@ import triton
import triton.testing
from sgl_kernel import dsv3_fused_a_gemm
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
# CI environment uses simplified parameters
if IS_CI:
num_tokens_vals = [1] # Only test 1 value in CI

11
sgl-kernel/benchmark/bench_dsv3_router_gemm.py
View File

@@ -1,11 +1,4 @@
import argparse
import os
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
import torch
import torch.nn.functional as F
@@ -13,6 +6,10 @@ import triton
import triton.testing
from sgl_kernel import dsv3_router_gemm
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
# CI environment uses simplified parameters
if IS_CI:
num_tokens_vals = [1] # Only test 1 value in CI

7
sgl-kernel/benchmark/bench_fp4_gemm.py
View File

@@ -10,12 +10,9 @@ from flashinfer.testing import bench_gpu_time_with_cupti
from sgl_kernel import cutlass_scaled_fp4_mm, scaled_fp4_quant
from sglang.srt.utils import get_device_capability, is_sm100_supported
from sglang.utils import is_in_ci
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max

8
sgl-kernel/benchmark/bench_fp8_blockwise_gemm.py
View File

@@ -9,6 +9,8 @@ import triton
from deep_gemm.utils.layout import get_mn_major_tma_aligned_tensor
from sgl_kernel import fp8_blockwise_scaled_mm
from sglang.utils import is_in_ci
# Optional vLLM import
try:
from vllm._custom_ops import cutlass_scaled_mm as vllm_scaled_mm
@@ -22,11 +24,7 @@ from sglang.srt.layers.quantization.fp8_kernel import (
w8a8_block_fp8_matmul_triton as w8a8_block_fp8_matmul,
)
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
def get_weight_shapes(args):

11
sgl-kernel/benchmark/bench_fp8_blockwise_group_gemm.py
View File

@@ -1,11 +1,4 @@
import argparse
import os
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
import random
from dataclasses import dataclass
from typing import List, Tuple
@@ -14,6 +7,10 @@ import deep_gemm
import torch
from sgl_kernel import fp8_blockwise_scaled_grouped_mm
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
def get_m_alignment_for_contiguous_layout():
return 128

7
sgl-kernel/benchmark/bench_fp8_gemm.py
View File

@@ -9,6 +9,7 @@ import triton
from sgl_kernel import fp8_scaled_mm as sgl_scaled_mm
from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
from sglang.utils import is_in_ci
# Optional vLLM import
try:
@@ -21,11 +22,7 @@ except ImportError:
vllm_scaled_fp8_quant = None
VLLM_AVAILABLE = False
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
# Weight Shapes are in the format
# ([K, N], TP_SPLIT_DIM)

8
sgl-kernel/benchmark/bench_int8_gemm.py
View File

@@ -7,6 +7,8 @@ import torch
import triton
from sgl_kernel import int8_scaled_mm
from sglang.utils import is_in_ci
# Optional vLLM import
try:
from vllm._custom_ops import cutlass_scaled_mm as vllm_scaled_mm
@@ -16,11 +18,7 @@ except ImportError:
vllm_scaled_mm = None
VLLM_AVAILABLE = False
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
def to_int8(tensor: torch.Tensor) -> torch.Tensor:

7
sgl-kernel/benchmark/bench_kimi_k2_moe_fused_gate.py
View File

@@ -8,12 +8,9 @@ import triton.language as tl
from sgl_kernel import kimi_k2_moe_fused_gate
from sglang.srt.layers.moe.topk import kimi_k2_biased_topk_impl
from sglang.utils import is_in_ci
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
def kimi_k2_biased_topk_torch_compile(scores, bias, topk, routed_scaling_factor):

8
sgl-kernel/benchmark/bench_moe_align_block_size.py
View File

@@ -7,6 +7,8 @@ import triton
import triton.language as tl
from sgl_kernel import moe_align_block_size as sgl_moe_align_block_size
from sglang.utils import is_in_ci
try:
from vllm import _custom_ops as ops
@@ -15,11 +17,7 @@ except ImportError:
ops = None
VLLM_AVAILABLE = False
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
USE_RANDOM_PERM = False

11
sgl-kernel/benchmark/bench_moe_ep_post_reorder.py
View File

@@ -1,15 +1,10 @@
import os
import torch
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
import triton
from sglang.srt.layers.moe.ep_moe.kernels import post_reorder_triton_kernel
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
# CI environment uses simplified parameters
if IS_CI:

7
sgl-kernel/benchmark/bench_moe_fused_gate.py
View File

@@ -8,12 +8,9 @@ import triton.language as tl
from sgl_kernel import moe_fused_gate
from sglang.srt.layers.moe.topk import biased_grouped_topk
from sglang.utils import is_in_ci
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
def biased_grouped_topk_org(scores, bias, num_expert_group, topk_group, topk):

8
sgl-kernel/benchmark/bench_moe_topk_sigmoid.py
View File

@@ -6,11 +6,9 @@ import torch
import triton
from sgl_kernel import topk_sigmoid
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
def torch_topk_sigmoid_native(

8
sgl-kernel/benchmark/bench_moe_topk_softmax.py
View File

@@ -6,6 +6,8 @@ import torch
import triton
from sgl_kernel import topk_softmax
from sglang.utils import is_in_ci
# Optional vLLM import
try:
from vllm import _custom_ops as vllm_custom_ops
@@ -15,11 +17,7 @@ except ImportError:
vllm_custom_ops = None
VLLM_AVAILABLE = False
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
def vllm_topk_softmax(gating_output, topk):

7
sgl-kernel/benchmark/bench_per_tensor_quant_fp8.py
View File

@@ -9,6 +9,7 @@ import triton
import triton.testing
from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
from sglang.utils import is_in_ci
# Optional imports
try:
@@ -23,11 +24,7 @@ from sglang.srt.utils import is_hip
_is_hip = is_hip()
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn

7
sgl-kernel/benchmark/bench_per_token_group_quant_8bit.py
View File

@@ -19,12 +19,9 @@ from sglang.srt.layers.quantization.fp8_kernel import (
)
from sglang.srt.utils import is_hip
from sglang.srt.utils.bench_utils import bench_kineto
from sglang.utils import is_in_ci
# 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

8
sgl-kernel/benchmark/bench_per_token_quant_fp8.py
View File

@@ -7,6 +7,8 @@ import triton
import triton.testing
from sgl_kernel import sgl_per_token_quant_fp8
from sglang.utils import is_in_ci
# Optional vLLM import
try:
from vllm import _custom_ops as ops
@@ -20,11 +22,7 @@ from sglang.srt.utils import is_hip
_is_hip = is_hip()
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn

8
sgl-kernel/benchmark/bench_qserve_w4a8_gemm.py
View File

@@ -11,11 +11,9 @@ from sgl_kernel import (
qserve_w4a8_per_group_gemm,
)
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
def to_int8(tensor: torch.Tensor) -> torch.Tensor:

8
sgl-kernel/benchmark/bench_rmsnorm.py
View File

@@ -13,6 +13,8 @@ import triton
import triton.testing
from sgl_kernel.utils import is_arch_support_pdl
from sglang.utils import is_in_ci
# Optional imports
try:
from flashinfer.norm import fused_add_rmsnorm, rmsnorm
@@ -31,11 +33,7 @@ except ImportError:
vllm_ops = None
VLLM_AVAILABLE = False
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
def str2int_list(arg: str) -> List[int]:

7
sgl-kernel/benchmark/bench_rotary_embedding.py
View File

@@ -12,12 +12,9 @@ from sgl_kernel.testing.rotary_embedding import (
)
from sglang.srt.utils.bench_utils import bench_kineto
from sglang.utils import is_in_ci
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
IS_CI = is_in_ci()
# CI environment uses simplified parameters
if IS_CI:

8
sgl-kernel/benchmark/bench_sum_scale.py
View File

@@ -6,11 +6,9 @@ import triton.language as tl
from sgl_kernel import moe_sum_reduce as moe_sum_reduce_cuda
from triton.testing import do_bench
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
@triton.jit

8
sgl-kernel/benchmark/bench_top_k_top_p_sampling.py
View File

@@ -7,11 +7,9 @@ import torch
import triton
import triton.testing
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
def torch_top_k_top_p_joint_sampling_from_probs(