fix

Signed-off-by: Dylan Chen <[email protected]>

Author: DylanChen-NV

tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py

@@ -121,6 +121,21 @@ def _(
         ret = mat_a.new_empty(shape, dtype=out_dtype)
         return ret
 
+    @torch.library.register_fake("trtllm::cuda_scaled_mm")
+    def _(
+        mat_a: torch.Tensor,
+        mat_b: torch.Tensor,
+        scale_a: torch.Tensor,
+        scale_b: torch.Tensor,
+        bias,
+        out_dtype,
+        userbuffers_id=False,
+    ):
+        shape = [i for i in mat_a.shape]
+        shape[-1] = mat_b.shape[-1]
+        ret = mat_a.new_empty(shape, dtype=out_dtype)
+        return ret
+
     @torch.library.register_fake("trtllm::cublas_mm")
     def _(mat_a, mat_b, bias, out_dtype):
         shape = list(mat_a.shape)