[Kernel] Reclassify kernel tests by ops group + move helpers out of the package (RFC #29630) (#32128)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

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

@@ -94,14 +94,14 @@ jobs:
               - "python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/**"
               - "python/sglang/srt/observability/**"
               - "python/sglang/kernels/ops/diffusion/**"
-              - "test/registered/jit/diffusion/**"
-              - "test/registered/jit/benchmark/diffusion/**"
+              - "test/registered/kernels/ops/diffusion/**"
+              - "test/registered/kernels/benchmark/diffusion/**"
               - "python/sglang/cli/**"
             jit_kernel:
               - ".github/workflows/pr-test.yml"
               - ".github/workflows/pr-test-jit-kernel.yml"
               - "python/pyproject.toml"
-              - "test/registered/jit/**"
+              - "test/registered/kernels/**"
               # 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/**"

.github/workflows/pr-test-amd-rocm720.yml

@@ -199,15 +199,15 @@ jobs:
               - ".github/workflows/pr-test-amd-rocm720.yml"
             jit_kernel:
               - "python/sglang/kernels/**"
-              - "test/registered/jit/**"
+              - "test/registered/kernels/**"
               - ".github/workflows/pr-test-amd-rocm720.yml"
             multimodal_gen:
               - "python/sglang/multimodal_gen/**/!(*.md|*.ipynb)"
               - "python/sglang/cli/**"
               - "python/sglang/srt/observability/**"
               - "python/sglang/kernels/ops/diffusion/**"
-              - "test/registered/jit/diffusion/**"
-              - "test/registered/jit/benchmark/diffusion/**"
+              - "test/registered/kernels/ops/diffusion/**"
+              - "test/registered/kernels/benchmark/diffusion/**"
               - "python/pyproject_rocm.toml"
               - "python/pyproject_other.toml"
 

.github/workflows/pr-test-amd.yml

@@ -187,15 +187,15 @@ jobs:
               - ".github/workflows/pr-test-amd.yml"
             jit_kernel:
               - "python/sglang/kernels/**"
-              - "test/registered/jit/**"
+              - "test/registered/kernels/**"
               - ".github/workflows/pr-test-amd.yml"
             multimodal_gen:
               - "python/sglang/multimodal_gen/**/!(*.md|*.ipynb)"
               - "python/sglang/cli/**"
               - "python/sglang/srt/observability/**"
               - "python/sglang/kernels/ops/diffusion/**"
-              - "test/registered/jit/diffusion/**"
-              - "test/registered/jit/benchmark/diffusion/**"
+              - "test/registered/kernels/ops/diffusion/**"
+              - "test/registered/kernels/benchmark/diffusion/**"
               - "python/pyproject_rocm.toml"
               - "python/pyproject_other.toml"
 

python/sglang/kernels/README.md

@@ -19,21 +19,28 @@ sglang/kernels/
   selector.py    # heuristic select_kernel() and cached get_kernel()
   fused_op.py    # BaseFusedOp: per-operator multi-backend contract
   ops/
-    <group>/     # one subpackage per operator group
+    <group>/     # one subpackage per operator group (see list below)
+  jit/           # shared JIT CUDA build/runtime infra: utils/, csrc/,
+                 # include/, __main__ (KERNEL_PATH resolves here)
 ```
 
-Groups populated in this phase: `activation`, `gemm`, `kvcache`, `layernorm`,
-`moe`, `quantization`. The remaining groups (`attention`, `communication`,
-`diffusion`, `grammar`, `mamba`, `memory`, `sampling`, `spatial`,
-`speculative`) are reserved package placeholders whose implementations still
-live in `sglang.kernels.jit` / `sgl_kernel` / `triton_ops` and will migrate in
-later phases.
+Operator groups (all populated): `activation`, `attention`, `communication`,
+`diffusion`, `embeddings`, `gemm`, `grammar`, `kv_canary`, `kvcache`,
+`layernorm`, `lplb`, `mamba`, `memory`, `model`, `moe`, `quantization`,
+`sampling`, `spatial`, `speculative`.
+
+As of the RFC #29630 finale (#32072) the legacy `sglang.jit_kernel` package has
+been **removed**: its shared build/runtime infra moved to `sglang.kernels.jit`
+and each JIT-backed operator into its group as
+`sglang.kernels.ops.<group>._jit_<op>`. Tests and benchmarks live under
+`test/registered/kernels/` (`ops/<group>/` for tests, `benchmark/<group>/` for
+benchmarks); shared test helpers are in `sglang.test.kernels`.
 
 ## How it works
 
-Implementations are not moved yet. Each `ops.<group>` function is a thin
-wrapper that forwards to a chosen backend, and every backend is described by a
-`KernelSpec` in the registry so alternatives can be inventoried and compared:
+Each `ops.<group>` function is a thin wrapper that forwards to a chosen
+backend, and every backend is described by a `KernelSpec` in the registry so
+alternatives can be inventoried and compared:
 
 - `register_kernel(KernelSpec(...))` records metadata only — an operator id
   (`"<group>.<name>"`), a backend, and an import path (`"module:attr"`). No
@@ -85,8 +92,8 @@ What this buys (see the
 
 - **Unified correctness testing** — a generic harness enumerates
   `available_backends()` and asserts each one matches `forward_native`
-  (`test/registered/kernels/test_fused_op_gpu_parity.py`); new backends are
-  picked up automatically.
+  (`test/registered/kernels/ops/layernorm/test_fused_op_gpu_parity.py`); new
+  backends are picked up automatically.
 - **One-switch debugging** — `SGLANG_FORCE_FUSED_OP_BACKEND=torch` (or
   `set_fused_op_backend(KernelBackend.TORCH)`) flips *every* fused op to its
   reference implementation for numerical-bug bisection.

python/sglang/kernels/ops/communication/mp.py

@@ -1,6 +1,6 @@
 """Multi-process / multi-GPU launching utilities (torchrun-based).
 
-Shared `multigpu_launch` helper that both `sglang.kernels.jit.tests.utils` and
+Shared `multigpu_launch` helper that both `sglang.test.kernels.utils` and
 `sglang.kernels.jit.benchmark.utils` build their domain-specific entry points on
 top of (`multigpu_pytest_main`, `multigpu_bench_main`).
 

python/sglang/test/kernels/kv_canary/_canary_helpers.py

@@ -5,21 +5,21 @@ from typing import Optional
 
 import torch
 
-from sglang.kernels.jit.tests.kv_canary._constants import (
+from sglang.kernels.ops.kv_canary import consts
+from sglang.kernels.ops.kv_canary.consts import splitmix64, splitmix64_mix3
+from sglang.kernels.ops.kv_canary.verify import VerifyPlan
+from sglang.kernels.ops.kv_canary.write import WritePlan
+from sglang.test.kernels.kv_canary._constants import (
     _I64_SIGN_BIT,
     _U64_MASK,
     DEFAULT_NUM_SLOTS,
     DEFAULT_RING_CAPACITY,
     DEFAULT_SLOT_STRIDE_BYTES,
 )
-from sglang.kernels.jit.tests.kv_canary._fixtures import (
+from sglang.test.kernels.kv_canary._fixtures import (
     make_real_kv_source,
     make_real_kv_sources,
 )
-from sglang.kernels.ops.kv_canary import consts
-from sglang.kernels.ops.kv_canary.consts import splitmix64, splitmix64_mix3
-from sglang.kernels.ops.kv_canary.verify import VerifyPlan
-from sglang.kernels.ops.kv_canary.write import WritePlan
 
 __all__ = [
     "FakeViolationLog",

python/sglang/test/kernels/kv_canary/_differential.py

@@ -5,12 +5,6 @@ from typing import Any, Callable, Iterator, Optional
 
 import torch
 
-from sglang.kernels.jit.tests.kv_canary._canary_helpers import (
-    FakeViolationLog,
-    assert_canary_buf_equal,
-    assert_canary_state_equal,
-    make_log_pair,
-)
 from sglang.kernels.ops.kv_canary import consts
 from sglang.kernels.ops.kv_canary.plan import launch_canary_plan_kernels
 from sglang.kernels.ops.kv_canary.plan_ref import (
@@ -30,6 +24,12 @@ from sglang.kernels.ops.kv_canary.write import WritePlan, launch_canary_write_ke
 from sglang.kernels.ops.kv_canary.write_ref import (
     launch_canary_write_kernel_torch_reference,
 )
+from sglang.test.kernels.kv_canary._canary_helpers import (
+    FakeViolationLog,
+    assert_canary_buf_equal,
+    assert_canary_state_equal,
+    make_log_pair,
+)
 
 _DEVICE = torch.device("cuda")
 

python/sglang/test/kernels/kv_canary/_fixtures.py

@@ -5,12 +5,12 @@ from typing import Literal, Optional
 
 import torch
 
-from sglang.kernels.jit.tests.kv_canary._constants import DEFAULT_NUM_SLOTS
 from sglang.kernels.ops.kv_canary.verify import (
     RealKvSource,
     VerifyPlan,
 )
 from sglang.kernels.ops.kv_canary.write import WritePlan
+from sglang.test.kernels.kv_canary._constants import DEFAULT_NUM_SLOTS
 
 _DEVICE = torch.device("cuda")
 

python/sglang/test/kernels/kv_canary/_fuzz_driver.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import random
 from typing import Any, Callable
 
-from sglang.kernels.jit.tests.kv_canary._differential import (
+from sglang.test.kernels.kv_canary._differential import (
     ShrinkResult,
     shrink_inputs,
 )

python/sglang/test/kernels/kv_canary/_invariants.py

@@ -11,10 +11,10 @@ from typing import Optional
 
 import torch
 
-from sglang.kernels.jit.tests.kv_canary._canary_helpers import FakeViolationLog
 from sglang.kernels.ops.kv_canary import consts
 from sglang.kernels.ops.kv_canary.verify import CanaryLaunchTag, VerifyPlan
 from sglang.kernels.ops.kv_canary.write import WritePlan
+from sglang.test.kernels.kv_canary._canary_helpers import FakeViolationLog
 
 
 class PlanInvariants:

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

@@ -7,7 +7,10 @@ import pytest
 import torch
 import triton
 
-from sglang.kernels.jit.tests.deepseek_v4.common import (
+from sglang.kernels.ops.attention.dsv4 import compress_forward
+from sglang.srt.utils import get_device
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.kernels.deepseek_v4.common import (
     LegacyContext,
     PagedContext,
     make_legacy_context,
@@ -15,9 +18,6 @@ from sglang.kernels.jit.tests.deepseek_v4.common import (
     make_state_pool,
     to_seq_extend,
 )
-from sglang.kernels.ops.attention.dsv4 import compress_forward
-from sglang.srt.utils import get_device
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_amd_ci(est_time=30, suite="nightly-amd-kernel-1-gpu", nightly=True)

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

@@ -7,7 +7,10 @@ import pytest
 import torch
 import triton
 
-from sglang.kernels.jit.tests.deepseek_v4.common import (
+from sglang.kernels.ops.attention.dsv4 import compress_forward
+from sglang.srt.utils import get_device
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.kernels.deepseek_v4.common import (
     LegacyContext,
     PagedContext,
     make_legacy_context,
@@ -15,9 +18,6 @@ from sglang.kernels.jit.tests.deepseek_v4.common import (
     make_state_pool,
     to_seq_extend,
 )
-from sglang.kernels.ops.attention.dsv4 import compress_forward
-from sglang.srt.utils import get_device
-from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_amd_ci(est_time=30, suite="nightly-amd-kernel-1-gpu", nightly=True)

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

@@ -89,16 +89,16 @@ from typing import Callable, Iterable, Literal, Optional
 import pytest
 import torch
 
-from sglang.kernels.jit.tests.deepseek_v4.common import (
-    make_legacy_context,
-    to_seq_extend,
-)
 from sglang.kernels.ops.attention.dsv4 import (
     CompressorDecodePlan,
     CompressorPrefillPlan,
     compress_forward,
 )
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.kernels.deepseek_v4.common import (
+    make_legacy_context,
+    to_seq_extend,
+)
 from sglang.utils import is_in_ci
 
 register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")