fix bmm fp8 (#4926)

195a09f5 · Yineng Zhang · GitHub · 9fccda31 · 195a09f5
Unverified Commit 195a09f5 authored Mar 30, 2025 by Yineng Zhang Committed by GitHub Mar 30, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

sgl-kernel/csrc/torch_extension.cc sgl-kernel/csrc/torch_extension.cc +4 -1

No files found.
--- a/sgl-kernel/csrc/torch_extension.cc
+++ b/sgl-kernel/csrc/torch_extension.cc
@@ -82,7 +82,10 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
  /*
   * From FlashInfer
   */
-  m.def("bmm_fp8", bmm_fp8);
+  m.def(
+      "bmm_fp8(Tensor A, Tensor B, Tensor! D, Tensor A_scale, Tensor B_scale, Tensor workspace_buffer, int "
+      "cublas_handle, int cuda_stream) -> ()");
+  m.impl("bmm_fp8", torch::kCUDA, &bmm_fp8);
  m.def("min_p_sampling_from_probs", min_p_sampling_from_probs);
  m.def("top_k_renorm_probs", top_k_renorm_probs);
  m.def("top_p_renorm_probs", top_p_renorm_probs);