fix fallback gemm

vllm-project · Jul 26, 2024 · d5f364b · d5f364b
1 parent 43414d6
commit d5f364b
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm89_int8_dispatch.cuh b/csrc/quantization/cutlass_w8a8/scaled_mm_c2x_sm89_int8_dispatch.cuh
@@ -12,15 +12,16 @@ namespace vllm {
 template <typename InType, typename OutType,
           template <typename, typename> typename Epilogue>
 struct sm89_int8_fallback_gemm {
+  // Shared mem requirement : 61440
   static_assert(std::is_same<InType, int8_t>());
-  using TileShape = typename cutlass::gemm::GemmShape<64, 128, 64>;
-  using WarpShape = typename cutlass::gemm::GemmShape<32, 64, 64>;
+  using TileShape = cutlass::gemm::GemmShape<32, 64, 128>;
+  using WarpShape = cutlass::gemm::GemmShape<16, 64, 64>;
   using InstructionShape = typename cutlass::gemm::GemmShape<16, 8, 32>;
-  using FP8MathOperator = typename cutlass::arch::OpMultiplyAdd;
+  static int32_t const MainLoopStages = 5;
+
   using Cutlass2xGemm =
       cutlass_2x_gemm<cutlass::arch::Sm89, enable_sm89_to_sm90, InType, OutType,
-                      Epilogue, TileShape, WarpShape, InstructionShape, 5,
-                      FP8MathOperator>;
+                      Epilogue, TileShape, WarpShape, InstructionShape, 5>;
 };
 
 struct sm89_int8_config_default {