[Piecewise CUDA Graph] Support INT8 (#14918)

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

@@ -28,6 +28,7 @@ from sglang.srt.utils import (
     set_weight_attrs,
     use_intel_amx_backend,
 )
+from sglang.srt.utils.patch_torch import register_fake_if_exists
 
 if TYPE_CHECKING:
     from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
@@ -39,6 +40,20 @@ _is_cpu = is_cpu()
 if _is_cuda:
     from sgl_kernel import int8_scaled_mm
 
+    @register_fake_if_exists("sgl_kernel::int8_scaled_mm")
+    def _int8_scaled_mm_abstract(
+        mat_a,
+        mat_b,
+        scales_a,
+        scales_b,
+        out_dtype,
+        bias=None,
+    ):
+        M = mat_a.shape[-2]
+        N = mat_b.shape[-1]
+        return mat_a.new_empty((M, N), dtype=out_dtype)
+
+
 logger = logging.getLogger(__name__)