1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
|
Fix use of undeclared identifier 'CHECK_NOSPARSE_CONTIGUOUS_CUDA' with USE='-flash'
Bug: https://github.com/pytorch/pytorch/issues/160826
--- a/aten/src/ATen/native/transformers/cuda/attention.cu
+++ b/aten/src/ATen/native/transformers/cuda/attention.cu
@@ -71,6 +71,7 @@
#include <ATen/native/transformers/cuda/sdp_utils.h>
#include <ATen/native/transformers/sdp_utils_cpp.h>
+#include <ATen/native/transformers/flash_api_common.h>
#ifdef USE_FLASH_ATTENTION
// FlashAttention Specific Imports
#include <ATen/native/transformers/cuda/flash_attn/flash_api.h>
--- a/aten/src/ATen/native/transformers/cuda/attention_backward.cu
+++ b/aten/src/ATen/native/transformers/cuda/attention_backward.cu
@@ -33,6 +33,7 @@
#include <ATen/ops/_scaled_dot_product_flash_attention_backward_native.h>
#endif
+#include <ATen/native/transformers/flash_api_common.h>
#ifdef USE_FLASH_ATTENTION
// FlashAttention Specific Imports
#include <ATen/native/transformers/cuda/flash_attn/flash_api.h>
--- /dev/null
+++ b/aten/src/ATen/native/transformers/flash_api_common.h
@@ -0,0 +1,28 @@
+#pragma once
+#include <cstdint>
+#include <limits>
+
+#include <ATen/core/Tensor.h>
+#include <c10/util/Exception.h>
+
+#define CHECK_NOSPARSE_CONTIGUOUS_CUDA(TENSOR) \
+ TORCH_CHECK(TENSOR.is_cuda(), #TENSOR " must be a CUDA tensor"); \
+ TORCH_CHECK(!TENSOR.is_sparse(), #TENSOR " must be a dense tensor"); \
+ TORCH_CHECK(TENSOR.is_contiguous());
+
+#define CHECK_NOSPARSE_LASTCONTIGUOUS_CUDA(TENSOR) \
+ TORCH_CHECK(TENSOR.is_cuda(), #TENSOR " must be a CUDA tensor"); \
+ TORCH_CHECK(!TENSOR.is_sparse(), #TENSOR " must be a dense tensor"); \
+ TORCH_CHECK( \
+ TENSOR.stride(-1) == 1, #TENSOR ": last dimension must be contiguous");
+
+#define CHECK_ALIGNED_PTR(PTR, ALIGNMENT) \
+ TORCH_CHECK( \
+ uint64_t(PTR) % ALIGNMENT == 0, #PTR " is not correctly aligned")
+
+#define ASSIGN_CHECK_OVERFLOW(A, B) \
+ { \
+ A = B; \
+ TORCH_CHECK( \
+ B < std::numeric_limits<decltype(A)>::max(), #B " overflows"); \
+ }
--- a/aten/src/ATen/native/transformers/hip/flash_attn/flash_api.h
+++ b/aten/src/ATen/native/transformers/hip/flash_attn/flash_api.h
@@ -4,28 +4,7 @@
#include <ATen/Context.h>
#include <ATen/core/Tensor.h>
#include <c10/util/Exception.h>
-
-#define CHECK_NOSPARSE_CONTIGUOUS_CUDA(TENSOR) \
- TORCH_CHECK(TENSOR.is_cuda(), #TENSOR " must be a CUDA tensor"); \
- TORCH_CHECK(!TENSOR.is_sparse(), #TENSOR " must be a dense tensor"); \
- TORCH_CHECK(TENSOR.is_contiguous());
-
-#define CHECK_NOSPARSE_LASTCONTIGUOUS_CUDA(TENSOR) \
- TORCH_CHECK(TENSOR.is_cuda(), #TENSOR " must be a CUDA tensor"); \
- TORCH_CHECK(!TENSOR.is_sparse(), #TENSOR " must be a dense tensor"); \
- TORCH_CHECK( \
- TENSOR.stride(-1) == 1, #TENSOR ": last dimension must be contiguous");
-
-#define CHECK_ALIGNED_PTR(PTR, ALIGNMENT) \
- TORCH_CHECK( \
- uint64_t(PTR) % ALIGNMENT == 0, #PTR " is not correctly aligned")
-
-#define ASSIGN_CHECK_OVERFLOW(A, B) \
- { \
- A = B; \
- TORCH_CHECK( \
- B < std::numeric_limits<decltype(A)>::max(), #B " overflows"); \
- }
+#include <ATen/native/transformers/flash_api_common.h>
namespace pytorch_flash {
|
