fix for compile

csrc/flash_attn/cuComplex.h

+#pragma once
+
+#include <hip/hip_complex.h>
+
+using cuFloatComplex = hipFloatComplex;
+using cuDoubleComplex = hipDoubleComplex;
+
+#ifndef make_cuFloatComplex
+#define make_cuFloatComplex make_hipFloatComplex
+#endif
+
+#ifndef make_cuDoubleComplex
+#define make_cuDoubleComplex make_hipDoubleComplex
+#endif
+
+__host__ __device__ inline float cuCrealf(cuFloatComplex z) { return z.x; }
+__host__ __device__ inline float cuCimagf(cuFloatComplex z) { return z.y; }
+__host__ __device__ inline double cuCreal(cuDoubleComplex z) { return z.x; }
+__host__ __device__ inline double cuCimag(cuDoubleComplex z) { return z.y; }

csrc/flash_attn/cuda.h

+#pragma once
+
+#include <hip/hip_runtime.h>

csrc/flash_attn/cuda/std/cassert

+#pragma once
+
+#include <cassert>

csrc/flash_attn/cuda/std/utility

+#pragma once
+
+#include <utility>
+
+namespace cuda {
+namespace std {
+using ::std::swap;
+}
+}

csrc/flash_attn/cuda_bf16.h

+#pragma once
+
+#include <hip/hip_bfloat16.h>
+
+using __nv_bfloat16 = hip_bfloat16;
+
+struct __nv_bfloat16_raw {
+  uint16_t x;
+  __host__ __device__ __nv_bfloat16_raw() = default;
+  __host__ __device__ explicit __nv_bfloat16_raw(__nv_bfloat16 const& v) : x(v.data) {}
+};

csrc/flash_attn/cuda_fp16.h

+#pragma once
+
+#include <hip/hip_fp16.h>

csrc/flash_attn/cuda_runtime.h

+#pragma once
+
+#include <hip/hip_runtime.h>

csrc/flash_attn/cuda_runtime_api.h

+#pragma once
+
+#include <hip/hip_runtime_api.h>

csrc/flash_attn/driver_types.h

+#pragma once
+
+#include "/opt/dtk/hip/include/hip/driver_types.h"
+
+typedef hipStream_t cudaStream_t;
+typedef hipError_t cudaError_t;
+
+#ifndef cudaSuccess
+#define cudaSuccess hipSuccess
+#endif
+
+#ifndef cudaEnableDefault
+#define cudaEnableDefault 0
+#endif
+
+enum cudaDriverEntryPointQueryResult {
+  cudaDriverEntryPointSuccess = 0,
+};

csrc/flash_attn/vector_types.h

+#pragma once
+
+#include <hip/hip_vector_types.h>

flash_attn/version.py

+version = '2.8.3'
+git_hash = 'c2a1b31'
+git_branch = 'master'
+abi = 'abi1'
+dtk = '2604'
+torch_version = '2.9'
+hcu_version = '2.8.3+das.opt0.dtk2604'

setup.py

@@ -901,7 +901,7 @@ if not SKIP_CUDA_BUILD:
                         # "-mllvm -full-vectorize-slp=true",
                         f"-DFLASH_ATTENTION_BF16_TYPE={bf16_type}",
                         # "-DHG_ROCM",
-                        # "-DFLASHATTENTION_DISABLE_BACKWARD",
+                        "-DFLASHATTENTION_DISABLE_BACKWARD",
                         # "-DBWDTRANS",
                         "-DBWDSEPARATE",
                         # "-DFLASHATTENTION_DISABLE_SPLITKV",
@@ -918,7 +918,7 @@ if not SKIP_CUDA_BUILD:
                         # "--ptxas-options=-v",
                         # "--ptxas-options=-O2",
                         # "-lineinfo",
-                        # "-DFLASHATTENTION_DISABLE_BACKWARD",
+                        "-DFLASHATTENTION_DISABLE_BACKWARD",
                         # "-DFLASHATTENTION_DISABLE_DROPOUT",
                         # "-DFLASHATTENTION_DISABLE_ALIBI",
                         "-DFLASHATTENTION_DISABLE_SOFTCAP",
@@ -944,8 +944,13 @@ if not SKIP_CUDA_BUILD:
                 Path(this_dir) / "csrc" / "flash_attn",
                 Path(this_dir) / "csrc" / "flash_attn" / "src",
                 Path(this_dir) / "csrc" / "cutlass" / "include",
+                Path("/usr/local/lib/python3.10/dist-packages/tilelang/3rdparty/cutlass/include"),
+            ]
         )
     )
+    ext_modules[-1].sources = [
+        src for src in ext_modules[-1].sources if "/flash_bwd" not in src and "/flash_bwd_attnmask" not in src
+    ]
 
 
 def get_package_version():