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

.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,

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 = (
+IS_CI = is_in_ci()
-    os.getenv("CI", "false").lower() == "true"
-    or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
-)
 
 if IS_CI:
     qweight_row_range = [128]

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 = (
+IS_CI = is_in_ci()
-    os.getenv("CI", "false").lower() == "true"
-    or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
-)
 
 NUM_BITS = 4
 GROUP_SIZE = 128

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 = (
+IS_CI = is_in_ci()
-    os.getenv("CI", "false").lower() == "true"
-    or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
-)
 
 SIZE_K = 4096
 SIZE_N = 4096

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 = (
+IS_CI = is_in_ci()
-    os.getenv("CI", "false").lower() == "true"
-    or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
-)
 
 SIZE_K = 4096
 SIZE_N = 4096

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 = (
+IS_CI = is_in_ci()
-    os.getenv("CI", "false").lower() == "true"
-    or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
-)
 
 SIZE_N = 4096
 NUM_BITS = 4

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 = (
+IS_CI = is_in_ci()
-    os.getenv("CI", "false").lower() == "true"
-    or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
-)
 
 
 def stack_and_dev(tensors):

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}"

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}"

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}"

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 = is_in_ci()
-IS_CI = (
-    os.getenv("CI", "false").lower() == "true"
-    or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
-)
 
 _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_]

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 (
+    return jit_kernel_is_in_ci()
-        os.getenv("CI", "false").lower() == "true"
-        or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
-    )
 
 
 def get_benchmark_range(full_range: List, ci_range: List) -> List:

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
 

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]
+BS_LIST = get_ci_test_range(BS_LIST, [1, 9, 256, 4109])
-N_Q_LIST = [8, 16]
+N_K_LIST = get_ci_test_range([2, 4], [2, 4])
-HEAD_DIM_LIST = [64, 128, 256, 512, 1024]
+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
 

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
 

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
 

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]
+BS_LIST = get_ci_test_range(BS_LIST, [1, 129, 2048, 2049])
-GQA_RATIO = [1, 4, 8]
+NUM_KV_HEADS_LIST = get_ci_test_range([1, 2, 8], [1, 8])
-ROPE_DIM_LIST = [64, 128, 256, 512]
+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("rope_dim", PARTIAL_ROPE_DIM_LIST)
-@pytest.mark.parametrize("head_dim", [64, 128, 256])
+@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.")

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_BS = get_ci_test_range([1, 7, 128], [1, 128])
-REPR_DIMS = [64, 128, 512, 1024, 96]
+REPR_DIMS = get_ci_test_range([64, 128, 512, 1024, 96], [64, 1024, 96])
 SMALL_CACHE = 4096
 
 

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(
+@pytest.mark.parametrize("batch_size", CORRECTNESS_BATCH_SIZES)
-    "batch_size", [1, 2, 8, 128, 256, 512, 1536, 2048, 4096, 11008, 16384]
+@pytest.mark.parametrize("dim", CORRECTNESS_DIMS)
-)
+@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize("dim", [32, 128, 256, 512, 1536, 2048, 4096, 8192])
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
 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("batch_size", DIFFUSERS_BATCH_SIZES)
-@pytest.mark.parametrize("dim", [32, 256, 512, 1536, 8192])
+@pytest.mark.parametrize("dim", DIFFUSERS_DIMS)
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+@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
 ):

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: