[Kernel] Classification cleanup: unify _jit_ naming, drop empty/model groups, add elementwise (RFC #29630) (#32148)

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

test/registered/kernels/benchmark/activation/bench_activation.py

@@ -6,16 +6,12 @@ from sgl_kernel import silu_and_mul as silu_and_mul_aot
 
 from sglang.kernels.jit.benchmark import marker
 from sglang.kernels.jit.benchmark.utils import create_random
-from sglang.kernels.ops.activation._jit_activation import (
-    gelu_and_mul as gelu_and_mul_jit,
-)
-from sglang.kernels.ops.activation._jit_activation import (
+from sglang.kernels.ops.activation.activation import gelu_and_mul as gelu_and_mul_jit
+from sglang.kernels.ops.activation.activation import (
     gelu_tanh_and_mul as gelu_tanh_and_mul_jit,
 )
-from sglang.kernels.ops.activation._jit_activation import relu2 as relu2_jit
-from sglang.kernels.ops.activation._jit_activation import (
-    silu_and_mul as silu_and_mul_jit,
-)
+from sglang.kernels.ops.activation.activation import relu2 as relu2_jit
+from sglang.kernels.ops.activation.activation import silu_and_mul as silu_and_mul_jit
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(

test/registered/kernels/benchmark/attention/bench_add_constant.py

@@ -7,7 +7,7 @@ from sglang.kernels.jit.benchmark.utils import (
     get_benchmark_range,
     run_benchmark_no_cudagraph,
 )
-from sglang.kernels.ops.attention.add_constant import (
+from sglang.kernels.ops.elementwise.add_constant import (
     _jit_add_constant_module,
     add_constant,
 )

test/registered/kernels/benchmark/attention/bench_hadamard.py

@@ -13,7 +13,7 @@ from sglang.kernels.jit.benchmark.utils import (
     get_benchmark_range,
     run_benchmark,
 )
-from sglang.kernels.ops.attention.hadamard import hadamard_transform
+from sglang.kernels.ops.quantization.hadamard import hadamard_transform
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(

test/registered/kernels/benchmark/diffusion/bench_norm_impls.py

@@ -17,10 +17,8 @@ from sglang.kernels.jit.benchmark.utils import DEFAULT_DEVICE
 from sglang.kernels.jit.utils import KERNEL_PATH
 from sglang.kernels.ops.diffusion.triton.norm import norm_infer, rms_norm_fn
 from sglang.kernels.ops.diffusion.triton.rmsnorm_onepass import triton_one_pass_rms_norm
-from sglang.kernels.ops.layernorm._jit_norm import (
-    fused_add_rmsnorm as jit_fused_add_rmsnorm,
-)
-from sglang.kernels.ops.layernorm._jit_norm import rmsnorm as jit_rmsnorm
+from sglang.kernels.ops.layernorm.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm
+from sglang.kernels.ops.layernorm.norm import rmsnorm as jit_rmsnorm
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.utils import is_in_ci
 

test/registered/kernels/benchmark/diffusion/bench_qknorm_rope.py

@@ -131,7 +131,7 @@ def clone_inputs(
 def split_qknorm_rope(inputs: dict[str, torch.Tensor | bool]) -> None:
     from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace
 
-    from sglang.kernels.ops.layernorm._jit_norm import fused_inplace_qknorm
+    from sglang.kernels.ops.layernorm.norm import fused_inplace_qknorm
 
     q = inputs["q"]
     k = inputs["k"]

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

@@ -10,10 +10,10 @@ import triton.testing
 
 from sglang.kernels.jit.benchmark import marker
 from sglang.kernels.jit.utils import get_jit_cuda_arch, is_hip_runtime
-from sglang.kernels.ops.gemm._jit_dsv3_fused_a_gemm import dsv3_fused_a_gemm
 from sglang.kernels.ops.gemm.cutedsl_dsv3_fused_a_gemm import (
     dsv3_fused_a_gemm as cutedsl_dsv3_fused_a_gemm,
 )
+from sglang.kernels.ops.gemm.dsv3_fused_a_gemm import dsv3_fused_a_gemm
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci
 

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

@@ -10,7 +10,7 @@ import torch.nn.functional as F
 from sglang.kernels.jit.benchmark import marker
 from sglang.kernels.jit.benchmark.utils import create_random
 from sglang.kernels.jit.utils import get_jit_cuda_arch, is_hip_runtime
-from sglang.kernels.ops.gemm._jit_dsv3_router_gemm import dsv3_router_gemm
+from sglang.kernels.ops.gemm.dsv3_router_gemm import dsv3_router_gemm
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(

test/registered/kernels/benchmark/kvcache/bench_set_mla_kv_buffer.py

@@ -21,9 +21,7 @@ from sglang.kernels.jit.benchmark.utils import (
     get_benchmark_range,
 )
 from sglang.kernels.jit.utils import is_arch_support_pdl
-from sglang.kernels.ops.kvcache._jit_set_mla_kv_buffer import (
-    set_mla_kv_buffer as jit_set,
-)
+from sglang.kernels.ops.kvcache.set_mla_kv_buffer import set_mla_kv_buffer as jit_set
 from sglang.srt.mem_cache.utils import set_mla_kv_buffer_kernel as sglang_triton_kernel
 from sglang.srt.mem_cache.utils import set_mla_kv_buffer_triton as sglang_wrapper
 from sglang.test.ci.ci_register import register_cuda_ci

test/registered/kernels/benchmark/layernorm/bench_norm.py

@@ -7,10 +7,8 @@ from flashinfer.norm import fused_add_rmsnorm as fi_fused_add_rmsnorm
 from flashinfer.norm import rmsnorm as fi_rmsnorm
 
 from sglang.kernels.jit.benchmark.utils import get_benchmark_range, run_benchmark
-from sglang.kernels.ops.layernorm._jit_norm import (
-    fused_add_rmsnorm as jit_fused_add_rmsnorm,
-)
-from sglang.kernels.ops.layernorm._jit_norm import rmsnorm as jit_rmsnorm
+from sglang.kernels.ops.layernorm.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm
+from sglang.kernels.ops.layernorm.norm import rmsnorm as jit_rmsnorm
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(

test/registered/kernels/benchmark/layernorm/bench_qknorm.py

@@ -2,7 +2,7 @@ import torch
 
 from sglang.kernels.jit.benchmark import marker
 from sglang.kernels.jit.benchmark.utils import create_random
-from sglang.kernels.ops.layernorm._jit_norm import fused_inplace_qknorm
+from sglang.kernels.ops.layernorm.norm import fused_inplace_qknorm
 from sglang.srt.utils import get_current_device_stream_fast
 from sglang.test.ci.ci_register import register_cuda_ci
 

test/registered/kernels/benchmark/layernorm/bench_qknorm_across_heads.py

@@ -7,7 +7,7 @@ import triton.testing
 from sgl_kernel import rmsnorm
 
 from sglang.kernels.jit.benchmark.utils import run_benchmark
-from sglang.kernels.ops.layernorm._jit_norm import fused_inplace_qknorm_across_heads
+from sglang.kernels.ops.layernorm.norm import fused_inplace_qknorm_across_heads
 from sglang.srt.utils import get_current_device_stream_fast
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.utils import is_in_ci

test/registered/kernels/benchmark/quantization/bench_per_tensor_quant_fp8.py

@@ -5,7 +5,7 @@ import triton
 import triton.testing
 
 from sglang.kernels.jit.benchmark.utils import get_benchmark_range, run_benchmark
-from sglang.kernels.ops.quantization._jit_per_tensor_quant_fp8 import (
+from sglang.kernels.ops.quantization.per_tensor_quant_fp8 import (
     per_tensor_quant_fp8,
 )
 from sglang.test.ci.ci_register import register_cuda_ci

test/registered/kernels/benchmark/quantization/bench_per_token_group_quant.py

@@ -1,20 +1,20 @@
 from sglang.kernels.jit.benchmark import marker
 from sglang.kernels.jit.benchmark.utils import create_empty, create_random
-
-# per_token_group_quant_8bit_v2 is DEPRECATED (no production call sites); the
-# kernel is kept only as the perf baseline for this benchmark.
-from sglang.kernels.ops.quantization._jit_per_token_group_quant import (
-    per_token_group_quant,
-)
-from sglang.kernels.ops.quantization._jit_per_token_group_quant_8bit_v2 import (
-    per_token_group_quant_8bit_v2,
-)
 from sglang.kernels.ops.quantization.fp8_kernel import (
     create_per_token_group_quant_fp8_output_scale,
     fp8_dtype,
     fp8_max,
     fp8_min,
 )
+
+# per_token_group_quant_8bit_v2 is DEPRECATED (no production call sites); the
+# kernel is kept only as the perf baseline for this benchmark.
+from sglang.kernels.ops.quantization.per_token_group_quant import (
+    per_token_group_quant,
+)
+from sglang.kernels.ops.quantization.per_token_group_quant_8bit_v2 import (
+    per_token_group_quant_8bit_v2,
+)
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(

test/registered/kernels/benchmark/quantization/bench_per_token_group_quant_8bit_v2.py

@@ -3,15 +3,15 @@ from sgl_kernel import sgl_per_token_group_quant_8bit
 
 from sglang.kernels.jit.benchmark import marker
 from sglang.kernels.jit.benchmark.utils import create_random
-from sglang.kernels.ops.quantization._jit_per_token_group_quant_8bit_v2 import (
-    per_token_group_quant_8bit_v2,
-)
 from sglang.kernels.ops.quantization.fp8_kernel import (
     create_per_token_group_quant_fp8_output_scale,
     fp8_dtype,
     fp8_max,
     fp8_min,
 )
+from sglang.kernels.ops.quantization.per_token_group_quant_8bit_v2 import (
+    per_token_group_quant_8bit_v2,
+)
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(

test/registered/kernels/benchmark/quantization/bench_per_token_group_quant_masked.py

@@ -4,21 +4,21 @@ import torch
 
 from sglang.kernels.jit.benchmark import marker
 from sglang.kernels.jit.benchmark.utils import create_empty, create_random
-
-# per_token_group_quant_8bit_v2 is DEPRECATED (no production call sites); the
-# kernel is kept only as the perf baseline for this benchmark.
-from sglang.kernels.ops.quantization._jit_per_token_group_quant import (
-    per_token_group_quant,
-)
-from sglang.kernels.ops.quantization._jit_per_token_group_quant_8bit_v2 import (
-    per_token_group_quant_8bit_v2,
-)
 from sglang.kernels.ops.quantization.fp8_kernel import (
     create_per_token_group_quant_fp8_output_scale,
     fp8_dtype,
     fp8_max,
     fp8_min,
 )
+
+# per_token_group_quant_8bit_v2 is DEPRECATED (no production call sites); the
+# kernel is kept only as the perf baseline for this benchmark.
+from sglang.kernels.ops.quantization.per_token_group_quant import (
+    per_token_group_quant,
+)
+from sglang.kernels.ops.quantization.per_token_group_quant_8bit_v2 import (
+    per_token_group_quant_8bit_v2,
+)
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(