Refactor JIT kernel CI to use run_suite.py registration system (#21239)

test/run_suite.py

@@ -41,6 +41,8 @@ PER_COMMIT_SUITES = {
         "stage-b-test-1-gpu-large",
         "stage-b-test-2-gpu-large",
         "stage-b-test-4-gpu-b200",
+        "stage-b-kernel-unit-1-gpu-large",
+        "stage-b-kernel-benchmark-1-gpu-large",
         "stage-c-test-4-gpu-h100",
         "stage-c-test-4-gpu-b200",
         "stage-c-test-4-gpu-gb200",
@@ -73,6 +75,7 @@ NIGHTLY_SUITES = {
         "nightly-8-gpu-h200-basic",  # Basic tests for large models on H200
         "nightly-8-gpu-b200-basic",  # Basic tests for large models on B200
         "nightly-8-gpu-common",  # Common tests that run on both H200 and B200
+        "nightly-kernel-1-gpu",
         # Eval and perf suites (2-gpu)
         "nightly-eval-text-2-gpu",
         "nightly-eval-vlm-2-gpu",
@@ -170,9 +173,11 @@ def run_a_suite(args):
     auto_partition_id = args.auto_partition_id
     auto_partition_size = args.auto_partition_size
 
-    # All tests (per-commit and nightly) are now in registered/
     # Use absolute paths so the script works from any working directory
     script_dir = os.path.dirname(os.path.abspath(__file__))
+    repo_root = os.path.dirname(script_dir)
+
+    # Registered tests under test/registered/
     files = [
         f
         for f in glob.glob(
@@ -180,7 +185,13 @@ def run_a_suite(args):
         )
         if not f.endswith("/conftest.py") and not f.endswith("/__init__.py")
     ]
-    # Strict: all registered files must have proper registration
+
+    # 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"))
+    files += glob.glob(os.path.join(jit_kernel_dir, "benchmark", "bench_*.py"))
+
+    # Strict: all discovered files must have proper registration
     sanity_check = True
 
     all_tests = collect_tests(files, sanity_check=sanity_check)