[Bugfix][Kernel] Register int8 scaled_mm as a custom op (#57331)

Signed-off-by: Kai <limkc.reply@gmail.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>

vllm/model_executor/kernels/linear/scaled_mm/triton.py

@@ -133,9 +133,7 @@ class TritonInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):
             x.contiguous(), i_s, i_zp, symmetric=symmetric
         )
 
-        out = triton_scaled_mm(
-            x_q, w_q, scale_a=x_s, scale_b=w_s, out_dtype=x.dtype, bias=bias
-        )
+        out = torch.ops.vllm.int8_scaled_mm_func(x_q, w_q, x_s, w_s, x.dtype, bias)
 
         if azp_adj is not None:
             # Asymmetric quantization: subtract the zero-point correction.
@@ -222,6 +220,44 @@ class TritonFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
         )
 
 
+def _int8_scaled_mm_func(
+    qx: torch.Tensor,
+    weight: torch.Tensor,
+    x_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    output_dtype: torch.dtype,
+    bias: torch.Tensor | None = None,
+) -> torch.Tensor:
+    return triton_scaled_mm(
+        qx,
+        weight,
+        scale_a=x_scale,
+        scale_b=weight_scale,
+        out_dtype=output_dtype,
+        bias=bias,
+    )
+
+
+def _int8_scaled_mm_fake(
+    qx: torch.Tensor,
+    weight: torch.Tensor,
+    x_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    output_dtype: torch.dtype,
+    bias: torch.Tensor | None = None,
+) -> torch.Tensor:
+    return torch.empty(
+        (qx.size(0), weight.size(1)), dtype=output_dtype, device=qx.device
+    )
+
+
+direct_register_custom_op(
+    "int8_scaled_mm_func",
+    _int8_scaled_mm_func,
+    fake_impl=_int8_scaled_mm_fake,
+)
+
+
 # TODO we should be able to change the type of block_size to GroupShape
 # after we resolve GroupShape compilation issue
 # https://github.com/vllm-project/vllm/issues/25270