diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 1735881f5..a1c355174 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -57,10 +57,7 @@ from sglang.srt.layers.quantization.fp8_utils import ( requant_weight_ue8m0_inplace, ) from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod -from sglang.srt.layers.quantization.marlin_utils_fp8 import ( - apply_fp8_marlin_linear, - prepare_fp8_layer_for_marlin, -) +from sglang.srt.layers.quantization.marlin_utils_fp8 import prepare_fp8_layer_for_marlin from sglang.srt.layers.quantization.unquant import ( UnquantizedFusedMoEMethod, UnquantizedLinearMethod, @@ -638,7 +635,7 @@ class Fp8LinearMethod(LinearMethodBase): bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: if self.use_marlin: - return apply_fp8_marlin_linear( + return torch.ops.sglang.apply_fp8_marlin_linear( input=x, weight=layer.weight, weight_scale=layer.weight_scale, diff --git a/python/sglang/srt/layers/quantization/marlin_utils_fp8.py b/python/sglang/srt/layers/quantization/marlin_utils_fp8.py index e699b6798..d5d008d5f 100644 --- a/python/sglang/srt/layers/quantization/marlin_utils_fp8.py +++ b/python/sglang/srt/layers/quantization/marlin_utils_fp8.py @@ -13,7 +13,7 @@ from sglang.srt.layers.quantization.marlin_utils import ( should_use_atomic_add_reduce, ) from sglang.srt.layers.quantization.utils import get_scalar_types -from sglang.srt.utils import is_cuda +from sglang.srt.utils import direct_register_custom_op, is_cuda _is_cuda = is_cuda() if _is_cuda: @@ -83,6 +83,30 @@ def apply_fp8_marlin_linear( return output.reshape(out_shape) +def fake_apply_fp8_marlin_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + workspace: torch.Tensor, + size_n: int, + size_k: int, + bias: Optional[torch.Tensor], + use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT, +) -> torch.Tensor: + + out_shape = input.shape[:-1] + (size_n,) + fake_output = torch.empty(out_shape, dtype=input.dtype, device=input.device) + return fake_output + + +direct_register_custom_op( + op_name="apply_fp8_marlin_linear", + op_func=apply_fp8_marlin_linear, + mutates_args=[], + fake_impl=fake_apply_fp8_marlin_linear, +) + + def prepare_fp8_layer_for_marlin( layer: torch.nn.Module, size_k_first: bool = True ) -> None: