diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 6f0f5f73f..29559a0b7 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -831,6 +831,79 @@ jobs: cd python/sglang/jit_kernel pytest tests/ + jit-kernel-unit-test-nightly: + needs: [check-changes] + if: | + github.event_name == 'schedule' && + needs.check-changes.outputs.jit_kernel == 'true' + runs-on: 1-gpu-runner + timeout-minutes: 240 + env: + RUNNER_LABELS: 1-gpu-runner + SGLANG_JIT_KERNEL_RUN_FULL_TESTS: "1" + steps: + - uses: actions/checkout@v4 + with: + ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} + + - name: Install dependencies + timeout-minutes: 20 + run: | + bash scripts/ci/cuda/ci_install_dependency.sh + + - name: Run full nightly test + timeout-minutes: 60 + run: | + cd python/sglang/jit_kernel + pytest tests/ + + jit-kernel-benchmark-test: + needs: [check-changes, call-gate] + # Skip for scheduled runs and when target_stage is set + if: | + github.event_name != 'schedule' && + inputs.test_parallel_dispatch != true && + !inputs.target_stage && + needs.check-changes.outputs.jit_kernel == 'true' + runs-on: 1-gpu-runner + timeout-minutes: 240 + env: + CI: true + RUNNER_LABELS: 1-gpu-runner + steps: + - uses: actions/checkout@v4 + with: + ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }} + + - name: Install dependencies + timeout-minutes: 20 + run: | + bash scripts/ci/cuda/ci_install_dependency.sh + + - name: Run benchmark smoke tests + timeout-minutes: 45 + run: | + cd python/sglang/jit_kernel/benchmark + echo "Running jit-kernel benchmark smoke tests in CI mode..." + + failures=() + + for bench_file in bench_*.py; do + echo "Testing $bench_file..." + if ! timeout 120 python3 "$bench_file"; then + failures+=("$bench_file") + fi + echo "Completed $bench_file" + echo "---" + done + + if [ ${#failures[@]} -ne 0 ]; then + echo "The following benchmark smoke tests failed: ${failures[*]}" + exit 1 + fi + + echo "All jit-kernel benchmark smoke tests completed successfully!" + # =============================================== primary ==================================================== stage-a-test-1: @@ -1695,6 +1768,8 @@ jobs: wait-for-stage-b, jit-kernel-unit-test, + jit-kernel-unit-test-nightly, + jit-kernel-benchmark-test, multimodal-gen-test-1-gpu, multimodal-gen-test-2-gpu, diff --git a/python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py b/python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py index 40c73f5df..a39027ea2 100644 --- a/python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py +++ b/python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py @@ -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] 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 f6cd8561e..d845e2bd0 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 @@ -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 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 b2f1d9e76..c6b7a7a59 100644 --- a/python/sglang/jit_kernel/benchmark/bench_awq_marlin_repack.py +++ b/python/sglang/jit_kernel/benchmark/bench_awq_marlin_repack.py @@ -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 diff --git a/python/sglang/jit_kernel/benchmark/bench_gptq_marlin.py b/python/sglang/jit_kernel/benchmark/bench_gptq_marlin.py index 77c1d61af..9ca606141 100644 --- a/python/sglang/jit_kernel/benchmark/bench_gptq_marlin.py +++ b/python/sglang/jit_kernel/benchmark/bench_gptq_marlin.py @@ -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 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 1253036f6..888e7fecb 100644 --- a/python/sglang/jit_kernel/benchmark/bench_gptq_marlin_repack.py +++ b/python/sglang/jit_kernel/benchmark/bench_gptq_marlin_repack.py @@ -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 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 5a0f480c8..d6d0649a5 100644 --- a/python/sglang/jit_kernel/benchmark/bench_moe_wna16_marlin.py +++ b/python/sglang/jit_kernel/benchmark/bench_moe_wna16_marlin.py @@ -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): diff --git a/python/sglang/jit_kernel/benchmark/bench_nvfp4_blockwise_moe.py b/python/sglang/jit_kernel/benchmark/bench_nvfp4_blockwise_moe.py index 54d37ddeb..3ae2e5cca 100644 --- a/python/sglang/jit_kernel/benchmark/bench_nvfp4_blockwise_moe.py +++ b/python/sglang/jit_kernel/benchmark/bench_nvfp4_blockwise_moe.py @@ -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}" diff --git a/python/sglang/jit_kernel/benchmark/bench_nvfp4_quant.py b/python/sglang/jit_kernel/benchmark/bench_nvfp4_quant.py index b1a7a24e6..e1a4eb7f5 100644 --- a/python/sglang/jit_kernel/benchmark/bench_nvfp4_quant.py +++ b/python/sglang/jit_kernel/benchmark/bench_nvfp4_quant.py @@ -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}" diff --git a/python/sglang/jit_kernel/benchmark/bench_nvfp4_scaled_mm.py b/python/sglang/jit_kernel/benchmark/bench_nvfp4_scaled_mm.py index 8c0bdc322..80d3b4f04 100644 --- a/python/sglang/jit_kernel/benchmark/bench_nvfp4_scaled_mm.py +++ b/python/sglang/jit_kernel/benchmark/bench_nvfp4_scaled_mm.py @@ -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}" 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 7edcf82c2..f0aadf145 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 @@ -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_] diff --git a/python/sglang/jit_kernel/benchmark/utils.py b/python/sglang/jit_kernel/benchmark/utils.py index 36af0e334..079a4489b 100644 --- a/python/sglang/jit_kernel/benchmark/utils.py +++ b/python/sglang/jit_kernel/benchmark/utils.py @@ -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: diff --git a/python/sglang/jit_kernel/tests/test_fused_add_rmsnorm.py b/python/sglang/jit_kernel/tests/test_fused_add_rmsnorm.py index 52c2dc612..fb17dcd39 100644 --- a/python/sglang/jit_kernel/tests/test_fused_add_rmsnorm.py +++ b/python/sglang/jit_kernel/tests/test_fused_add_rmsnorm.py @@ -3,6 +3,8 @@ import itertools import pytest import torch +from sglang.jit_kernel.utils import get_ci_test_range + def sglang_jit_fused_add_rmsnorm( input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float @@ -22,7 +24,11 @@ def flashinfer_fused_add_rmsnorm( BS_LIST = [2**n for n in range(0, 14)] BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)] -HIDDEN_SIZE_LIST = [512, 1024, 1536, 2048, 3072, 4096, 5120, 6144, 7168, 8192] +BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 4109]) +HIDDEN_SIZE_LIST = get_ci_test_range( + [512, 1024, 1536, 2048, 3072, 4096, 5120, 6144, 7168, 8192], + [512, 2048, 8192], +) DEVICE = "cuda" DTYPE = torch.bfloat16 diff --git a/python/sglang/jit_kernel/tests/test_qknorm.py b/python/sglang/jit_kernel/tests/test_qknorm.py index ee72e9ec6..72c574c71 100644 --- a/python/sglang/jit_kernel/tests/test_qknorm.py +++ b/python/sglang/jit_kernel/tests/test_qknorm.py @@ -4,6 +4,8 @@ import pytest import torch import triton +from sglang.jit_kernel.utils import get_ci_test_range + def sglang_aot_qknorm( q: torch.Tensor, @@ -61,9 +63,10 @@ def torch_impl_qknorm( BS_LIST = [2**n for n in range(0, 14)] BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)] -N_K_LIST = [2, 4] -N_Q_LIST = [8, 16] -HEAD_DIM_LIST = [64, 128, 256, 512, 1024] +BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 4109]) +N_K_LIST = get_ci_test_range([2, 4], [2, 4]) +N_Q_LIST = get_ci_test_range([8, 16], [8, 16]) +HEAD_DIM_LIST = get_ci_test_range([64, 128, 256, 512, 1024], [64, 256, 1024]) DEVICE = "cuda" DTYPE = torch.bfloat16 diff --git a/python/sglang/jit_kernel/tests/test_qknorm_across_heads.py b/python/sglang/jit_kernel/tests/test_qknorm_across_heads.py index f090d82ad..fb2123537 100644 --- a/python/sglang/jit_kernel/tests/test_qknorm_across_heads.py +++ b/python/sglang/jit_kernel/tests/test_qknorm_across_heads.py @@ -4,6 +4,8 @@ import pytest import torch import triton +from sglang.jit_kernel.utils import get_ci_test_range + def sglang_jit_qknorm_across_heads( q: torch.Tensor, @@ -46,7 +48,8 @@ def torch_impl_qknorm_across_heads( BS_LIST = [2**n for n in range(0, 14)] BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)] -HIDDEN_DIM_LIST = [512, 1024, 2048, 4096] +BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 4109]) +HIDDEN_DIM_LIST = get_ci_test_range([512, 1024, 2048, 4096], [512, 2048, 4096]) DEVICE = "cuda" DTYPE = torch.bfloat16 diff --git a/python/sglang/jit_kernel/tests/test_rmsnorm.py b/python/sglang/jit_kernel/tests/test_rmsnorm.py index 501124daf..d7bdb9543 100644 --- a/python/sglang/jit_kernel/tests/test_rmsnorm.py +++ b/python/sglang/jit_kernel/tests/test_rmsnorm.py @@ -4,6 +4,8 @@ import pytest import torch import triton +from sglang.jit_kernel.utils import get_ci_test_range + def sglang_jit_rmsnorm(input: torch.Tensor, weight: torch.Tensor) -> None: from sglang.jit_kernel.norm import rmsnorm @@ -19,7 +21,11 @@ def flashinfer_rmsnorm(input: torch.Tensor, weight: torch.Tensor) -> None: BS_LIST = [2**n for n in range(0, 14)] BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)] -HIDDEN_SIZE_LIST = [512, 1024, 1536, 2048, 3072, 4096, 5120, 6144, 7168, 8192] +BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 4109]) +HIDDEN_SIZE_LIST = get_ci_test_range( + [512, 1024, 1536, 2048, 3072, 4096, 5120, 6144, 7168, 8192], + [512, 2048, 8192], +) DEVICE = "cuda" DTYPE = torch.bfloat16 diff --git a/python/sglang/jit_kernel/tests/test_rope.py b/python/sglang/jit_kernel/tests/test_rope.py index 0aba5cf4c..fa2ae97b5 100644 --- a/python/sglang/jit_kernel/tests/test_rope.py +++ b/python/sglang/jit_kernel/tests/test_rope.py @@ -2,6 +2,8 @@ import pytest import torch import triton +from sglang.jit_kernel.utils import get_ci_test_range + DEVICE = "cuda" DTYPE = torch.bfloat16 MAX_SEQ_LEN = 131072 # common seq length @@ -87,11 +89,16 @@ def torch_impl_rope( BS_LIST = [2**x for x in range(12)] BS_LIST += [x + 1 for x in BS_LIST] # odd sizes to stress non-aligned paths -NUM_KV_HEADS_LIST = [1, 2, 8] -GQA_RATIO = [1, 4, 8] -ROPE_DIM_LIST = [64, 128, 256, 512] +BS_LIST = get_ci_test_range(BS_LIST, [1, 129, 2048, 2049]) +NUM_KV_HEADS_LIST = get_ci_test_range([1, 2, 8], [1, 8]) +GQA_RATIO = get_ci_test_range([1, 4, 8], [1, 8]) +ROPE_DIM_LIST = get_ci_test_range([64, 128, 256, 512], [64, 256]) IS_NEOX_LIST = [False, True] -DTYPE_LIST = [torch.bfloat16, torch.float16] +DTYPE_LIST = get_ci_test_range( + [torch.bfloat16, torch.float16], [torch.bfloat16, torch.float16] +) +PARTIAL_ROPE_DIM_LIST = get_ci_test_range([64, 80, 96, 128], [64, 96]) +HEAD_DIM_LIST = get_ci_test_range([64, 128, 256], [64, 256]) @pytest.mark.parametrize("batch_size", BS_LIST) @@ -151,8 +158,8 @@ def test_rope_position_dtypes(dtype: torch.dtype) -> None: @pytest.mark.parametrize("batch_size", BS_LIST) @pytest.mark.parametrize("is_neox", IS_NEOX_LIST) -@pytest.mark.parametrize("rope_dim", [64, 80, 96, 128]) -@pytest.mark.parametrize("head_dim", [64, 128, 256]) +@pytest.mark.parametrize("rope_dim", PARTIAL_ROPE_DIM_LIST) +@pytest.mark.parametrize("head_dim", HEAD_DIM_LIST) def test_partial_rope(batch_size: int, is_neox: bool, rope_dim: int, head_dim: int): if head_dim < rope_dim: pytest.skip("Invalid config: head_dim must be >= rope_dim.") diff --git a/python/sglang/jit_kernel/tests/test_store_cache.py b/python/sglang/jit_kernel/tests/test_store_cache.py index ea168d0b2..4f597a14a 100644 --- a/python/sglang/jit_kernel/tests/test_store_cache.py +++ b/python/sglang/jit_kernel/tests/test_store_cache.py @@ -4,10 +4,14 @@ import pytest import torch from sglang.jit_kernel.kvcache import can_use_store_cache, store_cache +from sglang.jit_kernel.utils import get_ci_test_range BS_LIST = [2**n for n in range(0, 15)] BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)] -HIDDEN_DIMS = [64, 128, 256, 512, 1024, 96, 98, 100] +BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 16399]) +HIDDEN_DIMS = get_ci_test_range( + [64, 128, 256, 512, 1024, 96, 98, 100], [64, 512, 1024, 98] +) CACHE_SIZE = 1024 * 1024 DTYPE = torch.bfloat16 DEVICE = "cuda" @@ -32,8 +36,8 @@ def test_store_cache(batch_size: int, element_dim: int) -> None: # Smaller subset for targeted tests below -REPR_BS = [1, 7, 128] -REPR_DIMS = [64, 128, 512, 1024, 96] +REPR_BS = get_ci_test_range([1, 7, 128], [1, 128]) +REPR_DIMS = get_ci_test_range([64, 128, 512, 1024, 96], [64, 1024, 96]) SMALL_CACHE = 4096 diff --git a/python/sglang/jit_kernel/tests/test_timestep_embedding.py b/python/sglang/jit_kernel/tests/test_timestep_embedding.py index 068363774..945fa1f13 100644 --- a/python/sglang/jit_kernel/tests/test_timestep_embedding.py +++ b/python/sglang/jit_kernel/tests/test_timestep_embedding.py @@ -12,6 +12,26 @@ except Exception: from sglang.jit_kernel.timestep_embedding import ( timestep_embedding as timestep_embedding_cuda, ) +from sglang.jit_kernel.utils import get_ci_test_range + +CORRECTNESS_BATCH_SIZES = get_ci_test_range( + [1, 2, 8, 128, 256, 512, 1536, 2048, 4096, 11008, 16384], + [1, 128, 2048, 16384], +) +CORRECTNESS_DIMS = get_ci_test_range( + [32, 128, 256, 512, 1536, 2048, 4096, 8192], + [32, 512, 8192], +) +DIFFUSERS_BATCH_SIZES = get_ci_test_range( + [1, 2, 8, 128, 256, 512, 1536, 2048, 16384], + [1, 512, 16384], +) +DIFFUSERS_DIMS = get_ci_test_range([32, 256, 512, 1536, 8192], [32, 512, 8192]) +DTYPES = get_ci_test_range( + [torch.float16, torch.bfloat16, torch.float32], + [torch.float16, torch.bfloat16], +) +SCALES = get_ci_test_range([1, 0.01], [1, 0.01]) def get_timestep_embedding_reference( @@ -47,11 +67,9 @@ def get_timestep_embedding_reference( return emb -@pytest.mark.parametrize( - "batch_size", [1, 2, 8, 128, 256, 512, 1536, 2048, 4096, 11008, 16384] -) -@pytest.mark.parametrize("dim", [32, 128, 256, 512, 1536, 2048, 4096, 8192]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("batch_size", CORRECTNESS_BATCH_SIZES) +@pytest.mark.parametrize("dim", CORRECTNESS_DIMS) +@pytest.mark.parametrize("dtype", DTYPES) def test_timestep_embedding_correctness_with_sgld(batch_size, dim, dtype): device = "cuda" t = torch.randint(low=0, high=1000, size=(batch_size,), device=device).to(dtype) @@ -64,12 +82,12 @@ def test_timestep_embedding_correctness_with_sgld(batch_size, dim, dtype): torch.testing.assert_close(torch_output, cuda_output, atol=1e-3, rtol=1e-3) -@pytest.mark.parametrize("batch_size", [1, 2, 8, 128, 256, 512, 1536, 2048, 16384]) -@pytest.mark.parametrize("dim", [32, 256, 512, 1536, 8192]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("batch_size", DIFFUSERS_BATCH_SIZES) +@pytest.mark.parametrize("dim", DIFFUSERS_DIMS) +@pytest.mark.parametrize("dtype", DTYPES) @pytest.mark.parametrize("flip_sin_to_cos", [False, True]) @pytest.mark.parametrize("downscale_freq_shift", [0, 1]) -@pytest.mark.parametrize("scale", [1, 0.01]) +@pytest.mark.parametrize("scale", SCALES) def test_timestep_embedding_correctness_with_diffusers( batch_size, dim, flip_sin_to_cos, downscale_freq_shift, scale, dtype ): diff --git a/python/sglang/jit_kernel/utils.py b/python/sglang/jit_kernel/utils.py index e8f1972ba..a1a35e0fe 100644 --- a/python/sglang/jit_kernel/utils.py +++ b/python/sglang/jit_kernel/utils.py @@ -11,6 +11,22 @@ if TYPE_CHECKING: from tvm_ffi import Module 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" + + +def get_ci_test_range(full_range: List[Any], ci_range: List[Any]) -> List[Any]: + if should_run_full_tests(): + return full_range + return ci_range if is_in_ci() else full_range def cache_once(fn: F) -> F: