[Kernel] Phase 4 (activation): move JIT activation ops into kernels.ops.activation (RFC #29630)

python/sglang/jit_kernel/__main__.py

@@ -7,10 +7,10 @@
 
 from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path
 
-from sglang.jit_kernel.utils import get_jit_cuda_arch, override_jit_cuda_arch
-from sglang.jit_kernel.utils.arch import get_default_target_flags
-from sglang.jit_kernel.utils.compile import DEFAULT_INCLUDE
-from sglang.jit_kernel.utils.deps import REGISTERED_DEPENDENCIES
+from sglang.kernels.jit import get_jit_cuda_arch, override_jit_cuda_arch
+from sglang.kernels.jit.arch import get_default_target_flags
+from sglang.kernels.jit.compile import DEFAULT_INCLUDE
+from sglang.kernels.jit.deps import REGISTERED_DEPENDENCIES
 
 
 def _clangd_major_version() -> int | None:

python/sglang/jit_kernel/activation.py

@@ -1,168 +1,10 @@
-from __future__ import annotations
+"""Compatibility shim (RFC #29630 Phase 4).
 
-from typing import TYPE_CHECKING, Optional
+JIT activation operators moved to ``sglang.kernels.ops.activation._jit_activation``.
+Re-exported here so existing ``sglang.jit_kernel.activation`` imports keep
+working; removed in Phase 5.
+"""
 
-import torch
+from sglang.kernels.ops.activation import _jit_activation as _impl
 
-from sglang.jit_kernel.utils import (
-    cache_once,
-    get_jit_cuda_arch,
-    is_arch_support_pdl,
-    is_hip_runtime,
-    load_jit,
-    make_cpp_args,
-)
-from sglang.srt.utils.custom_op import register_custom_op
-
-if TYPE_CHECKING:
-    from tvm_ffi.module import Module
-
-
-def _fast_math_flags() -> list[str]:
-    # Mirrors sgl-kernel's CMake policy: fast-math on SM90, precise on
-    # SM100+ (Blackwell needs bit-exact expf), off on HIP (clang rejects).
-    if is_hip_runtime():
-        return []
-    if get_jit_cuda_arch().major >= 10:
-        return []
-    return ["--use_fast_math"]
-
-
-@cache_once
-def _jit_activation_module(dtype: torch.dtype) -> Module:
-    args = make_cpp_args(dtype, is_arch_support_pdl())
-    return load_jit(
-        "activation",
-        *args,
-        cuda_files=["elementwise/activation.cuh"],
-        extra_cuda_cflags=_fast_math_flags(),
-        cuda_wrappers=[
-            ("run_activation", f"ActivationKernel<{args}>::run_activation"),
-            (
-                "run_activation_filtered",
-                f"ActivationKernel<{args}>::run_activation_filtered",
-            ),
-            (
-                "run_unary_activation",
-                f"ActivationKernel<{args}>::run_unary_activation",
-            ),
-        ],
-    )
-
-
-SUPPORTED_ACTIVATIONS = {"silu", "gelu", "gelu_tanh"}
-SUPPORTED_UNARY_ACTIVATIONS = {"relu2"}
-
-
-@register_custom_op(mutates_args=["out"])
-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)
-    input_2d = input.view(-1, hidden_size * 2)
-    out_2d = out.view(-1, hidden_size)
-    module.run_activation(input_2d, out_2d, op_name)
-
-
-@register_custom_op(mutates_args=["out"])
-def _run_activation_filtered_inplace(
-    op_name: str,
-    input: torch.Tensor,
-    out: torch.Tensor,
-    expert_ids: torch.Tensor,
-    expert_step: int,
-) -> None:
-    hidden_size = input.shape[-1] // 2
-    module = _jit_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)
-
-
-def run_activation(
-    op_name: str,
-    input: torch.Tensor,
-    out: Optional[torch.Tensor],
-    expert_ids: Optional[torch.Tensor] = None,
-    expert_step: int = 1,
-) -> torch.Tensor:
-    """Apply ``op_name`` activation followed by element-wise multiplication.
-
-    When ``expert_ids`` is provided, output rows are skipped for tokens whose
-    routed expert id is ``-1``. ``expert_step`` is 1 for per-token routing and
-    ``BLOCK_SIZE_M`` for sorted/TMA routing — i.e. ``expert_ids[token_id //
-    expert_step]`` is consulted before computing each row.
-    """
-    assert op_name in SUPPORTED_ACTIVATIONS, f"Unsupported activation: {op_name}"
-    hidden_size = input.shape[-1] // 2
-    if out is None:
-        out = input.new_empty(*input.shape[:-1], hidden_size)
-    if expert_ids is None:
-        _run_activation_inplace(op_name, input, out)
-    else:
-        _run_activation_filtered_inplace(op_name, input, out, expert_ids, expert_step)
-    return out
-
-
-@register_custom_op(mutates_args=["out"])
-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.run_unary_activation(input.view(-1, last), out.view(-1, last), op_name)
-
-
-def run_unary_activation(
-    op_name: str,
-    input: torch.Tensor,
-    out: Optional[torch.Tensor] = None,
-) -> torch.Tensor:
-    """Apply a standalone (non-gated) element-wise activation: ``out = act(input)``.
-
-    Unlike :func:`run_activation`, there is no gate/up split — ``input`` and
-    ``out`` share the same shape.
-    """
-    assert (
-        op_name in SUPPORTED_UNARY_ACTIVATIONS
-    ), f"Unsupported unary activation: {op_name}"
-    if out is None:
-        out = torch.empty_like(input)
-    _run_unary_activation_inplace(op_name, input, out)
-    return out
-
-
-def relu2(
-    input: torch.Tensor,
-    out: Optional[torch.Tensor] = None,
-) -> torch.Tensor:
-    """Squared ReLU: ``out = max(0, input) ** 2`` (element-wise)."""
-    return run_unary_activation("relu2", input, out)
-
-
-def silu_and_mul(
-    input: torch.Tensor,
-    out: Optional[torch.Tensor] = None,
-    expert_ids: Optional[torch.Tensor] = None,
-    expert_step: int = 1,
-) -> torch.Tensor:
-    return run_activation("silu", input, out, expert_ids, expert_step)
-
-
-def gelu_and_mul(
-    input: torch.Tensor,
-    out: Optional[torch.Tensor] = None,
-    expert_ids: Optional[torch.Tensor] = None,
-    expert_step: int = 1,
-) -> torch.Tensor:
-    return run_activation("gelu", input, out, expert_ids, expert_step)
-
-
-def gelu_tanh_and_mul(
-    input: torch.Tensor,
-    out: Optional[torch.Tensor] = None,
-    expert_ids: Optional[torch.Tensor] = None,
-    expert_step: int = 1,
-) -> torch.Tensor:
-    return run_activation("gelu_tanh", input, out, expert_ids, expert_step)
+globals().update({k: getattr(_impl, k) for k in dir(_impl) if not k.startswith("__")})

python/sglang/jit_kernel/add_constant.py

@@ -4,7 +4,7 @@
 
 import torch
 
-from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+from sglang.kernels.jit import cache_once, load_jit, make_cpp_args
 
 if TYPE_CHECKING:
     from tvm_ffi.module import Module

python/sglang/jit_kernel/all_reduce.py

@@ -7,14 +7,14 @@
 import tvm_ffi
 from tvm_ffi import Module
 
-from sglang.jit_kernel.utils import (
+from sglang.kernel_api_logging import debug_kernel_api
+from sglang.kernels.jit import (
     cache_once,
     is_arch_support_pdl,
     lazy_register_class,
     load_jit,
     make_cpp_args,
 )
-from sglang.kernel_api_logging import debug_kernel_api
 
 
 class AllReduceAlgo(enum.Enum):

python/sglang/jit_kernel/awq_dequantize.py

@@ -4,7 +4,7 @@
 
 import torch
 
-from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+from sglang.kernels.jit import cache_once, load_jit, make_cpp_args
 
 if TYPE_CHECKING:
     from tvm_ffi.module import Module

python/sglang/jit_kernel/awq_marlin_repack.py

@@ -4,8 +4,8 @@
 
 import torch
 
-from sglang.jit_kernel.utils import cache_once, load_jit
 from sglang.kernel_api_logging import debug_kernel_api
+from sglang.kernels.jit import cache_once, load_jit
 
 if TYPE_CHECKING:
     from tvm_ffi.module import Module

python/sglang/jit_kernel/benchmark/marker.py

@@ -21,7 +21,7 @@
 
 import torch
 
-from sglang.jit_kernel.utils import cache_once
+from sglang.kernels.jit import cache_once
 from sglang.utils import is_in_ci
 
 F = TypeVar("F", bound=Callable[..., "BenchResult"])

python/sglang/jit_kernel/clamp_position.py

@@ -4,7 +4,7 @@
 
 import torch
 
-from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+from sglang.kernels.jit import cache_once, load_jit, make_cpp_args
 
 if TYPE_CHECKING:
     from tvm_ffi.module import Module

python/sglang/jit_kernel/concat_mla.py

@@ -4,7 +4,7 @@
 
 import torch
 
-from sglang.jit_kernel.utils import cache_once, load_jit
+from sglang.kernels.jit import cache_once, load_jit
 
 if TYPE_CHECKING:
     from tvm_ffi.module import Module

python/sglang/jit_kernel/diffusion/causal_conv3d_cat_pad.py

@@ -4,7 +4,7 @@
 
 import torch
 
-from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+from sglang.kernels.jit import cache_once, load_jit, make_cpp_args
 from sglang.srt.utils.custom_op import register_custom_op
 
 if TYPE_CHECKING:

python/sglang/jit_kernel/diffusion/ltx2_qknorm_split_rope.py

@@ -4,7 +4,7 @@
 
 import torch
 
-from sglang.jit_kernel.utils import cache_once, load_jit
+from sglang.kernels.jit import cache_once, load_jit
 from sglang.srt.utils.custom_op import register_custom_op
 
 if TYPE_CHECKING:

python/sglang/jit_kernel/diffusion/norm_scale_shift_native.py

@@ -4,7 +4,7 @@
 
 import torch
 
-from sglang.jit_kernel.utils import cache_once, load_jit
+from sglang.kernels.jit import cache_once, load_jit
 
 if TYPE_CHECKING:
     from tvm_ffi.module import Module

python/sglang/jit_kernel/diffusion/qknorm_rope.py

@@ -5,7 +5,7 @@
 
 import torch
 
-from sglang.jit_kernel.utils import (
+from sglang.kernels.jit import (
     cache_once,
     is_arch_support_pdl,
     load_jit,

python/sglang/jit_kernel/diffusion/residual_gate_add.py

@@ -4,7 +4,7 @@
 
 import torch
 
-from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+from sglang.kernels.jit import cache_once, load_jit, make_cpp_args
 from sglang.srt.utils.custom_op import register_custom_op
 
 if TYPE_CHECKING:

python/sglang/jit_kernel/dsv32/elementwise.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from sglang.jit_kernel.utils import (
+from sglang.kernels.jit import (
     cache_once,
     is_arch_support_pdl,
     load_jit,

python/sglang/jit_kernel/dsv3_fused_a_gemm.py

@@ -11,13 +11,13 @@
 
 import torch
 
-from sglang.jit_kernel.utils import (
+from sglang.kernel_api_logging import debug_kernel_api
+from sglang.kernels.jit import (
     cache_once,
     is_arch_support_pdl,
     load_jit,
     make_cpp_args,
 )
-from sglang.kernel_api_logging import debug_kernel_api
 from sglang.srt.utils.common import direct_register_custom_op
 
 if TYPE_CHECKING:

python/sglang/jit_kernel/dsv3_router_gemm.py

@@ -11,13 +11,13 @@
 
 import torch
 
-from sglang.jit_kernel.utils import (
+from sglang.kernel_api_logging import debug_kernel_api
+from sglang.kernels.jit import (
     cache_once,
     is_arch_support_pdl,
     load_jit,
     make_cpp_args,
 )
-from sglang.kernel_api_logging import debug_kernel_api
 from sglang.srt.utils.custom_op import register_custom_op
 
 if TYPE_CHECKING:

python/sglang/jit_kernel/dsv4/attn.py

@@ -4,7 +4,7 @@
 import triton
 import triton.language as tl
 
-from sglang.jit_kernel.utils import (
+from sglang.kernels.jit import (
     cache_once,
     is_arch_support_pdl,
     is_hip_runtime,

python/sglang/jit_kernel/dsv4/compress.py

@@ -4,7 +4,7 @@
 
 import torch
 
-from sglang.jit_kernel.utils import (
+from sglang.kernels.jit import (
     cache_once,
     is_arch_support_pdl,
     load_jit,

python/sglang/jit_kernel/dsv4/compress_old.py

@@ -4,7 +4,7 @@
 
 import torch
 
-from sglang.jit_kernel.utils import (
+from sglang.kernels.jit import (
     cache_once,
     is_arch_support_pdl,
     load_jit,

python/sglang/jit_kernel/dsv4/elementwise.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from sglang.jit_kernel.utils import (
+from sglang.kernels.jit import (
     cache_once,
     is_arch_support_pdl,
     load_jit,