【Inference】Modify Layout assert (#74415)

Author: chang-wenbin

paddle/phi/kernels/fusion/cutlass/cutlass_extensions/gemm/kernel/fpA_intB_gemm.h

@@ -76,16 +76,6 @@ struct GemmFpAIntB {
   using LayoutC = typename Mma::LayoutC;
   using ElementScale = typename Mma::IteratorA::Element;
 
-  // NOTE: (changwenbin) Currently only A row major and B column major are
-  // supported. Other cases have not been verified yet.
-
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 750)
-  static_assert(
-      platform::is_same<LayoutA, layout::RowMajor>::value &&
-          platform::is_same<LayoutB, layout::ColumnMajor>::value,
-      "A must be row major and B must be col major in cuda_arch >= sm75");
-#endif
-
   static ComplexTransform const kTransformA = Mma::kTransformA;
   static ComplexTransform const kTransformB = Mma::kTransformA;
 
@@ -97,6 +87,20 @@ struct GemmFpAIntB {
   using InstructionShape = typename Mma::Policy::Operator::InstructionShape;
   using ArchTag = typename Mma::ArchTag;
 
+  // NOTE: (changwenbin) Ensure that the GEMM layout meets the requirements
+  // under different architectures.
+  static_assert(
+      (platform::is_same<ArchTag, arch::Sm75>::value &&
+       platform::is_same<LayoutA, layout::RowMajor>::value &&
+       platform::is_same<LayoutB, layout::ColumnMajor>::value) ||
+          (platform::is_same<ArchTag, arch::Sm80>::value &&
+           platform::is_same<LayoutA, layout::RowMajor>::value &&
+           platform::is_same<LayoutB, layout::ColumnMajor>::value) ||
+          (platform::is_same<ArchTag, arch::Sm70>::value &&
+           platform::is_same<LayoutA, layout::RowMajor>::value &&
+           platform::is_same<LayoutB, layout::RowMajor>::value),
+      "A must be row major and B must be col major in cuda_arch >= sm75");
+
   static int const kStages = Mma::kStages;
   static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
   static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;