diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 0da0aba82..28f7ad7d7 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -45,6 +45,7 @@ from sglang.srt.layers.quantization.utils import ( ) from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.utils.common import ( + direct_register_custom_op, get_bool_env_var, is_cuda, is_sm120_supported, @@ -100,7 +101,6 @@ except ImportError: logger = logging.getLogger(__name__) -@torch.library.custom_op("sglang::fp4_gemm", mutates_args=()) def _sglang_fp4_gemm( input: torch.Tensor, weight: torch.Tensor, @@ -119,7 +119,6 @@ def _sglang_fp4_gemm( 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, @@ -134,6 +133,14 @@ def _sglang_fp4_gemm_fake( return input.new_empty((M, N), dtype=out_dtype) +direct_register_custom_op( + op_name="fp4_gemm", + op_func=_sglang_fp4_gemm, + mutates_args=[], + fake_impl=_sglang_fp4_gemm_fake, +) + + if is_cuda() and (not is_sm120_supported()) and (fp4_quantize is not None): @register_fake_if_exists("sgl_kernel::scaled_fp4_quant")