[CI] Change torch #include to make it work with torch 2.1 Philox

Dao-AILab · Dec 7, 2024 · e782d28 · e782d28
1 parent 073afd5
commit e782d28
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 3 deletions.
diff --git a/csrc/flash_attn/flash_api.cpp b/csrc/flash_attn/flash_api.cpp
@@ -7,7 +7,8 @@
 #include <torch/nn/functional.h>
 #include <c10/cuda/CUDAGuard.h>
 #include <c10/cuda/CUDAStream.h>
-#include <ATen/cuda/CUDAGeneratorImpl.h>
+#include <ATen/cuda/CUDAGeneratorImpl.h>  // For at::Generator and at::PhiloxCudaState
+#include <ATen/cuda/CUDAGraphsUtils.cuh>  // For at::cuda::philox::unpack
 
 #include <cutlass/numeric_types.h>
 

diff --git a/csrc/flash_attn/src/flash.h b/csrc/flash_attn/src/flash.h
@@ -7,7 +7,7 @@
 #include <cuda.h>
 #include <vector>
 
-#include <ATen/cuda/PhiloxUtils.cuh> // For at::cuda::philox::unpack
+#include <ATen/cuda/CUDAGeneratorImpl.h> // For at::Generator and at::PhiloxCudaState
 
 constexpr int TOTAL_DIM = 0;
 constexpr int H_DIM = 1;

diff --git a/csrc/flash_attn/src/flash_fwd_kernel.h b/csrc/flash_attn/src/flash_fwd_kernel.h
@@ -4,6 +4,8 @@
 
 #pragma once
 
+#include <ATen/cuda/CUDAGraphsUtils.cuh> // For at::cuda::philox::unpack
+
 #include <cute/tensor.hpp>
 
 #include <cutlass/cutlass.h>

diff --git a/flash_attn/__init__.py b/flash_attn/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "2.7.1.post2"
+__version__ = "2.7.1.post3"
 
 from flash_attn.flash_attn_interface import (
     flash_attn_func,