[Fix] Register custom ops only if they exist (#13321)

python/sglang/srt/entrypoints/engine.py

@@ -683,7 +683,6 @@ class Engine(EngineBase):
 
 def _set_envs_and_config(server_args: ServerArgs):
     # Set global environments
-    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
     if "NCCL_CUMEM_ENABLE" not in os.environ or server_args.enable_symm_mem:
         os.environ["NCCL_CUMEM_ENABLE"] = str(int(server_args.enable_symm_mem))
     if (
@@ -756,10 +755,13 @@ def _set_envs_and_config(server_args: ServerArgs):
 
 
 def _init_tokenizer_manager(
-    server_args: ServerArgs, port_args: PortArgs
+    server_args: ServerArgs,
+    port_args: PortArgs,
+    TokenizerManagerClass: Optional[TokenizerManager] = None,
 ) -> TokenizerManager:
     # Launch tokenizer process
-    tokenizer_manager = TokenizerManager(server_args, port_args)
+    TokenizerManagerClass = TokenizerManagerClass or TokenizerManager
+    tokenizer_manager = TokenizerManagerClass(server_args, port_args)
 
     # Initialize templates
     template_manager = TemplateManager()

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

@@ -32,6 +32,7 @@ from sglang.srt.layers.quantization.marlin_utils import (
 from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
 from sglang.srt.layers.quantization.utils import get_scalar_types, replace_parameter
 from sglang.srt.layers.quantization.w8a8_int8 import npu_fused_experts
+from sglang.srt.utils.patch_torch import register_fake_if_exists
 
 if TYPE_CHECKING:
     from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
@@ -959,7 +960,7 @@ class AWQMoEAscendMethod(AWQMoEMethod):
 # Register fake implementations for torch.compile support
 if _is_cuda:
 
-    @torch.library.register_fake("sgl_kernel::awq_dequantize")
+    @register_fake_if_exists("sgl_kernel::awq_dequantize")
     def _(
         qweight,
         scales,
@@ -971,7 +972,7 @@ if _is_cuda:
         out_shape = qweight.shape[:-1] + (qweight.shape[-1] * 32 // num_bits,)
         return qweight.new_empty(out_shape, dtype=scales.dtype)
 
-    @torch.library.register_fake("sgl_kernel::awq_marlin_repack")
+    @register_fake_if_exists("sgl_kernel::awq_marlin_repack")
     def _(b_q_weight, size_k, size_n, num_bits):
         return b_q_weight.new_empty(
             (size_k // 16, size_n * (num_bits // 2)), dtype=b_q_weight.dtype

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

@@ -42,16 +42,16 @@ from sglang.srt.layers.quantization.utils import (
     replace_parameter,
     unpack_cols,
 )
+from sglang.srt.utils import is_cuda
+from sglang.srt.utils.patch_torch import register_fake_if_exists
 
 if TYPE_CHECKING:
     from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
     from sglang.srt.layers.moe.token_dispatcher import (
-        StandardDispatchOutput,
         CombineInput,
+        StandardDispatchOutput,
     )
 
-from sglang.srt.utils import is_cuda
-
 _is_cuda = is_cuda()
 
 if _is_cuda:
@@ -1099,21 +1099,21 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
 # Register fake implementations for torch.compile support
 if _is_cuda:
 
-    @torch.library.register_fake("sgl_kernel::gptq_gemm")
+    @register_fake_if_exists("sgl_kernel::gptq_gemm")
     def _(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, use_shuffle, bit):
         return a.new_empty((a.shape[0], b_q_weight.shape[-1]), dtype=a.dtype)
 
-    @torch.library.register_fake("sgl_kernel::gptq_marlin_repack")
+    @register_fake_if_exists("sgl_kernel::gptq_marlin_repack")
     def _(b_q_weight, perm, size_k, size_n, num_bits):
         return b_q_weight.new_empty(
             (size_k // 16, size_n * (num_bits // 2)), dtype=b_q_weight.dtype
         )
 
-    @torch.library.register_fake("sgl_kernel::gptq_shuffle")
+    @register_fake_if_exists("sgl_kernel::gptq_shuffle")
     def _(q_weight, q_perm, bit):
         return
 
-    @torch.library.register_fake("sgl_kernel::moe_wna16_marlin_gemm")
+    @register_fake_if_exists("sgl_kernel::moe_wna16_marlin_gemm")
     def _(
         a,
         c,

python/sglang/srt/managers/detokenizer_manager.py

@@ -27,7 +27,6 @@ import zmq
 from sglang.srt.managers.io_struct import (
     BatchEmbeddingOutput,
     BatchMultimodalDecodeReq,
-    BatchMultimodalOutput,
     BatchStrOutput,
     BatchTokenIDOutput,
     FreezeGCReq,
@@ -284,22 +283,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
         )
 
     def handle_multimodal_decode_req(self, recv_obj: BatchMultimodalDecodeReq):
-        outputs = self.tokenizer.detokenize(recv_obj)
-        return BatchMultimodalOutput(
-            rids=recv_obj.rids,
-            http_worker_ipcs=recv_obj.http_worker_ipcs,
-            finished_reasons=recv_obj.finished_reasons,
-            outputs=outputs,
-            prompt_tokens=recv_obj.prompt_tokens,
-            completion_tokens=recv_obj.completion_tokens,
-            cached_tokens=recv_obj.cached_tokens,
-            placeholder_tokens_idx=None,
-            placeholder_tokens_val=None,
-            queue_time=recv_obj.queue_time,
-            forward_entry_time=recv_obj.forward_entry_time,
-            prefill_launch_delay=recv_obj.prefill_launch_delay,
-            prefill_launch_latency=recv_obj.prefill_launch_latency,
-        )
+        raise NotImplementedError()
 
     def handle_freeze_gc_req(self, recv_req: FreezeGCReq):
         freeze_gc("Detokenizer Manager")

python/sglang/srt/utils/patch_torch.py

@@ -17,7 +17,7 @@ import torch
 from packaging import version
 from torch.multiprocessing import reductions
 
-from sglang.srt.utils import is_npu
+from sglang.srt.utils.common import is_npu
 
 _is_npu = is_npu()
 
@@ -88,3 +88,29 @@ def monkey_patch_torch_compile():
 
         af.auto_functionalized_v2._cacheable = True
         af.auto_functionalized._cacheable = True
+
+
+def register_fake_if_exists(op_name):
+    """
+    Decorator factory to conditionally register a fake for a custom op if it exists.
+    Parses op_name (e.g., 'sgl_kernel::gptq_gemm'), checks if the op exists via hasattr
+    on the namespace attribute of torch.ops. Registers the fake if present; otherwise,
+    returns the function unchanged.
+    Args:
+        op_name (str): Full operator name (e.g., 'sgl_kernel::gptq_gemm').
+    Returns:
+        callable: Decorator for the fake function.
+    Example:
+        @register_fake_if_exists('sgl_kernel::gptq_gemm')
+        def fake_gptq_gemm(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, use_shuffle, bit):
+            return a.new_empty((a.shape[0], b_q_weight.shape[-1]), dtype=a.dtype)
+    """
+
+    def decorator(func):
+        namespace, bare_op = op_name.split("::")
+        ops_namespace = getattr(torch.ops, namespace, None)
+        if ops_namespace and hasattr(ops_namespace, bare_op):
+            torch.library.register_fake(op_name, func)
+        return func
+
+    return decorator