[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
|
||||
```
|
||||
|
||||
Reference in New Issue
Block a user