[PCG]add piecewise cuda graph support for marlin linear (#20119)

Co-authored-by: undefined <zhouchen.arrebol@jd.com>

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 (
+from sglang.srt.layers.quantization.marlin_utils_fp8 import prepare_fp8_layer_for_marlin
-    apply_fp8_marlin_linear,
-    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,

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: