[CI] Move JIT kernel tests + benchmarks to test/registered/jit; add in-package guard (#27644)

.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"`
+- `test/registered/jit/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_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
+test/registered/jit/test_scale.py                     # NEW: Tests
-python/sglang/jit_kernel/benchmark/bench_scale.py     # NEW: Benchmark
+test/registered/jit/benchmark/bench_scale.py          # NEW: Benchmark
 ```

.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/<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 correctness tests use `test/registered/jit/**/test_*.py`.
-- JIT kernel benchmarks use `python/sglang/jit_kernel/benchmark/**/bench_*.py`.
+- JIT kernel benchmarks use `test/registered/jit/benchmark/**/bench_*.py`.
-- Those files are still executed by `test/run_suite.py`, but through dedicated kernel suites rather than `test/registered/`.
+- 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-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 `python/sglang/jit_kernel/benchmark/` |
+| `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
+- [ ] Placed in `test/registered/<category>/` (JIT kernel test/benchmark → `test/registered/jit/` or `test/registered/jit/benchmark/`)
-- [ ] JIT kernel work: files live in `python/sglang/jit_kernel/tests/` or `python/sglang/jit_kernel/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)

.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/**"
+              - "test/registered/jit/diffusion/**"
-              - "python/sglang/jit_kernel/benchmark/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.

.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/**"
+              - "test/registered/jit/diffusion/**"
-              - "python/sglang/jit_kernel/benchmark/diffusion/**"
+              - "test/registered/jit/benchmark/diffusion/**"
               - "python/pyproject_rocm.toml"
               - "python/pyproject_other.toml"
 

.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/**"
+              - "test/registered/jit/diffusion/**"
-              - "python/sglang/jit_kernel/benchmark/diffusion/**"
+              - "test/registered/jit/benchmark/diffusion/**"
               - "python/pyproject_rocm.toml"
               - "python/pyproject_other.toml"
 

.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:

.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

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
 

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
 

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`
+- `test/registered/jit/diffusion/test_qwen_image_modulation.py`
-- `python/sglang/jit_kernel/tests/diffusion/test_group_norm_silu.py`
+- `test/registered/jit/diffusion/test_group_norm_silu.py`
-- `python/sglang/jit_kernel/tests/diffusion/test_varlen_pack_pad.py`
+- `test/registered/jit/diffusion/test_varlen_pack_pad.py`
-- `python/sglang/jit_kernel/tests/diffusion/test_varlen_uspattn_equivalence.py`
+- `test/registered/jit/diffusion/test_varlen_uspattn_equivalence.py`
-- `python/sglang/jit_kernel/benchmark/diffusion/bench_qwen_image_modulation.py`
+- `test/registered/jit/benchmark/diffusion/bench_qwen_image_modulation.py`
-- `python/sglang/jit_kernel/benchmark/diffusion/bench_group_norm_silu.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`.
+- Validation: `test/registered/jit/diffusion/test_group_norm_silu.py`.
-- Microbench: `python/sglang/jit_kernel/benchmark/diffusion/bench_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**

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",
 }
 
 

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())

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:
+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: `python/sglang/jit_kernel/tests/test_*.py` → `base-b-kernel-unit-1-gpu-large`
+- Correctness tests: `test/registered/jit/test_*.py` → `base-b-kernel-unit-1-gpu-large`
-- Benchmarks: `python/sglang/jit_kernel/benchmark/bench_*.py` → `base-b-kernel-benchmark-1-gpu-large`
+- Benchmarks: `test/registered/jit/benchmark/bench_*.py` → `base-b-kernel-benchmark-1-gpu-large`
 
 ## Choosing a Suite
 

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

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"

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)

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)

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"
 )