Update Cutlass to V3.8-2

Differential Revision: D69890673

Author: jwfromm

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16i4bf16_rowwise.cu

@@ -99,9 +99,8 @@ at::Tensor bf16i4bf16_rowwise_impl(
                          // threadblocks in a
                          // cluster
   using CooperativeSchedule =
-      cutlass::gemm::KernelTmaWarpSpecializedCooperativeMixedInput;
-  using PongSchedule =
-      cutlass::gemm::KernelTmaWarpSpecializedPingpongMixedInput;
+      cutlass::gemm::KernelTmaWarpSpecializedCooperative;
+  using PongSchedule = cutlass::gemm::KernelTmaWarpSpecializedPingpong;
   using CooperativeEpilogueSchedule =
       cutlass::epilogue::TmaWarpSpecializedCooperative;
   using PongEpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized;

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/bf16i4bf16_rowwise_batched.cu

@@ -103,9 +103,8 @@ at::Tensor bf16i4bf16_rowwise_batched_impl(
                          // threadblocks in a
                          // cluster
   using CooperativeSchedule =
-      cutlass::gemm::KernelTmaWarpSpecializedCooperativeMixedInput;
-  using PongSchedule =
-      cutlass::gemm::KernelTmaWarpSpecializedPingpongMixedInput;
+      cutlass::gemm::KernelTmaWarpSpecializedCooperative;
+  using PongSchedule = cutlass::gemm::KernelTmaWarpSpecializedPingpong;
   using CooperativeEpilogueSchedule =
       cutlass::epilogue::TmaWarpSpecializedCooperative;
   using PongEpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized;

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16_rowwise_grouped.cu

@@ -67,17 +67,17 @@ struct GroupedGemmConfigs {
       conditional_t<PONG, PongEpilogueSchedule, CooperativeEpilogueSchedule>;
 
   // Implement rowwise scaling epilogue.
-  using XScale = cutlass::epilogue::fusion::Sm90ColBroadcastPtrArray<
+  using XScale = cutlass::epilogue::fusion::Sm90ColBroadcast<
       0,
       TileShape,
-      ElementComputeEpilogue,
+      ElementComputeEpilogue*,
       ElementComputeEpilogue,
       cute::Stride<cute::Int<1>, cute::Int<0>, cute::Int<0>>>;
 
-  using WScale = cutlass::epilogue::fusion::Sm90RowBroadcastPtrArray<
+  using WScale = cutlass::epilogue::fusion::Sm90RowBroadcast<
       0,
       TileShape,
-      ElementComputeEpilogue,
+      ElementComputeEpilogue*,
       ElementComputeEpilogue,
       cute::Stride<cute::Int<0>, cute::Int<1>, cute::Int<0>>>;
 

fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8i4bf16_rowwise.cu

@@ -93,9 +93,8 @@ at::Tensor f8i4bf16_rowwise_impl(
                          // threadblocks in a
                          // cluster
   using CooperativeSchedule =
-      cutlass::gemm::KernelTmaWarpSpecializedCooperativeMixedInput;
-  using PongSchedule =
-      cutlass::gemm::KernelTmaWarpSpecializedPingpongMixedInput;
+      cutlass::gemm::KernelTmaWarpSpecializedCooperative;
+  using PongSchedule = cutlass::gemm::KernelTmaWarpSpecializedPingpong;
   using CooperativeEpilogueSchedule =
       cutlass::epilogue::TmaWarpSpecializedCooperative;
   using PongEpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized;
@@ -260,7 +259,7 @@ at::Tensor dispatch_f8i4bf16_rowwise_kernel(
     return f8i4bf16_rowwise_impl<
         128,
         256,
-        64,
+        128,
         2,
         1,
         1,
@@ -271,7 +270,7 @@ at::Tensor dispatch_f8i4bf16_rowwise_kernel(
     return f8i4bf16_rowwise_impl<
         128,
         256,
-        64,
+        128,
         2,
         1,
         1,

fbgemm_gpu/experimental/gen_ai/src/quantize/quantize.cu

@@ -786,7 +786,7 @@ std::vector<at::Tensor> quantize_fp8_per_tensor(
   for (int i = 0; i < input.dim(); i++) {
     quantized_input_shape.push_back(input.size(i));
   }
-  std::vector<long int> scale_shape = {1};
+  std::vector<long int> scale_shape = {};
   input = input.cuda();
   at::Tensor quantized_input = torch::empty(
       quantized_input_shape,

fbgemm_gpu/experimental/gen_ai/test/quantize/quantize_test.py

@@ -1120,7 +1120,7 @@ def test_quantize_compile(self) -> None:
     @unittest.skipIf(
         not torch.version.cuda, "Skip on AMD: fast gemv op is not yet supported."
     )
-    def test_gemv(
+    def run_gemv(
         self, test_cases, gemv_op, atol, rtol, quantize_w=False, quantize_x=False
     ):
         for M, N, K in test_cases:
@@ -1150,7 +1150,7 @@ def test_bf16_gemv(self) -> None:
             (1, 7168, 8192),
             (1, 8192, 3584),
         ]
-        self.test_gemv(test_cases, torch.ops.fbgemm.bf16_fast_gemv, 9.0e-3, 9.0e-3)
+        self.run_gemv(test_cases, torch.ops.fbgemm.bf16_fast_gemv, 9.0e-3, 9.0e-3)
 
     @unittest.skipIf(
         not torch.version.cuda, "Skip on AMD: fast gemv op is not yet supported."
@@ -1164,7 +1164,7 @@ def test_bf16_fp8_gemv(self) -> None:
             (1, 7168, 8192),
             (1, 8192, 3584),
         ]
-        self.test_gemv(
+        self.run_gemv(
             test_cases,
             torch.ops.fbgemm.bf16fp8bf16_fast_gemv,
             1.0e-2,
@@ -1182,7 +1182,7 @@ def test_fp8_fp8_gemv(self) -> None:
             (1, 7168, 8192),
             (1, 8192, 3584),
         ]
-        self.test_gemv(
+        self.run_gemv(
             test_cases,
             torch.ops.fbgemm.fp8fp8bf16_fast_gemv,
             9.0e-2,