diff --git a/.claude/skills/add-sgl-kernel/SKILL.md b/.claude/skills/add-sgl-kernel/SKILL.md index 4c5dc7a52..662df875b 100644 --- a/.claude/skills/add-sgl-kernel/SKILL.md +++ b/.claude/skills/add-sgl-kernel/SKILL.md @@ -244,18 +244,15 @@ Create `sgl-kernel/benchmark/bench_scale.py`: ```python import itertools -import os import torch import triton import triton.testing import sgl_kernel +from sglang.utils import is_in_ci -IS_CI = ( - os.getenv("CI", "false").lower() == "true" - or os.getenv("GITHUB_ACTIONS", "false").lower() == "true" -) +IS_CI = is_in_ci() dtypes = [torch.float16] if IS_CI else [torch.float16, torch.bfloat16, torch.float32] sizes = [4096] if IS_CI else [2**n for n in range(10, 20)] # 1K … 512K diff --git a/benchmark/kernels/all_reduce/benchmark_torch_symm_mem.py b/benchmark/kernels/all_reduce/benchmark_torch_symm_mem.py index 030fd5bb2..5bdb7f5d6 100644 --- a/benchmark/kernels/all_reduce/benchmark_torch_symm_mem.py +++ b/benchmark/kernels/all_reduce/benchmark_torch_symm_mem.py @@ -44,12 +44,9 @@ from sglang.srt.distributed.parallel_state import ( initialize_model_parallel, set_torch_symm_mem_all_reduce, ) +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 torch_allreduce(torch_input: torch.Tensor, group: ProcessGroup) -> torch.Tensor: diff --git a/python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py b/python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py index a39027ea2..6a3fcbc93 100644 --- a/python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py +++ b/python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py @@ -5,7 +5,8 @@ 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 is_in_ci, run_benchmark +from sglang.jit_kernel.benchmark.utils import run_benchmark +from sglang.utils import is_in_ci try: from sgl_kernel import awq_dequantize as aot_awq_dequantize diff --git a/python/sglang/jit_kernel/benchmark/bench_awq_marlin_moe_repack.py b/python/sglang/jit_kernel/benchmark/bench_awq_marlin_moe_repack.py index 7ba7a1f2c..ac9faa517 100644 --- a/python/sglang/jit_kernel/benchmark/bench_awq_marlin_moe_repack.py +++ b/python/sglang/jit_kernel/benchmark/bench_awq_marlin_moe_repack.py @@ -7,8 +7,9 @@ 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 is_in_ci, run_benchmark +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" diff --git a/python/sglang/jit_kernel/benchmark/bench_awq_marlin_repack.py b/python/sglang/jit_kernel/benchmark/bench_awq_marlin_repack.py index e774e381c..8cf748fc0 100644 --- a/python/sglang/jit_kernel/benchmark/bench_awq_marlin_repack.py +++ b/python/sglang/jit_kernel/benchmark/bench_awq_marlin_repack.py @@ -7,8 +7,9 @@ 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 is_in_ci, run_benchmark +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" diff --git a/python/sglang/jit_kernel/benchmark/bench_concat_mla.py b/python/sglang/jit_kernel/benchmark/bench_concat_mla.py index 7e5c5c43a..7f2b542a0 100644 --- a/python/sglang/jit_kernel/benchmark/bench_concat_mla.py +++ b/python/sglang/jit_kernel/benchmark/bench_concat_mla.py @@ -6,9 +6,10 @@ import triton.testing from sgl_kernel import concat_mla_absorb_q as aot_absorb_q from sgl_kernel import concat_mla_k as aot_k -from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark +from sglang.jit_kernel.benchmark.utils import run_benchmark from sglang.jit_kernel.concat_mla import concat_mla_absorb_q as jit_absorb_q from sglang.jit_kernel.concat_mla import concat_mla_k as jit_k +from sglang.utils import is_in_ci IS_CI = is_in_ci() diff --git a/python/sglang/jit_kernel/benchmark/bench_fused_add_rmsnorm.py b/python/sglang/jit_kernel/benchmark/bench_fused_add_rmsnorm.py index ddce16b92..0d7b1ec23 100644 --- a/python/sglang/jit_kernel/benchmark/bench_fused_add_rmsnorm.py +++ b/python/sglang/jit_kernel/benchmark/bench_fused_add_rmsnorm.py @@ -5,8 +5,9 @@ import triton import triton.testing from flashinfer import fused_add_rmsnorm as fi_fused_add_rmsnorm -from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark +from sglang.jit_kernel.benchmark.utils import run_benchmark from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm +from sglang.utils import is_in_ci IS_CI = is_in_ci() diff --git a/python/sglang/jit_kernel/benchmark/bench_fused_norm_scale_shift.py b/python/sglang/jit_kernel/benchmark/bench_fused_norm_scale_shift.py index f6bf7f57f..24971231c 100644 --- a/python/sglang/jit_kernel/benchmark/bench_fused_norm_scale_shift.py +++ b/python/sglang/jit_kernel/benchmark/bench_fused_norm_scale_shift.py @@ -8,13 +8,14 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark_no_cudagraph +from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph from sglang.multimodal_gen.runtime.layers.layernorm import ( LayerNormScaleShift, RMSNormScaleShift, ScaleResidualLayerNormScaleShift, ScaleResidualRMSNormScaleShift, ) +from sglang.utils import is_in_ci if is_in_ci(): B_RANGE, S_RANGE, D_RANGE = [1], [128], [1024] diff --git a/python/sglang/jit_kernel/benchmark/bench_gptq_marlin.py b/python/sglang/jit_kernel/benchmark/bench_gptq_marlin.py index a6b4d0514..31881d785 100644 --- a/python/sglang/jit_kernel/benchmark/bench_gptq_marlin.py +++ b/python/sglang/jit_kernel/benchmark/bench_gptq_marlin.py @@ -3,10 +3,11 @@ import triton import triton.testing from sgl_kernel.scalar_type import scalar_types -from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark +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" diff --git a/python/sglang/jit_kernel/benchmark/bench_gptq_marlin_repack.py b/python/sglang/jit_kernel/benchmark/bench_gptq_marlin_repack.py index 027575b6c..3fe86a0e6 100644 --- a/python/sglang/jit_kernel/benchmark/bench_gptq_marlin_repack.py +++ b/python/sglang/jit_kernel/benchmark/bench_gptq_marlin_repack.py @@ -3,9 +3,10 @@ import triton import triton.testing from sgl_kernel.scalar_type import scalar_types -from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark +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" diff --git a/python/sglang/jit_kernel/benchmark/bench_moe_wna16_marlin.py b/python/sglang/jit_kernel/benchmark/bench_moe_wna16_marlin.py index 063ae908e..fa6e9b36d 100644 --- a/python/sglang/jit_kernel/benchmark/bench_moe_wna16_marlin.py +++ b/python/sglang/jit_kernel/benchmark/bench_moe_wna16_marlin.py @@ -3,10 +3,11 @@ import triton import triton.testing from sgl_kernel.scalar_type import scalar_types -from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark +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" diff --git a/python/sglang/jit_kernel/benchmark/bench_norm.py b/python/sglang/jit_kernel/benchmark/bench_norm.py index 2bccd672d..e58d9f5f9 100644 --- a/python/sglang/jit_kernel/benchmark/bench_norm.py +++ b/python/sglang/jit_kernel/benchmark/bench_norm.py @@ -6,9 +6,10 @@ import triton.testing from flashinfer.norm import fused_add_rmsnorm as fi_fused_add_rmsnorm from flashinfer.norm import rmsnorm as fi_rmsnorm -from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark +from sglang.jit_kernel.benchmark.utils import run_benchmark from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm +from sglang.utils import is_in_ci IS_CI = is_in_ci() diff --git a/python/sglang/jit_kernel/benchmark/bench_norm_impls.py b/python/sglang/jit_kernel/benchmark/bench_norm_impls.py index 5f3f74cd9..94b029597 100644 --- a/python/sglang/jit_kernel/benchmark/bench_norm_impls.py +++ b/python/sglang/jit_kernel/benchmark/bench_norm_impls.py @@ -15,12 +15,13 @@ from typing import Callable import torch import torch.nn.functional as F -from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE, is_in_ci +from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE from sglang.jit_kernel.diffusion.triton.norm import norm_infer, rms_norm_fn from sglang.jit_kernel.diffusion.triton.rmsnorm_onepass import triton_one_pass_rms_norm from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm from sglang.jit_kernel.utils import KERNEL_PATH +from sglang.utils import is_in_ci os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1") diff --git a/python/sglang/jit_kernel/benchmark/bench_per_token_group_quant_8bit.py b/python/sglang/jit_kernel/benchmark/bench_per_token_group_quant_8bit.py index f0aadf145..b4515df0a 100644 --- a/python/sglang/jit_kernel/benchmark/bench_per_token_group_quant_8bit.py +++ b/python/sglang/jit_kernel/benchmark/bench_per_token_group_quant_8bit.py @@ -5,10 +5,7 @@ import torch import triton 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.benchmark.utils import get_benchmark_range from sglang.jit_kernel.per_token_group_quant_8bit import ( per_token_group_quant_8bit as sglang_per_token_group_quant_8bit, ) @@ -20,6 +17,7 @@ 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 IS_CI = is_in_ci() diff --git a/python/sglang/jit_kernel/benchmark/bench_qknorm_across_heads.py b/python/sglang/jit_kernel/benchmark/bench_qknorm_across_heads.py index ab660f489..cbcff1031 100644 --- a/python/sglang/jit_kernel/benchmark/bench_qknorm_across_heads.py +++ b/python/sglang/jit_kernel/benchmark/bench_qknorm_across_heads.py @@ -6,9 +6,10 @@ import triton import triton.testing from sgl_kernel import rmsnorm -from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark +from sglang.jit_kernel.benchmark.utils import run_benchmark from sglang.jit_kernel.norm import fused_inplace_qknorm_across_heads from sglang.srt.utils import get_current_device_stream_fast +from sglang.utils import is_in_ci IS_CI = is_in_ci() diff --git a/python/sglang/jit_kernel/benchmark/bench_qwen_image_modulation.py b/python/sglang/jit_kernel/benchmark/bench_qwen_image_modulation.py index d16f9eda7..e0bc7dee0 100644 --- a/python/sglang/jit_kernel/benchmark/bench_qwen_image_modulation.py +++ b/python/sglang/jit_kernel/benchmark/bench_qwen_image_modulation.py @@ -3,13 +3,14 @@ from typing import Tuple import torch import triton.testing -from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark_no_cudagraph +from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph from sglang.jit_kernel.diffusion.triton.norm import norm_infer from sglang.jit_kernel.diffusion.triton.scale_shift import ( fuse_layernorm_scale_shift_gate_select01_kernel, fuse_residual_layernorm_scale_shift_gate_select01_kernel, fuse_scale_shift_gate_select01_kernel, ) +from sglang.utils import is_in_ci if is_in_ci(): B_RANGE, S_RANGE, D_RANGE = [1], [128], [3072] diff --git a/python/sglang/jit_kernel/benchmark/bench_renorm.py b/python/sglang/jit_kernel/benchmark/bench_renorm.py index 9c654ab7f..20cc9ae71 100644 --- a/python/sglang/jit_kernel/benchmark/bench_renorm.py +++ b/python/sglang/jit_kernel/benchmark/bench_renorm.py @@ -5,7 +5,8 @@ import torch import triton import triton.testing -from sglang.jit_kernel.benchmark.utils import is_in_ci, run_benchmark_no_cudagraph +from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph +from sglang.utils import is_in_ci def torch_top_k_renorm_probs(probs, top_k): diff --git a/python/sglang/jit_kernel/benchmark/utils.py b/python/sglang/jit_kernel/benchmark/utils.py index 079a4489b..c17ef4f9a 100644 --- a/python/sglang/jit_kernel/benchmark/utils.py +++ b/python/sglang/jit_kernel/benchmark/utils.py @@ -5,7 +5,7 @@ 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 +from sglang.utils import is_in_ci # Common constants DEFAULT_DTYPE = torch.bfloat16 @@ -13,11 +13,6 @@ DEFAULT_DEVICE = "cuda" DEFAULT_QUANTILES = [0.5, 0.2, 0.8] -def is_in_ci() -> bool: - """Check if running in CI environment.""" - return jit_kernel_is_in_ci() - - def get_benchmark_range(full_range: List, ci_range: List) -> List: """Return appropriate benchmark range based on CI environment.""" return ci_range if is_in_ci() else full_range diff --git a/python/sglang/jit_kernel/utils.py b/python/sglang/jit_kernel/utils.py index eb4748d71..fdcbc8e8a 100644 --- a/python/sglang/jit_kernel/utils.py +++ b/python/sglang/jit_kernel/utils.py @@ -7,6 +7,8 @@ from typing import TYPE_CHECKING, Any, Callable, List, Tuple, TypeAlias, TypeVar import torch +from sglang.utils import is_in_ci + if TYPE_CHECKING: from tvm_ffi import Module @@ -14,11 +16,6 @@ F = TypeVar("F", bound=Callable[..., Any]) _FULL_TEST_ENV_VAR = "SGLANG_JIT_KERNEL_RUN_FULL_TESTS" -def is_in_ci() -> bool: - ci_env_vars = ("SGLANG_IS_IN_CI", "CI", "GITHUB_ACTIONS") - return any(os.getenv(env_var, "false").lower() == "true" for env_var in ci_env_vars) - - def should_run_full_tests() -> bool: return os.getenv(_FULL_TEST_ENV_VAR, "false").lower() == "true" diff --git a/sgl-kernel/benchmark/bench_activation.py b/sgl-kernel/benchmark/bench_activation.py index 3caa5b936..43ec8d14d 100644 --- a/sgl-kernel/benchmark/bench_activation.py +++ b/sgl-kernel/benchmark/bench_activation.py @@ -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: diff --git a/sgl-kernel/benchmark/bench_awq_dequant.py b/sgl-kernel/benchmark/bench_awq_dequant.py index 6bd03ab8a..cb22ba07f 100644 --- a/sgl-kernel/benchmark/bench_awq_dequant.py +++ b/sgl-kernel/benchmark/bench_awq_dequant.py @@ -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( diff --git a/sgl-kernel/benchmark/bench_cutlass_mla.py b/sgl-kernel/benchmark/bench_cutlass_mla.py index 6947f309d..461d8862b 100644 --- a/sgl-kernel/benchmark/bench_cutlass_mla.py +++ b/sgl-kernel/benchmark/bench_cutlass_mla.py @@ -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: diff --git a/sgl-kernel/benchmark/bench_dsv3_fused_a_gemm.py b/sgl-kernel/benchmark/bench_dsv3_fused_a_gemm.py index bdf7f85de..43f0961f7 100644 --- a/sgl-kernel/benchmark/bench_dsv3_fused_a_gemm.py +++ b/sgl-kernel/benchmark/bench_dsv3_fused_a_gemm.py @@ -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 diff --git a/sgl-kernel/benchmark/bench_dsv3_router_gemm.py b/sgl-kernel/benchmark/bench_dsv3_router_gemm.py index 2daee279f..85ad88529 100644 --- a/sgl-kernel/benchmark/bench_dsv3_router_gemm.py +++ b/sgl-kernel/benchmark/bench_dsv3_router_gemm.py @@ -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 diff --git a/sgl-kernel/benchmark/bench_fp4_gemm.py b/sgl-kernel/benchmark/bench_fp4_gemm.py index 278b965c2..d5d358611 100755 --- a/sgl-kernel/benchmark/bench_fp4_gemm.py +++ b/sgl-kernel/benchmark/bench_fp4_gemm.py @@ -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 diff --git a/sgl-kernel/benchmark/bench_fp8_blockwise_gemm.py b/sgl-kernel/benchmark/bench_fp8_blockwise_gemm.py index 70766df94..29bd0a80a 100644 --- a/sgl-kernel/benchmark/bench_fp8_blockwise_gemm.py +++ b/sgl-kernel/benchmark/bench_fp8_blockwise_gemm.py @@ -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): diff --git a/sgl-kernel/benchmark/bench_fp8_blockwise_group_gemm.py b/sgl-kernel/benchmark/bench_fp8_blockwise_group_gemm.py index 19e425b52..063aca527 100644 --- a/sgl-kernel/benchmark/bench_fp8_blockwise_group_gemm.py +++ b/sgl-kernel/benchmark/bench_fp8_blockwise_group_gemm.py @@ -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 diff --git a/sgl-kernel/benchmark/bench_fp8_gemm.py b/sgl-kernel/benchmark/bench_fp8_gemm.py index 1dbedb66a..e37653d51 100644 --- a/sgl-kernel/benchmark/bench_fp8_gemm.py +++ b/sgl-kernel/benchmark/bench_fp8_gemm.py @@ -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) diff --git a/sgl-kernel/benchmark/bench_int8_gemm.py b/sgl-kernel/benchmark/bench_int8_gemm.py index 95f0f3bb8..722d89d7c 100644 --- a/sgl-kernel/benchmark/bench_int8_gemm.py +++ b/sgl-kernel/benchmark/bench_int8_gemm.py @@ -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: diff --git a/sgl-kernel/benchmark/bench_kimi_k2_moe_fused_gate.py b/sgl-kernel/benchmark/bench_kimi_k2_moe_fused_gate.py index 78b75231c..975a67a50 100644 --- a/sgl-kernel/benchmark/bench_kimi_k2_moe_fused_gate.py +++ b/sgl-kernel/benchmark/bench_kimi_k2_moe_fused_gate.py @@ -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): diff --git a/sgl-kernel/benchmark/bench_moe_align_block_size.py b/sgl-kernel/benchmark/bench_moe_align_block_size.py index 5dc4a0e4d..3debde2f4 100644 --- a/sgl-kernel/benchmark/bench_moe_align_block_size.py +++ b/sgl-kernel/benchmark/bench_moe_align_block_size.py @@ -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 diff --git a/sgl-kernel/benchmark/bench_moe_ep_post_reorder.py b/sgl-kernel/benchmark/bench_moe_ep_post_reorder.py index 2a617d72d..c3d9415a4 100644 --- a/sgl-kernel/benchmark/bench_moe_ep_post_reorder.py +++ b/sgl-kernel/benchmark/bench_moe_ep_post_reorder.py @@ -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: diff --git a/sgl-kernel/benchmark/bench_moe_fused_gate.py b/sgl-kernel/benchmark/bench_moe_fused_gate.py index cb5ac1760..508dcf1a9 100644 --- a/sgl-kernel/benchmark/bench_moe_fused_gate.py +++ b/sgl-kernel/benchmark/bench_moe_fused_gate.py @@ -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): diff --git a/sgl-kernel/benchmark/bench_moe_topk_sigmoid.py b/sgl-kernel/benchmark/bench_moe_topk_sigmoid.py index d34e68b98..cbc80a607 100644 --- a/sgl-kernel/benchmark/bench_moe_topk_sigmoid.py +++ b/sgl-kernel/benchmark/bench_moe_topk_sigmoid.py @@ -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( diff --git a/sgl-kernel/benchmark/bench_moe_topk_softmax.py b/sgl-kernel/benchmark/bench_moe_topk_softmax.py index e065981b8..4b2222405 100644 --- a/sgl-kernel/benchmark/bench_moe_topk_softmax.py +++ b/sgl-kernel/benchmark/bench_moe_topk_softmax.py @@ -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): diff --git a/sgl-kernel/benchmark/bench_per_tensor_quant_fp8.py b/sgl-kernel/benchmark/bench_per_tensor_quant_fp8.py index 43c2caddc..107123465 100644 --- a/sgl-kernel/benchmark/bench_per_tensor_quant_fp8.py +++ b/sgl-kernel/benchmark/bench_per_tensor_quant_fp8.py @@ -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 diff --git a/sgl-kernel/benchmark/bench_per_token_group_quant_8bit.py b/sgl-kernel/benchmark/bench_per_token_group_quant_8bit.py index 5ad6f9fc2..e4a8e7bbd 100644 --- a/sgl-kernel/benchmark/bench_per_token_group_quant_8bit.py +++ b/sgl-kernel/benchmark/bench_per_token_group_quant_8bit.py @@ -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 diff --git a/sgl-kernel/benchmark/bench_per_token_quant_fp8.py b/sgl-kernel/benchmark/bench_per_token_quant_fp8.py index 8db1869d1..dca014676 100644 --- a/sgl-kernel/benchmark/bench_per_token_quant_fp8.py +++ b/sgl-kernel/benchmark/bench_per_token_quant_fp8.py @@ -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 diff --git a/sgl-kernel/benchmark/bench_qserve_w4a8_gemm.py b/sgl-kernel/benchmark/bench_qserve_w4a8_gemm.py index 5827fa993..acef80f45 100644 --- a/sgl-kernel/benchmark/bench_qserve_w4a8_gemm.py +++ b/sgl-kernel/benchmark/bench_qserve_w4a8_gemm.py @@ -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: diff --git a/sgl-kernel/benchmark/bench_rmsnorm.py b/sgl-kernel/benchmark/bench_rmsnorm.py index d521ab05f..953b79792 100644 --- a/sgl-kernel/benchmark/bench_rmsnorm.py +++ b/sgl-kernel/benchmark/bench_rmsnorm.py @@ -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]: diff --git a/sgl-kernel/benchmark/bench_rotary_embedding.py b/sgl-kernel/benchmark/bench_rotary_embedding.py index 30614355f..bd3a900c7 100644 --- a/sgl-kernel/benchmark/bench_rotary_embedding.py +++ b/sgl-kernel/benchmark/bench_rotary_embedding.py @@ -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: diff --git a/sgl-kernel/benchmark/bench_sum_scale.py b/sgl-kernel/benchmark/bench_sum_scale.py index ad9621ee1..1c25ac151 100644 --- a/sgl-kernel/benchmark/bench_sum_scale.py +++ b/sgl-kernel/benchmark/bench_sum_scale.py @@ -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 diff --git a/sgl-kernel/benchmark/bench_top_k_top_p_sampling.py b/sgl-kernel/benchmark/bench_top_k_top_p_sampling.py index d84a34884..78a950362 100644 --- a/sgl-kernel/benchmark/bench_top_k_top_p_sampling.py +++ b/sgl-kernel/benchmark/bench_top_k_top_p_sampling.py @@ -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( diff --git a/test/registered/test_srt_backend.py b/test/registered/language/test_srt_backend.py similarity index 100% rename from test/registered/test_srt_backend.py rename to test/registered/language/test_srt_backend.py diff --git a/test/registered/test_hybrid_dp_ep_tp_mtp.py b/test/registered/moe/test_hybrid_dp_ep_tp_mtp.py similarity index 100% rename from test/registered/test_hybrid_dp_ep_tp_mtp.py rename to test/registered/moe/test_hybrid_dp_ep_tp_mtp.py