[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>
This commit is contained in:
Khoa Pham
2026-06-29 19:12:14 -07:00
committed by GitHub
co-authored by Cursor Claude Opus 4.8
parent bc8b3ab1f5
commit 3e16be2122
128 changed files with 274 additions and 188 deletions
+15 -12
View File
@@ -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
+6 -4
View File
@@ -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:**
+20 -15
View File
@@ -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.
---