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

2026-03-10 21:37:41 +08:00
parent 5a7c1b8ec6
commit 60cc06297e
20 changed files with 208 additions and 77 deletions
--- 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]
--- 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
--- 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
--- 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
--- 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
--- 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):
--- 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}"
--- 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}"
--- 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}"
--- 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_]
--- 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:
--- 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

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

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

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

--- 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.")
--- 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


--- 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
 ):
--- 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: