Port gfx12 native attention to ABI3 fork

.gitignore

@@ -5,6 +5,7 @@ __pycache__/
 
 # C extensions
 *.so
+*.pyd
 
 # Distribution / packaging
 .Python
@@ -174,3 +175,16 @@ cython_debug/
 .pypirc
 
 /wheelhouse
+csrc/cp_async_hip.cuh
+csrc/fused/fused.hip
+csrc/fused/pybind_hip.cpp
+csrc/numeric_conversion_hip.cuh
+csrc/qattn/qk_int_sv_gfx12_native.hip
+csrc/qattn/qk_int_sv_gfx12_native_*.hip
+csrc/reduction_utils_hip.cuh
+evaluate/
+*.pkl
+*.vscode/
+bench_qk_int8_pv_fp8_cuda_head96.py
+*test.py
+*unmask.py

bench/bench_qk_int8_pv_fp16_cuda.py

@@ -1,7 +1,8 @@
 import torch
 from flash_attn.utils.benchmark import benchmark_forward
 
-import sageattention._qattn_sm80 as qattn
+import sageattention._qattn_sm80
+qattn = torch.ops.sageattention_qattn_sm80
 
 import argparse
 

bench/bench_qk_int8_pv_fp8_cuda.py

@@ -1,7 +1,8 @@
 import torch
 from flash_attn.utils.benchmark import benchmark_forward
 
-import sageattention._qattn_sm89 as qattn
+import sageattention._qattn_sm89
+qattn = torch.ops.sageattention_qattn_sm89
 
 import argparse
 

bench/bench_qk_int8_pv_fp8_cuda_sm90.py

@@ -1,6 +1,9 @@
 import torch
 from flash_attn.utils.benchmark import benchmark_forward
-import sageattention._qattn_sm90 as qattn
+
+import sageattention._qattn_sm90
+qattn = torch.ops.sageattention_qattn_sm90
+
 import argparse
 
 parser = argparse.ArgumentParser(description='Benchmark QK Int8 PV FP8 SM90')

csrc/cp_async.cuh

@@ -33,7 +33,7 @@ enum class PrefetchMode {
   kPrefetch     // Fetch additional data from global memory to L2
 };
 
-#if (__CUDACC_VER_MAJOR__ >= 11)
+#if !defined(__HIP_PLATFORM_AMD__) && (__CUDACC_VER_MAJOR__ >= 11)
 #if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800))
 #define CP_ASYNC_ENABLED
 #endif
@@ -138,4 +138,4 @@ __device__ __forceinline__ void pred_load_128b(T* smem_ptr, const T* gmem_ptr, b
 #endif
 }
 
-} // namespace cp_async
+} // namespace cp_async

csrc/dispatch_utils.h

@@ -15,10 +15,11 @@
  */
 
 #pragma once
-#include <torch/extension.h>
 #include <cstdint>
 #include <sstream>
 #include <stdexcept>
+#include <torch/headeronly/core/ScalarType.h>
+#include <torch/headeronly/util/Exception.h>
 
 #define DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, ...)              \
   if (head_dim == 64) {                                         \
@@ -73,16 +74,16 @@
   }
 
 #define DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(pytorch_dtype, c_type, ...)                \
-  if (pytorch_dtype == at::ScalarType::Half) {                                          \
+  if (pytorch_dtype == torch::headeronly::ScalarType::Half) {                           \
     using c_type = half;                                                                \
     __VA_ARGS__                                                                         \
-  } else if (pytorch_dtype == at::ScalarType::BFloat16) {                               \
+  } else if (pytorch_dtype == torch::headeronly::ScalarType::BFloat16) {                \
     using c_type = nv_bfloat16;                                                         \
     __VA_ARGS__                                                                         \
   } else {                                                                              \
     std::ostringstream oss;                                                             \
-    oss << __PRETTY_FUNCTION__ << " failed to dispatch data type " << pytorch_dtype;    \
-    TORCH_CHECK(false, oss.str());                                                      \
+    oss << __PRETTY_FUNCTION__ << " failed to dispatch data type " << #pytorch_dtype;   \
+    STD_TORCH_CHECK(false, oss.str());                                                  \
   }
 
 #define DISPATCH_BLOCK_SIZE(block_size, BLOCK_SIZE, ...)        \