PaddlePaddle
diff --git a/‎paddle/phi/backends/dynload/cublas.h
Copy file name to clipboard
+12Lines changed: 12 additions & 0 deletions b/‎paddle/phi/backends/dynload/cublas.h
Copy file name to clipboard
+12Lines changed: 12 additions & 0 deletions
diff --git a/‎paddle/phi/kernels/funcs/blas/blas.h
Copy file name to clipboard
+14-14Lines changed: 14 additions & 14 deletions b/‎paddle/phi/kernels/funcs/blas/blas.h
Copy file name to clipboard
+14-14Lines changed: 14 additions & 14 deletions
@@ -17,6 +17,9 @@ limitations under the License. */
 #include <cublasXt.h>
 #include <cublas_v2.h>
 #include <cuda.h>
+#if CUDA_VERSION >= 12030
+#include <cublas_api.h>
+#endif
 
 #include <mutex>  // NOLINT
 #include <type_traits>
@@ -138,6 +141,15 @@ CUBLAS_BLAS_ROUTINE_EACH_R3(DECLARE_DYNAMIC_LOAD_CUBLAS_WRAP)
 CUBLAS_BLAS_ROUTINE_EACH_R4(DECLARE_DYNAMIC_LOAD_CUBLAS_WRAP)
 #endif
 
+#if CUDA_VERSION >= 12030
+#define CUBLAS_BLAS_ROUTINE_EACH_R5(__macro) \
+  __macro(cublasGemmStridedBatchedEx_64);    \
+  __macro(cublasGemmEx_64);                  \
+  __macro(cublasSgemmEx_64);
+
+CUBLAS_BLAS_ROUTINE_EACH_R5(DECLARE_DYNAMIC_LOAD_CUBLAS_WRAP)
+#endif
+
 #undef DECLARE_DYNAMIC_LOAD_CUBLAS_WRAP
 }  // namespace dynload
 }  // namespace phi
@@ -87,9 +87,9 @@ class Blas {
   template <typename T>
   void GEMM(CBLAS_TRANSPOSE transA,
             CBLAS_TRANSPOSE transB,
-            int M,
-            int N,
-            int K,
+            int64_t M,
+            int64_t N,
+            int64_t K,
             T alpha,
             const T* A,
             const T* B,
@@ -99,9 +99,9 @@ class Blas {
   template <typename T, typename U = T>
   void GEMM(CBLAS_TRANSPOSE transA,
             CBLAS_TRANSPOSE transB,
-            int M,
-            int N,
-            int K,
+            int64_t M,
+            int64_t N,
+            int64_t K,
             U alpha,
             const T* A,
             const T* B,
@@ -292,30 +292,30 @@ class Blas {
   template <typename T>
   void BatchedGEMM(CBLAS_TRANSPOSE transA,
                    CBLAS_TRANSPOSE transB,
-                   int M,
-                   int N,
-                   int K,
+                   int64_t M,
+                   int64_t N,
+                   int64_t K,
                    T alpha,
                    const T* A,
                    const T* B,
                    T beta,
                    T* C,
-                   int batchCount,
+                   int64_t batchCount,
                    int64_t strideA,
                    int64_t strideB) const;
 
   template <typename T, typename U = T>
   void BatchedGEMM(CBLAS_TRANSPOSE transA,
                    CBLAS_TRANSPOSE transB,
-                   int M,
-                   int N,
-                   int K,
+                   int64_t M,
+                   int64_t N,
+                   int64_t K,
                    U alpha,
                    const T* A,
                    const T* B,
                    U beta,
                    T* C,
-                   int batchCount,
+                   int64_t batchCount,
                    int64_t strideA,
                    int64_t strideB) const;