update and simplify CustomOp (#3249)

2025-02-01 18:56:44 +08:00
parent 17dbf976c5
commit 4eb4b401cc
8 changed files with 46 additions and 45 deletions
--- a/python/sglang/srt/layers/layernorm.py
+++ b/python/sglang/srt/layers/layernorm.py
@@ -29,14 +29,11 @@ if is_cuda_available():
        rmsnorm,
    )

-from vllm.model_executor.custom_op import CustomOp
-
-from sglang.srt.layers.custom_op_util import register_custom_op
+from sglang.srt.custom_op import CustomOp

 logger = logging.getLogger(__name__)


-@register_custom_op("sglang_rmsnorm")
 class RMSNorm(CustomOp):
    def __init__(
        self,
@@ -79,7 +76,6 @@ class RMSNorm(CustomOp):
            return x, residual


-@register_custom_op("sglang_gemma_rmsnorm")
 class GemmaRMSNorm(CustomOp):
    def __init__(
        self,