[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 (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    List,
-    Optional,
-    Tuple,
-    TypeAlias,
-    TypeVar,
-    Union,
-    overload,
-)
-
-import torch
-
-from sglang.srt.utils.common import direct_register_custom_op
+from typing import TYPE_CHECKING, Any, Callable, List, Tuple, TypeAlias, TypeVar, Union
 
 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