From a04999b3335c63ee1f694633e53d3c68c373ff7c Mon Sep 17 00:00:00 2001
From: zhyncs <me@zhyncs.com>
Date: Wed, 22 Jan 2025 04:04:31 -0800
Subject: [PATCH] fix return type of cuBLAS

---
 csrc/activation.cu                  | 2 +-
 include/flashinfer/gemm/bmm_fp8.cuh | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/csrc/activation.cu b/csrc/activation.cu
index dba033597..c997708db 100644
--- a/csrc/activation.cu
+++ b/csrc/activation.cu
@@ -15,7 +15,7 @@
  */
 #include <flashinfer/activation.cuh>
 
-#include "aot_extension_utils.h"
+#include "pytorch_extension_utils.h"
 
 using namespace flashinfer;
 
diff --git a/include/flashinfer/gemm/bmm_fp8.cuh b/include/flashinfer/gemm/bmm_fp8.cuh
index 853d803ce..777893416 100644
--- a/include/flashinfer/gemm/bmm_fp8.cuh
+++ b/include/flashinfer/gemm/bmm_fp8.cuh
@@ -45,7 +45,7 @@
 #else
 #define FLASHINFER_CUBLAS_CALL(EXPR)  \
   {                                   \
-    cudaError_t e = (EXPR);           \
+    cublasStatus_t e = (EXPR);        \
     if (e != CUBLAS_STATUS_SUCCESS) { \
       return e;                       \
     }                                 \