[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>

python/sglang/kernels/README.md

@@ -25,9 +25,9 @@ sglang/kernels/
 ```
 
 Operator groups (all populated): `activation`, `attention`, `communication`,
-`diffusion`, `embeddings`, `gemm`, `grammar`, `kv_canary`, `kvcache`,
-`layernorm`, `lplb`, `mamba`, `memory`, `model`, `moe`, `quantization`,
-`sampling`, `spatial`, `speculative`.
+`diffusion`, `elementwise`, `embeddings`, `gemm`, `grammar`, `kv_canary`,
+`kvcache`, `layernorm`, `lplb`, `mamba`, `memory`, `moe`, `quantization`,
+`sampling`, `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`

python/sglang/kernels/ops/__init__.py

@@ -19,6 +19,7 @@ _GROUPS = (
     "attention",
     "communication",
     "diffusion",
+    "elementwise",
     "embeddings",
     "gemm",
     "grammar",
@@ -29,11 +30,9 @@ _GROUPS = (
     "moe",
     "quantization",
     "sampling",
-    "spatial",
     "speculative",
     "lplb",
     "kv_canary",
-    "model",
 )
 
 for _group in _GROUPS:

python/sglang/kernels/ops/activation/__init__.py

@@ -31,7 +31,7 @@ _HIP = frozenset({CapabilityRequirement.HIP})
 # — the canonical OR-semantics case that a device-baked backend name couldn't.
 _CUDA_HIP = frozenset({CapabilityRequirement.CUDA, CapabilityRequirement.HIP})
 # JIT before AOT to match the production path (srt/layers/activation.py imports
-# from sglang.kernels.ops.activation._jit_activation on CUDA); auto-selection must not invert it.
+# from sglang.kernels.ops.activation.activation on CUDA); auto-selection must not invert it.
 _ACT_PRIORITY = (
     KernelBackend.JIT,
     KernelBackend.AOT,
@@ -82,7 +82,7 @@ class _GatedActivationOp(BaseFusedOp):
         expert_ids: Optional[torch.Tensor] = None,
         expert_step: int = 1,
     ) -> torch.Tensor:
-        import sglang.kernels.ops.activation._jit_activation as jit_activation
+        import sglang.kernels.ops.activation.activation as jit_activation
 
         return getattr(jit_activation, self.kernel_attr)(
             input, out, expert_ids, expert_step
@@ -183,7 +183,7 @@ class GeluTanhAndMulOp(_GatedActivationOp):
 class ReLU2Op(BaseFusedOp):
     """``out = relu(input) ** 2`` (single-input, not gated).
 
-    The real kernel is the CUDA JIT path (``sglang.kernels.ops.activation._jit_activation.relu2``,
+    The real kernel is the CUDA JIT path (``sglang.kernels.ops.activation.activation.relu2``,
     used in production on CUDA); elsewhere the torch reference runs.
     """
 
@@ -214,7 +214,7 @@ class ReLU2Op(BaseFusedOp):
     def forward_jit(
         self, input: torch.Tensor, out: Optional[torch.Tensor] = None
     ) -> torch.Tensor:
-        from sglang.kernels.ops.activation._jit_activation import relu2
+        from sglang.kernels.ops.activation.activation import relu2
 
         result = relu2(input)
         if out is None:

python/sglang/kernels/ops/activation/activation.py

@@ -29,7 +29,7 @@ def _fast_math_flags() -> list[str]:
 
 
 @cache_once
-def _jit_activation_module(dtype: torch.dtype) -> Module:
+def activation_module(dtype: torch.dtype) -> Module:
     args = make_cpp_args(dtype, is_arch_support_pdl())
     return load_jit(
         "activation",
@@ -59,7 +59,7 @@ def _run_activation_inplace(
     op_name: str, input: torch.Tensor, out: torch.Tensor
 ) -> None:
     hidden_size = input.shape[-1] // 2
-    module = _jit_activation_module(input.dtype)
+    module = activation_module(input.dtype)
     input_2d = input.view(-1, hidden_size * 2)
     out_2d = out.view(-1, hidden_size)
     module.run_activation(input_2d, out_2d, op_name)
@@ -74,7 +74,7 @@ def _run_activation_filtered_inplace(
     expert_step: int,
 ) -> None:
     hidden_size = input.shape[-1] // 2
-    module = _jit_activation_module(input.dtype)
+    module = activation_module(input.dtype)
     input_2d = input.view(-1, hidden_size * 2)
     out_2d = out.view(-1, hidden_size)
     module.run_activation_filtered(input_2d, out_2d, expert_ids, expert_step, op_name)
@@ -110,7 +110,7 @@ def _run_unary_activation_inplace(
     op_name: str, input: torch.Tensor, out: torch.Tensor
 ) -> None:
     last = input.shape[-1]
-    module = _jit_activation_module(input.dtype)
+    module = activation_module(input.dtype)
     module.run_unary_activation(input.view(-1, last), out.view(-1, last), op_name)
 
 

python/sglang/kernels/ops/attention/dsv4/compress_old.py

@@ -57,7 +57,7 @@ def _jit_compress_128_online_module(head_dim: int) -> Module:
 
 
 @cache_once
-def _jit_norm_rope_module(
+def norm_rope_module(
     dtype: torch.dtype,
     head_dim: int,
     rope_dim: int,
@@ -276,7 +276,7 @@ def compress_fused_norm_rope_inplace(
     plan: Union[CompressorDecodePlan, CompressorPrefillPlan],
 ) -> None:
     freq_cis = torch.view_as_real(freq_cis).flatten(-2)
-    module = _jit_norm_rope_module(kv.dtype, kv.shape[-1], freq_cis.shape[-1])
+    module = norm_rope_module(kv.dtype, kv.shape[-1], freq_cis.shape[-1])
     module.forward(
         kv,
         weight,
@@ -296,7 +296,7 @@ def fused_norm_rope_inplace(
     positions: torch.Tensor,
 ) -> None:
     freq_cis = torch.view_as_real(freq_cis).flatten(-2)
-    module = _jit_norm_rope_module(kv.dtype, kv.shape[-1], freq_cis.shape[-1])
+    module = norm_rope_module(kv.dtype, kv.shape[-1], freq_cis.shape[-1])
     module.forward(
         kv,
         weight,

python/sglang/kernels/ops/attention/log_scaling_tau.py

@@ -50,7 +50,7 @@ def apply_log_scaling_tau(x: torch.Tensor, tau: torch.Tensor) -> torch.Tensor:
         # Vectorized JIT kernel (16B loads, one row divide per vector) --
         # bit-identical output (same fp32-mul + bf16-round), ~2-3x the
         # scalar triton kernel below at every size.
-        from sglang.kernels.ops.model.inkling.inkling_row_scale import row_scale_bf16
+        from sglang.kernels.ops.attention.inkling_row_scale import row_scale_bf16
 
         x2d = torch.as_strided(x, (rows, inner), (x.stride(0), 1))
         return row_scale_bf16(x2d, tau.reshape(rows).float()).view(x.shape)

python/sglang/kernels/ops/diffusion/norm_scale_shift_native.py

@@ -59,7 +59,7 @@ def _row_bf16(t, device: torch.device):
 
 
 @cache_once
-def _jit_norm_scale_shift_module() -> Module:
+def norm_scale_shift_module() -> Module:
     return load_jit(
         "qwen_image_norm_scale_shift_native",
         cuda_files=["diffusion/norm_scale_shift.cuh"],
@@ -77,7 +77,7 @@ def _jit_norm_scale_shift_module() -> Module:
     )
 
 
-_module = _jit_norm_scale_shift_module
+_module = norm_scale_shift_module
 
 
 def try_fused_norm_scale_shift(x, weight, bias, scale, shift, norm_type, eps):

python/sglang/kernels/ops/elementwise/__init__.py

@@ -0,0 +1,11 @@
+"""Generic elementwise / fused-pointwise kernels.
+
+Home for cross-cutting pointwise kernels that do not belong to a single
+functional group: the fused-pointwise Triton collection (``elementwise``:
+softcap, sigmoid-mul, gated-activation and fused-rmsnorm variants shared
+across models) and the ``add_constant`` JIT reference kernel used by the
+developer guide. Individual functions register (or are imported) under the
+functional op id they logically belong to.
+"""
+
+__all__ = []

python/sglang/kernels/ops/gemm/__init__.py

@@ -59,7 +59,7 @@ register_kernel(
     KernelSpec(
         op="gemm.dsv3_fused_a_gemm",
         backend=KernelBackend.JIT,
-        target="sglang.kernels.ops.gemm._jit_dsv3_fused_a_gemm:dsv3_fused_a_gemm",
+        target="sglang.kernels.ops.gemm.dsv3_fused_a_gemm:dsv3_fused_a_gemm",
         capabilities=_CUDA,
         format_signature=FormatSignature(
             supported_dtypes=("bfloat16",),
@@ -72,7 +72,7 @@ register_kernel(
     KernelSpec(
         op="gemm.dsv3_router_gemm",
         backend=KernelBackend.JIT,
-        target="sglang.kernels.ops.gemm._jit_dsv3_router_gemm:dsv3_router_gemm",
+        target="sglang.kernels.ops.gemm.dsv3_router_gemm:dsv3_router_gemm",
         capabilities=_CUDA,
         format_signature=FormatSignature(
             supported_dtypes=("bfloat16",),

python/sglang/kernels/ops/gemm/dsv3_fused_a_gemm.py

@@ -25,7 +25,7 @@ if TYPE_CHECKING:
 
 
 @cache_once
-def _jit_dsv3_fused_a_gemm_module(hd_in: int, hd_out: int, use_pdl: bool) -> Module:
+def dsv3_fused_a_gemm_module(hd_in: int, hd_out: int, use_pdl: bool) -> Module:
     args = make_cpp_args(hd_in, hd_out, use_pdl)
     return load_jit(
         "dsv3_fused_a_gemm",
@@ -44,7 +44,7 @@ def _dsv3_fused_a_gemm_run(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Te
         device=mat_a.device,
         dtype=mat_a.dtype,
     )
-    module = _jit_dsv3_fused_a_gemm_module(
+    module = dsv3_fused_a_gemm_module(
         mat_a.shape[1], mat_b.shape[1], is_arch_support_pdl()
     )
     module.dsv3_fused_a_gemm(mat_a, mat_b, output)

python/sglang/kernels/ops/gemm/dsv3_router_gemm.py

@@ -25,7 +25,7 @@ if TYPE_CHECKING:
 
 
 @cache_once
-def _jit_dsv3_router_gemm_module(
+def dsv3_router_gemm_module(
     num_experts: int,
     hidden_dim: int,
     use_pdl: bool,
@@ -54,7 +54,7 @@ def _dsv3_router_gemm_custom_op(
     num_experts = router_weights.shape[0]
     hidden_dim = hidden_states.shape[1]
     out_float = output.dtype == torch.float32
-    module = _jit_dsv3_router_gemm_module(
+    module = dsv3_router_gemm_module(
         num_experts, hidden_dim, is_arch_support_pdl(), out_float
     )
     module.dsv3_router_gemm(hidden_states, router_weights, output)

python/sglang/kernels/ops/gemm/fused_a_gemm.py

@@ -2,7 +2,7 @@
 
 Dispatches to one of two interchangeable implementations via ``backend``:
 
-- ``"jit"``: runtime-compiled CUDA C++ (``sglang.kernels.ops.gemm._jit_dsv3_fused_a_gemm``).
+- ``"jit"``: runtime-compiled CUDA C++ (``sglang.kernels.ops.gemm.dsv3_fused_a_gemm``).
 - ``"cutedsl"``: CuTe DSL (``sglang.kernels.ops.gemm.cutedsl_dsv3_fused_a_gemm``).
 - ``"auto"``: CuTe DSL on SM120+, otherwise the JIT kernel.
 
@@ -69,9 +69,7 @@ def dsv3_fused_a_gemm(
         backend = _AUTO_BACKEND
 
     if backend == FusedAGemmBackend.JIT:
-        from sglang.kernels.ops.gemm._jit_dsv3_fused_a_gemm import (
-            dsv3_fused_a_gemm as impl,
-        )
+        from sglang.kernels.ops.gemm.dsv3_fused_a_gemm import dsv3_fused_a_gemm as impl
     else:
         from sglang.kernels.ops.gemm.cutedsl_dsv3_fused_a_gemm import (
             dsv3_fused_a_gemm as impl,

python/sglang/kernels/ops/kvcache/mla_buffer.py

@@ -116,10 +116,10 @@ def set_mla_kv_buffer_triton(
     Name retained for caller compatibility; the implementation is no longer
     Triton-only.
     """
-    from sglang.kernels.ops.kvcache._jit_set_mla_kv_buffer import (
+    from sglang.kernels.ops.kvcache.set_mla_kv_buffer import (
         can_use_set_mla_kv_buffer,
     )
-    from sglang.kernels.ops.kvcache._jit_set_mla_kv_buffer import (
+    from sglang.kernels.ops.kvcache.set_mla_kv_buffer import (
         set_mla_kv_buffer as jit_set_mla_kv_buffer,
     )
 

python/sglang/kernels/ops/kvcache/set_mla_kv_buffer.py

@@ -27,9 +27,7 @@ logger = logging.getLogger(__name__)
 
 
 @cache_once
-def _jit_set_mla_kv_buffer_module(
-    nope_bytes: int, rope_bytes: int, use_pdl: bool
-) -> Module:
+def set_mla_kv_buffer_module(nope_bytes: int, rope_bytes: int, use_pdl: bool) -> Module:
     args = make_cpp_args(nope_bytes, rope_bytes, use_pdl)
     return load_jit(
         f"set_mla_kv_buffer_{nope_bytes}_{rope_bytes}",
@@ -66,7 +64,7 @@ def can_use_set_mla_kv_buffer(nope_bytes: int, rope_bytes: int) -> bool:
         )
         return False
     try:
-        _jit_set_mla_kv_buffer_module(nope_bytes, rope_bytes, is_arch_support_pdl())
+        set_mla_kv_buffer_module(nope_bytes, rope_bytes, is_arch_support_pdl())
         return True
     except Exception as e:  # pragma: no cover - compile-time only
         logger.warning(
@@ -115,7 +113,5 @@ def set_mla_kv_buffer(
     if num_warps <= 0:
         num_warps = _pick_num_warps(n_loc)
 
-    module = _jit_set_mla_kv_buffer_module(
-        nope_bytes, rope_bytes, is_arch_support_pdl()
-    )
+    module = set_mla_kv_buffer_module(nope_bytes, rope_bytes, is_arch_support_pdl())
     module.set_mla_kv_buffer(buf, loc, src_nope, src_rope, num_warps)

python/sglang/kernels/ops/layernorm/__init__.py

@@ -113,7 +113,7 @@ class RMSNormOp(BaseFusedOp):
     ) -> torch.Tensor:
         import torch
 
-        from sglang.kernels.ops.layernorm._jit_norm import rmsnorm as jit_rmsnorm
+        from sglang.kernels.ops.layernorm.norm import rmsnorm as jit_rmsnorm
 
         if out is None:
             out = torch.empty_like(input)
@@ -227,7 +227,7 @@ class FusedAddRMSNormOp(BaseFusedOp):
         eps: float = 1e-6,
         enable_pdl: Optional[bool] = None,
     ) -> None:
-        from sglang.kernels.ops.layernorm._jit_norm import (
+        from sglang.kernels.ops.layernorm.norm import (
             fused_add_rmsnorm as jit_fused_add_rmsnorm,
         )
 
@@ -512,8 +512,6 @@ from sglang.kernels.spec import KernelSpec
 # Triton / TileLang kernels migrated from srt/layers top-level strays
 # (RFC #29630, Phase 2.5); registered for inventory.
 _PHASE25_KERNELS = [
-    ("elementwise", "fused_dual_residual_rmsnorm", "triton"),
-    ("elementwise", "fused_rmsnorm", "triton"),
     ("gemma4_fused_ops", "gemma4_fused_routing", "triton"),
     ("gemma4_fused_ops", "gemma_qkv_rmsnorm", "triton"),
     ("mhc_head", "fused_hc_head", "triton"),
@@ -527,3 +525,15 @@ for _mod, _fn, _bk in _PHASE25_KERNELS:
         )
     )
 del _mod, _fn, _bk
+
+# The fused-rmsnorm variants physically live in the shared fused-pointwise
+# collection (sglang.kernels.ops.elementwise.elementwise) but stay layernorm ops.
+for _fn in ("fused_dual_residual_rmsnorm", "fused_rmsnorm"):
+    register_kernel(
+        KernelSpec(
+            op=f"layernorm.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.elementwise.elementwise:{_fn}",
+        )
+    )
+del _fn

python/sglang/kernels/ops/model/__init__.py

@@ -1 +0,0 @@
-"""Model-specific JIT kernels (RFC #29630)."""

python/sglang/kernels/ops/model/inkling/__init__.py

@@ -1 +0,0 @@
-"""Inkling model-family JIT kernels."""

python/sglang/kernels/ops/quantization/__init__.py

@@ -56,7 +56,7 @@ register_kernel(
     KernelSpec(
         op="quantization.per_token_group_quant",
         backend=KernelBackend.JIT,
-        target="sglang.kernels.ops.quantization._jit_per_token_group_quant:per_token_group_quant",
+        target="sglang.kernels.ops.quantization.per_token_group_quant:per_token_group_quant",
         capabilities=_CUDA,
         format_signature=FormatSignature(
             supported_dtypes=("float8_e4m3fn", "int8"),

python/sglang/kernels/ops/quantization/fp8_kernel.py

@@ -59,7 +59,7 @@ if _is_cuda or _is_musa:
         per_token_group_quant,
         sgl_per_token_quant_fp8,
     )
-    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 as sgl_per_tensor_quant_fp8,
     )
 
@@ -543,7 +543,7 @@ def _run_per_token_group_quant_8bit_kernel(
     ``sglang_per_token_quant_fp8``.
     """
     if scale_ue8m0 and x_s.dtype == torch.float32 and not _is_musa:
-        from sglang.kernels.ops.quantization._jit_per_token_group_quant_8bit_v2 import (
+        from sglang.kernels.ops.quantization.per_token_group_quant_8bit_v2 import (
             per_token_group_quant_8bit_v2,
         )
 

python/sglang/kernels/ops/quantization/per_tensor_quant_fp8.py

@@ -12,7 +12,7 @@ if TYPE_CHECKING:
 
 
 @cache_once
-def _jit_per_tensor_quant_fp8_module(is_static: bool, dtype: torch.dtype) -> Module:
+def per_tensor_quant_fp8_module(is_static: bool, dtype: torch.dtype) -> Module:
     args = make_cpp_args(is_static, dtype)
     return load_jit(
         "per_tensor_quant_fp8",
@@ -41,7 +41,7 @@ def per_tensor_quant_fp8(
         output_s: Output scale tensor (float scalar or 1D tensor with 1 element)
         is_static: If True, assumes scale is pre-computed and skips absmax computation
     """
-    module = _jit_per_tensor_quant_fp8_module(is_static, input.dtype)
+    module = per_tensor_quant_fp8_module(is_static, input.dtype)
     module.per_tensor_quant_fp8(input.view(-1), output_q.view(-1), output_s.view(-1))
 
 

python/sglang/kernels/ops/quantization/per_token_group_quant_8bit_v2.py

@@ -1,4 +1,4 @@
-"""DEPRECATED: superseded by ``sglang.kernels.ops.quantization._jit_per_token_group_quant`` (the
+"""DEPRECATED: superseded by ``sglang.kernels.ops.quantization.per_token_group_quant`` (the
 default CUDA path). No sglang runtime code may call this kernel; it is kept
 only as the perf baseline for the per_token_group_quant benchmarks and its own
 bit-parity tests, and will be deleted once those move to torch references.

python/sglang/kernels/ops/spatial/__init__.py

@@ -1,49 +0,0 @@
-"""Spatial / green-context stream helpers."""
-
-from __future__ import annotations
-
-from typing import Optional
-
-from sglang.kernels.registry import register_kernel
-from sglang.kernels.selector import get_kernel
-from sglang.kernels.spec import FormatSignature, KernelBackend, KernelSpec
-
-register_kernel(
-    KernelSpec(
-        op="spatial.get_sm_available",
-        backend=KernelBackend.AOT,
-        target="sgl_kernel.spatial:get_sm_available",
-        format_signature=FormatSignature(
-            description="number of SMs available on device"
-        ),
-        description="Query available SM count (sgl_kernel wheel).",
-    )
-)
-register_kernel(
-    KernelSpec(
-        op="spatial.create_greenctx_stream_by_value",
-        backend=KernelBackend.AOT,
-        target="sgl_kernel.spatial:create_greenctx_stream_by_value",
-        format_signature=FormatSignature(
-            description="create two green-context streams partitioned by SM count"
-        ),
-        description="Green-context stream creation (sgl_kernel wheel).",
-    )
-)
-
-
-def get_sm_available(device_id: Optional[int] = None) -> int:
-    """Return the number of SMs available on ``device_id``."""
-    return get_kernel("spatial.get_sm_available", KernelBackend.AOT)(device_id)
-
-
-def create_greenctx_stream_by_value(
-    SM_a: int, SM_b: int, device_id: Optional[int] = None
-):
-    """Create two green-context streams partitioned by ``SM_a`` / ``SM_b``."""
-    return get_kernel("spatial.create_greenctx_stream_by_value", KernelBackend.AOT)(
-        SM_a, SM_b, device_id
-    )
-
-
-__all__ = ["get_sm_available", "create_greenctx_stream_by_value"]

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md

@@ -29,7 +29,7 @@ framework-specific optimization workflow.
 - `test/registered/jit/benchmark/diffusion/bench_qwen_image_modulation.py`
 - `test/registered/jit/benchmark/diffusion/bench_group_norm_silu.py`
 - `test/registered/jit/benchmark/diffusion/bench_residual_gate_add.py`
-- `python/sglang/kernels/ops/layernorm/_jit_norm.py`
+- `python/sglang/kernels/ops/layernorm/norm.py`
 - `python/sglang/multimodal_gen/runtime/platforms/cuda.py`
 - `python/sglang/multimodal_gen/runtime/layers/attention/selector.py`
 - `docs_new/docs/sglang-diffusion/attention_backends.mdx` (repo root)
@@ -137,7 +137,7 @@ framework-specific optimization workflow.
 **QK Norm Optimization**
 
 - Entry point: `apply_qk_norm` in `layernorm.py`.
-- Fast path: JIT fused inplace QK norm from `python/sglang/kernels/ops/layernorm/_jit_norm.py` via `fused_inplace_qknorm`.
+- Fast path: JIT fused inplace QK norm from `python/sglang/kernels/ops/layernorm/norm.py` via `fused_inplace_qknorm`.
 - Preconditions for fused path:
   - CUDA only.
   - `allow_inplace=True` and `q_eps == k_eps`.

python/sglang/multimodal_gen/runtime/layers/activation.py

@@ -19,7 +19,7 @@ _is_npu = current_platform.is_npu()
 _is_xpu = current_platform.is_xpu()
 
 if _is_cuda:
-    from sglang.kernels.ops.activation._jit_activation import silu_and_mul
+    from sglang.kernels.ops.activation.activation import silu_and_mul
 elif _is_hip or _is_xpu:
     from sgl_kernel import silu_and_mul
 

python/sglang/multimodal_gen/runtime/layers/layernorm.py

@@ -17,7 +17,7 @@ from sglang.kernels.ops.diffusion.qknorm_rope import (
 )
 from sglang.kernels.ops.diffusion.triton.rmsnorm_onepass import triton_one_pass_rms_norm
 from sglang.kernels.ops.diffusion.triton.scale_shift import fuse_scale_shift_kernel
-from sglang.kernels.ops.layernorm._jit_norm import (
+from sglang.kernels.ops.layernorm.norm import (
     can_use_fused_inplace_qknorm,
     fused_inplace_qknorm,
 )

python/sglang/srt/layers/activation.py

@@ -57,7 +57,7 @@ _is_xpu = is_xpu()
 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
 
 if _is_cuda:
-    from sglang.kernels.ops.activation._jit_activation import (
+    from sglang.kernels.ops.activation.activation import (
         gelu_and_mul,
         gelu_tanh_and_mul,
         relu2,

python/sglang/srt/layers/attention/dsa/dsa_indexer.py

@@ -343,7 +343,7 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor:
     elif _is_xpu:
         from sgl_kernel import hadamard_transform
     else:
-        from sglang.kernels.ops.attention.hadamard import hadamard_transform
+        from sglang.kernels.ops.quantization.hadamard import hadamard_transform
 
     hidden_size = x.size(-1)
     assert (

python/sglang/srt/layers/attention/vision.py

@@ -12,7 +12,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange
 
-from sglang.kernels.ops.layernorm._jit_norm import (
+from sglang.kernels.ops.layernorm.norm import (
     can_use_fused_inplace_qknorm as can_use_jit_qk_norm,
 )
 from sglang.srt.environ import envs

python/sglang/srt/layers/layernorm.py

@@ -147,10 +147,10 @@ if _is_cuda:
 
         _jit_rmsnorm_hf = None
 
-    from sglang.kernels.ops.layernorm._jit_norm import (
+    from sglang.kernels.ops.layernorm.norm import (
         fused_add_rmsnorm as _jit_fused_add_rmsnorm,
     )
-    from sglang.kernels.ops.layernorm._jit_norm import (
+    from sglang.kernels.ops.layernorm.norm import (
         is_supported_jit_fused_add_rmsnorm_hidden_size,
     )
 

python/sglang/srt/layers/moe/cutlass_moe.py

@@ -18,7 +18,7 @@ if _is_cuda:
         shuffle_rows,
     )
 
-    from sglang.kernels.ops.activation._jit_activation import silu_and_mul
+    from sglang.kernels.ops.activation.activation import silu_and_mul
 
 
 def cutlass_fused_experts_fp8(

python/sglang/srt/layers/moe/cutlass_w4a8_moe.py

@@ -18,7 +18,7 @@ if _is_cuda_alike:
     )
 
 if _is_cuda:
-    from sglang.kernels.ops.activation._jit_activation import silu_and_mul
+    from sglang.kernels.ops.activation.activation import silu_and_mul
 else:
     from sgl_kernel import silu_and_mul
 
@@ -35,7 +35,7 @@ from sglang.kernels.ops.moe.ep_moe_kernels import (
     silu_mul_dynamic_tensorwise_quant_for_cutlass_moe,
     silu_mul_static_tensorwise_quant_for_cutlass_moe,
 )
-from sglang.kernels.ops.quantization._jit_per_tensor_quant_fp8 import (
+from sglang.kernels.ops.quantization.per_tensor_quant_fp8 import (
     per_tensor_absmax_fp8,
     per_tensor_quant_fp8,
 )

python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py

@@ -11,7 +11,7 @@ _is_cuda = is_cuda()
 if _is_cuda:
     from sgl_kernel import moe_sum_reduce
 
-    from sglang.kernels.ops.activation._jit_activation import silu_and_mul
+    from sglang.kernels.ops.activation.activation import silu_and_mul
     from sglang.kernels.ops.moe.moe_wna16_marlin import moe_wna16_marlin_gemm
 
 

python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py

@@ -30,7 +30,7 @@ if is_sm120_supported():
     update_opt_flags_constraints({"is_persistent": False})
 
 if is_cuda():
-    from sglang.kernels.ops.activation._jit_activation import gelu_and_mul, silu_and_mul
+    from sglang.kernels.ops.activation.activation import gelu_and_mul, silu_and_mul
 else:
     from sgl_kernel import gelu_and_mul, silu_and_mul
 

python/sglang/srt/layers/moe/moe_runner/deep_gemm.py

@@ -56,7 +56,7 @@ _is_musa = is_musa()
 
 # Imported only for the SGLANG_OPT_FIX_MEGA_MOE_MEMORY=False fallback path.
 if not (_is_npu or _is_hip) and _is_cuda:
-    from sglang.kernels.ops.activation._jit_activation import (
+    from sglang.kernels.ops.activation.activation import (
         silu_and_mul as _legacy_silu_and_mul,
     )
 elif _is_musa:

python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py

@@ -60,7 +60,7 @@ _is_musa = is_musa()
 if _is_cuda:
     from sgl_kernel import moe_sum_reduce
 
-    from sglang.kernels.ops.activation._jit_activation import gelu_and_mul, silu_and_mul
+    from sglang.kernels.ops.activation.activation import gelu_and_mul, silu_and_mul
 elif _is_cpu and _is_cpu_amx_available:
     pass
 elif _is_hip:

python/sglang/srt/layers/moe/moe_runner/triton_utils/sigmoid_gate_topk_renorm.py

@@ -13,7 +13,7 @@ import triton
 import triton.language as tl
 
 from sglang.kernels.jit.utils import is_arch_support_pdl
-from sglang.kernels.ops.model.inkling.inkling_gate_topk_renorm import (
+from sglang.kernels.ops.moe.inkling_gate_topk_renorm import (
     inkling_gate_topk_renorm_v2,
 )
 from sglang.srt.environ import envs

python/sglang/srt/layers/quantization/gguf.py

@@ -51,7 +51,7 @@ if _is_cuda:
         ggml_mul_mat_vec_a8,
     )
 
-    from sglang.kernels.ops.activation._jit_activation import gelu_and_mul, silu_and_mul
+    from sglang.kernels.ops.activation.activation import gelu_and_mul, silu_and_mul
 elif _is_musa:
     from sgl_kernel import gelu_and_mul, moe_align_block_size, moe_sum, silu_and_mul
     from sgl_kernel.quantization import (

python/sglang/srt/models/deepseek_v2.py

@@ -213,8 +213,8 @@ if _use_aiter:
     pass
 
 if _is_cuda:
-    from sglang.kernels.ops.gemm._jit_dsv3_router_gemm import (
-        dsv3_router_gemm as _jit_dsv3_router_gemm,
+    from sglang.kernels.ops.gemm.dsv3_router_gemm import (
+        dsv3_router_gemm as dsv3_router_gemm,
     )
 elif _is_npu:
     from sglang.srt.hardware_backend.npu.modules.deepseek_v2_attention_mla_npu import (
@@ -512,7 +512,7 @@ class MoEGate(nn.Module):
                 and (self.weight.shape[0] == 256 or self.weight.shape[0] == 384)
                 and _device_sm >= 90
             ):
-                logits = _jit_dsv3_router_gemm(
+                logits = dsv3_router_gemm(
                     hidden_states, self.weight, out_dtype=torch.float32
                 )
 

python/sglang/srt/models/grok.py

@@ -22,7 +22,7 @@ import torch.nn.functional as F
 from torch import nn
 from transformers import PretrainedConfig
 
-from sglang.kernels.ops.layernorm.elementwise import (
+from sglang.kernels.ops.elementwise.elementwise import (
     fused_dual_residual_rmsnorm,
     fused_rmsnorm,
     gelu_and_mul_triton,

python/sglang/srt/models/inkling.py

@@ -625,7 +625,7 @@ class InklingCausalLLM(nn.Module):
             and sconv0 is not None
             and world in (4, 8)  # symm-mem multimem worlds, power-of-two
         ):
-            from sglang.kernels.ops.model.inkling.inkling_ar_fused import (
+            from sglang.kernels.ops.communication.inkling_ar_fused import (
                 compile_inkling_ar_sconv_norm,
             )
 
@@ -646,7 +646,7 @@ class InklingCausalLLM(nn.Module):
         # local/SWA layer (head_dim != 128) that never uses the prologue while
         # later full-attention layers do.
         if is_cuda() and envs.SGLANG_OPT_USE_INKLING_FUSED_ATTN_PROLOGUE.get():
-            from sglang.kernels.ops.model.inkling.inkling_attn_prologue import (
+            from sglang.kernels.ops.attention.inkling_attn_prologue import (
                 compile_inkling_attn_prologue,
             )
 

python/sglang/srt/models/inkling_common/attn.py

@@ -6,14 +6,14 @@ from functools import cache
 import torch
 from torch import nn
 
+from sglang.kernels.ops.attention.inkling_rel_proj import rel_proj_small_t
+from sglang.kernels.ops.attention.inkling_row_scale import row_compact_bf16
 from sglang.kernels.ops.attention.log_scaling_tau import (
     apply_log_scaling_tau as _apply_log_scaling_tau,
 )
 from sglang.kernels.ops.attention.score_mod import (
     relative_bias_score_mod as triton_relative_bias_score_mod,
 )
-from sglang.kernels.ops.model.inkling.inkling_rel_proj import rel_proj_small_t
-from sglang.kernels.ops.model.inkling.inkling_row_scale import row_compact_bf16
 from sglang.srt.environ import envs
 from sglang.srt.layers.linear import MergedColumnParallelLinear, RowParallelLinear
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
@@ -364,7 +364,7 @@ class InklingAttention(nn.Module):
 
     def _fused_attn_prologue_verify(self, q, k, v, forward_batch, log_scaling_tau=None):
         """Fused target-verify {k/v sconv + save_windows + qk-norm (+ KV store)}
-        (kernels/ops/model/inkling/inkling_attn_prologue.py); returns
+        (kernels/ops/attention/inkling_attn_prologue.py); returns
         ``(q, k, v, did_store)``.
 
         The fused kernel writes raw bf16 KV, so it only does the store when the
@@ -377,7 +377,7 @@ class InklingAttention(nn.Module):
         qk-norm stay fused either way. For the FA4 MXFP8 pool, the prologue can
         quantize Q and directly fill the fp8 K/V cache plus interleaved scale
         buffers, returning Q's per-token scales as ``q_descale``/``sfq``."""
-        from sglang.kernels.ops.model.inkling.inkling_attn_prologue import (
+        from sglang.kernels.ops.attention.inkling_attn_prologue import (
             inkling_attn_prologue_verify,
         )
         from sglang.srt.model_executor.forward_context import (
@@ -481,7 +481,7 @@ class InklingAttention(nn.Module):
         the backend store. Store gating (bf16 NHD / FA4 MXFP8 pools, SWA loc
         translation) is identical to the verify prologue. Returns
         (q, k, v, did_store, q_descale)."""
-        from sglang.kernels.ops.model.inkling.inkling_attn_prologue import (
+        from sglang.kernels.ops.attention.inkling_attn_prologue import (
             inkling_attn_prologue_extend,
         )
         from sglang.srt.model_executor.forward_context import (
@@ -615,7 +615,7 @@ class InklingAttention(nn.Module):
         kernel. Decode is one token/seq so the conv taps come from the working
         cache (no cross-token reads, no barrier). Returns
         (q, k, v, did_store, q_descale)."""
-        from sglang.kernels.ops.model.inkling.inkling_attn_prologue import (
+        from sglang.kernels.ops.attention.inkling_attn_prologue import (
             inkling_attn_prologue_decode,
         )
         from sglang.srt.model_executor.forward_context import (

python/sglang/srt/models/inkling_common/kernels/comm.py

@@ -94,7 +94,7 @@ def _ar_jit():
     lazy so importing comm.py doesn't pull in the JIT machinery)."""
     if not is_cuda():
         return None
-    from sglang.kernels.ops.model.inkling import inkling_all_reduce
+    from sglang.kernels.ops.communication import inkling_all_reduce
 
     return inkling_all_reduce
 
@@ -103,7 +103,7 @@ def _ar_jit():
 def _ar_fused_jit():
     if not is_cuda():
         return None
-    from sglang.kernels.ops.model.inkling import inkling_ar_fused
+    from sglang.kernels.ops.communication import inkling_ar_fused
 
     return inkling_ar_fused
 
@@ -241,7 +241,7 @@ def ar_sconv_norm_fusable(
     (attn-side: wo_ud AR -> attn_sconv -> mlp_norm; MoE-side: MoE AR ->
     mlp_sconv -> next attn_norm)
     can run as the single fused kernel
-    (kernels/ops/model/inkling/inkling_ar_fused.py). Must be
+    (kernels/ops/communication/inkling_ar_fused.py). Must be
     evaluated identically by the producing layer (MoE ``reduce=False``) and the
     consuming layer/tail -- it is a pure function of per-forward state."""
     if not is_cuda():
@@ -675,7 +675,7 @@ def all_gather_hidden(input: torch.Tensor, group: GroupCoordinator) -> torch.Ten
 def _ar_ssconv_jit():
     if not is_cuda():
         return None
-    from sglang.kernels.ops.model.inkling import inkling_ar_scattered_sconv
+    from sglang.kernels.ops.communication import inkling_ar_scattered_sconv
 
     return inkling_ar_scattered_sconv
 
@@ -689,7 +689,7 @@ def scattered_ar_sconv_fusable(
 ) -> bool:
     """True when an extend {reduce_scatter_hidden -> sconv(shard) ->
     all_gather_hidden} chain can run as the single fused v3/v3b-style kernel
-    (kernels/ops/model/inkling/inkling_ar_scattered_sconv.py). Pure function of
+    (kernels/ops/communication/inkling_ar_scattered_sconv.py). Pure function of
     per-forward
     state -- the producing layer (reduce=False) and the consuming site must
     evaluate it identically."""

python/sglang/srt/models/inkling_common/moe.py

@@ -10,7 +10,7 @@ from torch import nn
 from triton.language.extra import libdevice
 
 from sglang.kernels.jit.utils import is_arch_support_pdl
-from sglang.kernels.ops.model.inkling.inkling_gate_topk_renorm import (
+from sglang.kernels.ops.moe.inkling_gate_topk_renorm import (
     ensure_gate_gemv_fused_scratch,
     inkling_gate_gemv,
     inkling_gate_gemv_fused,

python/sglang/srt/models/qwen2_moe.py

@@ -27,7 +27,7 @@ import torch.nn.functional as F
 from torch import nn
 from transformers import PretrainedConfig
 
-from sglang.kernels.ops.layernorm.elementwise import fused_gate_sigmoid_mul_add
+from sglang.kernels.ops.elementwise.elementwise import fused_gate_sigmoid_mul_add
 from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo
 from sglang.srt.distributed import (
     get_pp_group,

python/sglang/srt/models/qwen3_5.py

@@ -27,7 +27,7 @@ from sglang.kernels.ops.attention.fla.layernorm_gated import RMSNorm as RMSNormG
 from sglang.kernels.ops.attention.triton_gdn_fused_proj import (
     fused_qkvzba_split_reshape_cat_contiguous,
 )
-from sglang.kernels.ops.layernorm.elementwise import fused_sigmoid_mul
+from sglang.kernels.ops.elementwise.elementwise import fused_sigmoid_mul
 
 # Configs
 from sglang.srt.configs.qwen3_5 import (

python/sglang/srt/models/utils.py

@@ -25,7 +25,7 @@ import triton
 import triton.language as tl
 
 from sglang.kernels.ops.attention.rope import FusedSetKVBufferArg
-from sglang.kernels.ops.layernorm._jit_norm import (
+from sglang.kernels.ops.layernorm.norm import (
     can_use_fused_inplace_qknorm,
     fused_inplace_qknorm,
 )