File size: 4,221 Bytes
fd6abd3 | 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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | // ============================================================================
// flash_attention.cpp - PyTorch C++ Extension
// ============================================================================
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <vector>
#include <cassert>
void launch_flash_attention_fwd(
const half* Q, const half* K, const half* V, half* O,
int batch_size, int num_heads, int N, int d, float scale
);
void launch_flash_attention_bwd(
const half* Q, const half* K, const half* V,
const half* O, const half* dO,
const float* L, const float* M,
half* dQ, half* dK, half* dV,
int batch_size, int num_heads, int N, int d, float scale
);
std::vector<at::Tensor> flash_attention_fwd(
at::Tensor Q, at::Tensor K, at::Tensor V, double scale_ = -1.0
) {
TORCH_CHECK(Q.is_cuda() && K.is_cuda() && V.is_cuda(), "Inputs must be CUDA tensors");
TORCH_CHECK(Q.dtype() == at::kHalf, "Inputs must be fp16");
int batch_size = Q.size(0);
int num_heads = Q.size(1);
int N = Q.size(2);
int d = Q.size(3);
float scale = (scale_ > 0) ? scale_ : 1.0f / std::sqrt(d);
auto O = at::empty_like(Q);
auto L = at::empty({batch_size, num_heads, N}, Q.options().dtype(at::kFloat));
auto M = at::empty({batch_size, num_heads, N}, Q.options().dtype(at::kFloat));
launch_flash_attention_fwd(
static_cast<const half*>(Q.data_ptr()),
static_cast<const half*>(K.data_ptr()),
static_cast<const half*>(V.data_ptr()),
static_cast<half*>(O.data_ptr()),
batch_size, num_heads, N, d, scale
);
return {O, L, M};
}
std::vector<at::Tensor> flash_attention_bwd(
at::Tensor Q, at::Tensor K, at::Tensor V,
at::Tensor O, at::Tensor dO,
at::Tensor L, at::Tensor M, double scale_ = -1.0
) {
TORCH_CHECK(Q.is_cuda() && dO.is_cuda(), "Inputs must be CUDA tensors");
int batch_size = Q.size(0);
int num_heads = Q.size(1);
int N = Q.size(2);
int d = Q.size(3);
float scale = (scale_ > 0) ? scale_ : 1.0f / std::sqrt(d);
auto dQ = at::empty_like(Q);
auto dK = at::empty_like(K);
auto dV = at::empty_like(V);
launch_flash_attention_bwd(
static_cast<const half*>(Q.data_ptr()),
static_cast<const half*>(K.data_ptr()),
static_cast<const half*>(V.data_ptr()),
static_cast<const half*>(O.data_ptr()),
static_cast<const half*>(dO.data_ptr()),
static_cast<const float*>(L.data_ptr()),
static_cast<const float*>(M.data_ptr()),
static_cast<half*>(dQ.data_ptr()),
static_cast<half*>(dK.data_ptr()),
static_cast<half*>(dV.data_ptr()),
batch_size, num_heads, N, d, scale
);
return {dQ, dK, dV};
}
class FlashAttentionFunction : public torch::autograd::Function<FlashAttentionFunction> {
public:
static torch::autograd::variable_list forward(
torch::autograd::AutogradContext* ctx,
torch::Tensor Q, torch::Tensor K, torch::Tensor V, double scale = -1.0
) {
auto outputs = flash_attention_fwd(Q, K, V, scale);
ctx->save_for_backward({Q, K, V, outputs[0], outputs[1], outputs[2]});
ctx->saved_data["scale"] = scale;
return {outputs[0]};
}
static torch::autograd::variable_list backward(
torch::autograd::AutogradContext* ctx,
torch::autograd::variable_list grad_outputs
) {
auto saved = ctx->get_saved_variables();
auto grads = flash_attention_bwd(
saved[0], saved[1], saved[2], saved[3],
grad_outputs[0], saved[4], saved[5],
ctx->saved_data["scale"].toDouble()
);
return {grads[0], grads[1], grads[2], torch::Tensor()};
}
};
torch::Tensor flash_attention(torch::Tensor Q, torch::Tensor K, torch::Tensor V, double scale = -1.0) {
return FlashAttentionFunction::apply({Q, K, V, scale})[0];
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("flash_attention_fwd", &flash_attention_fwd, "FlashAttention Forward");
m.def("flash_attention_bwd", &flash_attention_bwd, "FlashAttention Backward");
m.def("flash_attention", &flash_attention, "FlashAttention (autograd)");
} |