[Test] Consolidate kernel tests under plural kernels tree (#39966)

.claude/skills/write-sglang-test/SKILL.md

@@ -11,7 +11,7 @@ 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 shut the server down gracefully** — call `terminate_and_kill_process_tree(cls.process)`, never a bare `kill_process_tree`. SIGKILL alone skips the server's userspace cleanup and leaves its GPU memory charged to the dead process; the next class then OOMs while loading weights. Keep it defensive too: `hasattr`/null checks before accessing resources (e.g. `cls.process`) that `setUpClass` may not have finished allocating.
-3. **Place tests in `test/registered/<kind>/<subsystem>/`** — `<kind>` is `unit`, `kernel`, `e2e`, `accuracy`, `perf`, or `stress`; hardware belongs in registrations, not directory names
+3. **Place non-kernel tests in `test/registered/<kind>/<subsystem>/`** — `<kind>` is `unit`, `e2e`, `accuracy`, `perf`, or `stress`; kernel tests use `test/registered/kernels/{ops,benchmark}/<group>/`; hardware belongs in registrations, not directory names
 4. **Reuse server fixtures** — inherit from `DefaultServerBase` or write `setUpClass`/`tearDownClass` with `popen_launch_server`
 5. **Mock boundaries, not SGLang behavior** — mock slow or external dependencies only when the assertion still checks an observable result, state transition, or error. A test whose evidence is only `assert_called*` mirrors its mock and is not admissible. Launch a real server only when inference results or lifecycle behavior are the contract under test.
 
@@ -31,10 +31,8 @@ This skill covers **how to write and register tests**. For CI pipeline internals
 
 JIT kernel notes:
 - If the task is adding or updating code under `python/sglang/kernels/jit/`, prefer the `add-jit-kernel` skill first.
-- New JIT kernel correctness tests use `test/registered/kernel/jit/**/test_*.py`.
+- JIT kernel correctness tests use `test/registered/kernels/ops/<group>/test_*.py`.
-- New JIT kernel benchmarks use `test/registered/kernel/jit/benchmark/**/bench_*.py`.
+- JIT kernel benchmarks use `test/registered/kernels/benchmark/<group>/bench_*.py`.
-- `test/registered/jit/` also exists and still runs. It is a leftover from the kernel
-  reclassification (RFC #29630) that was never finished; do not add files there.
 - 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.
 
 ---
@@ -336,12 +334,12 @@ They are ordinary registered tests; only their stage differs:
 ```python
 from sglang.test.ci.ci_register import register_cuda_ci
 
-# Correctness tests in test/registered/kernel/jit/
+# Correctness tests in test/registered/kernels/ops/<group>/
 register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
 register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="8-gpu-h200")
 
-# Benchmarks in test/registered/kernel/jit/benchmark/
+# Benchmarks in test/registered/kernels/benchmark/<group>/
 register_cuda_ci(est_time=6, stage="base-b-kernel-benchmark", runner_config="1-gpu-large")
 
 # Optional nightly registration — same form, stage is just "nightly"
@@ -375,8 +373,9 @@ A `register_*_ci(...)` under `python/sglang/` is rejected by the
 
 **Decision rule** (see also `test/registered/README.md`):
 - CPU component logic, no server → `registered/unit/<subsystem>/`
-- JIT kernel correctness / benchmarks → `registered/kernel/jit/`
+- JIT kernel correctness → `registered/kernels/ops/<group>/`
-- Other accelerator operator correctness → `registered/kernel/<group>/`
+- JIT kernel benchmarks → `registered/kernels/benchmark/<group>/`
+- Other accelerator operator correctness → `registered/kernels/ops/<group>/`
 - Server needed → `registered/e2e/<subsystem>/`
 - Eval floor / performance contract → `registered/{accuracy,perf}/<family>/`
 - Local debugging → `manual/`
@@ -422,7 +421,7 @@ Before submitting a test:
 - [ ] Inherits from `CustomTestCase` (not `unittest.TestCase`)
 - [ ] Has `register_*_ci(...)` call at module level
 - [ ] Placed in `test/registered/<kind>/<subsystem>/`
-- [ ] JIT kernel work: test files live in `test/registered/kernel/jit/`; only test-only helpers stay under `python/sglang/kernels/jit/`
+- [ ] JIT kernel work: correctness tests live in `test/registered/kernels/ops/<group>/`, benchmarks live in `test/registered/kernels/benchmark/<group>/`, and only test-only helpers stay under `python/sglang/kernels/jit/`
 - [ ] 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)
 - [ ] `tearDownClass` is defensive — uses `hasattr`/null checks before accessing resources that may not have been allocated

.github/workflows/_pr-test-check-changes.yml

@@ -112,7 +112,6 @@ jobs:
               - "python/sglang/kernels/ops/diffusion/**"
               - "test/registered/kernels/ops/diffusion/**"
               - "test/registered/kernels/benchmark/diffusion/**"
-              - "test/registered/kernel/diffusion/**"
               - "test/registered/unit/diffusion/**"
               - "python/sglang/cli/**"
             jit_kernel:
@@ -120,7 +119,6 @@ jobs:
               - ".github/workflows/pr-test-jit-kernel.yml"
               - "python/pyproject.toml"
               - "test/registered/kernels/**"
-              - "test/registered/kernel/diffusion/**"
               # sglang.kernels is the migrated kernel namespace (RFC #29630 / #30044); the
               # base-b-kernel suites import it directly, so kernel edits must run them.
               - "python/sglang/kernels/!(*.md)"

python/sglang/multimodal_gen/runtime/layers/attention/backends/subblock_sparse/README.md

@@ -191,7 +191,7 @@ are comparable.
 | `../subblock_sparse_attn.py` | the `AttentionBackend`: schedule, gating, dense fallback |
 
 Tests: `test/unit/test_subblock_sparse_attention.py` and
-`test/registered/kernel/attention/test_subblock_sage_fp8_sm90.py`. The GPU
+`test/registered/kernels/ops/attention/test_subblock_sage_fp8_sm90.py`. The GPU
 test covers the native production dispatch. Running at a full block budget must
 reproduce dense attention up to the expected quantization error, pinning routing
 indices, ragged tails, scale domains and the softmax scale in one check.

python/sglang/srt/distributed/device_communicators/configs/custom_all_reduce_v2.py

@@ -1,9 +1,9 @@
 """Hand-tuned dispatch configs for the JIT custom all-reduce (v2).
 
 Thresholds and block counts come from sweeps of
-``test/registered/jit/benchmark/bench_custom_all_reduce.py`` on the listed
+``test/registered/kernels/benchmark/communication/bench_custom_all_reduce.py``
-GPUs; ``get_all_reduce_config`` picks the table for the current arch and
+on the listed GPUs; ``get_all_reduce_config`` picks the table for the current
-world size.
+arch and world size.
 """
 
 from functools import cache

scripts/lint/check_no_registered_tests_in_package.py

@@ -4,11 +4,12 @@ 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/
+test/registered/kernels/ops/ for kernel tests and
-for JIT kernel benchmarks) so they are not shipped in the wheel and are
+test/registered/kernels/benchmark/ for kernel benchmarks) so they are not
-collected by run_suite.py's registered glob. A registered file placed inside
+shipped in the wheel and are collected by run_suite.py's registered glob. A
-the package would be shipped to users AND silently dropped by run_suite.py
+registered file placed inside the package would be shipped to users AND
-(which no longer globs the package) -- it would never run in CI. This guard
+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
@@ -65,8 +66,8 @@ def main() -> int:
         )
         print(
             "  Registered tests and benchmarks must live under test/registered/\n"
-            "  (e.g. test/registered/jit/ for JIT kernel tests and\n"
+            "  (e.g. test/registered/kernels/ops/ for kernel tests and\n"
-            "  test/registered/jit/benchmark/ for JIT kernel benchmarks) so they\n"
+            "  test/registered/kernels/benchmark/ for kernel benchmarks) so they\n"
             "  are not shipped in the wheel and are collected by run_suite.py.\n"
         )
         for f in offenders:

scripts/lint/check_registered_tests.py

@@ -39,10 +39,11 @@ _MODERN_SHAPE = re.compile(r"^(.+)-test-(.+)$")
 # no suite any workflow invokes and the test silently never runs.
 _LEGACY_CUDA_PREFIXES = ("stress",)
 
-_TEST_KINDS = {"unit", "kernel", "e2e", "accuracy", "perf", "stress"}
+_TEST_KINDS = {"unit", "e2e", "accuracy", "perf", "stress"}
+_KERNEL_ROOT = "kernels"
 
 # Flat vendor trees. Vendor-only coverage fits no kind above: no XPU/NPU suite
-# carries the `-kernel-` infix `kernel` needs, and these launch device work.
+# carries the `-kernel-` infix the kernel tree needs, and these launch device work.
 _VENDOR_DIRS = {"amd", "mlx", "musa", "npu", "xpu"}
 
 
@@ -132,11 +133,22 @@ def taxonomy_errors(path: str, registries: list, tree: ast.AST) -> list[str]:
     relative_parts = parts[2:] if parts[:2] == ["test", "registered"] else []
     if relative_parts and relative_parts[0] in _VENDOR_DIRS:
         return []
+    if relative_parts and relative_parts[0] == _KERNEL_ROOT:
+        errors = []
+        if len(relative_parts) < 4 or relative_parts[1] not in {"ops", "benchmark"}:
+            errors.append(
+                f"{path}: kernel tests must live under "
+                "test/registered/kernels/{ops,benchmark}/<group>/"
+            )
+        if any("-kernel-" not in (r.effective_suite or "") for r in registries):
+            errors.append(f"{path}: kernel tests must use a *-kernel-* suite")
+        return errors
     if len(relative_parts) < 3 or relative_parts[0] not in _TEST_KINDS:
         return [
             f"{path}: registered tests must live under "
             "test/registered/<kind>/<subsystem>/; kind must be one of "
             + ", ".join(sorted(_TEST_KINDS))
+            + "; kernel tests use test/registered/kernels/{ops,benchmark}/<group>/"
         ]
 
     kind = relative_parts[0]
@@ -153,9 +165,6 @@ def taxonomy_errors(path: str, registries: list, tree: ast.AST) -> list[str]:
             errors.append(f"{path}: unit test est_time must be <= 60 seconds")
         if _contains_call(tree, "popen_launch_server"):
             errors.append(f"{path}: unit tests may not launch a server")
-    elif kind == "kernel":
-        if any("-kernel-" not in (r.effective_suite or "") for r in registries):
-            errors.append(f"{path}: kernel tests must use a *-kernel-* suite")
     elif kind in {"accuracy", "perf"}:
         invalid = [
             r

scripts/lint/test_check_registered_tests.py

@@ -0,0 +1,51 @@
+import ast
+import unittest
+from types import SimpleNamespace
+
+from scripts.lint.check_registered_tests import taxonomy_errors
+
+
+def _registry(suite: str):
+    return SimpleNamespace(effective_suite=suite, est_time=1)
+
+
+class TestRegisteredTestTaxonomy(unittest.TestCase):
+    def setUp(self):
+        self.tree = ast.parse("")
+        self.kernel_registry = [_registry("base-b-kernel-unit-test-1-gpu-large")]
+
+    def test_plural_kernel_ops_layout_is_accepted(self):
+        errors = taxonomy_errors(
+            "test/registered/kernels/ops/attention/test_example.py",
+            self.kernel_registry,
+            self.tree,
+        )
+        self.assertEqual(errors, [])
+
+    def test_plural_kernel_benchmark_layout_is_accepted(self):
+        errors = taxonomy_errors(
+            "test/registered/kernels/benchmark/attention/bench_example.py",
+            [_registry("base-b-kernel-benchmark-test-1-gpu-large")],
+            self.tree,
+        )
+        self.assertEqual(errors, [])
+
+    def test_singular_kernel_root_is_rejected(self):
+        errors = taxonomy_errors(
+            "test/registered/kernel/attention/test_example.py",
+            self.kernel_registry,
+            self.tree,
+        )
+        self.assertTrue(errors)
+
+    def test_kernel_group_is_required(self):
+        errors = taxonomy_errors(
+            "test/registered/kernels/ops/test_example.py",
+            self.kernel_registry,
+            self.tree,
+        )
+        self.assertTrue(errors)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/README.md

@@ -11,9 +11,9 @@ The CI pipeline runs in three sequential stages: **A** (pre-flight, ~3 min) →
 
 ## Folder Organization
 
-- `registered/`: CI test files, auto-discovered by `run_suite.py`. Most tests live here. JIT kernel tests are an exception (see below).
+- `registered/`: CI test files, including kernel tests and benchmarks, auto-discovered by `run_suite.py`.
 - `manual/`: Non-CI tests for local debugging or special setups.
-- `run_suite.py`: CI runner — scans `registered/` and JIT kernel directories.
+- `run_suite.py`: CI runner — scans `registered/` recursively.
 
 The system supports both [unittest](https://docs.python.org/3/library/unittest.html) and [pytest](https://docs.pytest.org/en/stable/). The launcher runs `python filename.py -f` with **failfast enabled by default**.
 
@@ -44,7 +44,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 test/registered/jit/test_add_constant.py
+python3 test/registered/kernels/ops/elementwise/test_add_constant.py
 
 # Run a suite
 python3 test/run_suite.py --hw cpu --suite base-a-test-cpu
@@ -72,13 +72,15 @@ 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.
 
-New and renamed tests use this layout:
+New and renamed non-kernel tests use this layout:
 
 ```text
 test/registered/<kind>/<subsystem>/test_*.py
 ```
 
-`<kind>` is one of `unit`, `kernel`, `e2e`, `accuracy`, `perf`, or `stress`.
+`<kind>` is one of `unit`, `e2e`, `accuracy`, `perf`, or `stress`. Kernel tests
+use `test/registered/kernels/{ops,benchmark}/<group>/`, retaining the established
+plural `kernels` root.
 Hardware is expressed by one or more `register_*_ci` calls, never by creating a
 new top-level hardware directory. The admission checker applies the layout and
 kind/suite contract incrementally while legacy paths are migrated.
@@ -87,13 +89,13 @@ Diffusion workflows also enter through `test/run_suite.py`; registered bridge
 files preserve their case-level pytest partitioning until the remaining
 diffusion cases are moved out of the package test-support tree.
 
-New JIT kernel correctness tests and benchmarks live under
+Kernel correctness tests and benchmarks use the established plural `kernels`
-`test/registered/kernel/jit/`; legacy `test/registered/jit/` files are migrated
+root and mirror the operator group under `python/sglang/kernels/ops/`. Helpers
-incrementally. Helpers stay alongside the kernel source under
+stay alongside the kernel source under `python/sglang/kernels/jit/` and are
-`python/sglang/kernels/jit/` and are imported by absolute path:
+imported by absolute path:
 
-- Correctness tests: `test/registered/kernel/jit/test_*.py` → `base-b-kernel-unit-test-1-gpu-large`
+- Correctness tests: `test/registered/kernels/ops/<group>/test_*.py` → `base-b-kernel-unit-test-1-gpu-large`
-- Benchmarks: `test/registered/kernel/jit/benchmark/bench_*.py` → `base-b-kernel-benchmark-test-1-gpu-large`
+- Benchmarks: `test/registered/kernels/benchmark/<group>/bench_*.py` → `base-b-kernel-benchmark-test-1-gpu-large`
 
 ## Choosing a Suite
 

test/registered/kernels/benchmark/attention/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 test/registered/jit/benchmark/bench_fused_qknorm_rope.py
+    python test/registered/kernels/benchmark/attention/bench_fused_qknorm_rope.py
 """
 
 import itertools

test/registered/kernels/benchmark/communication/bench_symm_mem_all_gather.py

@@ -11,9 +11,9 @@ Providers:
 Usage::
 
     # Benchmark on the default world sizes (2, 4, 8 GPUs):
-    python test/registered/jit/benchmark/bench_symm_mem_all_gather.py
+    python test/registered/kernels/benchmark/communication/bench_symm_mem_all_gather.py
     # Pick a specific world size (or comma-separated list):
-    python test/registered/jit/benchmark/bench_symm_mem_all_gather.py --num-gpu 8
+    python test/registered/kernels/benchmark/communication/bench_symm_mem_all_gather.py --num-gpu 8
 """
 
 from __future__ import annotations

test/registered/kernels/benchmark/gemm/bench_dsv3_fused_a_gemm.py

@@ -1,7 +1,7 @@
 """Benchmark for DeepSeek V3 fused QKV-A GEMM: CuTe DSL vs CUDA JIT vs torch.
 
 Run on SM90+ (Hopper or later):
-    python test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py
+    python test/registered/kernels/benchmark/gemm/bench_dsv3_fused_a_gemm.py
 """
 
 import torch

test/registered/kernels/ops/attention/test_deepseek_v4_compress_state_runtime_shapes.py

@@ -11,10 +11,10 @@
 #
 # Test command:
 #   python3 -m pytest -q \
-#     test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py
+#     test/registered/kernels/ops/attention/test_deepseek_v4_compress_state_runtime_shapes.py
 #
 # Runtime-shape benchmark command:
-#   python3 test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py \
+#   python3 test/registered/kernels/ops/attention/test_deepseek_v4_compress_state_runtime_shapes.py \
 #     --benchmark \
 #     --shape-source runtime \
 #     --warmup 20 \
@@ -22,7 +22,7 @@
 #     --csv /data00/eval_results/operator_bench/runtime_shape_bench.csv
 #
 # Synthetic Flash/Pro shape benchmark command:
-#   python3 test/registered/jit/test_deepseek_v4_compress_state_runtime_shapes.py \
+#   python3 test/registered/kernels/ops/attention/test_deepseek_v4_compress_state_runtime_shapes.py \
 #     --benchmark \
 #     --shape-source preset \
 #     --shape-presets all \

test/registered/kernels/ops/communication/test_symm_mem_all_gather.py

@@ -7,12 +7,12 @@ NCCL all-gather for a sweep of token counts, hidden widths, and the
 Usage::
 
     # Run on the default world sizes (2, 4, 8 GPUs):
-    python test/registered/jit/test_symm_mem_all_gather.py
+    python test/registered/kernels/ops/communication/test_symm_mem_all_gather.py
     # Pick a specific world size (or comma-separated list):
-    python test/registered/jit/test_symm_mem_all_gather.py --num-gpu 4
+    python test/registered/kernels/ops/communication/test_symm_mem_all_gather.py --num-gpu 4
-    python test/registered/jit/test_symm_mem_all_gather.py --num-gpu 2,4,8
+    python test/registered/kernels/ops/communication/test_symm_mem_all_gather.py --num-gpu 2,4,8
     # Extra pytest args (forwarded to each torchrun worker):
-    python test/registered/jit/test_symm_mem_all_gather.py -k 16384
+    python test/registered/kernels/ops/communication/test_symm_mem_all_gather.py -k 16384
 """
 
 from __future__ import annotations

test/registered/kernels/ops/kimi_k3/test_ar_fusion.py

@@ -10,7 +10,7 @@ semaphore window cycling.
 
 Usage::
 
-    python test/registered/jit/kimi_k3/test_ar_fusion.py   # relaunches under torchrun (8 GPUs)
+    python test/registered/kernels/ops/kimi_k3/test_ar_fusion.py   # relaunches under torchrun (8 GPUs)
 """
 
 from __future__ import annotations

test/registered/kernels/ops/qsa/test_qsa_strided_zero_fill.py

@@ -4,7 +4,7 @@ Poison the packed scratch with NaN, gather with the strided layout used by
 `_forward_trtllm_sparse`, and require that (a) valid rows are copied exactly and
 (b) every slot in [valid_count, stride) is zero, so the paged decode kernel can never
 multiply masked probabilities into stale NaN/Inf bytes. Also checks the compact
-(FA2 fallback) layout is unchanged. Intended for test/registered/kernel/qsa/.
+(FA2 fallback) layout is unchanged. Intended for test/registered/kernels/ops/qsa/.
 """
 
 import sys

test/registered/unit/README.md

@@ -2,7 +2,7 @@
 
 CPU-only component tests that do **not** launch a server, load model weights,
 or require an accelerator. GPU operator correctness belongs under
-`test/registered/kernel/<subsystem>/`.
+`test/registered/kernels/ops/<group>/`.
 
 ## Quick Start