chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Refactor] Clean up custom op (#15995)

Browse Source
This commit is contained in:
DarkSharpness
2026-01-01 21:41:56 +08:00
committed by GitHub
parent a3b1e8ef3d
commit f6f7af4068
23 changed files with 382 additions and 658 deletions
Download Patch File Download Diff File Expand all files Collapse all files

121
python/sglang/jit_kernel/utils.py
View File

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

134
python/sglang/srt/distributed/parallel_state.py
View File

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

6
python/sglang/srt/layers/dp_attention.py
View File

@@ -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)

20
python/sglang/srt/layers/elementwise.py
View File

@@ -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(

50
python/sglang/srt/layers/flashinfer_comm_fusion.py
View File

@@ -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:

9
python/sglang/srt/layers/layernorm.py
View File

@@ -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,

8
python/sglang/srt/layers/moe/ep_moe/layer.py
View File

@@ -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,

38
python/sglang/srt/layers/moe/fused_moe_triton/fused_marlin_moe.py
View File

@@ -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,
)

87
python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py
View File

@@ -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,

63
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
View File

@@ -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,
)

6
python/sglang/srt/layers/moe/moe_runner/marlin.py
View File

@@ -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,

38
python/sglang/srt/layers/moe/rocm_moe_utils.py
View File

@@ -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,

4
python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py
View File

@@ -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,

3
python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
View File

@@ -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,

42
python/sglang/srt/layers/quantization/fp8_kernel.py
View File

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

68
python/sglang/srt/layers/quantization/marlin_utils.py
View File

@@ -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,
)

46
python/sglang/srt/layers/quantization/modelopt_quant.py
View File

@@ -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,

47
python/sglang/srt/layers/quantization/mxfp4.py
View File

@@ -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__(

28
python/sglang/srt/layers/radix_attention.py
View File

@@ -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,
)

21
python/sglang/srt/models/qwen3_next.py
View File

@@ -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,
)

4
python/sglang/srt/models/utils.py
View File

@@ -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"])

6
python/sglang/srt/utils/common.py
View File

@@ -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,

191
python/sglang/srt/utils/custom_op.py Normal file
View File

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