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

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
+2 -5
View File
@@ -244,18 +244,15 @@ Create `sgl-kernel/benchmark/bench_scale.py`:
```python ```python
import itertools import itertools
import os
import torch import torch
import triton import triton
import triton.testing import triton.testing
import sgl_kernel import sgl_kernel
from sglang.utils import is_in_ci
IS_CI = ( IS_CI = is_in_ci()
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
dtypes = [torch.float16] if IS_CI else [torch.float16, torch.bfloat16, torch.float32] 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 sizes = [4096] if IS_CI else [2**n for n in range(10, 20)] # 1K … 512K
@@ -44,12 +44,9 @@ from sglang.srt.distributed.parallel_state import (
initialize_model_parallel, initialize_model_parallel,
set_torch_symm_mem_all_reduce, set_torch_symm_mem_all_reduce,
) )
from sglang.utils import is_in_ci
# CI environment detection IS_CI = is_in_ci()
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
def torch_allreduce(torch_input: torch.Tensor, group: ProcessGroup) -> torch.Tensor: def torch_allreduce(torch_input: torch.Tensor, group: ProcessGroup) -> torch.Tensor:
@@ -5,7 +5,8 @@ import triton
import triton.testing import triton.testing
from sglang.jit_kernel.awq_dequantize import awq_dequantize as jit_awq_dequantize 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: try:
from sgl_kernel import awq_dequantize as aot_awq_dequantize from sgl_kernel import awq_dequantize as aot_awq_dequantize
@@ -7,8 +7,9 @@ from sgl_kernel.scalar_type import scalar_types
from sglang.jit_kernel.awq_marlin_repack import ( from sglang.jit_kernel.awq_marlin_repack import (
awq_marlin_moe_repack as jit_awq_marlin_moe_repack, 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.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( AOT_AVAILABLE = hasattr(torch.ops.sgl_kernel, "awq_marlin_moe_repack") and hasattr(
torch.ops.sgl_kernel.awq_marlin_moe_repack, "default" torch.ops.sgl_kernel.awq_marlin_moe_repack, "default"
@@ -7,8 +7,9 @@ from sgl_kernel.scalar_type import scalar_types
from sglang.jit_kernel.awq_marlin_repack import ( from sglang.jit_kernel.awq_marlin_repack import (
awq_marlin_repack as jit_awq_marlin_repack, 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.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( AOT_AVAILABLE = hasattr(torch.ops.sgl_kernel, "awq_marlin_repack") and hasattr(
torch.ops.sgl_kernel.awq_marlin_repack, "default" torch.ops.sgl_kernel.awq_marlin_repack, "default"
@@ -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_absorb_q as aot_absorb_q
from sgl_kernel import concat_mla_k as aot_k 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_absorb_q as jit_absorb_q
from sglang.jit_kernel.concat_mla import concat_mla_k as jit_k 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() IS_CI = is_in_ci()
@@ -5,8 +5,9 @@ import triton
import triton.testing import triton.testing
from flashinfer import fused_add_rmsnorm as fi_fused_add_rmsnorm 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.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm
from sglang.utils import is_in_ci
IS_CI = is_in_ci() IS_CI = is_in_ci()
@@ -8,13 +8,14 @@ import torch
import triton import triton
import triton.testing 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 ( from sglang.multimodal_gen.runtime.layers.layernorm import (
LayerNormScaleShift, LayerNormScaleShift,
RMSNormScaleShift, RMSNormScaleShift,
ScaleResidualLayerNormScaleShift, ScaleResidualLayerNormScaleShift,
ScaleResidualRMSNormScaleShift, ScaleResidualRMSNormScaleShift,
) )
from sglang.utils import is_in_ci
if is_in_ci(): if is_in_ci():
B_RANGE, S_RANGE, D_RANGE = [1], [128], [1024] B_RANGE, S_RANGE, D_RANGE = [1], [128], [1024]
@@ -3,10 +3,11 @@ import triton
import triton.testing import triton.testing
from sgl_kernel.scalar_type import scalar_types 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.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.srt.layers.quantization.marlin_utils import marlin_make_workspace
from sglang.test.test_marlin_utils import marlin_quantize 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( AOT_AVAILABLE = hasattr(torch.ops.sgl_kernel, "gptq_marlin_gemm") and hasattr(
torch.ops.sgl_kernel.gptq_marlin_gemm, "default" torch.ops.sgl_kernel.gptq_marlin_gemm, "default"
@@ -3,9 +3,10 @@ import triton
import triton.testing import triton.testing
from sgl_kernel.scalar_type import scalar_types 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.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.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( AOT_AVAILABLE = hasattr(torch.ops.sgl_kernel, "gptq_marlin_repack") and hasattr(
torch.ops.sgl_kernel.gptq_marlin_repack, "default" torch.ops.sgl_kernel.gptq_marlin_repack, "default"
@@ -3,10 +3,11 @@ import triton
import triton.testing import triton.testing
from sgl_kernel.scalar_type import scalar_types 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.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.srt.layers.moe.fused_moe_triton import moe_align_block_size
from sglang.test.test_marlin_utils import marlin_quantize 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( AOT_AVAILABLE = hasattr(torch.ops.sgl_kernel, "moe_wna16_marlin_gemm") and hasattr(
torch.ops.sgl_kernel.moe_wna16_marlin_gemm, "default" torch.ops.sgl_kernel.moe_wna16_marlin_gemm, "default"
@@ -6,9 +6,10 @@ import triton.testing
from flashinfer.norm import fused_add_rmsnorm as fi_fused_add_rmsnorm from flashinfer.norm import fused_add_rmsnorm as fi_fused_add_rmsnorm
from flashinfer.norm import rmsnorm as fi_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 fused_add_rmsnorm as jit_fused_add_rmsnorm
from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm
from sglang.utils import is_in_ci
IS_CI = is_in_ci() IS_CI = is_in_ci()
@@ -15,12 +15,13 @@ from typing import Callable
import torch import torch
import torch.nn.functional as F 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.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.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 fused_add_rmsnorm as jit_fused_add_rmsnorm
from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm from sglang.jit_kernel.norm import rmsnorm as jit_rmsnorm
from sglang.jit_kernel.utils import KERNEL_PATH from sglang.jit_kernel.utils import KERNEL_PATH
from sglang.utils import is_in_ci
os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1") os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1")
@@ -5,10 +5,7 @@ import torch
import triton import triton
from sgl_kernel.test_utils import create_per_token_group_quant_test_data from sgl_kernel.test_utils import create_per_token_group_quant_test_data
from sglang.jit_kernel.benchmark.utils import ( from sglang.jit_kernel.benchmark.utils import get_benchmark_range
get_benchmark_range,
is_in_ci,
)
from sglang.jit_kernel.per_token_group_quant_8bit import ( from sglang.jit_kernel.per_token_group_quant_8bit import (
per_token_group_quant_8bit as sglang_per_token_group_quant_8bit, 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 import is_hip
from sglang.srt.utils.bench_utils import bench_kineto from sglang.srt.utils.bench_utils import bench_kineto
from sglang.utils import is_in_ci
IS_CI = is_in_ci() IS_CI = is_in_ci()
@@ -6,9 +6,10 @@ import triton
import triton.testing import triton.testing
from sgl_kernel import rmsnorm 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.jit_kernel.norm import fused_inplace_qknorm_across_heads
from sglang.srt.utils import get_current_device_stream_fast from sglang.srt.utils import get_current_device_stream_fast
from sglang.utils import is_in_ci
IS_CI = is_in_ci() IS_CI = is_in_ci()
@@ -3,13 +3,14 @@ from typing import Tuple
import torch import torch
import triton.testing 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.norm import norm_infer
from sglang.jit_kernel.diffusion.triton.scale_shift import ( from sglang.jit_kernel.diffusion.triton.scale_shift import (
fuse_layernorm_scale_shift_gate_select01_kernel, fuse_layernorm_scale_shift_gate_select01_kernel,
fuse_residual_layernorm_scale_shift_gate_select01_kernel, fuse_residual_layernorm_scale_shift_gate_select01_kernel,
fuse_scale_shift_gate_select01_kernel, fuse_scale_shift_gate_select01_kernel,
) )
from sglang.utils import is_in_ci
if is_in_ci(): if is_in_ci():
B_RANGE, S_RANGE, D_RANGE = [1], [128], [3072] B_RANGE, S_RANGE, D_RANGE = [1], [128], [3072]
@@ -5,7 +5,8 @@ import torch
import triton import triton
import triton.testing 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): def torch_top_k_renorm_probs(probs, top_k):
+1 -6
View File
@@ -5,7 +5,7 @@ from typing import Callable, List, Tuple
import torch import torch
import triton.testing 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 # Common constants
DEFAULT_DTYPE = torch.bfloat16 DEFAULT_DTYPE = torch.bfloat16
@@ -13,11 +13,6 @@ DEFAULT_DEVICE = "cuda"
DEFAULT_QUANTILES = [0.5, 0.2, 0.8] 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: def get_benchmark_range(full_range: List, ci_range: List) -> List:
"""Return appropriate benchmark range based on CI environment.""" """Return appropriate benchmark range based on CI environment."""
return ci_range if is_in_ci() else full_range return ci_range if is_in_ci() else full_range
+2 -5
View File
@@ -7,6 +7,8 @@ from typing import TYPE_CHECKING, Any, Callable, List, Tuple, TypeAlias, TypeVar
import torch import torch
from sglang.utils import is_in_ci
if TYPE_CHECKING: if TYPE_CHECKING:
from tvm_ffi import Module 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" _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: def should_run_full_tests() -> bool:
return os.getenv(_FULL_TEST_ENV_VAR, "false").lower() == "true" return os.getenv(_FULL_TEST_ENV_VAR, "false").lower() == "true"
+3 -5
View File
@@ -13,6 +13,8 @@ import triton
import triton.testing import triton.testing
from sgl_kernel import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul from sgl_kernel import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul
from sglang.utils import is_in_ci
# Optional vLLM import # Optional vLLM import
try: try:
from vllm import _custom_ops as vllm_ops from vllm import _custom_ops as vllm_ops
@@ -22,11 +24,7 @@ except ImportError:
vllm_ops = None vllm_ops = None
VLLM_AVAILABLE = False VLLM_AVAILABLE = False
# CI environment detection IS_CI = is_in_ci()
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
# gelu_quick is only available on HIP/ROCm platforms # gelu_quick is only available on HIP/ROCm platforms
try: try:
+3 -5
View File
@@ -7,6 +7,8 @@ import triton
import triton.testing import triton.testing
from sgl_kernel import awq_dequantize from sgl_kernel import awq_dequantize
from sglang.utils import is_in_ci
# Optional vLLM import # Optional vLLM import
try: try:
from vllm import _custom_ops as ops from vllm import _custom_ops as ops
@@ -16,11 +18,7 @@ except ImportError:
ops = None ops = None
VLLM_AVAILABLE = False VLLM_AVAILABLE = False
# CI environment detection IS_CI = is_in_ci()
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
def vllm_awq_dequantize( def vllm_awq_dequantize(
+2 -5
View File
@@ -8,12 +8,9 @@ import triton
from sgl_kernel import cutlass_mla_decode, cutlass_mla_get_workspace_size from sgl_kernel import cutlass_mla_decode, cutlass_mla_get_workspace_size
from sglang.srt.utils import get_device_capability from sglang.srt.utils import get_device_capability
from sglang.utils import is_in_ci
# CI environment detection IS_CI = is_in_ci()
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
# CI environment uses simplified parameters # CI environment uses simplified parameters
if IS_CI: if IS_CI:
@@ -1,11 +1,4 @@
import argparse 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
import torch.nn.functional as F import torch.nn.functional as F
@@ -13,6 +6,10 @@ import triton
import triton.testing import triton.testing
from sgl_kernel import dsv3_fused_a_gemm 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 # CI environment uses simplified parameters
if IS_CI: if IS_CI:
num_tokens_vals = [1] # Only test 1 value in CI num_tokens_vals = [1] # Only test 1 value in CI
@@ -1,11 +1,4 @@
import argparse 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
import torch.nn.functional as F import torch.nn.functional as F
@@ -13,6 +6,10 @@ import triton
import triton.testing import triton.testing
from sgl_kernel import dsv3_router_gemm from sgl_kernel import dsv3_router_gemm
from sglang.utils import is_in_ci
IS_CI = is_in_ci()
# CI environment uses simplified parameters # CI environment uses simplified parameters
if IS_CI: if IS_CI:
num_tokens_vals = [1] # Only test 1 value in CI num_tokens_vals = [1] # Only test 1 value in CI
+2 -5
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 sgl_kernel import cutlass_scaled_fp4_mm, scaled_fp4_quant
from sglang.srt.utils import get_device_capability, is_sm100_supported from sglang.srt.utils import get_device_capability, is_sm100_supported
from sglang.utils import is_in_ci
# CI environment detection IS_CI = is_in_ci()
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
FLOAT4_E2M1_MAX = 6.0 FLOAT4_E2M1_MAX = 6.0
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
@@ -9,6 +9,8 @@ import triton
from deep_gemm.utils.layout import get_mn_major_tma_aligned_tensor from deep_gemm.utils.layout import get_mn_major_tma_aligned_tensor
from sgl_kernel import fp8_blockwise_scaled_mm from sgl_kernel import fp8_blockwise_scaled_mm
from sglang.utils import is_in_ci
# Optional vLLM import # Optional vLLM import
try: try:
from vllm._custom_ops import cutlass_scaled_mm as vllm_scaled_mm 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, w8a8_block_fp8_matmul_triton as w8a8_block_fp8_matmul,
) )
# CI environment detection IS_CI = is_in_ci()
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
def get_weight_shapes(args): def get_weight_shapes(args):
@@ -1,11 +1,4 @@
import argparse 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 import random
from dataclasses import dataclass from dataclasses import dataclass
from typing import List, Tuple from typing import List, Tuple
@@ -14,6 +7,10 @@ import deep_gemm
import torch import torch
from sgl_kernel import fp8_blockwise_scaled_grouped_mm 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(): def get_m_alignment_for_contiguous_layout():
return 128 return 128
+2 -5
View File
@@ -9,6 +9,7 @@ import triton
from sgl_kernel import fp8_scaled_mm as sgl_scaled_mm 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.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
from sglang.utils import is_in_ci
# Optional vLLM import # Optional vLLM import
try: try:
@@ -21,11 +22,7 @@ except ImportError:
vllm_scaled_fp8_quant = None vllm_scaled_fp8_quant = None
VLLM_AVAILABLE = False VLLM_AVAILABLE = False
# CI environment detection IS_CI = is_in_ci()
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
# Weight Shapes are in the format # Weight Shapes are in the format
# ([K, N], TP_SPLIT_DIM) # ([K, N], TP_SPLIT_DIM)
+3 -5
View File
@@ -7,6 +7,8 @@ import torch
import triton import triton
from sgl_kernel import int8_scaled_mm from sgl_kernel import int8_scaled_mm
from sglang.utils import is_in_ci
# Optional vLLM import # Optional vLLM import
try: try:
from vllm._custom_ops import cutlass_scaled_mm as vllm_scaled_mm from vllm._custom_ops import cutlass_scaled_mm as vllm_scaled_mm
@@ -16,11 +18,7 @@ except ImportError:
vllm_scaled_mm = None vllm_scaled_mm = None
VLLM_AVAILABLE = False VLLM_AVAILABLE = False
# CI environment detection IS_CI = is_in_ci()
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
def to_int8(tensor: torch.Tensor) -> torch.Tensor: def to_int8(tensor: torch.Tensor) -> torch.Tensor:
@@ -8,12 +8,9 @@ import triton.language as tl
from sgl_kernel import kimi_k2_moe_fused_gate from sgl_kernel import kimi_k2_moe_fused_gate
from sglang.srt.layers.moe.topk import kimi_k2_biased_topk_impl from sglang.srt.layers.moe.topk import kimi_k2_biased_topk_impl
from sglang.utils import is_in_ci
# CI environment detection IS_CI = is_in_ci()
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
def kimi_k2_biased_topk_torch_compile(scores, bias, topk, routed_scaling_factor): def kimi_k2_biased_topk_torch_compile(scores, bias, topk, routed_scaling_factor):
@@ -7,6 +7,8 @@ import triton
import triton.language as tl import triton.language as tl
from sgl_kernel import moe_align_block_size as sgl_moe_align_block_size from sgl_kernel import moe_align_block_size as sgl_moe_align_block_size
from sglang.utils import is_in_ci
try: try:
from vllm import _custom_ops as ops from vllm import _custom_ops as ops
@@ -15,11 +17,7 @@ except ImportError:
ops = None ops = None
VLLM_AVAILABLE = False VLLM_AVAILABLE = False
# CI environment detection IS_CI = is_in_ci()
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
USE_RANDOM_PERM = False USE_RANDOM_PERM = False
@@ -1,15 +1,10 @@
import os
import torch import torch
# CI environment detection
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
import triton import triton
from sglang.srt.layers.moe.ep_moe.kernels import post_reorder_triton_kernel 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 # CI environment uses simplified parameters
if IS_CI: if IS_CI:
+2 -5
View File
@@ -8,12 +8,9 @@ import triton.language as tl
from sgl_kernel import moe_fused_gate from sgl_kernel import moe_fused_gate
from sglang.srt.layers.moe.topk import biased_grouped_topk from sglang.srt.layers.moe.topk import biased_grouped_topk
from sglang.utils import is_in_ci
# CI environment detection IS_CI = is_in_ci()
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
def biased_grouped_topk_org(scores, bias, num_expert_group, topk_group, topk): def biased_grouped_topk_org(scores, bias, num_expert_group, topk_group, topk):
@@ -6,11 +6,9 @@ import torch
import triton import triton
from sgl_kernel import topk_sigmoid from sgl_kernel import topk_sigmoid
# CI environment detection from sglang.utils import is_in_ci
IS_CI = (
os.getenv("CI", "false").lower() == "true" IS_CI = is_in_ci()
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
def torch_topk_sigmoid_native( def torch_topk_sigmoid_native(
@@ -6,6 +6,8 @@ import torch
import triton import triton
from sgl_kernel import topk_softmax from sgl_kernel import topk_softmax
from sglang.utils import is_in_ci
# Optional vLLM import # Optional vLLM import
try: try:
from vllm import _custom_ops as vllm_custom_ops from vllm import _custom_ops as vllm_custom_ops
@@ -15,11 +17,7 @@ except ImportError:
vllm_custom_ops = None vllm_custom_ops = None
VLLM_AVAILABLE = False VLLM_AVAILABLE = False
# CI environment detection IS_CI = is_in_ci()
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
def vllm_topk_softmax(gating_output, topk): def vllm_topk_softmax(gating_output, topk):
@@ -9,6 +9,7 @@ import triton
import triton.testing import triton.testing
from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8 from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
from sglang.utils import is_in_ci
# Optional imports # Optional imports
try: try:
@@ -23,11 +24,7 @@ from sglang.srt.utils import is_hip
_is_hip = is_hip() _is_hip = is_hip()
# CI environment detection IS_CI = is_in_ci()
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
@@ -19,12 +19,9 @@ from sglang.srt.layers.quantization.fp8_kernel import (
) )
from sglang.srt.utils import is_hip from sglang.srt.utils import is_hip
from sglang.srt.utils.bench_utils import bench_kineto from sglang.srt.utils.bench_utils import bench_kineto
from sglang.utils import is_in_ci
# CI environment detection IS_CI = is_in_ci()
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
_is_hip = is_hip() _is_hip = is_hip()
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
@@ -7,6 +7,8 @@ import triton
import triton.testing import triton.testing
from sgl_kernel import sgl_per_token_quant_fp8 from sgl_kernel import sgl_per_token_quant_fp8
from sglang.utils import is_in_ci
# Optional vLLM import # Optional vLLM import
try: try:
from vllm import _custom_ops as ops from vllm import _custom_ops as ops
@@ -20,11 +22,7 @@ from sglang.srt.utils import is_hip
_is_hip = is_hip() _is_hip = is_hip()
# CI environment detection IS_CI = is_in_ci()
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn
@@ -11,11 +11,9 @@ from sgl_kernel import (
qserve_w4a8_per_group_gemm, qserve_w4a8_per_group_gemm,
) )
# CI environment detection from sglang.utils import is_in_ci
IS_CI = (
os.getenv("CI", "false").lower() == "true" IS_CI = is_in_ci()
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
def to_int8(tensor: torch.Tensor) -> torch.Tensor: def to_int8(tensor: torch.Tensor) -> torch.Tensor:
+3 -5
View File
@@ -13,6 +13,8 @@ import triton
import triton.testing import triton.testing
from sgl_kernel.utils import is_arch_support_pdl from sgl_kernel.utils import is_arch_support_pdl
from sglang.utils import is_in_ci
# Optional imports # Optional imports
try: try:
from flashinfer.norm import fused_add_rmsnorm, rmsnorm from flashinfer.norm import fused_add_rmsnorm, rmsnorm
@@ -31,11 +33,7 @@ except ImportError:
vllm_ops = None vllm_ops = None
VLLM_AVAILABLE = False VLLM_AVAILABLE = False
# CI environment detection IS_CI = is_in_ci()
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
def str2int_list(arg: str) -> List[int]: def str2int_list(arg: str) -> List[int]:
@@ -12,12 +12,9 @@ from sgl_kernel.testing.rotary_embedding import (
) )
from sglang.srt.utils.bench_utils import bench_kineto from sglang.srt.utils.bench_utils import bench_kineto
from sglang.utils import is_in_ci
# CI environment detection IS_CI = is_in_ci()
IS_CI = (
os.getenv("CI", "false").lower() == "true"
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
# CI environment uses simplified parameters # CI environment uses simplified parameters
if IS_CI: if IS_CI:
+3 -5
View File
@@ -6,11 +6,9 @@ import triton.language as tl
from sgl_kernel import moe_sum_reduce as moe_sum_reduce_cuda from sgl_kernel import moe_sum_reduce as moe_sum_reduce_cuda
from triton.testing import do_bench from triton.testing import do_bench
# CI environment detection from sglang.utils import is_in_ci
IS_CI = (
os.getenv("CI", "false").lower() == "true" IS_CI = is_in_ci()
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
@triton.jit @triton.jit
@@ -7,11 +7,9 @@ import torch
import triton import triton
import triton.testing import triton.testing
# CI environment detection from sglang.utils import is_in_ci
IS_CI = (
os.getenv("CI", "false").lower() == "true" IS_CI = is_in_ci()
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
)
def torch_top_k_top_p_joint_sampling_from_probs( def torch_top_k_top_p_joint_sampling_from_probs(