diff --git a/.claude/skills/add-jit-kernel/SKILL.md b/.claude/skills/add-jit-kernel/SKILL.md index ee6fd15c4..66070f92e 100644 --- a/.claude/skills/add-jit-kernel/SKILL.md +++ b/.claude/skills/add-jit-kernel/SKILL.md @@ -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 ``` diff --git a/.claude/skills/write-sglang-test/SKILL.md b/.claude/skills/write-sglang-test/SKILL.md index 7f04c4bde..2f1f832f6 100644 --- a/.claude/skills/write-sglang-test/SKILL.md +++ b/.claude/skills/write-sglang-test/SKILL.md @@ -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//`** — 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//`** — 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//` - 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//`, 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//` (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) diff --git a/.github/workflows/_pr-test-check-changes.yml b/.github/workflows/_pr-test-check-changes.yml index c7b221170..dad1acfa8 100644 --- a/.github/workflows/_pr-test-check-changes.yml +++ b/.github/workflows/_pr-test-check-changes.yml @@ -93,14 +93,15 @@ jobs: - "python/sglang/multimodal_gen/**/!(*.md|*.ipynb)" - "python/sglang/srt/observability/**" - "python/sglang/jit_kernel/**" - - "python/sglang/jit_kernel/tests/diffusion/**" - - "python/sglang/jit_kernel/benchmark/diffusion/**" + - "test/registered/jit/diffusion/**" + - "test/registered/jit/benchmark/diffusion/**" - "python/sglang/cli/**" jit_kernel: - ".github/workflows/pr-test.yml" - ".github/workflows/pr-test-jit-kernel.yml" - "python/pyproject.toml" - "python/sglang/jit_kernel/**" + - "test/registered/jit/**" sgl_kernel: # Intentionally excludes ".github/workflows/pr-test-sgl-kernel.yml" — # see API-side detector below for rationale. diff --git a/.github/workflows/pr-test-amd-rocm720.yml b/.github/workflows/pr-test-amd-rocm720.yml index cef1df15a..c536369dc 100644 --- a/.github/workflows/pr-test-amd-rocm720.yml +++ b/.github/workflows/pr-test-amd-rocm720.yml @@ -177,14 +177,15 @@ jobs: - ".github/workflows/pr-test-amd-rocm720.yml" jit_kernel: - "python/sglang/jit_kernel/**" + - "test/registered/jit/**" - ".github/workflows/pr-test-amd-rocm720.yml" multimodal_gen: - "python/sglang/multimodal_gen/**/!(*.md|*.ipynb)" - "python/sglang/cli/**" - "python/sglang/srt/observability/**" - "python/sglang/jit_kernel/diffusion/**" - - "python/sglang/jit_kernel/tests/diffusion/**" - - "python/sglang/jit_kernel/benchmark/diffusion/**" + - "test/registered/jit/diffusion/**" + - "test/registered/jit/benchmark/diffusion/**" - "python/pyproject_rocm.toml" - "python/pyproject_other.toml" diff --git a/.github/workflows/pr-test-amd.yml b/.github/workflows/pr-test-amd.yml index 255d70019..0a569e352 100644 --- a/.github/workflows/pr-test-amd.yml +++ b/.github/workflows/pr-test-amd.yml @@ -175,14 +175,15 @@ jobs: - ".github/workflows/pr-test-amd.yml" jit_kernel: - "python/sglang/jit_kernel/**" + - "test/registered/jit/**" - ".github/workflows/pr-test-amd.yml" multimodal_gen: - "python/sglang/multimodal_gen/**/!(*.md|*.ipynb)" - "python/sglang/cli/**" - "python/sglang/srt/observability/**" - "python/sglang/jit_kernel/diffusion/**" - - "python/sglang/jit_kernel/tests/diffusion/**" - - "python/sglang/jit_kernel/benchmark/diffusion/**" + - "test/registered/jit/diffusion/**" + - "test/registered/jit/benchmark/diffusion/**" - "python/pyproject_rocm.toml" - "python/pyproject_other.toml" diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 97a51f40b..1887be7a3 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -388,7 +388,7 @@ jobs: caller_inputs: ${{ toJson(inputs) }} partitions: ${{ needs.check-changes.outputs.partitions }} run_timeout_minutes: '40' - extra_pytest_path: 'python/sglang/jit_kernel/tests/test_flash_attention_4.py' + extra_pytest_path: 'test/registered/jit/test_flash_attention_4.py' secrets: inherit call-multimodal-gen-tests: diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3a8958d40..54f328b0e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -99,6 +99,12 @@ repos: files: ^test/registered/.*\.py$ exclude: ^test/registered/.*/utils\.py$ pass_filenames: false + - id: check-no-registered-tests-in-package + name: reject CI-registered tests inside the sglang package + entry: python3 scripts/ci/check_no_registered_tests_in_package.py + language: system + files: ^python/sglang/.*\.py$ + pass_filenames: false - id: check-no-docs-changes name: reject changes under legacy docs/ entry: python3 scripts/ci/check_no_docs_changes.py diff --git a/docs/developer_guide/development_jit_kernel_guide.md b/docs/developer_guide/development_jit_kernel_guide.md index b09476e48..0059a82bd 100644 --- a/docs/developer_guide/development_jit_kernel_guide.md +++ b/docs/developer_guide/development_jit_kernel_guide.md @@ -262,7 +262,7 @@ Finally, import and use the kernel like a regular Python function: from sglang.jit_kernel.add_constant import add_constant ``` -For a complete, runnable example, refer to [test_add_constant.py](../../python/sglang/jit_kernel/tests/test_add_constant.py). +For a complete, runnable example, refer to [test_add_constant.py](../../test/registered/jit/test_add_constant.py). ## C++ Include Library Reference diff --git a/docs_new/docs/developer_guide/development_jit_kernel_guide.mdx b/docs_new/docs/developer_guide/development_jit_kernel_guide.mdx index 8b92fcd10..8c5623783 100644 --- a/docs_new/docs/developer_guide/development_jit_kernel_guide.mdx +++ b/docs_new/docs/developer_guide/development_jit_kernel_guide.mdx @@ -266,7 +266,7 @@ Finally, import and use the kernel like a regular Python function: from sglang.jit_kernel.add_constant import add_constant ``` -For a complete, runnable example, refer to [test_add_constant.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/tests/test_add_constant.py). +For a complete, runnable example, refer to [test_add_constant.py](https://github.com/sgl-project/sglang/blob/main/test/registered/jit/test_add_constant.py). ## C++ Include Library Reference diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md index 932a49e07..8b7739a30 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md @@ -19,12 +19,12 @@ framework-specific optimization workflow. - `python/sglang/jit_kernel/diffusion/triton/ltx2_rotary.py` - `python/sglang/jit_kernel/diffusion/triton/varlen_pack_pad.py` - `python/sglang/jit_kernel/diffusion/cutedsl/scale_residual_norm_scale_shift.py` -- `python/sglang/jit_kernel/tests/diffusion/test_qwen_image_modulation.py` -- `python/sglang/jit_kernel/tests/diffusion/test_group_norm_silu.py` -- `python/sglang/jit_kernel/tests/diffusion/test_varlen_pack_pad.py` -- `python/sglang/jit_kernel/tests/diffusion/test_varlen_uspattn_equivalence.py` -- `python/sglang/jit_kernel/benchmark/diffusion/bench_qwen_image_modulation.py` -- `python/sglang/jit_kernel/benchmark/diffusion/bench_group_norm_silu.py` +- `test/registered/jit/diffusion/test_qwen_image_modulation.py` +- `test/registered/jit/diffusion/test_group_norm_silu.py` +- `test/registered/jit/diffusion/test_varlen_pack_pad.py` +- `test/registered/jit/diffusion/test_varlen_uspattn_equivalence.py` +- `test/registered/jit/benchmark/diffusion/bench_qwen_image_modulation.py` +- `test/registered/jit/benchmark/diffusion/bench_group_norm_silu.py` - `python/sglang/jit_kernel/norm.py` - `python/sglang/multimodal_gen/runtime/platforms/cuda.py` - `python/sglang/multimodal_gen/runtime/layers/attention/selector.py` @@ -38,7 +38,7 @@ framework-specific optimization workflow. - Use cases: `x * (1 + scale) + shift`, `a * (k + b) + c`, and Qwen-style `(layernorm/residual layernorm) + scale/shift + gate select`. - Constraints: `x` must be CUDA and contiguous. `scale/shift` support 0D/1D/2D/3D/4D broadcast. 4D `[B, F, 1, C]` requires `L % F == 0`. - NPU fallback: `scale_shift.py` swaps to `npu_fallback` native path. -- Validation: `python/sglang/jit_kernel/tests/diffusion/test_qwen_image_modulation.py`. +- Validation: `test/registered/jit/diffusion/test_qwen_image_modulation.py`. 2. Norm + Scale/Shift fusion (CuTe DSL) - Kernels: `fused_norm_scale_shift`, `fused_scale_residual_norm_scale_shift` @@ -56,7 +56,7 @@ framework-specific optimization workflow. - `y = tanh(gate) * norm(x) + shift` - `y, y2 = tanh(gate) * norm(x) + shift`, then `y2 = norm(y) * (1 + scale)` - Constraints: same CuTe DSL envelope as the norm+scale/shift family in practice: contiguous last dim, fp16/bf16/fp32, and `D % 256 == 0`, `D <= 8192`. -- Validation: `python/sglang/jit_kernel/tests/diffusion/test_norm_tanh_mul_add_norm_scale.py` +- Validation: `test/registered/jit/diffusion/test_norm_tanh_mul_add_norm_scale.py` - Behavior: this is already a mainline fast path, so if Z-Image traces show the unfused chain, treat it as a missing or regressed existing optimization before proposing a new kernel. 4. Triton LayerNorm/RMSNorm fusion @@ -64,7 +64,7 @@ framework-specific optimization workflow. - Locations: `triton/norm.py`, `layernorm.py` - Use cases: fp32 RMSNorm with residual/dropout/rowscale/x1 branches, and inference-friendly `norm_infer`. - Constraints: last dim must be contiguous, and `N * element_size < 64KB`. -- Validation: `python/sglang/jit_kernel/tests/test_rmsnorm.py`. +- Validation: `test/registered/jit/test_rmsnorm.py`. 5. Triton one-pass RMSNorm (small hidden size fast path) - Kernel: `triton_one_pass_rms_norm` @@ -78,7 +78,7 @@ framework-specific optimization workflow. - Use case: GPT-J style RoPE when not Neox. - Constraints: `head_size` must be even. - NPU fallback: `npu_fallback.apply_rotary_embedding_native`. -- Validation: `python/sglang/jit_kernel/tests/test_rope.py`. +- Validation: `test/registered/jit/test_rope.py`. 7. LTX2 split RoPE fusion - Kernel: `apply_ltx2_split_rotary_emb` @@ -93,8 +93,8 @@ framework-specific optimization workflow. - Use case: `activation(group_norm(x))` when the activation is non-inplace `nn.SiLU` and the GroupNorm is affine. - Enablement: mainline uses `apply_group_norm_silu(...)` in HunyuanVideo VAE paths and LTX latent upsampler paths by default; there is no env toggle. The wrapper dispatches to Triton only when guards pass. - Constraints: CUDA inference path only; no grad, `x.requires_grad == False`, `nn.GroupNorm`, `nn.SiLU(inplace=False)`, affine norm with weight and bias. Unsupported cases fall back to native `activation(norm(x))`. -- Validation: `python/sglang/jit_kernel/tests/diffusion/test_group_norm_silu.py`. -- Microbench: `python/sglang/jit_kernel/benchmark/diffusion/bench_group_norm_silu.py`. +- Validation: `test/registered/jit/diffusion/test_group_norm_silu.py`. +- Microbench: `test/registered/jit/benchmark/diffusion/bench_group_norm_silu.py`. **Faster CUDA Kernel Usage Points** @@ -117,7 +117,7 @@ framework-specific optimization workflow. 4. Varlen USP attention pack/scatter - Locations: `runtime/layers/attention/layer.py`, `triton/varlen_pack_pad.py` - Behavior: masked `USPAttention.forward` can gather dense Q/K/V into packed `[total_valid, H, D]` rows with `fused_pack_qkv`, run varlen attention, then scatter back with `fused_scatter_to_padded`. -- Validation: `python/sglang/jit_kernel/tests/diffusion/test_varlen_pack_pad.py` and `test_varlen_uspattn_equivalence.py`. +- Validation: `test/registered/jit/diffusion/test_varlen_pack_pad.py` and `test_varlen_uspattn_equivalence.py`. - Workflow rule: if a masked attention trace spends time in Python/advanced indexing pack or scatter, first check whether this fused varlen path should have engaged. **QK Norm Optimization** @@ -130,7 +130,7 @@ framework-specific optimization workflow. - `can_use_fused_inplace_qknorm(head_dim, dtype)` returns true. - Supported head dims: `64, 128, 256, 512, 1024`. - Behavior: Fused path operates on `q` and `k` in place after reshaping to `[B, -1, head_dim]`. If preconditions fail, fall back to per-tensor RMSNorm. -- Validation: `python/sglang/jit_kernel/tests/test_qknorm.py` and `python/sglang/jit_kernel/tests/test_qknorm_across_heads.py`. +- Validation: `test/registered/jit/test_qknorm.py` and `test/registered/jit/test_qknorm_across_heads.py`. **QK Norm + RoPE Optimization** @@ -145,7 +145,7 @@ framework-specific optimization workflow. - `can_use_fused_inplace_qknorm_rope(head_dim, rope_dim, is_neox, dtype)` returns true. - Supported head dims: `64, 128, 256`. - Behavior: `apply_qk_norm_rope` prefers the fused JIT kernel when all guards pass; otherwise it falls back to `apply_qk_norm(...)` plus `apply_flashinfer_rope_qk_inplace(...)`. -- Validation: `python/sglang/jit_kernel/tests/diffusion/test_qknorm_rope.py`. +- Validation: `test/registered/jit/diffusion/test_qknorm_rope.py`. - Workflow rule: treat LTX2 traces that miss the generic fused path as an enablement/shape-guard issue first, and check the separate LTX2 split-RoPE path before proposing new attention-prep kernels. **Nunchaku Fused GELU MLP** diff --git a/scripts/ci/check_no_docs_changes.py b/scripts/ci/check_no_docs_changes.py index 9ab5f32ae..b3c02cb5c 100755 --- a/scripts/ci/check_no_docs_changes.py +++ b/scripts/ci/check_no_docs_changes.py @@ -17,6 +17,9 @@ LEGACY_DOCS_ALLOWLIST = { "docs/_static/css/custom_log.css", "docs/_static/js/deprecation_banner.js", "docs/conf.py", + # Has relative links into the source tree that the offline lychee check + # validates, so it must be updated when the linked source files move. + "docs/developer_guide/development_jit_kernel_guide.md", } diff --git a/scripts/ci/check_no_registered_tests_in_package.py b/scripts/ci/check_no_registered_tests_in_package.py new file mode 100755 index 000000000..a1e6a30ff --- /dev/null +++ b/scripts/ci/check_no_registered_tests_in_package.py @@ -0,0 +1,81 @@ +#!/usr/bin/env python3 +""" +Pre-commit hook: reject CI-registered tests that live inside the importable +`sglang` package (python/sglang/). + +Registered tests and benchmarks must live under test/registered/ (e.g. +test/registered/jit/ for JIT kernel tests and test/registered/jit/benchmark/ +for JIT kernel benchmarks) so they are not shipped in the wheel and are +collected by run_suite.py's registered glob. A registered file placed inside +the package would be shipped to users AND silently dropped by run_suite.py +(which no longer globs the package) -- it would never run in CI. This guard +turns that silent skip into a hard failure. + +Reuses ut_parse_one_file() from ci_register.py (AST-based) so the registry +detection matches run_suite.py's collect_tests() exactly. +""" + +import glob +import importlib.util +import os +import sys + +# Markers whose mere presence in the source is worth an AST parse. Anything +# without one of these strings cannot register a test, so we skip parsing it. +_MARKERS = ( + "register_cuda_ci", + "register_amd_ci", + "register_cpu_ci", + "register_npu_ci", + "register_xpu_ci", + "register_musa_ci", +) + + +def main() -> int: + # Import ci_register directly to avoid pulling in all of sglang. + spec = importlib.util.spec_from_file_location( + "ci_register", + os.path.join("python", "sglang", "test", "ci", "ci_register.py"), + ) + ci_register = importlib.util.module_from_spec(spec) + spec.loader.exec_module(ci_register) + + offenders = [] + for f in sorted(glob.glob("python/sglang/**/*.py", recursive=True)): + try: + with open(f, "r", encoding="utf-8") as fh: + source = fh.read() + except (OSError, UnicodeDecodeError): + continue + if not any(marker in source for marker in _MARKERS): + continue + try: + registries, _has_main_entry = ci_register.ut_parse_one_file(f) + except Exception: + # A malformed register call still indicates a misplaced test. + offenders.append(f) + continue + if registries: + offenders.append(f) + + if offenders: + print( + "ERROR: CI-registered test(s)/benchmark(s) found inside the sglang package:" + ) + print( + " Registered tests and benchmarks must live under test/registered/\n" + " (e.g. test/registered/jit/ for JIT kernel tests and\n" + " test/registered/jit/benchmark/ for JIT kernel benchmarks) so they\n" + " are not shipped in the wheel and are collected by run_suite.py.\n" + ) + for f in offenders: + print(f" {f}") + print() + return 1 + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/test/README.md b/test/README.md index 414e02b40..ddc526f75 100644 --- a/test/README.md +++ b/test/README.md @@ -45,7 +45,7 @@ python3 test/registered/core/test_srt_endpoint.py python3 test/registered/core/test_srt_endpoint.py TestSRTEndpoint.test_simple_decode # Single JIT kernel test -python3 python/sglang/jit_kernel/tests/test_add_constant.py +python3 test/registered/jit/test_add_constant.py # Run a suite python3 test/run_suite.py --hw cpu --suite base-a-test-cpu @@ -73,9 +73,9 @@ 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 files live outside `test/registered/` but still use registration: -- Correctness tests: `python/sglang/jit_kernel/tests/test_*.py` → `base-b-kernel-unit-1-gpu-large` -- Benchmarks: `python/sglang/jit_kernel/benchmark/bench_*.py` → `base-b-kernel-benchmark-1-gpu-large` +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` ## Choosing a Suite diff --git a/python/sglang/jit_kernel/benchmark/bench_activation.py b/test/registered/jit/benchmark/bench_activation.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_activation.py rename to test/registered/jit/benchmark/bench_activation.py diff --git a/python/sglang/jit_kernel/benchmark/bench_add_constant.py b/test/registered/jit/benchmark/bench_add_constant.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_add_constant.py rename to test/registered/jit/benchmark/bench_add_constant.py diff --git a/python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py b/test/registered/jit/benchmark/bench_awq_dequantize.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py rename to test/registered/jit/benchmark/bench_awq_dequantize.py diff --git a/python/sglang/jit_kernel/benchmark/bench_clamp_position.py b/test/registered/jit/benchmark/bench_clamp_position.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_clamp_position.py rename to test/registered/jit/benchmark/bench_clamp_position.py diff --git a/python/sglang/jit_kernel/benchmark/bench_concat_mla.py b/test/registered/jit/benchmark/bench_concat_mla.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_concat_mla.py rename to test/registered/jit/benchmark/bench_concat_mla.py diff --git a/python/sglang/jit_kernel/benchmark/bench_custom_all_reduce.py b/test/registered/jit/benchmark/bench_custom_all_reduce.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_custom_all_reduce.py rename to test/registered/jit/benchmark/bench_custom_all_reduce.py diff --git a/python/sglang/jit_kernel/benchmark/bench_dsv4_fp4_indexer.py b/test/registered/jit/benchmark/bench_dsv4_fp4_indexer.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_dsv4_fp4_indexer.py rename to test/registered/jit/benchmark/bench_dsv4_fp4_indexer.py diff --git a/python/sglang/jit_kernel/benchmark/bench_fused_qknorm_rope.py b/test/registered/jit/benchmark/bench_fused_qknorm_rope.py similarity index 99% rename from python/sglang/jit_kernel/benchmark/bench_fused_qknorm_rope.py rename to test/registered/jit/benchmark/bench_fused_qknorm_rope.py index 5ae05f78b..543be0489 100644 --- a/python/sglang/jit_kernel/benchmark/bench_fused_qknorm_rope.py +++ b/test/registered/jit/benchmark/bench_fused_qknorm_rope.py @@ -5,7 +5,7 @@ Measures throughput (µs) for fused_qk_norm_rope across typical LLM configurations (head_dim × num_heads × num_tokens). Run: - python python/sglang/jit_kernel/benchmark/bench_fused_qknorm_rope.py + python test/registered/jit/benchmark/bench_fused_qknorm_rope.py """ import itertools diff --git a/python/sglang/jit_kernel/benchmark/bench_hadamard.py b/test/registered/jit/benchmark/bench_hadamard.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_hadamard.py rename to test/registered/jit/benchmark/bench_hadamard.py diff --git a/python/sglang/jit_kernel/benchmark/bench_hicache.py b/test/registered/jit/benchmark/bench_hicache.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_hicache.py rename to test/registered/jit/benchmark/bench_hicache.py diff --git a/python/sglang/jit_kernel/benchmark/bench_hisparse.py b/test/registered/jit/benchmark/bench_hisparse.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_hisparse.py rename to test/registered/jit/benchmark/bench_hisparse.py diff --git a/python/sglang/jit_kernel/benchmark/bench_mla_kv_pack_quantize_fp8.py b/test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_mla_kv_pack_quantize_fp8.py rename to test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py diff --git a/python/sglang/jit_kernel/benchmark/bench_mxfp8_moe.py b/test/registered/jit/benchmark/bench_mxfp8_moe.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_mxfp8_moe.py rename to test/registered/jit/benchmark/bench_mxfp8_moe.py diff --git a/python/sglang/jit_kernel/benchmark/bench_ngram_compute_decode.py b/test/registered/jit/benchmark/bench_ngram_compute_decode.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_ngram_compute_decode.py rename to test/registered/jit/benchmark/bench_ngram_compute_decode.py diff --git a/python/sglang/jit_kernel/benchmark/bench_ngram_update_token_table.py b/test/registered/jit/benchmark/bench_ngram_update_token_table.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_ngram_update_token_table.py rename to test/registered/jit/benchmark/bench_ngram_update_token_table.py diff --git a/python/sglang/jit_kernel/benchmark/bench_norm.py b/test/registered/jit/benchmark/bench_norm.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_norm.py rename to test/registered/jit/benchmark/bench_norm.py diff --git a/python/sglang/jit_kernel/benchmark/bench_nvfp4_blockwise_moe.py b/test/registered/jit/benchmark/bench_nvfp4_blockwise_moe.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_nvfp4_blockwise_moe.py rename to test/registered/jit/benchmark/bench_nvfp4_blockwise_moe.py diff --git a/python/sglang/jit_kernel/benchmark/bench_nvfp4_quant.py b/test/registered/jit/benchmark/bench_nvfp4_quant.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_nvfp4_quant.py rename to test/registered/jit/benchmark/bench_nvfp4_quant.py diff --git a/python/sglang/jit_kernel/benchmark/bench_nvfp4_scaled_mm.py b/test/registered/jit/benchmark/bench_nvfp4_scaled_mm.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_nvfp4_scaled_mm.py rename to test/registered/jit/benchmark/bench_nvfp4_scaled_mm.py diff --git a/python/sglang/jit_kernel/benchmark/bench_per_tensor_quant_fp8.py b/test/registered/jit/benchmark/bench_per_tensor_quant_fp8.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_per_tensor_quant_fp8.py rename to test/registered/jit/benchmark/bench_per_tensor_quant_fp8.py diff --git a/python/sglang/jit_kernel/benchmark/bench_per_token_group_quant_8bit.py b/test/registered/jit/benchmark/bench_per_token_group_quant_8bit.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_per_token_group_quant_8bit.py rename to test/registered/jit/benchmark/bench_per_token_group_quant_8bit.py diff --git a/python/sglang/jit_kernel/benchmark/bench_qknorm.py b/test/registered/jit/benchmark/bench_qknorm.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_qknorm.py rename to test/registered/jit/benchmark/bench_qknorm.py diff --git a/python/sglang/jit_kernel/benchmark/bench_qknorm_across_heads.py b/test/registered/jit/benchmark/bench_qknorm_across_heads.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_qknorm_across_heads.py rename to test/registered/jit/benchmark/bench_qknorm_across_heads.py diff --git a/python/sglang/jit_kernel/benchmark/bench_renorm.py b/test/registered/jit/benchmark/bench_renorm.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_renorm.py rename to test/registered/jit/benchmark/bench_renorm.py diff --git a/python/sglang/jit_kernel/benchmark/bench_resolve_future_token_ids.py b/test/registered/jit/benchmark/bench_resolve_future_token_ids.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_resolve_future_token_ids.py rename to test/registered/jit/benchmark/bench_resolve_future_token_ids.py diff --git a/python/sglang/jit_kernel/benchmark/bench_rope.py b/test/registered/jit/benchmark/bench_rope.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_rope.py rename to test/registered/jit/benchmark/bench_rope.py diff --git a/python/sglang/jit_kernel/benchmark/bench_set_mla_kv_buffer.py b/test/registered/jit/benchmark/bench_set_mla_kv_buffer.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_set_mla_kv_buffer.py rename to test/registered/jit/benchmark/bench_set_mla_kv_buffer.py diff --git a/python/sglang/jit_kernel/benchmark/bench_store_cache.py b/test/registered/jit/benchmark/bench_store_cache.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_store_cache.py rename to test/registered/jit/benchmark/bench_store_cache.py diff --git a/python/sglang/jit_kernel/benchmark/bench_tp_qknorm.py b/test/registered/jit/benchmark/bench_tp_qknorm.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/bench_tp_qknorm.py rename to test/registered/jit/benchmark/bench_tp_qknorm.py diff --git a/python/sglang/jit_kernel/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py b/test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py similarity index 98% rename from python/sglang/jit_kernel/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py rename to test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py index 35c727dec..7ba1ffe87 100644 --- a/python/sglang/jit_kernel/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py +++ b/test/registered/jit/benchmark/diffusion/bench_diffusion_nvfp4_scaled_mm.py @@ -12,6 +12,7 @@ import sgl_kernel import torch from sglang.jit_kernel.benchmark.utils import DEFAULT_DTYPE +from sglang.jit_kernel.utils import KERNEL_PATH from sglang.test.ci.ci_register import register_cuda_ci from sglang.utils import is_in_ci @@ -25,7 +26,9 @@ SCRIPT_DIR = Path(__file__).resolve().parent REPO_ROOT = ( Path(os.environ["SGLANG_NVFP4_REPO_ROOT"]) if os.environ.get("SGLANG_NVFP4_REPO_ROOT") - else Path(__file__).resolve().parents[5] + # Anchor on the installed jit_kernel package (python/sglang/jit_kernel) so + # this stays correct regardless of where the benchmark file lives. + else KERNEL_PATH.parents[2] ) DEFAULT_OUTPUT_DIR = REPO_ROOT / "outputs" / "nvfp4_benchmarks" DEFAULT_SHAPE_LIBRARY = SCRIPT_DIR / "diffusion_nvfp4_shapes.json" diff --git a/python/sglang/jit_kernel/benchmark/diffusion/bench_fused_norm_scale_shift.py b/test/registered/jit/benchmark/diffusion/bench_fused_norm_scale_shift.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/diffusion/bench_fused_norm_scale_shift.py rename to test/registered/jit/benchmark/diffusion/bench_fused_norm_scale_shift.py diff --git a/python/sglang/jit_kernel/benchmark/diffusion/bench_group_norm_silu.py b/test/registered/jit/benchmark/diffusion/bench_group_norm_silu.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/diffusion/bench_group_norm_silu.py rename to test/registered/jit/benchmark/diffusion/bench_group_norm_silu.py diff --git a/python/sglang/jit_kernel/benchmark/diffusion/bench_norm_impls.py b/test/registered/jit/benchmark/diffusion/bench_norm_impls.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/diffusion/bench_norm_impls.py rename to test/registered/jit/benchmark/diffusion/bench_norm_impls.py diff --git a/python/sglang/jit_kernel/benchmark/diffusion/bench_qknorm_rope.py b/test/registered/jit/benchmark/diffusion/bench_qknorm_rope.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/diffusion/bench_qknorm_rope.py rename to test/registered/jit/benchmark/diffusion/bench_qknorm_rope.py diff --git a/python/sglang/jit_kernel/benchmark/diffusion/bench_qwen_image_modulation.py b/test/registered/jit/benchmark/diffusion/bench_qwen_image_modulation.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/diffusion/bench_qwen_image_modulation.py rename to test/registered/jit/benchmark/diffusion/bench_qwen_image_modulation.py diff --git a/python/sglang/jit_kernel/benchmark/diffusion/diffusion_nvfp4_shapes.json b/test/registered/jit/benchmark/diffusion/diffusion_nvfp4_shapes.json similarity index 100% rename from python/sglang/jit_kernel/benchmark/diffusion/diffusion_nvfp4_shapes.json rename to test/registered/jit/benchmark/diffusion/diffusion_nvfp4_shapes.json diff --git a/python/sglang/jit_kernel/benchmark/kv_canary/bench_plan.py b/test/registered/jit/benchmark/kv_canary/bench_plan.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/kv_canary/bench_plan.py rename to test/registered/jit/benchmark/kv_canary/bench_plan.py diff --git a/python/sglang/jit_kernel/benchmark/kv_canary/bench_scatter_req_token_ids.py b/test/registered/jit/benchmark/kv_canary/bench_scatter_req_token_ids.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/kv_canary/bench_scatter_req_token_ids.py rename to test/registered/jit/benchmark/kv_canary/bench_scatter_req_token_ids.py diff --git a/python/sglang/jit_kernel/benchmark/kv_canary/bench_verify.py b/test/registered/jit/benchmark/kv_canary/bench_verify.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/kv_canary/bench_verify.py rename to test/registered/jit/benchmark/kv_canary/bench_verify.py diff --git a/python/sglang/jit_kernel/benchmark/kv_canary/bench_write.py b/test/registered/jit/benchmark/kv_canary/bench_write.py similarity index 100% rename from python/sglang/jit_kernel/benchmark/kv_canary/bench_write.py rename to test/registered/jit/benchmark/kv_canary/bench_write.py diff --git a/python/sglang/jit_kernel/tests/deepseek_v4/test_c128_v2.py b/test/registered/jit/deepseek_v4/test_c128_v2.py similarity index 99% rename from python/sglang/jit_kernel/tests/deepseek_v4/test_c128_v2.py rename to test/registered/jit/deepseek_v4/test_c128_v2.py index 5ec3462c0..2493074d6 100644 --- a/python/sglang/jit_kernel/tests/deepseek_v4/test_c128_v2.py +++ b/test/registered/jit/deepseek_v4/test_c128_v2.py @@ -7,7 +7,6 @@ import pytest import torch import triton -from sglang.jit_kernel.benchmark.bench_activation import register_cuda_ci from sglang.jit_kernel.dsv4 import compress_forward from sglang.jit_kernel.tests.deepseek_v4.common import ( LegacyContext, @@ -17,6 +16,7 @@ from sglang.jit_kernel.tests.deepseek_v4.common import ( make_state_pool, to_seq_extend, ) +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="nightly-kernel-1-gpu", nightly=True) diff --git a/python/sglang/jit_kernel/tests/deepseek_v4/test_c4_v2.py b/test/registered/jit/deepseek_v4/test_c4_v2.py similarity index 99% rename from python/sglang/jit_kernel/tests/deepseek_v4/test_c4_v2.py rename to test/registered/jit/deepseek_v4/test_c4_v2.py index 7f8e7c06d..cf5d99c07 100644 --- a/python/sglang/jit_kernel/tests/deepseek_v4/test_c4_v2.py +++ b/test/registered/jit/deepseek_v4/test_c4_v2.py @@ -7,7 +7,6 @@ import pytest import torch import triton -from sglang.jit_kernel.benchmark.bench_activation import register_cuda_ci from sglang.jit_kernel.dsv4 import compress_forward from sglang.jit_kernel.tests.deepseek_v4.common import ( LegacyContext, @@ -17,6 +16,7 @@ from sglang.jit_kernel.tests.deepseek_v4.common import ( make_state_pool, to_seq_extend, ) +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="nightly-kernel-1-gpu", nightly=True) diff --git a/python/sglang/jit_kernel/tests/deepseek_v4/test_fp4_indexer.py b/test/registered/jit/deepseek_v4/test_fp4_indexer.py similarity index 100% rename from python/sglang/jit_kernel/tests/deepseek_v4/test_fp4_indexer.py rename to test/registered/jit/deepseek_v4/test_fp4_indexer.py diff --git a/python/sglang/jit_kernel/tests/diffusion/test_diffusion_modelopt_fp8_scaled_mm.py b/test/registered/jit/diffusion/test_diffusion_modelopt_fp8_scaled_mm.py similarity index 100% rename from python/sglang/jit_kernel/tests/diffusion/test_diffusion_modelopt_fp8_scaled_mm.py rename to test/registered/jit/diffusion/test_diffusion_modelopt_fp8_scaled_mm.py diff --git a/python/sglang/jit_kernel/tests/diffusion/test_diffusion_nvfp4_scaled_mm.py b/test/registered/jit/diffusion/test_diffusion_nvfp4_scaled_mm.py similarity index 100% rename from python/sglang/jit_kernel/tests/diffusion/test_diffusion_nvfp4_scaled_mm.py rename to test/registered/jit/diffusion/test_diffusion_nvfp4_scaled_mm.py diff --git a/python/sglang/jit_kernel/tests/diffusion/test_flydsl_fused_norm.py b/test/registered/jit/diffusion/test_flydsl_fused_norm.py similarity index 100% rename from python/sglang/jit_kernel/tests/diffusion/test_flydsl_fused_norm.py rename to test/registered/jit/diffusion/test_flydsl_fused_norm.py diff --git a/python/sglang/jit_kernel/tests/diffusion/test_fused_norm_scale_shift.py b/test/registered/jit/diffusion/test_fused_norm_scale_shift.py similarity index 100% rename from python/sglang/jit_kernel/tests/diffusion/test_fused_norm_scale_shift.py rename to test/registered/jit/diffusion/test_fused_norm_scale_shift.py diff --git a/python/sglang/jit_kernel/tests/diffusion/test_group_norm_silu.py b/test/registered/jit/diffusion/test_group_norm_silu.py similarity index 100% rename from python/sglang/jit_kernel/tests/diffusion/test_group_norm_silu.py rename to test/registered/jit/diffusion/test_group_norm_silu.py diff --git a/python/sglang/jit_kernel/tests/diffusion/test_qknorm_rope.py b/test/registered/jit/diffusion/test_qknorm_rope.py similarity index 100% rename from python/sglang/jit_kernel/tests/diffusion/test_qknorm_rope.py rename to test/registered/jit/diffusion/test_qknorm_rope.py diff --git a/python/sglang/jit_kernel/tests/diffusion/test_qwen_image_modulation.py b/test/registered/jit/diffusion/test_qwen_image_modulation.py similarity index 100% rename from python/sglang/jit_kernel/tests/diffusion/test_qwen_image_modulation.py rename to test/registered/jit/diffusion/test_qwen_image_modulation.py diff --git a/python/sglang/jit_kernel/tests/diffusion/test_varlen_pack_pad.py b/test/registered/jit/diffusion/test_varlen_pack_pad.py similarity index 100% rename from python/sglang/jit_kernel/tests/diffusion/test_varlen_pack_pad.py rename to test/registered/jit/diffusion/test_varlen_pack_pad.py diff --git a/python/sglang/jit_kernel/tests/diffusion/test_varlen_uspattn_equivalence.py b/test/registered/jit/diffusion/test_varlen_uspattn_equivalence.py similarity index 100% rename from python/sglang/jit_kernel/tests/diffusion/test_varlen_uspattn_equivalence.py rename to test/registered/jit/diffusion/test_varlen_uspattn_equivalence.py diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_const_sync.py b/test/registered/jit/kv_canary/test_const_sync.py similarity index 88% rename from python/sglang/jit_kernel/tests/kv_canary/test_const_sync.py rename to test/registered/jit/kv_canary/test_const_sync.py index c5fd4799a..d6c800dcd 100644 --- a/python/sglang/jit_kernel/tests/kv_canary/test_const_sync.py +++ b/test/registered/jit/kv_canary/test_const_sync.py @@ -3,14 +3,20 @@ from __future__ import annotations import re from pathlib import Path +import sglang.jit_kernel from sglang.jit_kernel.kv_canary import consts from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=5, suite="base-b-kernel-unit-1-gpu-large") +# Resolve the kernel source against the installed jit_kernel package rather +# than this file's location, so the test stays correct wherever it lives. _CONSTS_CUH: Path = ( - Path(__file__).resolve().parents[2] / "csrc" / "kv_canary" / "consts.cuh" + Path(sglang.jit_kernel.__file__).resolve().parent + / "csrc" + / "kv_canary" + / "consts.cuh" ) diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_kernel_config.py b/test/registered/jit/kv_canary/test_kernel_config.py similarity index 100% rename from python/sglang/jit_kernel/tests/kv_canary/test_kernel_config.py rename to test/registered/jit/kv_canary/test_kernel_config.py diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_pipeline_e2e.py b/test/registered/jit/kv_canary/test_pipeline_e2e.py similarity index 100% rename from python/sglang/jit_kernel/tests/kv_canary/test_pipeline_e2e.py rename to test/registered/jit/kv_canary/test_pipeline_e2e.py diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_plan_fuzz.py b/test/registered/jit/kv_canary/test_plan_fuzz.py similarity index 100% rename from python/sglang/jit_kernel/tests/kv_canary/test_plan_fuzz.py rename to test/registered/jit/kv_canary/test_plan_fuzz.py diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_plan_hand.py b/test/registered/jit/kv_canary/test_plan_hand.py similarity index 100% rename from python/sglang/jit_kernel/tests/kv_canary/test_plan_hand.py rename to test/registered/jit/kv_canary/test_plan_hand.py diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_scatter_req_token_ids.py b/test/registered/jit/kv_canary/test_scatter_req_token_ids.py similarity index 100% rename from python/sglang/jit_kernel/tests/kv_canary/test_scatter_req_token_ids.py rename to test/registered/jit/kv_canary/test_scatter_req_token_ids.py diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_utils.py b/test/registered/jit/kv_canary/test_utils.py similarity index 100% rename from python/sglang/jit_kernel/tests/kv_canary/test_utils.py rename to test/registered/jit/kv_canary/test_utils.py diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_verify_fuzz.py b/test/registered/jit/kv_canary/test_verify_fuzz.py similarity index 100% rename from python/sglang/jit_kernel/tests/kv_canary/test_verify_fuzz.py rename to test/registered/jit/kv_canary/test_verify_fuzz.py diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_verify_hand.py b/test/registered/jit/kv_canary/test_verify_hand.py similarity index 100% rename from python/sglang/jit_kernel/tests/kv_canary/test_verify_hand.py rename to test/registered/jit/kv_canary/test_verify_hand.py diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_write_fuzz.py b/test/registered/jit/kv_canary/test_write_fuzz.py similarity index 100% rename from python/sglang/jit_kernel/tests/kv_canary/test_write_fuzz.py rename to test/registered/jit/kv_canary/test_write_fuzz.py diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_write_hand.py b/test/registered/jit/kv_canary/test_write_hand.py similarity index 100% rename from python/sglang/jit_kernel/tests/kv_canary/test_write_hand.py rename to test/registered/jit/kv_canary/test_write_hand.py diff --git a/python/sglang/jit_kernel/tests/test_activation.py b/test/registered/jit/test_activation.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_activation.py rename to test/registered/jit/test_activation.py diff --git a/python/sglang/jit_kernel/tests/test_add_constant.py b/test/registered/jit/test_add_constant.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_add_constant.py rename to test/registered/jit/test_add_constant.py diff --git a/python/sglang/jit_kernel/tests/test_awq_dequantize.py b/test/registered/jit/test_awq_dequantize.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_awq_dequantize.py rename to test/registered/jit/test_awq_dequantize.py diff --git a/python/sglang/jit_kernel/tests/test_awq_marlin_moe_repack.py b/test/registered/jit/test_awq_marlin_moe_repack.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_awq_marlin_moe_repack.py rename to test/registered/jit/test_awq_marlin_moe_repack.py diff --git a/python/sglang/jit_kernel/tests/test_awq_marlin_repack.py b/test/registered/jit/test_awq_marlin_repack.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_awq_marlin_repack.py rename to test/registered/jit/test_awq_marlin_repack.py diff --git a/python/sglang/jit_kernel/tests/test_clamp_position.py b/test/registered/jit/test_clamp_position.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_clamp_position.py rename to test/registered/jit/test_clamp_position.py diff --git a/python/sglang/jit_kernel/tests/test_concat_mla.py b/test/registered/jit/test_concat_mla.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_concat_mla.py rename to test/registered/jit/test_concat_mla.py diff --git a/python/sglang/jit_kernel/tests/test_custom_all_reduce.py b/test/registered/jit/test_custom_all_reduce.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_custom_all_reduce.py rename to test/registered/jit/test_custom_all_reduce.py diff --git a/python/sglang/jit_kernel/tests/test_cutedsl_gdn.py b/test/registered/jit/test_cutedsl_gdn.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_cutedsl_gdn.py rename to test/registered/jit/test_cutedsl_gdn.py diff --git a/python/sglang/jit_kernel/tests/test_flash_attention_4.py b/test/registered/jit/test_flash_attention_4.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_flash_attention_4.py rename to test/registered/jit/test_flash_attention_4.py diff --git a/python/sglang/jit_kernel/tests/test_fused_add_rmsnorm.py b/test/registered/jit/test_fused_add_rmsnorm.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_fused_add_rmsnorm.py rename to test/registered/jit/test_fused_add_rmsnorm.py diff --git a/python/sglang/jit_kernel/tests/test_fused_metadata_copy.py b/test/registered/jit/test_fused_metadata_copy.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_fused_metadata_copy.py rename to test/registered/jit/test_fused_metadata_copy.py diff --git a/python/sglang/jit_kernel/tests/test_fused_store_index_cache.py b/test/registered/jit/test_fused_store_index_cache.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_fused_store_index_cache.py rename to test/registered/jit/test_fused_store_index_cache.py diff --git a/python/sglang/jit_kernel/tests/test_fused_verify_triton_gdn.py b/test/registered/jit/test_fused_verify_triton_gdn.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_fused_verify_triton_gdn.py rename to test/registered/jit/test_fused_verify_triton_gdn.py diff --git a/python/sglang/jit_kernel/tests/test_gptq_marlin.py b/test/registered/jit/test_gptq_marlin.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_gptq_marlin.py rename to test/registered/jit/test_gptq_marlin.py diff --git a/python/sglang/jit_kernel/tests/test_gptq_marlin_repack.py b/test/registered/jit/test_gptq_marlin_repack.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_gptq_marlin_repack.py rename to test/registered/jit/test_gptq_marlin_repack.py diff --git a/python/sglang/jit_kernel/tests/test_grouped_topk.py b/test/registered/jit/test_grouped_topk.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_grouped_topk.py rename to test/registered/jit/test_grouped_topk.py diff --git a/python/sglang/jit_kernel/tests/test_hadamard_jit.py b/test/registered/jit/test_hadamard_jit.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_hadamard_jit.py rename to test/registered/jit/test_hadamard_jit.py diff --git a/python/sglang/jit_kernel/tests/test_hicache.py b/test/registered/jit/test_hicache.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_hicache.py rename to test/registered/jit/test_hicache.py diff --git a/python/sglang/jit_kernel/tests/test_hisparse.py b/test/registered/jit/test_hisparse.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_hisparse.py rename to test/registered/jit/test_hisparse.py diff --git a/python/sglang/jit_kernel/tests/test_mla_kv_pack_quantize_fp8.py b/test/registered/jit/test_mla_kv_pack_quantize_fp8.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_mla_kv_pack_quantize_fp8.py rename to test/registered/jit/test_mla_kv_pack_quantize_fp8.py diff --git a/python/sglang/jit_kernel/tests/test_moe_align_block_size.py b/test/registered/jit/test_moe_align_block_size.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_moe_align_block_size.py rename to test/registered/jit/test_moe_align_block_size.py diff --git a/python/sglang/jit_kernel/tests/test_moe_lora_align_block_size.py b/test/registered/jit/test_moe_lora_align_block_size.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_moe_lora_align_block_size.py rename to test/registered/jit/test_moe_lora_align_block_size.py diff --git a/python/sglang/jit_kernel/tests/test_moe_wna16_marlin.py b/test/registered/jit/test_moe_wna16_marlin.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_moe_wna16_marlin.py rename to test/registered/jit/test_moe_wna16_marlin.py diff --git a/python/sglang/jit_kernel/tests/test_mxfp8_moe.py b/test/registered/jit/test_mxfp8_moe.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_mxfp8_moe.py rename to test/registered/jit/test_mxfp8_moe.py diff --git a/python/sglang/jit_kernel/tests/test_ngram_embedding.py b/test/registered/jit/test_ngram_embedding.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_ngram_embedding.py rename to test/registered/jit/test_ngram_embedding.py diff --git a/python/sglang/jit_kernel/tests/test_nvfp4_blockwise_moe.py b/test/registered/jit/test_nvfp4_blockwise_moe.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_nvfp4_blockwise_moe.py rename to test/registered/jit/test_nvfp4_blockwise_moe.py diff --git a/python/sglang/jit_kernel/tests/test_nvfp4_gemm.py b/test/registered/jit/test_nvfp4_gemm.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_nvfp4_gemm.py rename to test/registered/jit/test_nvfp4_gemm.py diff --git a/python/sglang/jit_kernel/tests/test_nvfp4_quant.py b/test/registered/jit/test_nvfp4_quant.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_nvfp4_quant.py rename to test/registered/jit/test_nvfp4_quant.py diff --git a/python/sglang/jit_kernel/tests/test_per_tensor_quant_fp8.py b/test/registered/jit/test_per_tensor_quant_fp8.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_per_tensor_quant_fp8.py rename to test/registered/jit/test_per_tensor_quant_fp8.py diff --git a/python/sglang/jit_kernel/tests/test_per_token_group_quant_8bit.py b/test/registered/jit/test_per_token_group_quant_8bit.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_per_token_group_quant_8bit.py rename to test/registered/jit/test_per_token_group_quant_8bit.py diff --git a/python/sglang/jit_kernel/tests/test_pos_enc.py b/test/registered/jit/test_pos_enc.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_pos_enc.py rename to test/registered/jit/test_pos_enc.py diff --git a/python/sglang/jit_kernel/tests/test_qknorm.py b/test/registered/jit/test_qknorm.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_qknorm.py rename to test/registered/jit/test_qknorm.py diff --git a/python/sglang/jit_kernel/tests/test_qknorm_across_heads.py b/test/registered/jit/test_qknorm_across_heads.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_qknorm_across_heads.py rename to test/registered/jit/test_qknorm_across_heads.py diff --git a/python/sglang/jit_kernel/tests/test_renorm.py b/test/registered/jit/test_renorm.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_renorm.py rename to test/registered/jit/test_renorm.py diff --git a/python/sglang/jit_kernel/tests/test_resolve_future_token_ids.py b/test/registered/jit/test_resolve_future_token_ids.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_resolve_future_token_ids.py rename to test/registered/jit/test_resolve_future_token_ids.py diff --git a/python/sglang/jit_kernel/tests/test_rmsnorm.py b/test/registered/jit/test_rmsnorm.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_rmsnorm.py rename to test/registered/jit/test_rmsnorm.py diff --git a/python/sglang/jit_kernel/tests/test_rmsnorm_hf.py b/test/registered/jit/test_rmsnorm_hf.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_rmsnorm_hf.py rename to test/registered/jit/test_rmsnorm_hf.py diff --git a/python/sglang/jit_kernel/tests/test_rope.py b/test/registered/jit/test_rope.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_rope.py rename to test/registered/jit/test_rope.py diff --git a/python/sglang/jit_kernel/tests/test_set_mla_kv_buffer.py b/test/registered/jit/test_set_mla_kv_buffer.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_set_mla_kv_buffer.py rename to test/registered/jit/test_set_mla_kv_buffer.py diff --git a/python/sglang/jit_kernel/tests/test_store_cache.py b/test/registered/jit/test_store_cache.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_store_cache.py rename to test/registered/jit/test_store_cache.py diff --git a/python/sglang/jit_kernel/tests/test_timestep_embedding.py b/test/registered/jit/test_timestep_embedding.py similarity index 100% rename from python/sglang/jit_kernel/tests/test_timestep_embedding.py rename to test/registered/jit/test_timestep_embedding.py diff --git a/python/sglang/jit_kernel/tests/test_tp_qknorm.py b/test/registered/jit/test_tp_qknorm.py similarity index 97% rename from python/sglang/jit_kernel/tests/test_tp_qknorm.py rename to test/registered/jit/test_tp_qknorm.py index 854d984d1..5c5a99759 100644 --- a/python/sglang/jit_kernel/tests/test_tp_qknorm.py +++ b/test/registered/jit/test_tp_qknorm.py @@ -10,8 +10,7 @@ import torch.distributed as dist import triton from sglang.jit_kernel.all_reduce import fused_parallel_qknorm -from sglang.jit_kernel.tests.test_custom_all_reduce import multiprocess_test -from sglang.jit_kernel.tests.utils import multiprocess_main +from sglang.jit_kernel.tests.utils import multiprocess_main, multiprocess_test from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci( diff --git a/test/registered/unit/mem_cache/test_hicache_nixl_storage.py b/test/registered/unit/mem_cache/test_hicache_nixl_storage.py index a6ea710a8..05a034790 100644 --- a/test/registered/unit/mem_cache/test_hicache_nixl_storage.py +++ b/test/registered/unit/mem_cache/test_hicache_nixl_storage.py @@ -2,7 +2,7 @@ from sglang.test.ci.ci_register import register_cuda_ci -register_cuda_ci(est_time=30, stage="base-a", runner_config="1-gpu-small") +register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small") import os import shutil diff --git a/test/run_suite.py b/test/run_suite.py index 1f4bfc952..1f935973a 100644 --- a/test/run_suite.py +++ b/test/run_suite.py @@ -298,15 +298,6 @@ def run_a_suite(args): and not f.endswith("/cpu/utils.py") ] - # JIT kernel tests and benchmarks (live alongside kernel source) - jit_kernel_dir = os.path.join(repo_root, "python", "sglang", "jit_kernel") - files += glob.glob( - os.path.join(jit_kernel_dir, "tests", "**", "test_*.py"), recursive=True - ) - files += glob.glob( - os.path.join(jit_kernel_dir, "benchmark", "**", "bench_*.py"), recursive=True - ) - # Strict: all discovered files must have proper registration sanity_check = True