[CI] Migrate JIT tests to runner config registration (#29066)

Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-29 19:12:14 -07:00
parent bc8b3ab1f5
commit 3e16be2122
128 changed files with 274 additions and 188 deletions
@@ -435,31 +435,34 @@ if torch.cuda.get_device_capability()[0] < 9:

 JIT kernel correctness tests and benchmarks live under `test/registered/jit/` and `test/registered/jit/benchmark/` (NOT inside the `sglang` package -- a `register_*_ci(...)` call anywhere under `python/sglang/` is rejected by the `check-no-registered-tests-in-package` pre-commit hook). Only their test-only helpers (e.g. `benchmark/marker.py`) stay alongside the kernel source under `python/sglang/jit_kernel/` and are imported by absolute path. **CI does not run `pytest` in those directories directly.** The unified runner `test/run_suite.py` discovers every `test_*.py` and `bench_*.py` under `test/registered/`, collects `register_*_ci(...)` calls by **statically parsing each file's AST**, and executes the selected suite. Every test file must register at least one CUDA entry or the collector fails its sanity check.

- **PR / per-commit CUDA suites** (see `test/run_suite.py` → `PER_COMMIT_SUITES`): JIT unit tests use `base-b-kernel-unit-1-gpu-large` on H100 and `base-b-kernel-unit-1-gpu-b200` on B200/SM100 paths (see `.github/workflows/pr-test-jit-kernel.yml`). Multi-GPU JIT tests use `base-b-kernel-unit-8-gpu-h200`.
+- **PR / per-commit CUDA suites** (see `test/run_suite.py` → `PER_COMMIT_SUITES`): JIT unit tests use `base-b-kernel-unit-test-1-gpu-large` on H100 and `base-b-kernel-unit-test-4-gpu-b200` on B200/SM100 paths (see `.github/workflows/pr-test-jit-kernel.yml`). Multi-GPU JIT tests use `base-b-kernel-unit-test-8-gpu-h200`.
 - **Nightly kernel suite**: `nightly-kernel-1-gpu` with `--nightly` — typically used with `SGLANG_JIT_KERNEL_RUN_FULL_TESTS=1` in CI for expanded parameter grids (see `python/sglang/jit_kernel/utils.py` → `should_run_full_tests` / `get_ci_test_range`). Wired in `.github/workflows/nightly-test-nvidia.yml` (e.g. `python3 run_suite.py --hw cuda --suite nightly-kernel-1-gpu --nightly --continue-on-error`).

-Registration pattern (module level, **literal** `est_time` and `suite` strings — required for AST parsing):
+Registration pattern (module level, **literal** `est_time`, `stage`, and `runner_config` values — required for AST parsing):

 ```python
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 # Optional B200/SM100 registration for tests that cover Blackwell-specific code paths
-# register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-b200")
+# register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
 # Optional second registration: same file also listed under the nightly kernel suite
+# (nightly suites use the legacy single-string suite=, not stage/runner_config)
 # register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
 ```

-Keep `est_time` and `suite` as literal values. `run_suite.py` collects them from the file AST, so computed values and helper wrappers can break CI discovery.
+CI generates the suite name as `{stage}-test-{runner_config}`, so `stage="base-b-kernel-unit", runner_config="1-gpu-large"` becomes the `base-b-kernel-unit-test-1-gpu-large` suite you pass to `run_suite.py` below — don't put the `-test-` infix in `register_cuda_ci`. The single-string `suite=` form is only for nightly/stress/weekly suites.
+
+Keep `est_time`, `stage`, `runner_config`, and `suite` as literal values. `run_suite.py` collects them from the file AST, so computed values and helper wrappers can break CI discovery.

 Use `register_cuda_ci(..., disabled="reason")` if the file must stay in-tree but should be skipped in CI (e.g. multi-GPU only).

 **Run like CI** (from repo root):

 ```bash
-(cd test && python3 run_suite.py --hw cuda --suite base-b-kernel-unit-1-gpu-large)
+(cd test && python3 run_suite.py --hw cuda --suite base-b-kernel-unit-test-1-gpu-large)
 # For B200/SM100-specific coverage:
-(cd test && python3 run_suite.py --hw cuda --suite base-b-kernel-unit-1-gpu-b200)
+(cd test && python3 run_suite.py --hw cuda --suite base-b-kernel-unit-test-4-gpu-b200)
 ```

 For fast iteration you can still run `pytest` on a single file locally; CI coverage is via `run_suite.py`.
@@ -472,7 +475,7 @@ import torch
 from sglang.jit_kernel.scale import scale
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")


@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
@@ -517,7 +520,7 @@ if __name__ == "__main__":

 ## Step 5: Add a benchmark (required)

-Benchmarks are `bench_*.py` files under `test/registered/jit/benchmark/`. They are picked up by the same `run_suite.py` machinery as unit tests. Register them for **`base-b-kernel-benchmark-1-gpu-large`** (PR JIT benchmark job: `python3 run_suite.py --hw cuda --suite base-b-kernel-benchmark-1-gpu-large`).
+Benchmarks are `bench_*.py` files under `test/registered/jit/benchmark/`. They are picked up by the same `run_suite.py` machinery as unit tests. Register them for **`base-b-kernel-benchmark-test-1-gpu-large`** (PR JIT benchmark job: `python3 run_suite.py --hw cuda --suite base-b-kernel-benchmark-test-1-gpu-large`).

 Benchmarks use the project's own `marker` framework (in `python/sglang/jit_kernel/benchmark/marker.py`) — **do not** use `triton.testing.perf_report` / `triton.testing.do_bench` directly. The marker framework provides (public names: `benchmark`, `parametrize`, `do_bench`, `skip`, `BenchResult`, `BenchSkip`):

@@ -544,7 +547,7 @@ from sglang.jit_kernel.benchmark.utils import create_random
 from sglang.jit_kernel.scale import scale as jit_scale
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=6, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(est_time=6, stage="base-b-kernel-benchmark", runner_config="1-gpu-large")


@torch.compile()
@@ -598,14 +601,14 @@ python test/registered/jit/benchmark/bench_scale.py
 Run the benchmark suite the way CI does:

 ```bash
-cd test && python3 run_suite.py --hw cuda --suite base-b-kernel-benchmark-1-gpu-large
+cd test && python3 run_suite.py --hw cuda --suite base-b-kernel-benchmark-test-1-gpu-large
 ```

 ---

 ## Troubleshooting

- **`No CI registry found in ...` from `run_suite.py`**: add a module-level `register_cuda_ci(...)` with literal `est_time` and `suite` (and optional `nightly=True`); starred args and non-literal values break AST collection
+- **`No CI registry found in ...` from `run_suite.py`**: add a module-level `register_cuda_ci(...)` with literal `est_time`, `stage`, and `runner_config` (and optional `nightly=True`); starred args and non-literal values break AST collection
 - **JIT compilation fails**: ensure the `.cuh` file is under `python/sglang/jit_kernel/csrc/`; reduce template argument combinations
 - **CUDA crash / illegal memory access**: `CUDA_LAUNCH_BLOCKING=1`; `compute-sanitizer --tool memcheck python ...`
 - **Unstable benchmark results**: `marker.do_bench` uses CUDA-graph-based timing by default; set `use_cuda_graph=False` only if the kernel can't be captured. Make sure `graph_clone_args` covers every *read* tensor — reusing a single buffer keeps it L2-hot and skews results
@@ -277,10 +277,10 @@ Large suites are split across matrix jobs using the **LPT (Longest Processing Ti
 | `base-b-test-1-gpu-large` | 14 | `1-gpu-h100` | dynamic (3 or 14) |
 | `base-b-test-2-gpu-large` | 4 | `2-gpu-h100` | — |
 | `base-b-test-4-gpu-b200` | 1 (no matrix) | `4-gpu-b200` | — |
-| `base-b-kernel-unit-1-gpu-large` | 1 (no matrix) | `1-gpu-h100` | — |
-| `base-b-kernel-unit-1-gpu-b200` | 1 (no matrix) | `4-gpu-b200` | — |
-| `base-b-kernel-unit-8-gpu-h200` | 1 (no matrix) | `8-gpu-h200` | — |
-| `base-b-kernel-benchmark-1-gpu-large` | 1 (no matrix) | `1-gpu-h100` | — |
+| `base-b-kernel-unit-test-1-gpu-large` | 1 (no matrix) | `1-gpu-h100` | — |
+| `base-b-kernel-unit-test-4-gpu-b200` | 1 (no matrix) | `4-gpu-b200` | — |
+| `base-b-kernel-unit-test-8-gpu-h200` | 1 (no matrix) | `8-gpu-h200` | — |
+| `base-b-kernel-benchmark-test-1-gpu-large` | 1 (no matrix) | `1-gpu-h100` | — |
 | `base-c-test-4-gpu-h100` | 3 | `4-gpu-h100` | — |
 | `base-c-test-8-gpu-h200` | 4 | `8-gpu-h200` | — |
 | `base-c-test-8-gpu-h20` | 2 | `8-gpu-h20` | — |
@@ -290,6 +290,8 @@ Large suites are split across matrix jobs using the **LPT (Longest Processing Ti
 | `base-c-test-8-gpu-b200` | registered only | `8-gpu-b200` | — |
 | `base-c-test-4-gpu-gb200` | registered only | `4-gpu-gb200` | — |

+> **Suite names are generated**, not hand-written: each comes from a test's `register_*_ci(stage=..., runner_config=...)` as `{stage}-test-{runner_config}`, and `runner_config` maps to the `Runner` column via `scripts/ci/runner_configs.yml`.
+>
 > **Note**: Kernel suites (`base-b-kernel-*`) run via `pr-test-jit-kernel.yml` and `pr-test-sgl-kernel.yml`, not the main `pr-test.yml`. `base-c-test-8-gpu-b200` is registered in `test/run_suite.py` but not wired to PR CI. The GB200 job is currently commented out in `pr-test.yml` until a company-owned runner is provisioned. Multimodal diffusion uses `python/sglang/multimodal_gen/test/run_suite.py`, not `test/run_suite.py`.

 **Workflow usage:**
@@ -50,8 +50,13 @@ Defined in `python/sglang/test/test_utils.py`:

 ### Naming Conventions

- **Suite**: `base-{a,b,c}-test-{gpu_count}-gpu-{hardware}` (e.g., `base-b-test-1-gpu-small`)
- **CI runner**: `{gpu_count}-gpu-{hardware}` (e.g., `1-gpu-5090`, `4-gpu-h100`, `8-gpu-h200`)
+A per-commit suite name is **generated** from registration metadata as `{stage}-test-{runner_config}` — you don't hand-write it:
+
+- **`stage`** — the CI stage (e.g. `base-b`, `base-b-kernel-unit`, `base-c`).
+- **`runner_config`** — a runner-pool key from `scripts/ci/runner_configs.yml`, which maps it to the physical runner label (so `1-gpu-large` runs on `1-gpu-h100`). AMD/NPU use their own keys (e.g. `amd`).
+- **Suite** — `register_cuda_ci(stage="base-b", runner_config="1-gpu-small")` → `base-b-test-1-gpu-small`, the name you pass to `run_suite.py --suite`. The `-test-` is just the connector; never put it in `register_*_ci`.
+
+> Legacy single-string `suite=` is only for suites that don't fit that shape — nightly/stress/weekly and some AMD/CPU/NPU pools (e.g. `suite="nightly-kernel-1-gpu", nightly=True`). Per-commit tests always use `stage=` + `runner_config=`.

 ### All CI Suites

@@ -65,10 +70,10 @@ Defined in `python/sglang/test/test_utils.py`:
 | `base-b-test-1-gpu-large` | `1-gpu-h100` | Tests that need H100-class memory or kernels (e.g. FA3) |
 | `base-b-test-2-gpu-large` | `2-gpu-h100` | Two-GPU correctness and parallelism (TP/PP) on H100 |
 | `base-b-test-4-gpu-b200` | `4-gpu-b200` | Early Blackwell coverage (SM100+ paths) on four GPUs |
-| `base-b-kernel-unit-1-gpu-large` | `1-gpu-h100` | JIT kernel correctness tests under `test/registered/jit/` |
-| `base-b-kernel-unit-1-gpu-b200` | `4-gpu-b200` | JIT kernel correctness tests for Blackwell / SM100-specific paths |
-| `base-b-kernel-unit-8-gpu-h200` | `8-gpu-h200` | Multi-GPU JIT kernel correctness tests under `test/registered/jit/` |
-| `base-b-kernel-benchmark-1-gpu-large` | `1-gpu-h100` | JIT kernel benchmark files under `test/registered/jit/benchmark/` |
+| `base-b-kernel-unit-test-1-gpu-large` | `1-gpu-h100` | JIT kernel correctness tests under `test/registered/jit/` |
+| `base-b-kernel-unit-test-4-gpu-b200` | `4-gpu-b200` | JIT kernel correctness tests for Blackwell / SM100-specific paths |
+| `base-b-kernel-unit-test-8-gpu-h200` | `8-gpu-h200` | Multi-GPU JIT kernel correctness tests under `test/registered/jit/` |
+| `base-b-kernel-benchmark-test-1-gpu-large` | `1-gpu-h100` | JIT kernel benchmark files under `test/registered/jit/benchmark/` |
 | `base-c-test-4-gpu-h100` | `4-gpu-h100` | Large 4-GPU H100 integration and scaling tests |
 | `base-c-test-8-gpu-h200` | `8-gpu-h200` | Large 8-GPU H200 runs for big models and parallelism |
 | `base-c-test-8-gpu-h20` | `8-gpu-h20` | Large 8-GPU H20 runs for big models |
@@ -132,9 +137,9 @@ Use the lightest suite that meets your test's needs:
 - **No GPU required** → `base-a-test-cpu`
 - **Most small GPU tests** → `base-b-test-1-gpu-small` (default choice)
 - **Need H100 memory or Hopper features** → `base-b-test-1-gpu-large`
- **JIT kernel correctness** → `base-b-kernel-unit-1-gpu-large`
- **JIT kernel correctness for B200 / SM100 paths** → `base-b-kernel-unit-1-gpu-b200`
- **JIT kernel benchmarks** → `base-b-kernel-benchmark-1-gpu-large`
+- **JIT kernel correctness** → `base-b-kernel-unit-test-1-gpu-large`
+- **JIT kernel correctness for B200 / SM100 paths** → `base-b-kernel-unit-test-4-gpu-b200`
+- **JIT kernel benchmarks** → `base-b-kernel-benchmark-test-1-gpu-large`
 - **Multi-GPU** → only when the test actually needs multiple GPUs

 ---
@@ -353,19 +358,19 @@ JIT kernel files live outside `test/registered/` but still use registration:
 from sglang.test.ci.ci_register import register_cuda_ci

 # Correctness tests in test/registered/jit/
-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-b200")
-register_cuda_ci(est_time=120, suite="base-b-kernel-unit-8-gpu-h200")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="8-gpu-h200")

 # Benchmarks in test/registered/jit/benchmark/
-register_cuda_ci(est_time=6, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(est_time=6, stage="base-b-kernel-benchmark", runner_config="1-gpu-large")

-# Optional nightly registration
+# Optional nightly registration — nightly suites use the legacy single-string suite=
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
 register_cuda_ci(est_time=120, suite="nightly-kernel-8-gpu-h200", nightly=True)
 ```

-Keep `est_time` and `suite` as **literal values** — `run_suite.py` collects them by AST parsing
+The `stage` + `runner_config` calls generate suites like `base-b-kernel-unit-test-1-gpu-large`; nightly keeps the legacy `suite=` string. Keep `est_time`, `stage`, `runner_config`, and `suite` as **literal values** — `run_suite.py` collects them by AST parsing.

 ---

@@ -76,7 +76,7 @@ jobs:
        timeout-minutes: 30
        run: |
          cd test/
-          python3 run_suite.py --hw cuda --suite base-b-kernel-unit-1-gpu-large
+          python3 run_suite.py --hw cuda --suite base-b-kernel-unit-test-1-gpu-large

  jit-kernel-multigpu-unit-test:
    if: |
@@ -114,7 +114,7 @@ jobs:
        timeout-minutes: 45
        run: |
          cd test/
-          python3 run_suite.py --hw cuda --suite base-b-kernel-unit-8-gpu-h200
+          python3 run_suite.py --hw cuda --suite base-b-kernel-unit-test-8-gpu-h200

  jit-kernel-benchmark-test:
    if: |
@@ -154,7 +154,7 @@ jobs:
        timeout-minutes: 45
        run: |
          cd test/
-          python3 run_suite.py --hw cuda --suite base-b-kernel-benchmark-1-gpu-large
+          python3 run_suite.py --hw cuda --suite base-b-kernel-benchmark-test-1-gpu-large

  jit-kernel-b200-test:
    if: |
@@ -194,4 +194,4 @@ jobs:
        timeout-minutes: 30
        run: |
          cd test/
-          python3 run_suite.py --hw cuda --suite base-b-kernel-unit-1-gpu-b200
+          python3 run_suite.py --hw cuda --suite base-b-kernel-unit-test-4-gpu-b200
@@ -73,8 +73,8 @@ Parameters: `est_time` (seconds), `stage` + `runner_config` (target stage and ru
 Keep `est_time`, `stage`, `runner_config` as **literal values** — `run_suite.py` collects them by AST parsing.

 JIT kernel correctness tests and benchmarks live under `test/registered/jit/`, same as other registered tests (their helpers stay alongside the kernel source under `python/sglang/jit_kernel/` and are imported by absolute path):
- Correctness tests: `test/registered/jit/test_*.py` → `base-b-kernel-unit-1-gpu-large`
- Benchmarks: `test/registered/jit/benchmark/bench_*.py` → `base-b-kernel-benchmark-1-gpu-large`
+- Correctness tests: `test/registered/jit/test_*.py` → `base-b-kernel-unit-test-1-gpu-large`
+- Benchmarks: `test/registered/jit/benchmark/bench_*.py` → `base-b-kernel-benchmark-test-1-gpu-large`

 ## Choosing a Suite

@@ -85,8 +85,8 @@ Use the lightest suite that meets your test's needs. Full suite tables are in th
 | No GPU required | `base-a-test-cpu` |
 | Small GPU (fits 5090, 32GB) | `base-b-test-1-gpu-small` (most tests go here) |
 | Large GPU memory or Hopper features | `base-b-test-1-gpu-large` |
-| JIT kernel correctness | `base-b-kernel-unit-1-gpu-large` |
-| JIT kernel benchmarks | `base-b-kernel-benchmark-1-gpu-large` |
+| JIT kernel correctness | `base-b-kernel-unit-test-1-gpu-large` |
+| JIT kernel benchmarks | `base-b-kernel-benchmark-test-1-gpu-large` |
 | Multi-GPU (2/4/8) | `base-b-test-2-gpu-large`, `base-c-test-*` |
 | Long-running or experimental | `nightly-*` suites |

@@ -15,7 +15,7 @@ from sglang.test.ci.ci_register import register_cuda_ci

 # CuteDSL prefill kernel only exists on Blackwell. Single-GPU kernel-unit
 # suite is the right slot (matches existing jit_kernel test_*.py pattern).
-register_cuda_ci(est_time=60, suite="base-b-kernel-unit-1-gpu-b200")
+register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="4-gpu-b200")

 if not (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10):
    pytest.skip(
@@ -17,7 +17,7 @@ from sglang.test.ci.ci_register import register_cuda_ci

 # CuteDSL prefill kernel only exists on Blackwell. Single-GPU kernel-unit suite,
 # same slot as the GDN prefill test.
-register_cuda_ci(est_time=60, suite="base-b-kernel-unit-1-gpu-b200")
+register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="4-gpu-b200")

 if not (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10):
    pytest.skip(
@@ -12,7 +12,9 @@ from sglang.jit_kernel.benchmark import marker
 from sglang.jit_kernel.benchmark.utils import create_random
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=30, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=30, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)


@torch.compile
@@ -10,7 +10,9 @@ from sglang.jit_kernel.benchmark.utils import (
 )
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=15, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=15, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 CONSTANT = 7
 SIZE_LIST = get_benchmark_range(
@@ -9,7 +9,9 @@ from sglang.jit_kernel.benchmark.utils import run_benchmark
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci

-register_cuda_ci(est_time=5, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 try:
    from sgl_kernel import awq_dequantize as aot_awq_dequantize
@@ -13,8 +13,10 @@ from sglang.jit_kernel.clamp_position import clamp_position_cuda
 from sglang.srt.utils import get_compiler_backend
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=13, suite="base-b-kernel-benchmark-1-gpu-large")
-register_amd_ci(est_time=16, suite="jit-kernel-unit-test-amd")
+register_cuda_ci(
+    est_time=13, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)
+register_amd_ci(est_time=16, stage="jit-kernel-unit", runner_config="amd")

 SIZE_LIST = get_benchmark_range(
    full_range=[2**n for n in range(4, 16)],
@@ -12,7 +12,9 @@ from sglang.jit_kernel.concat_mla import concat_mla_k as jit_k
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci

-register_cuda_ci(est_time=6, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=6, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 IS_CI = is_in_ci()

@@ -33,7 +33,8 @@ from sglang.test.ci.ci_register import register_cuda_ci

 register_cuda_ci(
    est_time=120,
-    suite="base-b-kernel-benchmark-1-gpu-large",
+    stage="base-b-kernel-benchmark",
+    runner_config="1-gpu-large",
    disabled="requires multi-GPU, self-skips in CI",
 )

@@ -10,7 +10,9 @@ from sglang.jit_kernel.benchmark.utils import get_benchmark_range
 from sglang.srt.utils import is_sm100_supported
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=5, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 try:
    import deep_gemm
@@ -20,7 +20,9 @@ from sglang.jit_kernel.fused_qknorm_rope import (
 )
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=6, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=6, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 try:
    from sgl_kernel import fused_qk_norm_rope as fused_qk_norm_rope_aot
@@ -16,7 +16,9 @@ from sglang.jit_kernel.benchmark.utils import (
 from sglang.jit_kernel.hadamard import hadamard_transform
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=5, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 # AOT kernel: might not be available in all environments.
 # This is used for performance baseline comparison.
@@ -33,7 +33,9 @@ from sglang.jit_kernel.hicache import (
 )
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=29, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=29, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 DISABLE_TORCH = os.environ.get("DISABLE_TORCH", "0") == "1"
 PAGE_SIZE = 1
@@ -9,7 +9,9 @@ from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE, DEFAULT_DTYPE
 from sglang.jit_kernel.hisparse import load_cache_to_device_buffer_mla
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=12, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=12, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 DEVICE = DEFAULT_DEVICE
 DTYPE = DEFAULT_DTYPE
@@ -20,7 +20,9 @@ from sglang.jit_kernel.mla_kv_pack_quantize_fp8 import (
 from sglang.jit_kernel.utils import is_arch_support_pdl
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=15, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=15, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)


@triton.jit
@@ -13,7 +13,9 @@ from sglang.jit_kernel.mxfp8 import (
 )
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=5, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)


 def is_sm100_supported(device=None) -> bool:
@@ -13,7 +13,9 @@ from sglang.jit_kernel.ngram_embedding import (
 )
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=15, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=15, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 NE_N = 8
 NE_K = 2
@@ -13,7 +13,9 @@ from sglang.jit_kernel.ngram_embedding import (
 )
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=15, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=15, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 MAX_CONTEXT_LEN = 4096
 BATCH_SIZE_LIST = get_benchmark_range(
@@ -11,7 +11,9 @@ 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.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=30, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=30, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)


 DTYPE = torch.bfloat16
@@ -15,7 +15,9 @@ from sglang.jit_kernel.nvfp4 import (
 from sglang.srt.utils import is_sm100_supported
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=5, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 FLOAT4_E2M1_MAX = 6.0
 FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
@@ -10,7 +10,9 @@ from sglang.jit_kernel.nvfp4 import scaled_fp4_quant
 from sglang.srt.utils import is_sm100_supported
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=5, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 FLOAT4_E2M1_MAX = 6.0
 FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
@@ -10,7 +10,9 @@ from sglang.jit_kernel.nvfp4 import cutlass_scaled_fp4_mm, scaled_fp4_quant
 from sglang.srt.utils import is_sm100_supported, is_sm120_supported
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=5, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 FLOAT4_E2M1_MAX = 6.0
 FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
@@ -18,7 +18,9 @@ from sglang.jit_kernel.benchmark.utils import (
 from sglang.jit_kernel.dsv4.online_c128_mtp import _jit_online_c128_mtp_module
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=10, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=10, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 HEAD_DIM = 512
 STATE_DIM = HEAD_DIM * 3
@@ -8,7 +8,9 @@ from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
 from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=5, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 try:
    from vllm import _custom_ops as ops
@@ -20,7 +20,9 @@ from sglang.srt.utils.bench_utils import bench_kineto
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci

-register_cuda_ci(est_time=13, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=13, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 IS_CI = is_in_ci()

@@ -14,7 +14,9 @@ from sglang.srt.layers.quantization.fp8_kernel import (
 )
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=6, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=6, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 G = 128
 HIDDEN = 4096
@@ -6,7 +6,9 @@ from sglang.jit_kernel.norm import fused_inplace_qknorm
 from sglang.srt.utils import get_current_device_stream_fast
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=10, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=10, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 alt_stream = torch.cuda.Stream()

@@ -12,7 +12,9 @@ from sglang.srt.utils import get_current_device_stream_fast
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci

-register_cuda_ci(est_time=12, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=12, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 IS_CI = is_in_ci()

@@ -9,7 +9,9 @@ from sglang.jit_kernel.benchmark.utils import run_benchmark_no_cudagraph
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci

-register_cuda_ci(est_time=5, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)


 def torch_top_k_renorm_probs(probs, top_k):
@@ -13,8 +13,10 @@ from sglang.jit_kernel.resolve_future_token_ids import resolve_future_token_ids_
 from sglang.srt.utils import get_compiler_backend
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=10, suite="base-b-kernel-benchmark-1-gpu-large")
-register_amd_ci(est_time=10, suite="jit-kernel-unit-test-amd")
+register_cuda_ci(
+    est_time=10, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)
+register_amd_ci(est_time=10, stage="jit-kernel-unit", runner_config="amd")

 SIZE_LIST = get_benchmark_range(
    full_range=[2**n for n in range(4, 16)],  # 16 … 32K elements
@@ -12,7 +12,9 @@ from sglang.jit_kernel.benchmark.utils import (
 )
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=6, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=6, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 MAX_SEQ_LEN = 131072
 ROPE_BASE = 10000.0
@@ -26,7 +26,9 @@ from sglang.srt.mem_cache.utils import set_mla_kv_buffer_kernel as sglang_triton
 from sglang.srt.mem_cache.utils import set_mla_kv_buffer_triton as sglang_wrapper
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=9, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=9, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)


 def _triton_baseline(kv_buffer, loc, cache_k_nope, cache_k_rope):
@@ -9,7 +9,9 @@ from sglang.jit_kernel.benchmark.utils import (
 from sglang.jit_kernel.kvcache import store_cache
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=9, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=9, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)


@torch.compile()
@@ -41,7 +41,8 @@ from sglang.test.ci.ci_register import register_cuda_ci

 register_cuda_ci(
    est_time=120,
-    suite="base-b-kernel-benchmark-1-gpu-large",
+    stage="base-b-kernel-benchmark",
+    runner_config="1-gpu-large",
    disabled="requires multi-GPU, self-skips in CI",
 )

@@ -18,7 +18,8 @@ from sglang.utils import is_in_ci

 register_cuda_ci(
    est_time=120,
-    suite="base-b-kernel-benchmark-1-gpu-large",
+    stage="base-b-kernel-benchmark",
+    runner_config="1-gpu-large",
    disabled="standalone diffusion NVFP4 benchmark",
 )

@@ -20,7 +20,8 @@ from sglang.utils import is_in_ci

 register_cuda_ci(
    est_time=17,
-    suite="base-b-kernel-benchmark-1-gpu-large",
+    stage="base-b-kernel-benchmark",
+    runner_config="1-gpu-large",
    disabled="Temporarily skipped to unblock flashinfer upgrade. Ref: https://github.com/sgl-project/sglang/actions/runs/23735552939/job/69139238979?pr=21422",
 )

@@ -16,7 +16,8 @@ from sglang.utils import is_in_ci

 register_cuda_ci(
    est_time=45,
-    suite="base-b-kernel-benchmark-1-gpu-large",
+    stage="base-b-kernel-benchmark",
+    runner_config="1-gpu-large",
    disabled="standalone benchmark",
 )

@@ -24,7 +24,8 @@ from sglang.utils import is_in_ci

 register_cuda_ci(
    est_time=120,
-    suite="base-b-kernel-benchmark-1-gpu-large",
+    stage="base-b-kernel-benchmark",
+    runner_config="1-gpu-large",
    disabled="self-skips in CI, standalone tool",
 )

@@ -13,7 +13,9 @@ from sglang.jit_kernel.benchmark.utils import (
 )
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=13, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=13, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 MAX_SEQ_LEN = 131072
 ROPE_BASE = 10000.0
@@ -12,7 +12,9 @@ from sglang.jit_kernel.diffusion.triton.scale_shift import (
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci

-register_cuda_ci(est_time=13, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=13, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 if is_in_ci():
    B_RANGE, S_RANGE, D_RANGE = [1], [128], [3072]
@@ -17,7 +17,9 @@ from sglang.srt.layers.attention.minimax_sparse_ops.decode.flash_with_topk_idx i
 )
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=8, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=8, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 BLOCK_SIZE = 128
 TOPK = 16
@@ -12,7 +12,9 @@ from sglang.jit_kernel.minimax_qknorm_rope import (
 )
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=6, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=6, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 HEAD_DIM, ROTARY_DIM, BASE, EPS, MAXPOS = 128, 64, 5_000_000, 1e-6, 131072
 NQ, NK = 64, 4
@@ -7,7 +7,9 @@ from sglang.jit_kernel.benchmark import marker
 from sglang.jit_kernel.minimax_store_kv_index import store_kv_index
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=6, suite="base-b-kernel-benchmark-1-gpu-large")
+register_cuda_ci(
+    est_time=6, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)

 HEAD_DIM = 128
 NUM_KV_HEADS = 1
@@ -18,7 +18,7 @@ from sglang.jit_kernel.tests.deepseek_v4.common import (
 )
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=30, suite="nightly-kernel-1-gpu", nightly=True)
 register_amd_ci(est_time=30, suite="nightly-amd-kernel-1-gpu", nightly=True)

@@ -18,7 +18,7 @@ from sglang.jit_kernel.tests.deepseek_v4.common import (
 )
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=30, suite="nightly-kernel-1-gpu", nightly=True)
 register_amd_ci(est_time=30, suite="nightly-amd-kernel-1-gpu", nightly=True)

@@ -21,7 +21,7 @@ from sglang.srt.layers.deepseek_v4_rope import (
 )
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=60, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=60, suite="nightly-kernel-1-gpu", nightly=True)

 HEAD_DIM = 128
@@ -14,7 +14,7 @@ from sglang.srt.layers.quantization.fp8_utils import (
 )
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=20, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=80, suite="nightly-kernel-1-gpu", nightly=True)

 DEVICE = "cuda"
@@ -17,7 +17,7 @@ from sglang.srt.layers.quantization.modelopt_quant import pad_nvfp4_weight
 from sglang.test.ci.ci_register import register_cuda_ci

 # B200-only correctness coverage for diffusion NVFP4 scaled mm.
-register_cuda_ci(est_time=15, suite="base-b-kernel-unit-1-gpu-b200")
+register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="4-gpu-b200")

 DEVICE = "cuda"
 DTYPE = torch.bfloat16
@@ -6,7 +6,7 @@ import torch.nn.functional as F

 from sglang.test.ci.ci_register import register_amd_ci

-register_amd_ci(est_time=30, suite="jit-kernel-unit-test-amd")
+register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd")

 DEVICE = "cuda"
 D = 5120
@@ -13,7 +13,7 @@ from sglang.jit_kernel.diffusion.cutedsl.scale_residual_norm_scale_shift import
 )
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=28, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=28, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)

 DEVICE = "cuda"
@@ -9,7 +9,7 @@ from sglang.jit_kernel.diffusion.group_norm_silu import apply_group_norm_silu
 from sglang.jit_kernel.diffusion.triton.group_norm_silu import triton_group_norm_silu
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=8, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=8, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
 register_amd_ci(est_time=15, suite="nightly-amd-kernel-1-gpu", nightly=True)

@@ -8,7 +8,7 @@ import triton
 from sglang.jit_kernel.utils import get_ci_test_range
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=44, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=44, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=176, suite="nightly-kernel-1-gpu", nightly=True)

 DEVICE = "cuda"
@@ -12,7 +12,7 @@ from sglang.jit_kernel.diffusion.triton.scale_shift import (
 from sglang.jit_kernel.utils import get_ci_test_range
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=15, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
 register_amd_ci(est_time=30, suite="nightly-amd-kernel-1-gpu", nightly=True)

@@ -15,7 +15,7 @@ from sglang.jit_kernel.diffusion.triton.varlen_pack_pad import (
 from sglang.jit_kernel.utils import get_ci_test_range
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=60, suite="nightly-kernel-1-gpu", nightly=True)
 register_amd_ci(est_time=15, suite="nightly-amd-kernel-1-gpu", nightly=True)

@@ -29,7 +29,7 @@ from sglang.multimodal_gen.runtime.layers.attention.layer import (
 )
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=15, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=60, suite="nightly-kernel-1-gpu", nightly=True)

 DEVICE = "cuda"
@@ -7,8 +7,8 @@ import sglang.jit_kernel
 from sglang.jit_kernel.kv_canary import consts
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=5, suite="base-b-kernel-unit-1-gpu-large")
-register_amd_ci(est_time=5, suite="jit-kernel-unit-test-amd")
+register_cuda_ci(est_time=5, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_amd_ci(est_time=5, stage="jit-kernel-unit", runner_config="amd")


 # Resolve the kernel source against the installed jit_kernel package rather
@@ -35,8 +35,8 @@ from sglang.jit_kernel.tests.kv_canary._fixtures import (
 )
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=60, suite="base-b-kernel-unit-1-gpu-large")
-register_amd_ci(est_time=60, suite="jit-kernel-unit-test-amd")
+register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_amd_ci(est_time=60, stage="jit-kernel-unit", runner_config="amd")

 _DEVICE = torch.device("cuda")

@@ -40,8 +40,8 @@ from sglang.jit_kernel.tests.kv_canary._fixtures import (
 )
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
-register_amd_ci(est_time=30, suite="jit-kernel-unit-test-amd")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd")


 _DEVICE = torch.device("cuda")
@@ -22,8 +22,8 @@ from sglang.jit_kernel.tests.kv_canary._fuzz_driver import (
 from sglang.jit_kernel.tests.kv_canary._invariants import PlanInvariants
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
-register_amd_ci(est_time=30, suite="jit-kernel-unit-test-amd")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd")


 _DEVICE = torch.device("cuda")
@@ -22,8 +22,8 @@ from sglang.jit_kernel.tests.kv_canary._fixtures import (
 from sglang.jit_kernel.tests.kv_canary._invariants import PlanInvariants
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
-register_amd_ci(est_time=30, suite="jit-kernel-unit-test-amd")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd")


 _DEVICE = torch.device("cuda")
@@ -13,8 +13,8 @@ from sglang.jit_kernel.kv_canary.scatter_req_token_ids import (
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.test_utils import CustomTestCase

-register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
-register_amd_ci(est_time=10, suite="jit-kernel-unit-test-amd")
+register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_amd_ci(est_time=10, stage="jit-kernel-unit", runner_config="amd")


 _DEVICE = torch.device("cuda")
@@ -8,7 +8,7 @@ from sglang.jit_kernel.benchmark.kv_canary.utils import (
 )
 from sglang.test.ci.ci_register import register_cpu_ci

-register_cpu_ci(est_time=3, suite="base-a-test-cpu")
+register_cpu_ci(est_time=3, stage="base-a", runner_config="cpu")


 def test_fast_matrix_cases_include_e2e_decode_and_chunked_prefill_scenarios() -> None:
@@ -31,8 +31,8 @@ from sglang.jit_kernel.tests.kv_canary._fuzz_driver import (
 from sglang.jit_kernel.tests.kv_canary._invariants import VerifyInvariants
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
-register_amd_ci(est_time=30, suite="jit-kernel-unit-test-amd")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd")


 _DEVICE = torch.device("cuda")
@@ -53,8 +53,8 @@ from sglang.jit_kernel.tests.kv_canary._hand_oracle import (
 )
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
-register_amd_ci(est_time=30, suite="jit-kernel-unit-test-amd")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd")


 _DEVICE = torch.device("cuda")
@@ -31,8 +31,8 @@ from sglang.jit_kernel.tests.kv_canary._fuzz_driver import (
 from sglang.jit_kernel.tests.kv_canary._invariants import WriteInvariants
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
-register_amd_ci(est_time=30, suite="jit-kernel-unit-test-amd")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd")


 _DEVICE = torch.device("cuda")
@@ -51,8 +51,8 @@ from sglang.jit_kernel.tests.kv_canary._hand_oracle import (
 )
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
-register_amd_ci(est_time=30, suite="jit-kernel-unit-test-amd")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd")


 _DEVICE = torch.device("cuda")
@@ -12,8 +12,8 @@ import torch
 from sglang.jit_kernel.minimax_decode_topk import minimax_decode_topk
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=40, suite="base-b-kernel-unit-1-gpu-large")
-register_cuda_ci(est_time=40, suite="base-b-kernel-unit-1-gpu-b200")
+register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
 register_amd_ci(est_time=15, suite="nightly-amd-kernel-1-gpu", nightly=True)


@@ -25,7 +25,7 @@ from sglang.srt.layers.attention.minimax_sparse_ops.decode.topk_sparse import (
 )
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=25, suite="base-b-kernel-unit-1-gpu-b200")
+register_cuda_ci(est_time=25, stage="base-b-kernel-unit", runner_config="4-gpu-b200")

 dev = "cuda"

@@ -14,8 +14,8 @@ from sglang.jit_kernel.minimax_qknorm_rope import (
 )
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-b200")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200")

 dev = "cuda"
 HEAD_DIM, ROTARY_DIM, BASE, EPS = 128, 64, 5_000_000, 1e-6
@@ -11,8 +11,8 @@ import torch
 from sglang.jit_kernel.minimax_store_kv_index import store_kv_index
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-b200")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
 register_amd_ci(est_time=10, suite="nightly-amd-kernel-1-gpu", nightly=True)

 dev = "cuda"
@@ -12,9 +12,9 @@ from sglang.jit_kernel.activation import (
 from sglang.jit_kernel.utils import get_ci_test_range
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=20, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=30, suite="nightly-kernel-1-gpu", nightly=True)
-register_amd_ci(est_time=20, suite="jit-kernel-unit-test-amd")
+register_amd_ci(est_time=20, stage="jit-kernel-unit", runner_config="amd")


 OPS = SUPPORTED_ACTIVATIONS
@@ -6,7 +6,7 @@ import torch
 from sglang.jit_kernel.add_constant import add_constant
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=45, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=180, suite="nightly-kernel-1-gpu", nightly=True)


@@ -25,7 +25,7 @@ import torch.distributed as dist
 from sglang.srt.environ import envs
 from sglang.test.ci.ci_register import register_amd_ci

-register_amd_ci(est_time=120, suite="sgl-kernel-unit-test-2-gpu-amd")
+register_amd_ci(est_time=120, stage="sgl-kernel-unit", runner_config="2-gpu-amd")


 def get_open_port():
@@ -22,7 +22,7 @@ import torch.distributed as dist

 from sglang.test.ci.ci_register import register_amd_ci

-register_amd_ci(est_time=120, suite="sgl-kernel-unit-test-2-gpu-amd")
+register_amd_ci(est_time=120, stage="sgl-kernel-unit", runner_config="2-gpu-amd")


 def get_open_port():
@@ -7,7 +7,7 @@ import torch
 from sglang.jit_kernel.awq_dequantize import awq_dequantize as jit_awq_dequantize
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=9, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=9, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)

 try:
@@ -11,7 +11,7 @@ from sglang.jit_kernel.awq_marlin_repack import (
 from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)


@@ -12,7 +12,7 @@ from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_marlin_utils import get_weight_perm, marlin_weights

-register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)


@@ -6,9 +6,9 @@ import torch
 from sglang.jit_kernel.clamp_position import clamp_position_cuda
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_cuda_ci(est_time=12, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=12, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
-register_amd_ci(est_time=12, suite="jit-kernel-unit-test-amd")
+register_amd_ci(est_time=12, stage="jit-kernel-unit", runner_config="amd")


 def _reference_clamp_position(seq_lens):
@@ -7,7 +7,7 @@ import triton

 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=17, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=17, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)


@@ -46,7 +46,8 @@ from sglang.test.ci.ci_register import register_cuda_ci

 register_cuda_ci(
    est_time=300,
-    suite="base-b-kernel-unit-8-gpu-h200",
+    stage="base-b-kernel-unit",
+    runner_config="8-gpu-h200",
 )
 register_cuda_ci(
    est_time=300,
@@ -29,7 +29,7 @@ try:
 except ImportError:
    TRITON_AVAILABLE = False

-register_cuda_ci(est_time=5, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=5, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)


@@ -101,7 +101,7 @@ from sglang.jit_kernel.tests.deepseek_v4.common import (
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.utils import is_in_ci

-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_amd_ci(est_time=25, suite="nightly-amd-kernel-1-gpu", nightly=True)

 Mode = Literal["decode", "prefill"]
@@ -15,7 +15,7 @@ from sglang.jit_kernel.flash_attention import flash_attn_with_kvcache
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase

-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")

 # FA3 only_qv path is SM90 (Hopper) only — skip on pre-Hopper and on
 # Blackwell+ (sm100+) where FA3 is not built.
@@ -14,8 +14,8 @@ from einops import rearrange, repeat
 from sglang.jit_kernel.flash_attention import flash_attn_varlen_func
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=120, suite="base-b-kernel-unit-1-gpu-large")
-register_cuda_ci(est_time=120, suite="base-b-kernel-unit-1-gpu-b200")
+register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
 register_cuda_ci(est_time=900, suite="nightly-kernel-1-gpu", nightly=True)

 # Skip this test on Hopper machine
@@ -7,7 +7,7 @@ import torch
 from sglang.jit_kernel.utils import get_ci_test_range
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)


@@ -16,7 +16,7 @@ import torch

 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=100, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=100, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=400, suite="nightly-kernel-1-gpu", nightly=True)

 # =============================================================================
@@ -7,7 +7,7 @@ import torch
 from sglang.srt.models.utils import fused_qk_gemma_rmsnorm_with_gate
 from sglang.test.ci.ci_register import register_amd_ci

-register_amd_ci(est_time=20, suite="jit-kernel-unit-test-amd")
+register_amd_ci(est_time=20, stage="jit-kernel-unit", runner_config="amd")


 def reference_qk_gemma_rmsnorm_with_gate(
@@ -48,7 +48,7 @@ try:
 except ImportError:
    _is_fp8_fnuz = False

-register_cuda_ci(est_time=24, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=24, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
 register_amd_ci(est_time=24, suite="nightly-amd-kernel-1-gpu", nightly=True)

@@ -26,7 +26,7 @@ try:
 except ImportError:
    KERNELS_AVAILABLE = False

-register_cuda_ci(est_time=6, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=6, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
 register_amd_ci(est_time=10, suite="nightly-amd-kernel-1-gpu", nightly=True)

@@ -23,7 +23,7 @@ from sglang.test.test_marlin_utils import (
    marlin_quantize,
 )

-register_cuda_ci(est_time=13, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=13, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)

 MNK_FACTORS = [
@@ -13,7 +13,7 @@ from sglang.srt.layers.quantization.utils import (
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_marlin_utils import get_weight_perm, marlin_weights

-register_cuda_ci(est_time=16, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=16, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)

 MARLIN_K_CHUNKS = [128]
@@ -9,7 +9,7 @@ from sglang.jit_kernel.utils import get_ci_test_range
 from sglang.srt.layers.moe.topk import biased_grouped_topk_impl
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)


@@ -16,7 +16,7 @@ from sglang.jit_kernel.hadamard import (
 )
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=128, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=128, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=512, suite="nightly-kernel-1-gpu", nightly=True)

 # Exact M×N Hadamard matrices (±1 entries) copied from
@@ -16,7 +16,7 @@ from sglang.srt.mem_cache.pool_host.common import (
 from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)

 pytestmark = pytest.mark.skipif(
@@ -11,8 +11,8 @@ from sglang.jit_kernel.hisparse import (
 from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

-register_amd_ci(est_time=30, suite="stage-b-test-1-gpu-small-amd")
-register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
+register_amd_ci(est_time=30, stage="stage-b", runner_config="1-gpu-small-amd")
+register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)

 pytestmark = pytest.mark.skipif(
@@ -17,7 +17,7 @@ import torch
 from sglang.jit_kernel.kpool_topk_transform import fast_kpool_topk_transform_fused
 from sglang.test.ci.ci_register import register_cuda_ci

-register_cuda_ci(est_time=60, suite="base-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large")


 def _ref_torch_kpool_transform_impl(
--- a/Show More
+++ b/Show More