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

2026-09-18 07:37:48 +08:00
parent b98a2d1096
commit 7bc9152447
53 changed files with 113 additions and 53 deletions
@@ -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
-`test/registered/kernel/jit/`; legacy `test/registered/jit/` files are migrated
-incrementally. Helpers stay alongside the kernel source under
-`python/sglang/kernels/jit/` and are imported by absolute path:
+Kernel correctness tests and benchmarks use the established plural `kernels`
+root and mirror the operator group under `python/sglang/kernels/ops/`. Helpers
+stay alongside the kernel source under `python/sglang/kernels/jit/` and are
+imported by absolute path:

- Correctness tests: `test/registered/kernel/jit/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`
+- Correctness tests: `test/registered/kernels/ops/<group>/test_*.py` → `base-b-kernel-unit-test-1-gpu-large`
+- Benchmarks: `test/registered/kernels/benchmark/<group>/bench_*.py` → `base-b-kernel-benchmark-test-1-gpu-large`

 ## Choosing a Suite

@@ -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
@@ -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
@@ -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
@@ -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 \
@@ -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/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 4
+    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
@@ -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
@@ -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
@@ -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