[CI] Move JIT kernel tests + benchmarks to test/registered/jit; add in-package guard (#27644)
This commit is contained in:
Liangsheng Yin
@@ -433,7 +433,7 @@ if torch.cuda.get_device_capability()[0] < 9:
|
||||
|
||||
## Step 4: Write tests (required)
|
||||
|
||||
JIT kernel tests live under `python/sglang/jit_kernel/tests/`. **CI does not run `pytest` in that directory directly.** The unified runner `test/run_suite.py` discovers every `test_*.py` there (and every `bench_*.py` under `benchmark/`), 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.
|
||||
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`.
|
||||
- **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`).
|
||||
@@ -464,7 +464,7 @@ Use `register_cuda_ci(..., disabled="reason")` if the file must stay in-tree but
|
||||
|
||||
For fast iteration you can still run `pytest` on a single file locally; CI coverage is via `run_suite.py`.
|
||||
|
||||
Create `python/sglang/jit_kernel/tests/test_scale.py`:
|
||||
Create `test/registered/jit/test_scale.py`:
|
||||
|
||||
```python
|
||||
import pytest
|
||||
@@ -517,7 +517,7 @@ if __name__ == "__main__":
|
||||
|
||||
## Step 5: Add a benchmark (required)
|
||||
|
||||
Benchmarks are `bench_*.py` files under `python/sglang/jit_kernel/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-1-gpu-large`** (PR JIT benchmark job: `python3 run_suite.py --hw cuda --suite base-b-kernel-benchmark-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:
|
||||
|
||||
@@ -531,7 +531,7 @@ Benchmarks use the project's own `marker` framework (in `python/sglang/jit_kerne
|
||||
- **`utils.create_random(*shape)` / `utils.create_empty(*shape)`** — shorthand for `torch.randn` / `torch.empty` with `DEFAULT_DTYPE` (`bfloat16`) and `DEFAULT_DEVICE` (`"cuda"`). Override via the `dtype=` / `device=` kwargs.
|
||||
- **`utils.get_benchmark_range(full_range, ci_range)`** — returns the smaller `ci_range` under CI (`is_in_ci()`), the `full_range` locally. Use this so PR CI stays fast while local sweeps stay broad.
|
||||
|
||||
Create `python/sglang/jit_kernel/benchmark/bench_scale.py`:
|
||||
Create `test/registered/jit/benchmark/bench_scale.py`:
|
||||
|
||||
```python
|
||||
import torch
|
||||
@@ -593,7 +593,7 @@ if __name__ == "__main__":
|
||||
Run locally:
|
||||
|
||||
```bash
|
||||
python python/sglang/jit_kernel/benchmark/bench_scale.py
|
||||
python test/registered/jit/benchmark/bench_scale.py
|
||||
```
|
||||
|
||||
Run the benchmark suite the way CI does:
|
||||
@@ -617,7 +617,7 @@ cd test && python3 run_suite.py --hw cuda --suite base-b-kernel-benchmark-1-gpu-
|
||||
## References
|
||||
|
||||
- `docs/developer_guide/development_jit_kernel_guide.md`
|
||||
- `test/run_suite.py` — suite names, discovery of `jit_kernel/tests/` and `jit_kernel/benchmark/`, execution entrypoint for CI
|
||||
- `test/run_suite.py` — suite names, discovery of `test/registered/`, execution entrypoint for CI
|
||||
- `python/sglang/test/ci/ci_register.py` — `register_cuda_ci` and AST registration rules
|
||||
- `python/sglang/jit_kernel/utils.py` — `cache_once`, `load_jit`, `make_cpp_args`, `should_run_full_tests`, `get_ci_test_range`
|
||||
- `python/sglang/jit_kernel/include/sgl_kernel/tensor.h` — `TensorMatcher`, `SymbolicSize/DType/Device`
|
||||
@@ -635,14 +635,14 @@ cd test && python3 run_suite.py --hw cuda --suite base-b-kernel-benchmark-1-gpu-
|
||||
- `python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh` — real example using `runtime::get_blocks_per_sm` + persistent kernel pattern
|
||||
- `python/sglang/jit_kernel/benchmark/marker.py` — `mark_benchmark`, `mark_args`, `do_bench`, `BenchResult`
|
||||
- `python/sglang/jit_kernel/benchmark/utils.py` — `create_random` / `create_empty` / `get_benchmark_range` helpers and `DEFAULT_DTYPE` / `DEFAULT_DEVICE`
|
||||
- `python/sglang/jit_kernel/benchmark/bench_qknorm.py` — real example: multi-axis `mark_args` + `memory_args="all"`
|
||||
- `python/sglang/jit_kernel/benchmark/bench_store_cache.py` — real example: scoped `memory_args` + selective `graph_clone_args`
|
||||
- `test/registered/jit/benchmark/bench_qknorm.py` — real example: multi-axis `mark_args` + `memory_args="all"`
|
||||
- `test/registered/jit/benchmark/bench_store_cache.py` — real example: scoped `memory_args` + selective `graph_clone_args`
|
||||
|
||||
## Summary of Files Created
|
||||
|
||||
```
|
||||
python/sglang/jit_kernel/csrc/elementwise/scale.cuh # NEW: CUDA kernel
|
||||
python/sglang/jit_kernel/scale.py # NEW: Python wrapper
|
||||
python/sglang/jit_kernel/tests/test_scale.py # NEW: Tests
|
||||
python/sglang/jit_kernel/benchmark/bench_scale.py # NEW: Benchmark
|
||||
test/registered/jit/test_scale.py # NEW: Tests
|
||||
test/registered/jit/benchmark/bench_scale.py # NEW: Benchmark
|
||||
```
|
||||
|
||||
@@ -11,15 +11,15 @@ This skill covers **how to write and register tests**. For CI pipeline internals
|
||||
|
||||
1. **Always use `CustomTestCase`** — never raw `unittest.TestCase`. It ensures `tearDownClass` runs even when `setUpClass` fails, preventing resource leaks in CI.
|
||||
2. **`tearDownClass` must be defensive** — use `hasattr`/null checks before accessing resources (e.g. `cls.process`) that `setUpClass` may not have finished allocating.
|
||||
3. **Place tests in `test/registered/<category>/`** — except JIT kernel tests and benchmarks, which live in `python/sglang/jit_kernel/tests/` and `python/sglang/jit_kernel/benchmark/` (nested subfolders are allowed)
|
||||
3. **Place tests in `test/registered/<category>/`** — including JIT kernel tests and benchmarks, which live in `test/registered/jit/` and `test/registered/jit/benchmark/` (nested subfolders are allowed)
|
||||
4. **Reuse server fixtures** — inherit from `DefaultServerBase` or write `setUpClass`/`tearDownClass` with `popen_launch_server`
|
||||
5. **Prefer mock over real server** — when testing logic that doesn't need a server / engine launch (middleware, request routing, config validation, argument parsing), use `unittest.mock.patch` / `MagicMock` and place tests in `test/registered/unit/`. Only launch a real server when the test genuinely needs inference results or server lifecycle behavior.
|
||||
|
||||
JIT kernel exception:
|
||||
JIT kernel notes:
|
||||
- If the task is adding or updating code under `python/sglang/jit_kernel/`, prefer the `add-jit-kernel` skill first.
|
||||
- JIT kernel correctness tests use `python/sglang/jit_kernel/tests/**/test_*.py`.
|
||||
- JIT kernel benchmarks use `python/sglang/jit_kernel/benchmark/**/bench_*.py`.
|
||||
- Those files are still executed by `test/run_suite.py`, but through dedicated kernel suites rather than `test/registered/`.
|
||||
- JIT kernel correctness tests use `test/registered/jit/**/test_*.py`.
|
||||
- JIT kernel benchmarks use `test/registered/jit/benchmark/**/bench_*.py`.
|
||||
- Those files are executed by `test/run_suite.py` through dedicated kernel suites (`base-b-kernel-*`); a `register_*_ci(...)` call placed under `python/sglang/` is rejected by the `check-no-registered-tests-in-package` pre-commit hook.
|
||||
|
||||
---
|
||||
|
||||
@@ -65,10 +65,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 `python/sglang/jit_kernel/tests/` |
|
||||
| `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 `python/sglang/jit_kernel/tests/` |
|
||||
| `base-b-kernel-benchmark-1-gpu-large` | `1-gpu-h100` | JIT kernel benchmark files under `python/sglang/jit_kernel/benchmark/` |
|
||||
| `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-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 |
|
||||
@@ -352,12 +352,12 @@ JIT kernel files live outside `test/registered/` but still use registration:
|
||||
```python
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
# Correctness tests in python/sglang/jit_kernel/tests/
|
||||
# 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")
|
||||
|
||||
# Benchmarks in python/sglang/jit_kernel/benchmark/
|
||||
# Benchmarks in test/registered/jit/benchmark/
|
||||
register_cuda_ci(est_time=6, suite="base-b-kernel-benchmark-1-gpu-large")
|
||||
|
||||
# Optional nightly registration
|
||||
@@ -393,7 +393,7 @@ python/sglang/jit_kernel/
|
||||
|
||||
**Decision rule** (see also `test/registered/README.md`):
|
||||
- Component logic, no server → `registered/unit/`
|
||||
- JIT kernel correctness / benchmarks → `python/sglang/jit_kernel/tests/` or `python/sglang/jit_kernel/benchmark/`
|
||||
- JIT kernel correctness / benchmarks → `test/registered/jit/` or `test/registered/jit/benchmark/`
|
||||
- Other kernel correctness → `registered/kernels/`
|
||||
- Server needed → `registered/<category>/`
|
||||
- Local debugging → `manual/`
|
||||
@@ -437,8 +437,8 @@ Before submitting a test:
|
||||
|
||||
- [ ] Inherits from `CustomTestCase` (not `unittest.TestCase`)
|
||||
- [ ] Has `register_*_ci(...)` call at module level
|
||||
- [ ] Placed in `test/registered/<category>/`, unless this is a JIT kernel test/benchmark
|
||||
- [ ] JIT kernel work: files live in `python/sglang/jit_kernel/tests/` or `python/sglang/jit_kernel/benchmark/`
|
||||
- [ ] Placed in `test/registered/<category>/` (JIT kernel test/benchmark → `test/registered/jit/` or `test/registered/jit/benchmark/`)
|
||||
- [ ] JIT kernel work: test files live in `test/registered/jit/`; only test-only helpers stay under `python/sglang/jit_kernel/`
|
||||
- [ ] Backend-independent tests: `register_cuda_ci` only + smallest model
|
||||
- [ ] Logic that doesn't need a server / engine launch → unit test in `registered/unit/` (see Unit Tests section)
|
||||
- [ ] `setUpClass` launches server, `tearDownClass` kills it (if server-based)
|
||||
|
||||
Reference in New Issue
Block a user