add cutlass support for blackwell fp8 gemm (#13798)

2025-03-04 07:55:07 -08:00
parent b3cf368d79
commit f89978ad7c
11 changed files with 272 additions and 65 deletions
--- a/csrc/sparse/cutlass/sparse_scaled_mm_c3x.cuh
+++ b/csrc/sparse/cutlass/sparse_scaled_mm_c3x.cuh
@@ -65,12 +65,7 @@ struct cutlass_sparse_3x_gemm {
      typename std::conditional<std::is_same_v<ElementAB, int8_t>, int32_t,
                                float>::type;

-  using EpilogueDescriptor =
-      cutlass::epilogue::collective::detail::EpilogueDescriptor<
-          TileShape, cutlass::epilogue::collective::EpilogueTileAuto, ElementD,
-          ElementD, EpilogueSchedule>;
-
-  using Epilogue = Epilogue_<ElementAcc, ElementD, EpilogueDescriptor>;
+  using Epilogue = Epilogue_<ElementAcc, ElementD, TileShape>;

  using ElementC = void;
  using LayoutC = cutlass::layout::RowMajor;