[Refactor] Clean up custom op (#15995)

python/sglang/jit_kernel/utils.py

@@ -1,25 +1,9 @@
 from __future__ import annotations
 
 import functools
-import inspect
 import pathlib
 from functools import lru_cache
-from typing import (
+from typing import TYPE_CHECKING, Any, Callable, List, Tuple, TypeAlias, TypeVar, Union
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    List,
-    Optional,
-    Tuple,
-    TypeAlias,
-    TypeVar,
-    Union,
-    overload,
-)
-
-import torch
-
-from sglang.srt.utils.common import direct_register_custom_op
 
 if TYPE_CHECKING:
     from tvm_ffi import Module
@@ -176,106 +160,3 @@ def is_arch_support_pdl() -> bool:
 
     device = torch.cuda.current_device()
     return torch.cuda.get_device_capability(device)[0] >= 9
-
-
-def fake_inplace_impl(*args, **kwargs) -> None:
-    pass
-
-
-@overload
-def register_jit_op(
-    fn: F,
-    *,
-    op_name: Optional[str] = None,
-    out_list: Optional[List[int]] = None,
-    out_args: Optional[List[str]] = None,
-    fake_impl: Optional[Callable] = fake_inplace_impl,
-) -> F: ...
-
-
-@overload
-def register_jit_op(
-    *,
-    op_name: Optional[str] = None,
-    out_list: Optional[List[int]] = None,
-    out_args: Optional[List[str]] = None,
-    fake_impl: Optional[Callable] = fake_inplace_impl,
-) -> Callable[[F], F]: ...
-
-
-# Real implementation
-def register_jit_op(
-    fn=None,
-    *,
-    op_name: Optional[str] = None,
-    out_list: Optional[List[int]] = None,
-    out_args: Optional[List[str]] = None,
-    fake_impl: Optional[Callable] = fake_inplace_impl,
-) -> Any:
-    """
-    A decorator to register a JIT custom operator.
-
-    Example usage:
-    ```python
-    @register_jit_op(op_name="my_op", out_list=[0])
-    def my_inplace_op(x: torch.Tensor) -> None:
-        x.add_(1)
-
-    def fake_impl(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
-        return x + y
-
-    @register_jit_op(op_name="my_op2", out_args=["x"], fake_impl=fake_impl)
-    def my_op(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
-        return x.add_(y)
-    ```
-
-    :param fn: The function to be registered as a JIT custom operator.
-               If None, return a decorator.
-    :type fn: Callable
-    :param op_name: The name of the operator. If None, use the function name
-    :type op_name: Optional[str]
-    :param out_list: A list of argument indices that are mutated in-place.
-    :type out_list: Optional[List[int]]
-    :param out_args: A list of argument names that are mutated in-place.
-    :type out_args: Optional[List[str]]
-    :param fake_impl: A fake implementation for the operator, used for
-                      torch.compile compatibility.
-                      By default, a no-op function is used, which suits
-                      for most in-place operations.
-    :type fake_impl: Optional[Callable]
-    :return: The registered JIT custom operator, or a decorator.
-             NOTE: the real register will occur at the first call of the function.
-    :rtype: Callable
-    """
-
-    def decorator(fn):
-        real_impl = None
-        resolved_name = op_name or fn.__name__
-
-        @functools.wraps(fn)
-        def wrapper(*args, **kwargs):
-            nonlocal real_impl
-            if real_impl is None:
-                if not hasattr(torch.ops.sglang, resolved_name):
-                    signature = inspect.signature(fn)
-                    mutates_args = []
-                    param_names = list(signature.parameters.keys())
-                    for id in out_list or []:
-                        mutates_args.append(param_names[id])
-                    for name in out_args or []:
-                        mutates_args.append(name)
-                    mutates_args = list(set(mutates_args))
-                    direct_register_custom_op(
-                        op_name=resolved_name,
-                        op_func=fn,
-                        mutates_args=mutates_args,
-                        fake_impl=fake_impl,
-                    )
-                real_impl = getattr(torch.ops.sglang, resolved_name)
-            return real_impl(*args, **kwargs)
-
-        return wrapper
-
-    if fn is not None:
-        return decorator(fn)
-    return decorator

python/sglang/srt/distributed/parallel_state.py

@@ -41,7 +41,6 @@ from torch.distributed import Backend, ProcessGroup
 
 from sglang.srt.environ import envs
 from sglang.srt.utils import (
-    direct_register_custom_op,
     get_bool_env_var,
     get_current_device_stream_fast,
     get_int_env_var,
@@ -52,14 +51,12 @@ from sglang.srt.utils import (
     is_npu,
     is_shm_available,
     is_xpu,
-    supports_custom_op,
 )
+from sglang.srt.utils.custom_op import register_custom_op
 
 _is_npu = is_npu()
 _is_cpu = is_cpu()
 _is_xpu = is_xpu()
-_supports_custom_op = supports_custom_op()
-
 
 TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])
 
@@ -127,87 +124,46 @@ def _register_group(group: "GroupCoordinator") -> None:
     _groups[group.unique_name] = weakref.ref(group)
 
 
-if _supports_custom_op:
+@register_custom_op(mutates_args=["tensor"])
+def inplace_all_reduce(tensor: torch.Tensor, group_name: str) -> None:
+    assert group_name in _groups, f"Group {group_name} is not found."
+    group = _groups[group_name]()
+    if group is None:
+        raise ValueError(f"Group {group_name} is destroyed.")
+    group._all_reduce_in_place(tensor)
 
-    def inplace_all_reduce(tensor: torch.Tensor, group_name: str) -> None:
-        assert group_name in _groups, f"Group {group_name} is not found."
-        group = _groups[group_name]()
-        if group is None:
-            raise ValueError(f"Group {group_name} is destroyed.")
-        group._all_reduce_in_place(tensor)
 
-    def inplace_all_reduce_fake(tensor: torch.Tensor, group_name: str) -> None:
+@register_custom_op(out_shape="tensor")
-        return
+def outplace_all_reduce(
+    tensor: torch.Tensor, group_name: str, outplace_all_reduce_method: str
+) -> torch.Tensor:
+    assert group_name in _groups, f"Group {group_name} is not found."
+    group = _groups[group_name]()
+    if group is None:
+        raise ValueError(f"Group {group_name} is destroyed.")
+    return group._all_reduce_out_place(tensor, outplace_all_reduce_method)
 
-    direct_register_custom_op(
-        op_name="inplace_all_reduce",
-        op_func=inplace_all_reduce,
-        mutates_args=["tensor"],
-        fake_impl=inplace_all_reduce_fake,
-    )
 
-    def outplace_all_reduce(
+@register_custom_op(mutates_args=["output"])
-        tensor: torch.Tensor, group_name: str, outplace_all_reduce_method: str
+def reg_all_gather_into_tensor(
-    ) -> torch.Tensor:
+    output: torch.Tensor, input: torch.Tensor, group_name: str
-        assert group_name in _groups, f"Group {group_name} is not found."
+) -> None:
-        group = _groups[group_name]()
+    assert group_name in _groups, f"Group {group_name} is not found."
-        if group is None:
+    group = _groups[group_name]()
-            raise ValueError(f"Group {group_name} is destroyed.")
+    if group is None:
-        return group._all_reduce_out_place(tensor, outplace_all_reduce_method)
+        raise ValueError(f"Group {group_name} is destroyed.")
+    group._all_gather_into_tensor(output, input)
 
-    def outplace_all_reduce_fake(
-        tensor: torch.Tensor, group_name: str, outplace_all_reduce_method: str
-    ) -> torch.Tensor:
-        return torch.empty_like(tensor)
 
-    direct_register_custom_op(
+@register_custom_op(mutates_args=["output"])
-        op_name="outplace_all_reduce",
+def reg_reduce_scatter_tensor(
-        op_func=outplace_all_reduce,
+    output: torch.Tensor, input: torch.Tensor, group_name: str
-        mutates_args=[],
+) -> None:
-        fake_impl=outplace_all_reduce_fake,
+    assert group_name in _groups, f"Group {group_name} is not found."
-    )
+    group = _groups[group_name]()
-
+    if group is None:
-    def reg_all_gather_into_tensor(
+        raise ValueError(f"Group {group_name} is destroyed.")
-        output: torch.Tensor, input: torch.Tensor, group_name: str
+    group._reduce_scatter_tensor(output, input)
-    ) -> None:
-        assert group_name in _groups, f"Group {group_name} is not found."
-        group = _groups[group_name]()
-        if group is None:
-            raise ValueError(f"Group {group_name} is destroyed.")
-        group._all_gather_into_tensor(output, input)
-
-    def reg_all_gather_into_tensor_fake(
-        output: torch.Tensor, input: torch.Tensor, group_name: str
-    ) -> None:
-        pass
-
-    direct_register_custom_op(
-        op_name="reg_all_gather_into_tensor",
-        op_func=reg_all_gather_into_tensor,
-        mutates_args=["output"],
-        fake_impl=reg_all_gather_into_tensor_fake,
-    )
-
-    def reg_reduce_scatter_tensor(
-        output: torch.Tensor, input: torch.Tensor, group_name: str
-    ) -> None:
-        assert group_name in _groups, f"Group {group_name} is not found."
-        group = _groups[group_name]()
-        if group is None:
-            raise ValueError(f"Group {group_name} is destroyed.")
-        group._reduce_scatter_tensor(output, input)
-
-    def reg_reduce_scatter_tensor_fake(
-        output: torch.Tensor, input: torch.Tensor, group_name: str
-    ) -> None:
-        pass
-
-    direct_register_custom_op(
-        op_name="reg_reduce_scatter_tensor",
-        op_func=reg_reduce_scatter_tensor,
-        mutates_args=["output"],
-        fake_impl=reg_reduce_scatter_tensor_fake,
-    )
 
 
 class GroupCoordinator:
@@ -590,10 +546,6 @@ class GroupCoordinator:
                 torch.distributed.all_reduce(input_, group=self.device_group)
             return input_
 
-        if not _supports_custom_op:
-            self._all_reduce_in_place(input_)
-            return input_
-
         if self.hpu_communicator is not None and not self.hpu_communicator.disabled:
             return self.hpu_communicator.all_reduce(input_)
 
@@ -636,13 +588,13 @@ class GroupCoordinator:
         ):
             outplace_all_reduce_method = "torch_symm_mem"
         if outplace_all_reduce_method is not None:
-            return torch.ops.sglang.outplace_all_reduce(
+            return outplace_all_reduce(
                 input_,
                 group_name=self.unique_name,
                 outplace_all_reduce_method=outplace_all_reduce_method,
             )
         else:
-            torch.ops.sglang.inplace_all_reduce(input_, group_name=self.unique_name)
+            inplace_all_reduce(input_, group_name=self.unique_name)
             return input_
 
     def _all_reduce_out_place(
@@ -698,12 +650,10 @@ class GroupCoordinator:
         return output
 
     def reduce_scatter_tensor(self, output: torch.Tensor, input: torch.Tensor):
-        if _is_npu or not supports_custom_op():
+        if _is_npu:
             self._reduce_scatter_tensor(output, input)
         else:
-            torch.ops.sglang.reg_reduce_scatter_tensor(
+            reg_reduce_scatter_tensor(output, input, group_name=self.unique_name)
-                output, input, group_name=self.unique_name
-            )
 
     def reduce_scatter(
         self,
@@ -764,12 +714,10 @@ class GroupCoordinator:
             )
 
     def all_gather_into_tensor(self, output: torch.Tensor, input: torch.Tensor):
-        if _is_npu or _is_xpu or not _supports_custom_op:
+        if _is_npu or _is_xpu:
             self._all_gather_into_tensor(output, input)
         else:
-            torch.ops.sglang.reg_all_gather_into_tensor(
+            reg_all_gather_into_tensor(output, input, group_name=self.unique_name)
-                output, input, group_name=self.unique_name
-            )
 
     def cp_all_gather_into_tensor_async(
         self, output: torch.Tensor, input: torch.Tensor, stream=None

python/sglang/srt/layers/dp_attention.py

@@ -467,9 +467,9 @@ def _dp_gather_via_all_reduce(
         not local_tokens.dtype.is_floating_point
         and get_tensor_model_parallel_world_size() <= NUM_GPUS_PER_NODE
     ):
-        torch.ops.sglang.inplace_all_reduce(
+        from sglang.srt.distributed.parallel_state import inplace_all_reduce
-            global_tokens, group_name=get_tp_group().unique_name
+
-        )
+        inplace_all_reduce(global_tokens, group_name=get_tp_group().unique_name)
 
     else:
         global_tokens[:] = tensor_model_parallel_all_reduce(global_tokens)

python/sglang/srt/layers/elementwise.py

@@ -4,7 +4,8 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.utils import direct_register_custom_op, is_hip
+from sglang.srt.utils import is_hip
+from sglang.srt.utils.custom_op import register_custom_op
 
 _is_hip = is_hip()
 
@@ -358,6 +359,7 @@ def experts_combine_kernel(
     tl.store(out_hidden_states + start_index_mlp + offsets, combined_x, mask=mask)
 
 
+@register_custom_op(out_shape="mlp_hidden_states")
 def experts_combine_triton(
     moe_hidden_states: torch.Tensor,
     mlp_hidden_states: torch.Tensor,
@@ -401,22 +403,6 @@ def experts_combine_triton(
     return out_hidden_states
 
 
-def experts_combine_triton_fake(
-    moe_hidden_states: torch.Tensor,
-    mlp_hidden_states: torch.Tensor,
-    output_buffer: Optional[torch.Tensor] = None,
-) -> torch.Tensor:
-    return torch.empty_like(mlp_hidden_states)
-
-
-direct_register_custom_op(
-    op_name="experts_combine_triton",
-    op_func=experts_combine_triton,
-    mutates_args=[],
-    fake_impl=experts_combine_triton_fake,
-)
-
-
 # gelu on first half of vector
 @triton.jit
 def gelu_and_mul_kernel(

python/sglang/srt/layers/flashinfer_comm_fusion.py

@@ -5,11 +5,8 @@ import torch
 import torch.distributed as dist
 
 from sglang.srt.distributed import get_tensor_model_parallel_world_size
-from sglang.srt.utils import (
+from sglang.srt.utils import is_flashinfer_available
-    direct_register_custom_op,
+from sglang.srt.utils.custom_op import register_custom_op
-    is_flashinfer_available,
-    supports_custom_op,
-)
 
 logger = logging.getLogger(__name__)
 
@@ -123,6 +120,25 @@ def ensure_workspace_initialized(
     return _workspace_manager.initialized
 
 
+def fake_flashinfer_allreduce_residual_rmsnorm(
+    input_tensor: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float = 1e-6,
+    max_token_num: int = 16384,
+    use_oneshot: Optional[bool] = None,
+    trigger_completion_at_end: bool = False,
+    fp32_acc: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    residual_out = torch.empty_like(residual)
+    norm_out = torch.empty_like(input_tensor)
+    return norm_out, residual_out
+
+
+@register_custom_op(
+    mutates_args=["input_tensor", "residual", "weight"],
+    fake_impl=fake_flashinfer_allreduce_residual_rmsnorm,
+)
 def flashinfer_allreduce_residual_rmsnorm(
     input_tensor: torch.Tensor,
     residual: torch.Tensor,
@@ -202,30 +218,6 @@ def flashinfer_allreduce_residual_rmsnorm(
     return norm_out, residual_out
 
 
-def fake_flashinfer_allreduce_residual_rmsnorm(
-    input_tensor: torch.Tensor,
-    residual: torch.Tensor,
-    weight: torch.Tensor,
-    eps: float = 1e-6,
-    max_token_num: int = 16384,
-    use_oneshot: Optional[bool] = None,
-    trigger_completion_at_end: bool = False,
-    fp32_acc: bool = False,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    residual_out = torch.empty_like(residual)
-    norm_out = torch.empty_like(input_tensor)
-    return norm_out, residual_out
-
-
-if supports_custom_op():
-    direct_register_custom_op(
-        "flashinfer_allreduce_residual_rmsnorm",
-        flashinfer_allreduce_residual_rmsnorm,
-        mutates_args=["input_tensor", "residual", "weight"],
-        fake_impl=fake_flashinfer_allreduce_residual_rmsnorm,
-    )
-
-
 def cleanup_flashinfer_workspace():
     global _workspace_manager
     if _workspace_manager is not None:

python/sglang/srt/layers/layernorm.py

@@ -35,7 +35,6 @@ from sglang.srt.utils import (
     is_hip,
     is_npu,
     is_xpu,
-    supports_custom_op,
 )
 
 _is_cuda = is_cuda()
@@ -266,14 +265,8 @@ class RMSNorm(MultiPlatformOp):
                 flashinfer_allreduce_residual_rmsnorm,
             )
 
-            fused_op = (
-                torch.ops.sglang.flashinfer_allreduce_residual_rmsnorm
-                if supports_custom_op()
-                else flashinfer_allreduce_residual_rmsnorm
-            )
-
             if get_tensor_model_parallel_world_size() > 1:
-                fused_result = fused_op(
+                fused_result = flashinfer_allreduce_residual_rmsnorm(
                     input_tensor=x,
                     residual=residual,
                     weight=self.weight,

python/sglang/srt/layers/moe/ep_moe/layer.py

@@ -16,7 +16,11 @@ from sglang.srt.layers.moe import (
     get_moe_a2a_backend,
     get_moe_runner_backend,
 )
-from sglang.srt.layers.moe.fused_moe_triton.layer import FlashInferFusedMoE, FusedMoE
+from sglang.srt.layers.moe.fused_moe_triton.layer import (
+    FlashInferFusedMoE,
+    FusedMoE,
+    moe_forward_piecewise_cuda_graph_impl,
+)
 from sglang.srt.layers.moe.token_dispatcher.deepep import (
     DeepEPLLCombineInput,
     DeepEPNormalCombineInput,
@@ -156,7 +160,7 @@ class DeepEPMoE(FusedMoE):
             assert TopKOutputChecker.format_is_standard(
                 topk_output
             ), "Only standard topk output is supported for piecewise cuda graph"
-            return torch.ops.sglang.moe_forward_piecewise_cuda_graph_impl(
+            return moe_forward_piecewise_cuda_graph_impl(
                 hidden_states,
                 topk_output.topk_weights,
                 topk_output.topk_ids,

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

@@ -2,7 +2,8 @@ from typing import Optional
 
 import torch
 
-from sglang.srt.utils import direct_register_custom_op, is_cuda
+from sglang.srt.utils import is_cuda
+from sglang.srt.utils.custom_op import register_custom_op
 
 _is_cuda = is_cuda()
 
@@ -20,6 +21,7 @@ def get_scalar_type(num_bits: int, has_zp: bool):
         return scalar_types.uint4b8 if num_bits == 4 else scalar_types.uint8b128
 
 
+@register_custom_op(out_shape="hidden_states")
 def fused_marlin_moe(
     hidden_states: torch.Tensor,
     w1: torch.Tensor,
@@ -214,37 +216,3 @@ def fused_marlin_moe(
         routed_scaling_factor,
     )
     return output
-
-
-def fused_marlin_moe_fake(
-    hidden_states: torch.Tensor,
-    w1: torch.Tensor,
-    w2: torch.Tensor,
-    w1_scale: torch.Tensor,
-    w2_scale: torch.Tensor,
-    gating_output: torch.Tensor,
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    global_num_experts: int = -1,
-    expert_map: Optional[torch.Tensor] = None,
-    g_idx1: Optional[torch.Tensor] = None,
-    g_idx2: Optional[torch.Tensor] = None,
-    sort_indices1: Optional[torch.Tensor] = None,
-    sort_indices2: Optional[torch.Tensor] = None,
-    w1_zeros: Optional[torch.Tensor] = None,
-    w2_zeros: Optional[torch.Tensor] = None,
-    workspace: Optional[torch.Tensor] = None,
-    num_bits: int = 8,
-    is_k_full: bool = True,
-    inplace: bool = False,
-    routed_scaling_factor: Optional[float] = None,
-) -> torch.Tensor:
-    return torch.empty_like(hidden_states)
-
-
-direct_register_custom_op(
-    op_name="fused_marlin_moe",
-    op_func=fused_marlin_moe,
-    mutates_args=[],
-    fake_impl=fused_marlin_moe_fake,
-)

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

@@ -16,12 +16,12 @@ import triton.language as tl
 from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
 from sglang.srt.utils import (
     cpu_has_amx_support,
-    direct_register_custom_op,
     get_bool_env_var,
     is_cpu,
     is_cuda,
     is_hip,
 )
+from sglang.srt.utils.custom_op import register_custom_op
 
 from .fused_moe_triton_config import get_config_dtype_str, try_get_optimal_moe_config
 from .fused_moe_triton_kernels import (
@@ -59,6 +59,7 @@ elif _is_hip:
 padding_size = 128 if bool(int(os.getenv("SGLANG_MOE_PADDING", "0"))) else 0
 
 
+@register_custom_op(mutates_args=["hidden_states"])
 def inplace_fused_experts(
     hidden_states: torch.Tensor,
     w1: torch.Tensor,
@@ -119,45 +120,7 @@ def inplace_fused_experts(
     )
 
 
-def inplace_fused_experts_fake(
+@register_custom_op(out_shape="hidden_states")
-    hidden_states: torch.Tensor,
-    w1: torch.Tensor,
-    w2: torch.Tensor,
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    b1: Optional[torch.Tensor] = None,
-    b2: Optional[torch.Tensor] = None,
-    activation: str = "silu",
-    is_gated: bool = True,
-    apply_router_weight_on_input: bool = False,
-    use_fp8_w8a8: bool = False,
-    use_int8_w8a8: bool = False,
-    use_int8_w8a16: bool = False,
-    use_int4_w4a16: bool = False,
-    per_channel_quant: bool = False,
-    w1_scale: Optional[torch.Tensor] = None,
-    w2_scale: Optional[torch.Tensor] = None,
-    w1_zp: Optional[torch.Tensor] = None,
-    w2_zp: Optional[torch.Tensor] = None,
-    a1_scale: Optional[torch.Tensor] = None,
-    a2_scale: Optional[torch.Tensor] = None,
-    block_shape: Optional[List[int]] = None,
-    routed_scaling_factor: Optional[float] = None,
-    gemm1_alpha: Optional[float] = None,
-    gemm1_limit: Optional[float] = None,
-    filter_expert: bool = True,
-) -> None:
-    pass
-
-
-direct_register_custom_op(
-    op_name="inplace_fused_experts",
-    op_func=inplace_fused_experts,
-    mutates_args=["hidden_states"],
-    fake_impl=inplace_fused_experts_fake,
-)
-
-
 def outplace_fused_experts(
     hidden_states: torch.Tensor,
     w1: torch.Tensor,
@@ -219,46 +182,6 @@ def outplace_fused_experts(
     )
 
 
-def outplace_fused_experts_fake(
-    hidden_states: torch.Tensor,
-    w1: torch.Tensor,
-    w2: torch.Tensor,
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    b1: Optional[torch.Tensor] = None,
-    b2: Optional[torch.Tensor] = None,
-    activation: str = "silu",
-    is_gated: bool = True,
-    apply_router_weight_on_input: bool = False,
-    use_fp8_w8a8: bool = False,
-    use_int8_w8a8: bool = False,
-    use_int8_w8a16: bool = False,
-    use_int4_w4a16: bool = False,
-    per_channel_quant: bool = False,
-    w1_scale: Optional[torch.Tensor] = None,
-    w2_scale: Optional[torch.Tensor] = None,
-    w1_zp: Optional[torch.Tensor] = None,
-    w2_zp: Optional[torch.Tensor] = None,
-    a1_scale: Optional[torch.Tensor] = None,
-    a2_scale: Optional[torch.Tensor] = None,
-    block_shape: Optional[List[int]] = None,
-    no_combine: bool = False,
-    routed_scaling_factor: Optional[float] = None,
-    gemm1_alpha: Optional[float] = None,
-    gemm1_limit: Optional[float] = None,
-    filter_expert: bool = True,
-) -> torch.Tensor:
-    return torch.empty_like(hidden_states)
-
-
-direct_register_custom_op(
-    op_name="outplace_fused_experts",
-    op_func=outplace_fused_experts,
-    mutates_args=[],
-    fake_impl=outplace_fused_experts_fake,
-)
-
-
 def fused_experts(
     hidden_states: torch.Tensor,
     w1: torch.Tensor,
@@ -287,7 +210,7 @@ def fused_experts(
     )
     if moe_runner_config.inplace:
         assert not moe_runner_config.no_combine, "no combine + inplace makes no sense"
-        torch.ops.sglang.inplace_fused_experts(
+        inplace_fused_experts(
             hidden_states,
             w1,
             w2,
@@ -317,7 +240,7 @@ def fused_experts(
         )
         return hidden_states
     else:
-        return torch.ops.sglang.outplace_fused_experts(
+        return outplace_fused_experts(
             hidden_states,
             w1,
             w2,

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

@@ -67,7 +67,7 @@ from sglang.srt.utils import (
     next_power_of_2,
     round_up,
 )
-from sglang.srt.utils.common import direct_register_custom_op
+from sglang.srt.utils.custom_op import register_custom_op
 
 if is_flashinfer_available():
     from flashinfer import fp4_quantize
@@ -909,7 +909,7 @@ class FusedMoE(torch.nn.Module):
             assert TopKOutputChecker.format_is_standard(
                 topk_output
             ), "Only standard topk output is supported for piecewise cuda graph"
-            return torch.ops.sglang.moe_forward_piecewise_cuda_graph_impl(
+            return moe_forward_piecewise_cuda_graph_impl(
                 hidden_states,
                 topk_output.topk_weights,
                 topk_output.topk_ids,
@@ -1081,7 +1081,7 @@ class FlashInferFusedMoE(FusedMoE):
             assert TopKOutputChecker.format_is_standard(
                 topk_output
             ), "Only standard topk output is supported for piecewise cuda graph"
-            return torch.ops.sglang.moe_forward_piecewise_cuda_graph_impl(
+            return moe_forward_piecewise_cuda_graph_impl(
                 hidden_states,
                 topk_output.topk_weights,
                 topk_output.topk_ids,
@@ -1210,16 +1210,14 @@ class FlashInferFP4MoE(FusedMoE):
         ), "Only bypassed topk output is supported for flashinfer fp4 moe"
 
         if is_in_piecewise_cuda_graph():
-            return (
+            return flashinfer_fp4_moe_forward_piecewise_cuda_graph_impl(
-                torch.ops.sglang.flashinfer_fp4_moe_forward_piecewise_cuda_graph_impl(
+                hidden_states,
-                    hidden_states,
+                topk_output.router_logits,
-                    topk_output.router_logits,
+                topk_output.topk_config.top_k,
-                    topk_output.topk_config.top_k,
+                topk_output.topk_config.topk_group,
-                    topk_output.topk_config.topk_group,
+                topk_output.topk_config.num_expert_group,
-                    topk_output.topk_config.num_expert_group,
+                topk_output.topk_config.correction_bias,
-                    topk_output.topk_config.correction_bias,
+                self.layer_id,
-                    self.layer_id,
-                )
             )
         else:
             return self.forward_impl(hidden_states, topk_output)
@@ -1317,6 +1315,7 @@ class FlashInferFP4MoE(FusedMoE):
         return result
 
 
+@register_custom_op(out_shape="hidden_states")
 def moe_forward_piecewise_cuda_graph_impl(
     hidden_states: torch.Tensor,
     topk_weights: torch.Tensor,
@@ -1333,16 +1332,7 @@ def moe_forward_piecewise_cuda_graph_impl(
     return moe_layer.forward_impl(hidden_states, topk_output)
 
 
-def moe_forward_piecewise_cuda_graph_impl_fake(
+@register_custom_op(out_shape="hidden_states")
-    hidden_states: torch.Tensor,
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    router_logits: torch.Tensor,
-    layer_id: int,
-) -> torch.Tensor:
-    return torch.empty_like(hidden_states)
-
-
 def flashinfer_fp4_moe_forward_piecewise_cuda_graph_impl(
     hidden_states: torch.Tensor,
     router_logits: torch.Tensor,
@@ -1365,30 +1355,3 @@ def flashinfer_fp4_moe_forward_piecewise_cuda_graph_impl(
     forward_context = get_forward_context()
     moe_layer = forward_context.moe_layers[layer_id]
     return moe_layer.forward_impl(hidden_states, topk_output)
-
-
-def flashinfer_fp4_moe_forward_piecewise_cuda_graph_impl_fake(
-    hidden_states: torch.Tensor,
-    router_logits: torch.Tensor,
-    top_k: int,
-    topk_group: Optional[int],
-    num_expert_group: Optional[int],
-    correction_bias: Optional[torch.Tensor],
-    layer_id: int,
-) -> torch.Tensor:
-    return torch.empty_like(hidden_states)
-
-
-direct_register_custom_op(
-    op_name="moe_forward_piecewise_cuda_graph_impl",
-    op_func=moe_forward_piecewise_cuda_graph_impl,
-    mutates_args=[],
-    fake_impl=moe_forward_piecewise_cuda_graph_impl_fake,
-)
-
-direct_register_custom_op(
-    op_name="flashinfer_fp4_moe_forward_piecewise_cuda_graph_impl",
-    op_func=flashinfer_fp4_moe_forward_piecewise_cuda_graph_impl,
-    mutates_args=[],
-    fake_impl=flashinfer_fp4_moe_forward_piecewise_cuda_graph_impl_fake,
-)

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

@@ -80,9 +80,7 @@ def fused_experts_none_to_marlin(
     runner_config: MoeRunnerConfig,
 ) -> StandardCombineInput:
     global MARLIN_MOE_WORKSPACE
-    from sglang.srt.layers.moe.fused_moe_triton.fused_marlin_moe import (  # noqa
+    from sglang.srt.layers.moe.fused_moe_triton.fused_marlin_moe import fused_marlin_moe
-        fused_marlin_moe,
-    )
     from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
     from sglang.srt.layers.quantization.marlin_utils import marlin_make_workspace
 
@@ -99,7 +97,7 @@ def fused_experts_none_to_marlin(
             hidden_states.device, max_blocks_per_sm=4
         )
 
-    output = torch.ops.sglang.fused_marlin_moe(
+    output = fused_marlin_moe(
         hidden_states=hidden_states,
         w1=quant_info.w13_qweight,
         w2=quant_info.w2_qweight,

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

@@ -6,7 +6,8 @@ from typing import Optional
 
 import torch
 
-from sglang.srt.utils import direct_register_custom_op, get_bool_env_var, is_hip
+from sglang.srt.utils import get_bool_env_var, is_hip
+from sglang.srt.utils.custom_op import register_custom_op
 
 _is_hip = is_hip()
 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
@@ -19,7 +20,10 @@ class ActivationMethod(IntEnum):
     GELU = 1
 
 
-def rocm_aiter_asm_moe_tkw1_impl(
+# NOTE: for non _use_aiter case, use lazy registration to avoid overhead
+# (registration may not be trigger actually, since it will not be called)
+@register_custom_op(out_shape="hidden_states", eager=_use_aiter)
+def rocm_aiter_asm_moe_tkw1(
     hidden_states: torch.Tensor,
     w1: torch.Tensor,
     w2: torch.Tensor,
@@ -57,34 +61,6 @@ def rocm_aiter_asm_moe_tkw1_impl(
     )
 
 
-def rocm_aiter_asm_moe_tkw1_fake(
-    hidden_states: torch.Tensor,
-    w1: torch.Tensor,
-    w2: torch.Tensor,
-    topk_weights: torch.Tensor,
-    topk_ids: torch.Tensor,
-    fc1_scale: Optional[torch.Tensor] = None,
-    fc2_scale: Optional[torch.Tensor] = None,
-    fc1_smooth_scale: Optional[torch.Tensor] = None,
-    fc2_smooth_scale: Optional[torch.Tensor] = None,
-    a16: bool = False,
-    per_tensor_quant_scale: Optional[torch.Tensor] = None,
-    expert_mask: Optional[torch.Tensor] = None,
-    activation_method: int = ActivationMethod.SILU.value,
-) -> torch.Tensor:
-    return torch.empty_like(hidden_states)
-
-
-if _use_aiter:
-
-    direct_register_custom_op(
-        op_name="rocm_aiter_asm_moe_tkw1",
-        op_func=rocm_aiter_asm_moe_tkw1_impl,
-        mutates_args=[],
-        fake_impl=rocm_aiter_asm_moe_tkw1_fake,
-    )
-
-
 def rocm_fused_experts_tkw1(
     hidden_states: torch.Tensor,
     w1: torch.Tensor,
@@ -121,7 +97,7 @@ def rocm_fused_experts_tkw1(
             "Only support topk=1 when" " `apply_router_weight_on_input` is True"
         )
 
-        return torch.ops.sglang.rocm_aiter_asm_moe_tkw1(
+        return rocm_aiter_asm_moe_tkw1(
             hidden_states,
             w1,
             w2,

python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py

@@ -1115,7 +1115,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
         layer: torch.nn.Module,
         dispatch_output: StandardDispatchOutput,
     ) -> CombineInput:
-        from sglang.srt.layers.moe.fused_moe_triton.fused_marlin_moe import (  # noqa
+        from sglang.srt.layers.moe.fused_moe_triton.fused_marlin_moe import (
             fused_marlin_moe,
         )
         from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
@@ -1139,7 +1139,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
             if expert_map is not None:
                 global_num_experts = self.moe_runner_config.num_experts
 
-        output = torch.ops.sglang.fused_marlin_moe(
+        output = fused_marlin_moe(
             x,
             layer.w13_weight_packed,
             layer.w2_weight_packed,

python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py

@@ -18,6 +18,7 @@ from sglang.srt.layers.quantization.compressed_tensors.schemes import (
 from sglang.srt.layers.quantization.modelopt_quant import (
     FLASHINFER_FP4_GEMM_BACKEND,
     enable_flashinfer_fp4_gemm,
+    fp4_gemm,
     fp4_quantize,
 )
 from sglang.srt.layers.quantization.utils import swizzle_blockscale
@@ -153,7 +154,7 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
             w = layer.weight_packed.T
             w_blockscale = layer.weight_scale.T
 
-        out = torch.ops.sglang.fp4_gemm(
+        out = fp4_gemm(
             x_fp4,
             w,
             x_blockscale,

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

@@ -26,7 +26,6 @@ import triton.language as tl
 from sglang.srt.layers import deep_gemm_wrapper
 from sglang.srt.utils import (
     ceil_align,
-    direct_register_custom_op,
     get_bool_env_var,
     get_device_core_count,
     get_device_name,
@@ -34,8 +33,8 @@ from sglang.srt.utils import (
     is_cuda,
     is_hip,
     log_info_on_rank0,
-    supports_custom_op,
 )
+from sglang.srt.utils.custom_op import register_custom_op
 
 _is_hip = is_hip()
 _is_cuda = is_cuda()
@@ -90,32 +89,16 @@ else:
     fp8_max = torch.finfo(fp8_dtype).max
 fp8_min = -fp8_max
 
-if supports_custom_op():
 
-    def deep_gemm_fp8_fp8_bf16_nt(
+@register_custom_op(mutates_args=["C"])
-        A: torch.Tensor,
+def deep_gemm_fp8_fp8_bf16_nt(
-        As: torch.Tensor,
+    A: torch.Tensor,
-        B: torch.Tensor,
+    As: torch.Tensor,
-        Bs: torch.Tensor,
+    B: torch.Tensor,
-        C: torch.Tensor,
+    Bs: torch.Tensor,
-    ) -> None:
+    C: torch.Tensor,
-        deep_gemm_wrapper.gemm_nt_f8f8bf16((A, As), (B, Bs), C)
+) -> None:
-
+    deep_gemm_wrapper.gemm_nt_f8f8bf16((A, As), (B, Bs), C)
-    def deep_gemm_fp8_fp8_bf16_nt_fake(
-        A: torch.Tensor,
-        As: torch.Tensor,
-        B: torch.Tensor,
-        Bs: torch.Tensor,
-        C: torch.Tensor,
-    ) -> None:
-        return
-
-    direct_register_custom_op(
-        op_name="deep_gemm_fp8_fp8_bf16_nt",
-        op_func=deep_gemm_fp8_fp8_bf16_nt,
-        mutates_args=["C"],
-        fake_impl=deep_gemm_fp8_fp8_bf16_nt_fake,
-    )
 
 
 @triton.jit
@@ -1081,10 +1064,7 @@ def w8a8_block_fp8_matmul_deepgemm(
     # Deepgemm only supports output tensor type as bfloat16
     assert C.dtype == torch.bfloat16 and deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
 
-    if supports_custom_op():
+    deep_gemm_fp8_fp8_bf16_nt(A, As, B, Bs, C)
-        torch.ops.sglang.deep_gemm_fp8_fp8_bf16_nt(A, As, B, Bs, C)
-    else:
-        deep_gemm_wrapper.gemm_nt_f8f8bf16((A, As), (B, Bs), C)
 
     return C
 

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

@@ -26,6 +26,7 @@ from sglang.srt.layers.quantization.utils import (
     unpack_cols,
 )
 from sglang.srt.utils import get_device_capability, is_cuda
+from sglang.srt.utils.custom_op import register_custom_op
 
 if TYPE_CHECKING:
     from sglang.srt.layers.linear import LinearBase
@@ -38,7 +39,6 @@ try:
 except ImportError:
     ops = None
 
-from sglang.srt.utils import direct_register_custom_op
 
 _is_cuda = is_cuda()
 
@@ -509,7 +509,7 @@ def apply_gptq_marlin_linear(
             is_zp_float=False,
         )
     else:
-        output = torch.ops.sglang.unified_apply_gptq_marlin_gemm_with_wtype(
+        output = unified_apply_gptq_marlin_gemm_with_wtype(
             input=reshaped_x,
             weight=weight,
             weight_scale=weight_scale,
@@ -578,7 +578,7 @@ def apply_awq_marlin_linear(
             is_zp_float=False,
         )
     else:
-        output = torch.ops.sglang.unified_apply_gptq_marlin_gemm(
+        output = unified_apply_gptq_marlin_gemm(
             input=reshaped_x,
             weight=weight,
             weight_scale=weight_scale,
@@ -860,6 +860,26 @@ class MarlinLinearMethod(LinearMethodBase):
         return output
 
 
+def fake_unified_apply_gptq_marlin_gemm(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    weight_zp: torch.Tensor,
+    g_idx: torch.Tensor,
+    g_idx_sort_indices: torch.Tensor,
+    workspace: torch.Tensor,
+    output_size_per_partition: int,
+    input_size_per_partition: int,
+    use_atomic_add: bool,
+    use_fp32_reduce: bool,
+    is_zp_float: bool,
+) -> torch.Tensor:
+    return input.new_empty(
+        (input.shape[0], output_size_per_partition), dtype=input.dtype
+    )
+
+
+@register_custom_op(fake_impl=fake_unified_apply_gptq_marlin_gemm)
 def unified_apply_gptq_marlin_gemm(
     input: torch.Tensor,
     weight: torch.Tensor,
@@ -896,7 +916,7 @@ def unified_apply_gptq_marlin_gemm(
     )
 
 
-def fake_unified_apply_gptq_marlin_gemm(
+def fake_unified_apply_gptq_marlin_gemm_with_wtype(
     input: torch.Tensor,
     weight: torch.Tensor,
     weight_scale: torch.Tensor,
@@ -904,8 +924,10 @@ def fake_unified_apply_gptq_marlin_gemm(
     g_idx: torch.Tensor,
     g_idx_sort_indices: torch.Tensor,
     workspace: torch.Tensor,
+    wtype_id: int,
     output_size_per_partition: int,
     input_size_per_partition: int,
+    is_k_full: bool,
     use_atomic_add: bool,
     use_fp32_reduce: bool,
     is_zp_float: bool,
@@ -915,14 +937,7 @@ def fake_unified_apply_gptq_marlin_gemm(
     )
 
 
-direct_register_custom_op(
+@register_custom_op(fake_impl=fake_unified_apply_gptq_marlin_gemm_with_wtype)
-    op_name="unified_apply_gptq_marlin_gemm",
-    op_func=unified_apply_gptq_marlin_gemm,
-    mutates_args=[],
-    fake_impl=fake_unified_apply_gptq_marlin_gemm,
-)
-
-
 def unified_apply_gptq_marlin_gemm_with_wtype(
     input: torch.Tensor,
     weight: torch.Tensor,
@@ -966,32 +981,3 @@ def unified_apply_gptq_marlin_gemm_with_wtype(
         use_fp32_reduce=use_fp32_reduce,
         is_zp_float=is_zp_float,
     )
-
-
-def fake_unified_apply_gptq_marlin_gemm_with_wtype(
-    input: torch.Tensor,
-    weight: torch.Tensor,
-    weight_scale: torch.Tensor,
-    weight_zp: torch.Tensor,
-    g_idx: torch.Tensor,
-    g_idx_sort_indices: torch.Tensor,
-    workspace: torch.Tensor,
-    wtype_id: int,
-    output_size_per_partition: int,
-    input_size_per_partition: int,
-    is_k_full: bool,
-    use_atomic_add: bool,
-    use_fp32_reduce: bool,
-    is_zp_float: bool,
-) -> torch.Tensor:
-    return input.new_empty(
-        (input.shape[0], output_size_per_partition), dtype=input.dtype
-    )
-
-
-direct_register_custom_op(
-    op_name="unified_apply_gptq_marlin_gemm_with_wtype",
-    op_func=unified_apply_gptq_marlin_gemm_with_wtype,
-    mutates_args=[],
-    fake_impl=fake_unified_apply_gptq_marlin_gemm_with_wtype,
-)

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

@@ -53,6 +53,7 @@ from sglang.srt.utils.common import (
     is_sm120_supported,
     next_power_of_2,
 )
+from sglang.srt.utils.custom_op import register_custom_op
 from sglang.srt.utils.patch_torch import register_fake_if_exists
 
 if TYPE_CHECKING:
@@ -73,15 +74,13 @@ except ImportError:
     fp4_quantize = None
 
 try:
-    from flashinfer import mm_fp4 as fp4_gemm
+    from flashinfer import mm_fp4 as flashinfer_fp4_gemm
     from flashinfer import reorder_rows_for_gated_act_gemm, shuffle_matrix_sf_a
 
     enable_flashinfer_fp4_gemm = True
 except ImportError:
     if is_cuda():
-        from sgl_kernel import cutlass_scaled_fp4_mm as fp4_gemm
+        from sgl_kernel import cutlass_scaled_fp4_mm as cutlass_fp4_gemm
-    else:
-        fp4_gemm = None
     enable_flashinfer_fp4_gemm = False
     reorder_rows_for_gated_act_gemm = None
     shuffle_matrix_a = None
@@ -103,8 +102,22 @@ except ImportError:
 logger = logging.getLogger(__name__)
 
 
-@torch.library.custom_op("sglang::fp4_gemm", mutates_args=())
+def _sglang_fp4_gemm_fake(
-def _sglang_fp4_gemm(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    alpha: torch.Tensor,
+    out_dtype: torch.dtype,
+    out_features: int,
+) -> torch.Tensor:
+    M = input.shape[-2]
+    N = int(out_features)
+    return input.new_empty((M, N), dtype=out_dtype)
+
+
+@register_custom_op(fake_impl=_sglang_fp4_gemm_fake)
+def fp4_gemm(
     input: torch.Tensor,
     weight: torch.Tensor,
     input_sf: torch.Tensor,
@@ -115,26 +128,11 @@ def _sglang_fp4_gemm(
 ) -> torch.Tensor:
     backend = FLASHINFER_FP4_GEMM_BACKEND if FLASHINFER_FP4_GEMM_BACKEND else "cutlass"
     if enable_flashinfer_fp4_gemm:
-        return fp4_gemm(
+        return flashinfer_fp4_gemm(
             input, weight, input_sf, weight_sf, alpha, out_dtype, backend=backend
         )
     else:
-        return fp4_gemm(input, weight, input_sf, weight_sf, alpha, out_dtype)
+        return cutlass_fp4_gemm(input, weight, input_sf, weight_sf, alpha, out_dtype)
-
-
-@torch.library.register_fake("sglang::fp4_gemm")
-def _sglang_fp4_gemm_fake(
-    input,
-    weight,
-    input_sf,
-    weight_sf,
-    alpha,
-    out_dtype,
-    out_features: int,
-):
-    M = input.shape[-2]
-    N = int(out_features)
-    return input.new_empty((M, N), dtype=out_dtype)
 
 
 if is_cuda() and (not is_sm120_supported()) and (fp4_quantize is not None):
@@ -1229,7 +1227,7 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
         backend = (
             FLASHINFER_FP4_GEMM_BACKEND if FLASHINFER_FP4_GEMM_BACKEND else "cutlass"
         )
-        out = torch.ops.sglang.fp4_gemm(
+        out = fp4_gemm(
             x_fp4,
             w,
             x_scale_interleaved,

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

@@ -38,7 +38,6 @@ from sglang.srt.layers.quantization.base_config import (
 from sglang.srt.layers.quantization.utils import is_layer_skipped
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
-    direct_register_custom_op,
     is_cuda,
     is_flashinfer_available,
     is_gfx95_supported,
@@ -51,6 +50,7 @@ from sglang.srt.utils import (
     round_up,
     set_weight_attrs,
 )
+from sglang.srt.utils.custom_op import register_custom_op
 
 _is_sm100_supported = is_cuda() and is_sm100_supported()
 has_triton_kernels = is_triton_kernels_available()
@@ -118,7 +118,16 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps):
     return quant_tensor, InFlexData(), scale
 
 
-def _dequant_mxfp4(
+def _dequant_mxfp4_fake(
+    x: torch.Tensor, scale: torch.Tensor, float_dtype: torch.dtype
+) -> torch.Tensor:
+    return torch.empty(
+        (*x.shape[:-1], x.shape[-1] * 2), dtype=float_dtype, device=x.device
+    )
+
+
+@register_custom_op(fake_impl=_dequant_mxfp4_fake)
+def dequant_mxfp4(
     x: torch.Tensor, scale: torch.Tensor, float_dtype: torch.dtype
 ) -> torch.Tensor:
     try:
@@ -133,15 +142,8 @@ def _dequant_mxfp4(
     return mx.dq_mxfp4(x, scale, float_dtype)
 
 
-def _dequant_mxfp4_fake(
+@register_custom_op(out_shape="x")
-    x: torch.Tensor, scale: torch.Tensor, float_dtype: torch.dtype
+def quant_dequant_mxfp4(
-) -> torch.Tensor:
-    return torch.empty(
-        (*x.shape[:-1], x.shape[-1] * 2), dtype=float_dtype, device=x.device
-    )
-
-
-def _quant_dequant_mxfp4(
     x: torch.Tensor, scale_calculation_mode: str = "even"
 ) -> torch.Tensor:
     try:
@@ -156,29 +158,6 @@ def _quant_dequant_mxfp4(
     return mx.qdq_mxfp4(x, scale_calculation_mode)
 
 
-def _quant_dequant_mxfp4_fake(
-    x: torch.Tensor, scale_calculation_mode: str = "even"
-) -> torch.Tensor:
-    return torch.empty_like(x)
-
-
-direct_register_custom_op(
-    op_name="dequant_mxfp4",
-    op_func=_dequant_mxfp4,
-    mutates_args=[],
-    fake_impl=_dequant_mxfp4_fake,
-)
-dequant_mxfp4 = torch.ops.sglang.dequant_mxfp4
-
-direct_register_custom_op(
-    op_name="quant_dequant_mxfp4",
-    op_func=_quant_dequant_mxfp4,
-    mutates_args=[],
-    fake_impl=_quant_dequant_mxfp4_fake,
-)
-quant_dequant_mxfp4 = torch.ops.sglang.quant_dequant_mxfp4
-
-
 class Mxfp4Config(QuantizationConfig):
 
     def __init__(

python/sglang/srt/layers/radix_attention.py

@@ -21,7 +21,7 @@ import torch
 from torch import nn
 
 from sglang.srt.compilation.piecewise_context_manager import get_forward_context
-from sglang.srt.utils import direct_register_custom_op
+from sglang.srt.utils.custom_op import register_custom_op
 
 if TYPE_CHECKING:
     from sglang.srt.layers.quantization.base_config import QuantizationConfig
@@ -115,7 +115,7 @@ class RadixAttention(nn.Module):
                 output = q.new_empty((q.shape[0], self.tp_q_head_num * self.v_head_dim))
             else:
                 output = torch.empty_like(q)
-            torch.ops.sglang.unified_attention_with_output(
+            unified_attention_with_output(
                 q, k, v, output, save_kv_cache, self.layer_id, **kwargs
             )
             return output
@@ -131,6 +131,7 @@ class RadixAttention(nn.Module):
             )
 
 
+@register_custom_op(mutates_args=["output"])
 def unified_attention_with_output(
     query: torch.Tensor,
     key: torch.Tensor,
@@ -165,26 +166,3 @@ def unified_attention_with_output(
 
     output.view(ret.shape).copy_(ret)
     return
-
-
-def unified_attention_with_output_fake(
-    query: torch.Tensor,
-    key: torch.Tensor,
-    value: torch.Tensor,
-    output: torch.Tensor,
-    save_kv_cache: bool,
-    layer_id: int,
-    *,
-    q_rope: Optional[torch.Tensor] = None,
-    k_rope: Optional[torch.Tensor] = None,
-    sinks: Optional[torch.Tensor] = None,
-) -> None:
-    return
-
-
-direct_register_custom_op(
-    op_name="unified_attention_with_output",
-    op_func=unified_attention_with_output,
-    mutates_args=["output"],
-    fake_impl=unified_attention_with_output_fake,
-)

python/sglang/srt/models/qwen3_next.py

@@ -49,6 +49,7 @@ from sglang.srt.utils import (
     make_layers,
     set_weight_attrs,
 )
+from sglang.srt.utils.custom_op import register_custom_op
 
 logger = logging.getLogger(__name__)
 _is_cuda = is_cuda()
@@ -59,7 +60,6 @@ import triton
 import triton.language as tl
 
 from sglang.srt.compilation.piecewise_context_manager import get_forward_context
-from sglang.srt.utils import direct_register_custom_op
 
 
 @triton.jit
@@ -371,7 +371,7 @@ class Qwen3GatedDeltaNet(nn.Module):
     ):
         output = torch.empty_like(hidden_states)
         if forward_batch.forward_mode.is_extend() and get_forward_context() is not None:
-            torch.ops.sglang.gdn_with_output(
+            gdn_with_output(
                 hidden_states,
                 output,
                 self.layer_id,
@@ -1046,6 +1046,7 @@ class Qwen3NextForCausalLM(nn.Module):
 EntryClass = Qwen3NextForCausalLM
 
 
+@register_custom_op(mutates_args=["output"])
 def gdn_with_output(
     hidden_states: torch.Tensor,
     output: torch.Tensor,
@@ -1064,19 +1065,3 @@ def gdn_with_output(
 
     output.view(ret.shape).copy_(ret)
     return
-
-
-def gdn_with_output_fake(
-    hidden_states: torch.Tensor,
-    output: torch.Tensor,
-    layer_id: int,
-) -> None:
-    return
-
-
-direct_register_custom_op(
-    op_name="gdn_with_output",
-    op_func=gdn_with_output,
-    mutates_args=["output"],
-    fake_impl=gdn_with_output_fake,
-)

python/sglang/srt/models/utils.py

@@ -22,12 +22,12 @@ import numpy as np
 import torch
 
 from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm, fused_inplace_qknorm
-from sglang.jit_kernel.utils import register_jit_op
 from sglang.srt.environ import envs
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.utils import is_cuda
+from sglang.srt.utils.custom_op import register_custom_op
 
 if TYPE_CHECKING:
     from sglang.srt.layers.layernorm import RMSNorm
@@ -265,4 +265,4 @@ def apply_qk_norm(
 
 
 # Register the inplace op
-fused_inplace_qknorm = register_jit_op(fused_inplace_qknorm, out_args=["q", "k"])
+fused_inplace_qknorm = register_custom_op(fused_inplace_qknorm, mutates_args=["q", "k"])

python/sglang/srt/utils/common.py

@@ -1972,12 +1972,6 @@ def get_compiler_backend(mode=None) -> str:
 sglang_lib = Library("sglang", "FRAGMENT")  # noqa
 
 
-# Some backends use pytorch version < 2.4.0 which doesn't
-# support `torch.library.custom_op`.
-def supports_custom_op() -> bool:
-    return hasattr(torch.library, "custom_op")
-
-
 def direct_register_custom_op(
     op_name: str,
     op_func: Callable,

python/sglang/srt/utils/custom_op.py

@@ -0,0 +1,191 @@
+from __future__ import annotations
+
+import inspect
+from typing import Any, Callable, List, Optional, TypeVar, Union, overload
+
+import torch
+
+F = TypeVar("F", bound=Callable)
+
+
+@overload
+def register_custom_op(
+    fn: F,
+    *,
+    op_name: Optional[str] = None,
+    mutates_args: Optional[List[str]] = None,
+    out_shape: Optional[Union[int, str]] = None,
+    eager: bool = True,
+) -> F: ...
+
+
+@overload
+def register_custom_op(
+    fn: F,
+    *,
+    op_name: Optional[str] = None,
+    mutates_args: Optional[List[str]] = None,
+    fake_impl: Optional[Callable],
+    eager: bool = True,
+) -> F: ...
+
+
+@overload
+def register_custom_op(
+    *,
+    op_name: Optional[str] = None,
+    mutates_args: Optional[List[str]] = None,
+    out_shape: Optional[Union[int, str]] = None,
+    eager: bool = True,
+) -> Callable[[F], F]: ...
+
+
+@overload
+def register_custom_op(
+    *,
+    op_name: Optional[str] = None,
+    mutates_args: Optional[List[str]] = None,
+    fake_impl: Optional[Callable],
+    eager: bool = True,
+) -> Callable[[F], F]: ...
+
+
+# Real implementation
+def register_custom_op(
+    fn: Optional[Callable] = None,
+    *,
+    op_name: Optional[str] = None,
+    mutates_args: Optional[List[str]] = None,
+    eager: bool = True,
+    **extra_kwargs,
+) -> Any:
+    """
+    A decorator to register a custom operator.
+
+    Example usage:
+    ```python
+    # inplace operator, out_shape is None by default
+    @register_custom_op(mutates_args=["x"])
+    def add_1_(x: torch.Tensor) -> None:
+        x.add_(1)
+
+    # operator with output, out_shape indicates the position of output
+    @register_custom_op(mutates_args=["x"], out_shape=0)
+    def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        return x.add_(y)
+    ```
+
+    :param fn: The function to be registered as a custom operator.
+               If None, return a decorator.
+    :type fn: Callable
+    :param op_name: The name of the operator. If None, use the function name
+    :type op_name: Optional[str]
+    :param mutates_args: A list of argument names that are mutated in-place.
+    :type mutates_args: List[str]
+    :param out_shape: The position (int for positional, str for keyword) of the output-shape tensor.
+                      It is used to generate a fake implementation for torch.compile compatibility.
+                      If the operator is inplace and has no output, set to None.
+    :type out_shape: Optional[List[Union[int, str]]]
+    :param fake_impl: A fake implementation for the operator.
+                      Only one of `out_shape` or `fake_impl` should be provided.
+    :type fake_impl: Optional[Callable]
+    :param eager: Whether to register the operator eagerly.
+                  If False, the registration will be deferred until the first call.
+                  If you met any issue with torch.compile, try to set eager=True.
+                  Currently, to avoid misuse, we set eager=True by default.
+    :type eager: bool
+    :return: The registered JIT custom operator, or a decorator.
+             NOTE: the real register will occur at the first call of the function.
+    :rtype: Callable
+    """
+    extra_kwarg_keys = set(extra_kwargs.keys())
+    expected_kwarg_keys = set({"out_shape", "fake_impl"})
+    assert (
+        expected_kwarg_keys >= extra_kwarg_keys
+    ), f"Unexpected extra kwargs: {extra_kwarg_keys - expected_kwarg_keys}"
+
+    has_out_shape = "out_shape" in extra_kwargs
+    has_fake_impl = "fake_impl" in extra_kwargs
+    assert not (
+        has_out_shape and has_fake_impl
+    ), "Only one of `out_shape` or `fake_impl` should be provided."
+    # Assume inplace if neither out_shape nor fake_impl is provided
+    if not (has_out_shape or has_fake_impl):
+        extra_kwargs["out_shape"] = None
+
+    def decorator(op_func: Callable) -> Callable:
+        wrapper = CustomOpWrapper(
+            op_name=op_name or op_func.__name__,
+            op_func=op_func,
+            mutates_args=mutates_args or [],
+            **extra_kwargs,
+        )
+        return wrapper.real_impl if eager else wrapper
+
+    if fn is not None:
+        return decorator(fn)
+    return decorator
+
+
+class CustomOpWrapper:
+    def __init__(
+        self,
+        op_name: str,
+        op_func: Callable,
+        mutates_args: List[str],
+        **extra_kwargs,
+    ):
+        self.op_name = op_name
+        self.op_func = op_func
+        self.mutates_args = mutates_args
+        self.extra_kwargs = extra_kwargs
+        self._impl: Optional[Callable] = None
+
+    def __call__(self, *args, **kwargs):
+        return self.real_impl(*args, **kwargs)
+
+    @property
+    def real_impl(self) -> Callable:
+        if self._impl is None:
+            if not hasattr(torch.ops.sglang, self.op_name):
+                from sglang.srt.utils.common import direct_register_custom_op
+
+                # NOTE(dark): if torch compile fail here, mark the decorator as eager
+                # lazy registration does not work with torch compile
+                direct_register_custom_op(
+                    op_name=self.op_name,
+                    op_func=self.op_func,
+                    mutates_args=self.mutates_args,
+                    fake_impl=self.fake_impl,
+                )
+            self._impl = getattr(torch.ops.sglang, self.op_name)
+            assert self._impl is not None
+        return self._impl
+
+    @property
+    def fake_impl(self) -> Callable:
+        if "fake_impl" in self.extra_kwargs:
+            return self.extra_kwargs["fake_impl"]
+        assert "out_shape" in self.extra_kwargs
+        signature = inspect.signature(self.op_func)
+        out_shape = self.extra_kwargs["out_shape"]
+        # check out_shape in signature
+
+        def fake_impl(*args, **kwargs):
+            if out_shape is None:
+                return None
+            bound = signature.bind(*args, **kwargs)
+            bound.apply_defaults()
+            try:
+                return torch.empty_like(
+                    bound.args[out_shape]
+                    if isinstance(out_shape, int)
+                    else bound.arguments[out_shape]
+                )
+            except (IndexError, KeyError):
+                raise RuntimeError(
+                    f"Cannot find output argument at position `{out_shape}` for "
+                    f"custom operator `{self.op_name}` with signature `{signature}`."
+                )
+
+        return fake_impl