[CI] Don't include <ATen/cuda/CUDAGraphsUtils.cuh>

csrc/flash_attn/flash_api.cpp

@@ -8,7 +8,7 @@
 #include <c10/cuda/CUDAGuard.h>
 #include <c10/cuda/CUDAStream.h>
 #include <ATen/cuda/CUDAGeneratorImpl.h>  // For at::Generator and at::PhiloxCudaState
-#include <ATen/cuda/CUDAGraphsUtils.cuh>  // For at::cuda::philox::unpack
+#include "philox_unpack.cuh"  // For at::cuda::philox::unpack
 
 #include <cutlass/numeric_types.h>
 

csrc/flash_attn/src/flash_fwd_kernel.h

@@ -4,7 +4,7 @@
 
 #pragma once
 
-#include <ATen/cuda/CUDAGraphsUtils.cuh> // For at::cuda::philox::unpack
+#include "philox_unpack.cuh" // For at::cuda::philox::unpack
 
 #include <cute/tensor.hpp>
 

csrc/flash_attn/src/philox_unpack.cuh

@@ -0,0 +1,4 @@
+// This is purely so that it works with torch 2.1. For torch 2.2+ we can include ATen/cuda/PhiloxUtils.cuh
+
+#pragma once
+#include <ATen/cuda/detail/UnpackRaw.cuh>

flash_attn/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.7.1.post3"
+__version__ = "2.7.1.post4"
 
 from flash_attn.flash_attn_interface import (
     flash_attn_func,

setup.py

@@ -114,7 +114,7 @@ def check_if_rocm_home_none(global_option: str) -> None:
 
 
 def append_nvcc_threads(nvcc_extra_args):
-    nvcc_threads = os.getenv("NVCC_THREADS") or "4"
+    nvcc_threads = os.getenv("NVCC_THREADS") or "2"
     return nvcc_extra_args + ["--threads", nvcc_threads]