diff --git a/python/sglang/jit_kernel/utils.py b/python/sglang/jit_kernel/utils.py index 07bd75906..4de953497 100644 --- a/python/sglang/jit_kernel/utils.py +++ b/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 diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 8e3429dec..39af99f24 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/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: - return +@register_custom_op(out_shape="tensor") +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( - 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) +@register_custom_op(mutates_args=["output"]) +def reg_all_gather_into_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._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( - op_name="outplace_all_reduce", - op_func=outplace_all_reduce, - mutates_args=[], - fake_impl=outplace_all_reduce_fake, - ) - - def reg_all_gather_into_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._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, - ) +@register_custom_op(mutates_args=["output"]) +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) 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( - output, input, group_name=self.unique_name - ) + reg_reduce_scatter_tensor(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( - output, input, group_name=self.unique_name - ) + reg_all_gather_into_tensor(output, input, group_name=self.unique_name) def cp_all_gather_into_tensor_async( self, output: torch.Tensor, input: torch.Tensor, stream=None diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index 1d79819e5..6c2db5f26 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/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( - global_tokens, group_name=get_tp_group().unique_name - ) + from sglang.srt.distributed.parallel_state import inplace_all_reduce + + inplace_all_reduce(global_tokens, group_name=get_tp_group().unique_name) else: global_tokens[:] = tensor_model_parallel_all_reduce(global_tokens) diff --git a/python/sglang/srt/layers/elementwise.py b/python/sglang/srt/layers/elementwise.py index a3e930064..dd07fc485 100644 --- a/python/sglang/srt/layers/elementwise.py +++ b/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( diff --git a/python/sglang/srt/layers/flashinfer_comm_fusion.py b/python/sglang/srt/layers/flashinfer_comm_fusion.py index 21b0b20f8..ecc89bb6d 100644 --- a/python/sglang/srt/layers/flashinfer_comm_fusion.py +++ b/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 ( - direct_register_custom_op, - is_flashinfer_available, - supports_custom_op, -) +from sglang.srt.utils import is_flashinfer_available +from sglang.srt.utils.custom_op import register_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: diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 8655b02a3..8e9d14910 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/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, diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index a525637ec..53c25b069 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/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, diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py index a9ea71e33..1f1c3e709 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py +++ b/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, -) diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py index fa809e982..a1885fade 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py +++ b/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( - 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, -) - - +@register_custom_op(out_shape="hidden_states") 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, diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 7922d5f5b..839463518 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/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 ( - torch.ops.sglang.flashinfer_fp4_moe_forward_piecewise_cuda_graph_impl( - hidden_states, - topk_output.router_logits, - topk_output.topk_config.top_k, - topk_output.topk_config.topk_group, - topk_output.topk_config.num_expert_group, - topk_output.topk_config.correction_bias, - self.layer_id, - ) + return flashinfer_fp4_moe_forward_piecewise_cuda_graph_impl( + hidden_states, + topk_output.router_logits, + topk_output.topk_config.top_k, + topk_output.topk_config.topk_group, + topk_output.topk_config.num_expert_group, + topk_output.topk_config.correction_bias, + 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( - 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) - - +@register_custom_op(out_shape="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, -) diff --git a/python/sglang/srt/layers/moe/moe_runner/marlin.py b/python/sglang/srt/layers/moe/moe_runner/marlin.py index ea35b875f..45104dd27 100644 --- a/python/sglang/srt/layers/moe/moe_runner/marlin.py +++ b/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 - fused_marlin_moe, - ) + from sglang.srt.layers.moe.fused_moe_triton.fused_marlin_moe import 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, diff --git a/python/sglang/srt/layers/moe/rocm_moe_utils.py b/python/sglang/srt/layers/moe/rocm_moe_utils.py index efa6bb1bb..70b1dd757 100644 --- a/python/sglang/srt/layers/moe/rocm_moe_utils.py +++ b/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, diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py index b06bb0cc1..c5e5a11fc 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/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, diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py index 26ef087d6..89a77efc5 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py +++ b/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, diff --git a/python/sglang/srt/layers/quantization/fp8_kernel.py b/python/sglang/srt/layers/quantization/fp8_kernel.py index 3f73c0a90..df231f1a0 100644 --- a/python/sglang/srt/layers/quantization/fp8_kernel.py +++ b/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( - A: torch.Tensor, - As: torch.Tensor, - B: torch.Tensor, - Bs: torch.Tensor, - C: torch.Tensor, - ) -> 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, - ) +@register_custom_op(mutates_args=["C"]) +def deep_gemm_fp8_fp8_bf16_nt( + A: torch.Tensor, + As: torch.Tensor, + B: torch.Tensor, + Bs: torch.Tensor, + C: torch.Tensor, +) -> None: + deep_gemm_wrapper.gemm_nt_f8f8bf16((A, As), (B, Bs), C) @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(): - 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) + deep_gemm_fp8_fp8_bf16_nt(A, As, B, Bs, C) return C diff --git a/python/sglang/srt/layers/quantization/marlin_utils.py b/python/sglang/srt/layers/quantization/marlin_utils.py index 96b2efd25..cde7463c2 100644 --- a/python/sglang/srt/layers/quantization/marlin_utils.py +++ b/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( - op_name="unified_apply_gptq_marlin_gemm", - op_func=unified_apply_gptq_marlin_gemm, - mutates_args=[], - fake_impl=fake_unified_apply_gptq_marlin_gemm, -) - - +@register_custom_op(fake_impl=fake_unified_apply_gptq_marlin_gemm_with_wtype) 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, -) diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 79a1c7ef2..b1092aec6 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/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 - else: - fp4_gemm = None + from sgl_kernel import cutlass_scaled_fp4_mm as cutlass_fp4_gemm 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( +def _sglang_fp4_gemm_fake( + 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) - - -@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) + return cutlass_fp4_gemm(input, weight, input_sf, weight_sf, alpha, 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, diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index 84424458c..274d7fe35 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/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( - 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 - ) - - -def _quant_dequant_mxfp4( +@register_custom_op(out_shape="x") +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__( diff --git a/python/sglang/srt/layers/radix_attention.py b/python/sglang/srt/layers/radix_attention.py index 3110cbbb7..199b2aa36 100644 --- a/python/sglang/srt/layers/radix_attention.py +++ b/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, -) diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py index b38a83d57..28f36f661 100644 --- a/python/sglang/srt/models/qwen3_next.py +++ b/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, -) diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index 19bb72442..3ebd82448 100644 --- a/python/sglang/srt/models/utils.py +++ b/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"]) diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 6259e9c48..1aa46107b 100644 --- a/python/sglang/srt/utils/common.py +++ b/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, diff --git a/python/sglang/srt/utils/custom_op.py b/python/sglang/srt/utils/custom_op.py new file mode 100644 index 000000000..9b4bd90d6 --- /dev/null +++ b/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