diff --git a/.claude/skills/add-jit-kernel/SKILL.md b/.claude/skills/add-jit-kernel/SKILL.md index ce6116a54..1bbe3c83e 100644 --- a/.claude/skills/add-jit-kernel/SKILL.md +++ b/.claude/skills/add-jit-kernel/SKILL.md @@ -414,12 +414,40 @@ 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. + +- **PR / per-commit CUDA suites** (see `test/run_suite.py` → `PER_COMMIT_SUITES`): JIT unit tests use `stage-b-kernel-unit-1-gpu-large` (see `.github/workflows/pr-test-jit-kernel.yml`: `python3 run_suite.py --hw cuda --suite stage-b-kernel-unit-1-gpu-large`). +- **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): + +```python +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, suite="stage-b-kernel-unit-1-gpu-large") +# Optional second registration: same file also listed under the nightly kernel suite +# register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True) +``` + +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 stage-b-kernel-unit-1-gpu-large +``` + +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`: ```python import pytest 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="stage-b-kernel-unit-1-gpu-large") @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) @@ -464,6 +492,8 @@ 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 **`stage-b-kernel-benchmark-1-gpu-large`** (PR JIT benchmark job: `python3 run_suite.py --hw cuda --suite stage-b-kernel-benchmark-1-gpu-large`). + Create `python/sglang/jit_kernel/benchmark/bench_scale.py`: ```python @@ -480,7 +510,9 @@ from sglang.jit_kernel.benchmark.utils import ( run_benchmark, ) 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="stage-b-kernel-benchmark-1-gpu-large") SIZE_LIST = get_benchmark_range( full_range=[2**n for n in range(10, 20)], # 1K … 512K elements @@ -519,16 +551,23 @@ if __name__ == "__main__": benchmark.run(print_data=True) ``` -Run: +Run locally: ```bash python python/sglang/jit_kernel/benchmark/bench_scale.py ``` +Run the benchmark suite the way CI does: + +```bash +cd test && python3 run_suite.py --hw cuda --suite stage-b-kernel-benchmark-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 - **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**: `run_benchmark` uses CUDA-graph-based timing by default @@ -538,7 +577,9 @@ python python/sglang/jit_kernel/benchmark/bench_scale.py ## References - `docs/developer_guide/development_jit_kernel_guide.md` -- `python/sglang/jit_kernel/utils.py` — `cache_once`, `load_jit`, `make_cpp_args` +- `test/run_suite.py` — suite names, discovery of `jit_kernel/tests/` and `jit_kernel/benchmark/`, 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` - `python/sglang/jit_kernel/include/sgl_kernel/utils.cuh` — type aliases, `LaunchKernel`, `SGL_DEVICE` - `python/sglang/jit_kernel/include/sgl_kernel/vec.cuh` — `AlignedVector`