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

Author: tridao

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>
 

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;

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>

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,