Skip to content

GitLab

"ppocr/git@developer.sourcefind.cn:wangsen/paddle_dbnet.git" did not exist on "f8a3c3f9205d8cdaa7dd637a4005008138516723"
Commit 35d589fa authored by robotcator's avatar robotcator
Browse files

Merge branch 'main' of github.com:robotcator/flash-attention into workflow

parents 10d07459 52fb4b72
Hide whitespace changes
Inline Side-by-side
Showing with 391 additions and 52 deletions
csrc/flash_attn/fmha_api.cpp
View file @ 35d589fa
...@@ -28,6 +28,7 @@ ...@@ -28,6 +28,7 @@
#include <torch/extension.h> #include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h> #include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include "fmha.h" #include "fmha.h"
...@@ -45,9 +46,9 @@ void set_params_fprop(FMHA_fprop_params &params, ...@@ -45,9 +46,9 @@ void set_params_fprop(FMHA_fprop_params &params,
const at::Tensor q, const at::Tensor q,
const at::Tensor k, const at::Tensor k,
const at::Tensor v, const at::Tensor v,
at::Tensor out,
void *cu_seqlens_q_d, void *cu_seqlens_q_d,
void *cu_seqlens_k_d, void *cu_seqlens_k_d,
void *o_packed_d,
void *o_tmp_d, void *o_tmp_d,
void *s_d, void *s_d,
void *softmax_lse_d, void *softmax_lse_d,
...@@ -73,10 +74,12 @@ void set_params_fprop(FMHA_fprop_params &params, ...@@ -73,10 +74,12 @@ void set_params_fprop(FMHA_fprop_params &params,
params.q_head_stride_in_elts = q.stride(1); params.q_head_stride_in_elts = q.stride(1);
params.k_head_stride_in_elts = k.stride(1); params.k_head_stride_in_elts = k.stride(1);
params.v_head_stride_in_elts = v.stride(1); params.v_head_stride_in_elts = v.stride(1);
params.o_ptr = o_packed_d; params.o_ptr = out.data_ptr();
params.o_row_stride_in_elts = h * d; params.o_row_stride_in_elts = out.stride(0);
params.o_head_stride_in_elts = d; params.o_head_stride_in_elts = out.stride(1);
params.o_tmp_ptr = o_tmp_d; params.o_tmp_ptr = o_tmp_d;
params.o_tmp_row_stride_in_elts = h * d;
params.o_tmp_head_stride_in_elts = d;
params.cu_seqlens_q = static_cast<int *>(cu_seqlens_q_d); params.cu_seqlens_q = static_cast<int *>(cu_seqlens_q_d);
params.cu_seqlens_k = static_cast<int *>(cu_seqlens_k_d); params.cu_seqlens_k = static_cast<int *>(cu_seqlens_k_d);
...@@ -127,12 +130,12 @@ void set_params_dgrad(FMHA_dgrad_params &params, ...@@ -127,12 +130,12 @@ void set_params_dgrad(FMHA_dgrad_params &params,
const at::Tensor q, const at::Tensor q,
const at::Tensor k, const at::Tensor k,
const at::Tensor v, const at::Tensor v,
const at::Tensor out,
at::Tensor dq, at::Tensor dq,
at::Tensor dk, at::Tensor dk,
at::Tensor dv, at::Tensor dv,
void *cu_seqlens_q_d, void *cu_seqlens_q_d,
void *cu_seqlens_k_d, void *cu_seqlens_k_d,
void *o_packed_d,
void *dq_tmp_d, void *dq_tmp_d,
void *do_packed_d, void *do_packed_d,
void *softmax_lse_d, void *softmax_lse_d,
...@@ -143,10 +146,9 @@ void set_params_dgrad(FMHA_dgrad_params &params, ...@@ -143,10 +146,9 @@ void set_params_dgrad(FMHA_dgrad_params &params,
set_params_fprop(params, set_params_fprop(params,
b, seqlen_q, seqlen_k, h, d, b, seqlen_q, seqlen_k, h, d,
q, k, v, q, k, v, out,
cu_seqlens_q_d, cu_seqlens_q_d,
cu_seqlens_k_d, cu_seqlens_k_d,
o_packed_d,
dq_tmp_d, // Reusing the o_tmp_ptr variable to store dq_tmp dq_tmp_d, // Reusing the o_tmp_ptr variable to store dq_tmp
nullptr, nullptr,
softmax_lse_d, softmax_lse_d,
...@@ -174,6 +176,7 @@ std::vector<at::Tensor> ...@@ -174,6 +176,7 @@ std::vector<at::Tensor>
mha_fwd(const at::Tensor &q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i mha_fwd(const at::Tensor &q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
const at::Tensor &k, // total_k x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i const at::Tensor &k, // total_k x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
const at::Tensor &v, // total_k x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i const at::Tensor &v, // total_k x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
at::Tensor &out, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
const at::Tensor &cu_seqlens_q, // b+1 const at::Tensor &cu_seqlens_q, // b+1
const at::Tensor &cu_seqlens_k, // b+1 const at::Tensor &cu_seqlens_k, // b+1
const int max_seqlen_q_, const int max_seqlen_q_,
...@@ -198,18 +201,21 @@ mha_fwd(const at::Tensor &q, // total_q x num_heads x head_size, total_q ...@@ -198,18 +201,21 @@ mha_fwd(const at::Tensor &q, // total_q x num_heads x head_size, total_q
TORCH_CHECK(q_dtype == torch::kFloat16 || (is_sm8x && q_dtype == torch::kBFloat16)); TORCH_CHECK(q_dtype == torch::kFloat16 || (is_sm8x && q_dtype == torch::kBFloat16));
TORCH_CHECK(k.dtype() == q_dtype); TORCH_CHECK(k.dtype() == q_dtype);
TORCH_CHECK(v.dtype() == q_dtype); TORCH_CHECK(v.dtype() == q_dtype);
TORCH_CHECK(out.dtype() == q_dtype);
TORCH_CHECK(cu_seqlens_q.dtype() == torch::kInt32); TORCH_CHECK(cu_seqlens_q.dtype() == torch::kInt32);
TORCH_CHECK(cu_seqlens_k.dtype() == torch::kInt32); TORCH_CHECK(cu_seqlens_k.dtype() == torch::kInt32);
TORCH_CHECK(q.is_cuda()); TORCH_CHECK(q.is_cuda());
TORCH_CHECK(k.is_cuda()); TORCH_CHECK(k.is_cuda());
TORCH_CHECK(v.is_cuda()); TORCH_CHECK(v.is_cuda());
TORCH_CHECK(out.is_cuda());
TORCH_CHECK(cu_seqlens_q.is_cuda()); TORCH_CHECK(cu_seqlens_q.is_cuda());
TORCH_CHECK(cu_seqlens_k.is_cuda()); TORCH_CHECK(cu_seqlens_k.is_cuda());
TORCH_CHECK(q.stride(-1) == 1); TORCH_CHECK(q.stride(-1) == 1);
TORCH_CHECK(k.stride(-1) == 1); TORCH_CHECK(k.stride(-1) == 1);
TORCH_CHECK(v.stride(-1) == 1); TORCH_CHECK(v.stride(-1) == 1);
TORCH_CHECK(out.stride(-1) == 1);
TORCH_CHECK(cu_seqlens_k.is_contiguous()); TORCH_CHECK(cu_seqlens_k.is_contiguous());
TORCH_CHECK(cu_seqlens_k.is_contiguous()); TORCH_CHECK(cu_seqlens_k.is_contiguous());
...@@ -226,6 +232,7 @@ mha_fwd(const at::Tensor &q, // total_q x num_heads x head_size, total_q ...@@ -226,6 +232,7 @@ mha_fwd(const at::Tensor &q, // total_q x num_heads x head_size, total_q
CHECK_SHAPE(q, total_q, num_heads, head_size); CHECK_SHAPE(q, total_q, num_heads, head_size);
CHECK_SHAPE(k, total_k, num_heads, head_size); CHECK_SHAPE(k, total_k, num_heads, head_size);
CHECK_SHAPE(v, total_k, num_heads, head_size); CHECK_SHAPE(v, total_k, num_heads, head_size);
CHECK_SHAPE(out, total_q, num_heads, head_size);
CHECK_SHAPE(cu_seqlens_q, batch_size + 1); CHECK_SHAPE(cu_seqlens_q, batch_size + 1);
CHECK_SHAPE(cu_seqlens_k, batch_size + 1); CHECK_SHAPE(cu_seqlens_k, batch_size + 1);
...@@ -240,9 +247,12 @@ mha_fwd(const at::Tensor &q, // total_q x num_heads x head_size, total_q ...@@ -240,9 +247,12 @@ mha_fwd(const at::Tensor &q, // total_q x num_heads x head_size, total_q
int max_seqlen_q = ((max_seqlen_q_ + 16 - 1) / 16) * 16; int max_seqlen_q = ((max_seqlen_q_ + 16 - 1) / 16) * 16;
bool loop = max_seqlen_k > blocksize_c; bool loop = max_seqlen_k > blocksize_c;
// Otherwise the kernel will be launched from cuda:0 device
at::cuda::CUDAGuard device_guard{q.get_device()};
auto opts = q.options(); auto opts = q.options();
auto o = torch::empty({ total_q, num_heads, head_size }, opts); // auto o = torch::empty({ total_q, num_heads, head_size }, opts);
at::Tensor o_tmp; at::Tensor o_tmp;
if (loop) { o_tmp = torch::empty({total_q, num_heads, head_size}, opts.dtype(at::kFloat)); } if (loop) { o_tmp = torch::empty({total_q, num_heads, head_size}, opts.dtype(at::kFloat)); }
...@@ -254,7 +264,7 @@ mha_fwd(const at::Tensor &q, // total_q x num_heads x head_size, total_q ...@@ -254,7 +264,7 @@ mha_fwd(const at::Tensor &q, // total_q x num_heads x head_size, total_q
if (return_softmax) { s = torch::empty({ batch_size, num_heads, max_seqlen_q, max_seqlen_k }, opts); } if (return_softmax) { s = torch::empty({ batch_size, num_heads, max_seqlen_q, max_seqlen_k }, opts); }
if( zero_tensors ) { if( zero_tensors ) {
o.zero_(); out.zero_();
softmax_lse.fill_(-std::numeric_limits<float>::infinity()); softmax_lse.fill_(-std::numeric_limits<float>::infinity());
if (return_softmax) {s.zero_();} if (return_softmax) {s.zero_();}
} }
...@@ -268,10 +278,9 @@ mha_fwd(const at::Tensor &q, // total_q x num_heads x head_size, total_q ...@@ -268,10 +278,9 @@ mha_fwd(const at::Tensor &q, // total_q x num_heads x head_size, total_q
max_seqlen_k, max_seqlen_k,
num_heads, num_heads,
head_size, head_size,
q, k, v, q, k, v, out,
cu_seqlens_q.data_ptr(), cu_seqlens_q.data_ptr(),
cu_seqlens_k.data_ptr(), cu_seqlens_k.data_ptr(),
o.data_ptr(),
loop ? o_tmp.data_ptr() : nullptr, loop ? o_tmp.data_ptr() : nullptr,
return_softmax ? s.data_ptr() : nullptr, return_softmax ? s.data_ptr() : nullptr,
softmax_lse.data_ptr(), softmax_lse.data_ptr(),
...@@ -293,7 +302,7 @@ mha_fwd(const at::Tensor &q, // total_q x num_heads x head_size, total_q ...@@ -293,7 +302,7 @@ mha_fwd(const at::Tensor &q, // total_q x num_heads x head_size, total_q
run_fmha_fp16_sm80(launch_params, /*configure=*/false); run_fmha_fp16_sm80(launch_params, /*configure=*/false);
std::vector<at::Tensor> result = {o, softmax_lse}; std::vector<at::Tensor> result = {softmax_lse};
if (return_softmax) {result.push_back(s);} if (return_softmax) {result.push_back(s);}
return result; return result;
} }
...@@ -395,6 +404,9 @@ mha_bwd(const at::Tensor &dout, // total_q x num_heads, x head_size ...@@ -395,6 +404,9 @@ mha_bwd(const at::Tensor &dout, // total_q x num_heads, x head_size
int max_seqlen_q = ((max_seqlen_q_ + 16 - 1) / 16) * 16; int max_seqlen_q = ((max_seqlen_q_ + 16 - 1) / 16) * 16;
bool loop = max_seqlen_k > blocksize_c; bool loop = max_seqlen_k > blocksize_c;
// Otherwise the kernel will be launched from cuda:0 device
at::cuda::CUDAGuard device_guard{q.get_device()};
// It's possible the softmax_lse_ from the fwd has a different length since blocksize_c could be different. // It's possible the softmax_lse_ from the fwd has a different length since blocksize_c could be different.
auto softmax_lse = softmax_lse_.index({torch::indexing::Slice(), torch::indexing::Slice(), torch::indexing::Slice(torch::indexing::None, max_seqlen_q)}).contiguous(); auto softmax_lse = softmax_lse_.index({torch::indexing::Slice(), torch::indexing::Slice(), torch::indexing::Slice(torch::indexing::None, max_seqlen_q)}).contiguous();
...@@ -418,11 +430,10 @@ mha_bwd(const at::Tensor &dout, // total_q x num_heads, x head_size ...@@ -418,11 +430,10 @@ mha_bwd(const at::Tensor &dout, // total_q x num_heads, x head_size
max_seqlen_k, max_seqlen_k,
num_heads, num_heads,
head_size, head_size,
q, k, v, q, k, v, out,
dq, dk, dv, dq, dk, dv,
cu_seqlens_q.data_ptr(), cu_seqlens_q.data_ptr(),
cu_seqlens_k.data_ptr(), cu_seqlens_k.data_ptr(),
out.data_ptr(),
loop ? dq_tmp.data_ptr() : nullptr, loop ? dq_tmp.data_ptr() : nullptr,
dout.data_ptr(), dout.data_ptr(),
softmax_lse.data_ptr(), softmax_lse.data_ptr(),
...@@ -541,10 +552,9 @@ mha_fwd_block(const at::Tensor &q, // total_q x num_heads x head_size, t ...@@ -541,10 +552,9 @@ mha_fwd_block(const at::Tensor &q, // total_q x num_heads x head_size, t
max_seqlen_k, max_seqlen_k,
num_heads, num_heads,
head_size, head_size,
q, k, v, q, k, v, o,
cu_seqlens_q.data_ptr(), cu_seqlens_q.data_ptr(),
cu_seqlens_k.data_ptr(), cu_seqlens_k.data_ptr(),
o.data_ptr(),
loop ? o_tmp.data_ptr() : nullptr, loop ? o_tmp.data_ptr() : nullptr,
return_softmax ? s.data_ptr() : nullptr, return_softmax ? s.data_ptr() : nullptr,
softmax_lse.data_ptr(), softmax_lse.data_ptr(),
...@@ -686,11 +696,10 @@ mha_bwd_block(const at::Tensor &dout, // total x num_heads, x head_size ...@@ -686,11 +696,10 @@ mha_bwd_block(const at::Tensor &dout, // total x num_heads, x head_size
max_seqlen_k, max_seqlen_k,
num_heads, num_heads,
head_size, head_size,
q, k, v, q, k, v, out,
dq, dk, dv, dq, dk, dv,
cu_seqlens_q.data_ptr(), cu_seqlens_q.data_ptr(),
cu_seqlens_k.data_ptr(), cu_seqlens_k.data_ptr(),
out.data_ptr(),
loop ? dq_tmp.data_ptr() : nullptr, loop ? dq_tmp.data_ptr() : nullptr,
dout.data_ptr(), dout.data_ptr(),
softmax_lse.data_ptr(), softmax_lse.data_ptr(),
......
csrc/flash_attn/src/fmha.h
View file @ 35d589fa
...@@ -81,6 +81,8 @@ struct FMHA_fprop_params : public Qkv_params { ...@@ -81,6 +81,8 @@ struct FMHA_fprop_params : public Qkv_params {
// size_t o_stride_in_bytes; // size_t o_stride_in_bytes;
uint32_t o_row_stride_in_elts; uint32_t o_row_stride_in_elts;
uint32_t o_head_stride_in_elts; uint32_t o_head_stride_in_elts;
uint32_t o_tmp_row_stride_in_elts;
uint32_t o_tmp_head_stride_in_elts;
// The pointer to the O_tmp matrix, which holds O intermediate value during // The pointer to the O_tmp matrix, which holds O intermediate value during
// the loop; // the loop;
......
csrc/flash_attn/src/fmha_dgrad_fp16_kernel_loop.sm80.cu
View file @ 35d589fa
...@@ -2,6 +2,7 @@ ...@@ -2,6 +2,7 @@
*/ */
#include "static_switch.h" #include "static_switch.h"
#include "fp16_switch.h"
#include "fmha.h" #include "fmha.h"
#include "fmha_dgrad_kernel_1xN_loop.h" #include "fmha_dgrad_kernel_1xN_loop.h"
...@@ -52,8 +53,8 @@ void run_fmha_dgrad_fp16_sm80_loop_(const FMHA_dgrad_params &params, cudaStream_ ...@@ -52,8 +53,8 @@ void run_fmha_dgrad_fp16_sm80_loop_(const FMHA_dgrad_params &params, cudaStream_
} }
void run_fmha_dgrad_fp16_sm80(const FMHA_dgrad_params &params, cudaStream_t stream) { void run_fmha_dgrad_fp16_sm80(const FMHA_dgrad_params &params, cudaStream_t stream) {
BOOL_SWITCH(params.is_bf16, IsBf16Const, [&] { // work around for MSVC issue
using elem_type = std::conditional<IsBf16Const, __nv_bfloat16, __half>::type; FP16_SWITCH(params.is_bf16, [&] {
auto dprops = at::cuda::getCurrentDeviceProperties(); auto dprops = at::cuda::getCurrentDeviceProperties();
if (params.d == 16) { if (params.d == 16) {
if( params.seqlen_k == 128 ) { if( params.seqlen_k == 128 ) {
......
csrc/flash_attn/src/fmha_fprop_fp16_kernel.sm80.cu
View file @ 35d589fa
...@@ -29,6 +29,7 @@ ...@@ -29,6 +29,7 @@
#include <cuda_bf16.h> #include <cuda_bf16.h>
#include "static_switch.h" #include "static_switch.h"
#include "fp16_switch.h"
#include "fmha.h" #include "fmha.h"
#include "fmha_fprop_kernel_1xN.h" #include "fmha_fprop_kernel_1xN.h"
...@@ -83,8 +84,7 @@ void run_fmha_fp16_sm80_loop_(Launch_params<FMHA_fprop_params> &launch_params, ...@@ -83,8 +84,7 @@ void run_fmha_fp16_sm80_loop_(Launch_params<FMHA_fprop_params> &launch_params,
void run_fmha_fp16_sm80(Launch_params<FMHA_fprop_params> &launch_params, void run_fmha_fp16_sm80(Launch_params<FMHA_fprop_params> &launch_params,
const bool configure) { const bool configure) {
BOOL_SWITCH(launch_params.params.is_bf16, IsBf16Const, [&] { FP16_SWITCH(launch_params.params.is_bf16, [&] {
using elem_type = std::conditional<IsBf16Const, __nv_bfloat16, __half>::type;
auto dprops = at::cuda::getCurrentDeviceProperties(); auto dprops = at::cuda::getCurrentDeviceProperties();
if (launch_params.params.d == 16) { if (launch_params.params.d == 16) {
if( launch_params.params.seqlen_k == 128 ) { if( launch_params.params.seqlen_k == 128 ) {
......
csrc/flash_attn/src/fmha_fprop_kernel_1xN.h
View file @ 35d589fa
...@@ -259,7 +259,7 @@ inline __device__ void device_1xN_(const Params &params, const int bidb, const i ...@@ -259,7 +259,7 @@ inline __device__ void device_1xN_(const Params &params, const int bidb, const i
Gmem_tile_q gmem_q(params.q_ptr, params.q_row_stride_in_elts, params.q_head_stride_in_elts, binfo, tidx, true); Gmem_tile_q gmem_q(params.q_ptr, params.q_row_stride_in_elts, params.q_head_stride_in_elts, binfo, tidx, true);
// Allocate the global memory tile loader for O. // Allocate the global memory tile loader for O.
Gmem_tile_o gmem_o(params.o_ptr, params.o_row_stride_in_elts, params.o_head_stride_in_elts, binfo, tidx); Gmem_tile_o gmem_o(params.o_ptr, params.o_row_stride_in_elts, params.o_head_stride_in_elts, binfo, tidx);
Gmem_tile_o_tmp gmem_o_tmp(params.o_tmp_ptr, params.o_row_stride_in_elts, params.o_head_stride_in_elts, binfo, tidx); Gmem_tile_o_tmp gmem_o_tmp(params.o_tmp_ptr, params.o_tmp_row_stride_in_elts, params.o_tmp_head_stride_in_elts, binfo, tidx);
// Allocate the global memory tile loader for S. // Allocate the global memory tile loader for S.
Gmem_tile_s gmem_s(params, binfo, tidx); Gmem_tile_s gmem_s(params, binfo, tidx);
Gmem_softmax_sum gmem_softmax_lse(params.softmax_lse_ptr, params, tidx); Gmem_softmax_sum gmem_softmax_lse(params.softmax_lse_ptr, params, tidx);
......
csrc/flash_attn/src/fp16_switch.h 0 → 100644
View file @ 35d589fa
// Inspired by https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h
// and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h
// modified from static_switch.h
// because MSVC cannot handle std::conditional with constexpr variable
#pragma once
/// @param COND - a boolean expression to switch by
/// @param ... - code to execute for true and false
///
/// Usage:
/// ```
/// FP16_SWITCH(flag, [&] {
/// some_function(...);
/// });
/// ```
#define FP16_SWITCH(COND, ...) \
[&] { \
if (COND) { \
using elem_type = __nv_bfloat16; \
return __VA_ARGS__(); \
} else { \
using elem_type = __half; \
return __VA_ARGS__(); \
} \
}()
\ No newline at end of file
csrc/flash_attn/src/static_switch.h
View file @ 35d589fa
...@@ -22,4 +22,4 @@ ...@@ -22,4 +22,4 @@
constexpr bool CONST_NAME = false; \ constexpr bool CONST_NAME = false; \
return __VA_ARGS__(); \ return __VA_ARGS__(); \
} \ } \
}() }()
\ No newline at end of file
flash_attn/bert_padding.py
View file @ 35d589fa
# Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py # Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py
import numpy as np
import torch import torch
import torch.nn.functional as F import torch.nn.functional as F
...@@ -15,7 +13,7 @@ class IndexFirstAxis(torch.autograd.Function): ...@@ -15,7 +13,7 @@ class IndexFirstAxis(torch.autograd.Function):
ctx.save_for_backward(indices) ctx.save_for_backward(indices)
assert input.ndim >= 2 assert input.ndim >= 2
ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]
second_dim = np.prod(other_shape) second_dim = other_shape.numel()
# TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
# return input[indices] # return input[indices]
return torch.gather(rearrange(input, 'b ... -> b (...)'), 0, return torch.gather(rearrange(input, 'b ... -> b (...)'), 0,
...@@ -71,7 +69,7 @@ class IndexFirstAxisResidual(torch.autograd.Function): ...@@ -71,7 +69,7 @@ class IndexFirstAxisResidual(torch.autograd.Function):
ctx.save_for_backward(indices) ctx.save_for_backward(indices)
assert input.ndim >= 2 assert input.ndim >= 2
ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]
second_dim = np.prod(other_shape) second_dim = other_shape.numel()
# TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
output = input[indices] output = input[indices]
# We don't want to reshape input (b ... -> b (...)) since it could change the channel_last # We don't want to reshape input (b ... -> b (...)) since it could change the channel_last
......
flash_attn/flash_attn_interface.py
View file @ 35d589fa
import torch import torch
import torch.nn as nn import torch.nn as nn
import torch.nn.functional as F
import flash_attn_cuda import flash_attn_cuda
...@@ -14,11 +15,11 @@ def _get_block_size(device, head_dim, is_dropout): ...@@ -14,11 +15,11 @@ def _get_block_size(device, head_dim, is_dropout):
return 256 if (torch.cuda.get_device_capability(device) == (8, 0) and not is_dropout) else 128 return 256 if (torch.cuda.get_device_capability(device) == (8, 0) and not is_dropout) else 128
def _flash_attn_forward(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p, def _flash_attn_forward(q, k, v, out, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
softmax_scale, causal, return_softmax): dropout_p, softmax_scale, causal, return_softmax, generator=None):
out, softmax_lse, *rest = flash_attn_cuda.fwd( softmax_lse, *rest = flash_attn_cuda.fwd(
q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p, softmax_scale, q, k, v, out, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p,
False, causal, return_softmax, None softmax_scale, False, causal, return_softmax, generator
) )
# if out.isnan().any() or softmax_lse.isnan().any(): # if out.isnan().any() or softmax_lse.isnan().any():
# breakpoint() # breakpoint()
...@@ -27,10 +28,11 @@ def _flash_attn_forward(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_s ...@@ -27,10 +28,11 @@ def _flash_attn_forward(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_s
def _flash_attn_backward(dout, q, k, v, out, softmax_lse, dq, dk, dv, cu_seqlens_q, cu_seqlens_k, def _flash_attn_backward(dout, q, k, v, out, softmax_lse, dq, dk, dv, cu_seqlens_q, cu_seqlens_k,
max_seqlen_q, max_seqlen_k, dropout_p, softmax_scale, causal): max_seqlen_q, max_seqlen_k, dropout_p, softmax_scale, causal,
generator=None):
softmax_d = flash_attn_cuda.bwd( softmax_d = flash_attn_cuda.bwd(
dout, q, k, v, out, softmax_lse, dq, dk, dv, cu_seqlens_q, cu_seqlens_k, dout, q, k, v, out, softmax_lse, dq, dk, dv, cu_seqlens_q, cu_seqlens_k,
max_seqlen_q, max_seqlen_k, dropout_p, softmax_scale, False, causal, None) max_seqlen_q, max_seqlen_k, dropout_p, softmax_scale, False, causal, generator)
# if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any(): # if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any():
# breakpoint() # breakpoint()
return dq, dk, dv, softmax_d return dq, dk, dv, softmax_d
...@@ -45,8 +47,9 @@ class FlashAttnQKVPackedFunc(torch.autograd.Function): ...@@ -45,8 +47,9 @@ class FlashAttnQKVPackedFunc(torch.autograd.Function):
if softmax_scale is None: if softmax_scale is None:
softmax_scale = qkv.shape[-1] ** (-0.5) softmax_scale = qkv.shape[-1] ** (-0.5)
out, softmax_lse, S_dmask = _flash_attn_forward( out, softmax_lse, S_dmask = _flash_attn_forward(
qkv[:, 0], qkv[:, 1], qkv[:, 2], cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, qkv[:, 0], qkv[:, 1], qkv[:, 2], torch.empty_like(qkv[:, 0]), cu_seqlens, cu_seqlens,
dropout_p, softmax_scale, causal=causal, return_softmax=return_softmax max_seqlen, max_seqlen, dropout_p, softmax_scale, causal=causal,
return_softmax=return_softmax
) )
ctx.save_for_backward(qkv, out, softmax_lse, cu_seqlens, rng_state) ctx.save_for_backward(qkv, out, softmax_lse, cu_seqlens, rng_state)
ctx.dropout_p = dropout_p ctx.dropout_p = dropout_p
...@@ -82,8 +85,8 @@ class FlashAttnKVPackedFunc(torch.autograd.Function): ...@@ -82,8 +85,8 @@ class FlashAttnKVPackedFunc(torch.autograd.Function):
if softmax_scale is None: if softmax_scale is None:
softmax_scale = q.shape[-1] ** (-0.5) softmax_scale = q.shape[-1] ** (-0.5)
out, softmax_lse, S_dmask = _flash_attn_forward( out, softmax_lse, S_dmask = _flash_attn_forward(
q, kv[:, 0], kv[:, 1], cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, q, kv[:, 0], kv[:, 1], torch.empty_like(q), cu_seqlens_q, cu_seqlens_k, max_seqlen_q,
dropout_p, softmax_scale, causal=causal, return_softmax=return_softmax max_seqlen_k, dropout_p, softmax_scale, causal=causal, return_softmax=return_softmax
) )
ctx.save_for_backward(q, kv, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state) ctx.save_for_backward(q, kv, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state)
ctx.dropout_p = dropout_p ctx.dropout_p = dropout_p
...@@ -121,7 +124,7 @@ class FlashAttnFunc(torch.autograd.Function): ...@@ -121,7 +124,7 @@ class FlashAttnFunc(torch.autograd.Function):
if softmax_scale is None: if softmax_scale is None:
softmax_scale = q.shape[-1] ** (-0.5) softmax_scale = q.shape[-1] ** (-0.5)
out, softmax_lse, S_dmask = _flash_attn_forward( out, softmax_lse, S_dmask = _flash_attn_forward(
q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, q, k, v, torch.empty_like(q), cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
dropout_p, softmax_scale, causal=causal, return_softmax=return_softmax dropout_p, softmax_scale, causal=causal, return_softmax=return_softmax
) )
ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state) ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state)
...@@ -148,6 +151,85 @@ class FlashAttnFunc(torch.autograd.Function): ...@@ -148,6 +151,85 @@ class FlashAttnFunc(torch.autograd.Function):
return dq, dk, dv, None, None, None, None, None, None, None, None return dq, dk, dv, None, None, None, None, None, None, None, None
class FlashAttnQKVPackedSplitFunc(torch.autograd.Function):
@staticmethod
def forward(ctx, qkv, cu_seqlens, max_seqlen0, max_seqlen1, batch_size0, dropout_p,
softmax_scale, causal, return_softmax):
# Save rng_state because the backward pass will regenerate the dropout mask
if dropout_p > 0:
rng_state0 = torch.cuda.get_rng_state()
generator1 = torch.Generator(device='cuda')
rng_state1 = generator1.get_state()
else:
rng_state0, generator1, rng_state1 = None, None, None
if softmax_scale is None:
softmax_scale = qkv.shape[-1] ** (-0.5)
out = torch.empty_like(qkv[:, 0])
_, softmax_lse0, S_dmask0 = _flash_attn_forward(
qkv[:, 0], qkv[:, 1], qkv[:, 2], out, cu_seqlens[:batch_size0 + 1],
cu_seqlens[:batch_size0 + 1], max_seqlen0, max_seqlen0, dropout_p, softmax_scale,
causal=causal, return_softmax=return_softmax
)
s = torch.cuda.Stream()
with torch.cuda.stream(s):
_, softmax_lse1, S_dmask1 = _flash_attn_forward(
qkv[:, 0], qkv[:, 1], qkv[:, 2], out, cu_seqlens[batch_size0:],
cu_seqlens[batch_size0:], max_seqlen1, max_seqlen1, dropout_p, softmax_scale,
causal=causal, return_softmax=return_softmax, generator=generator1
)
torch.cuda.current_stream().wait_stream(s)
ctx.save_for_backward(qkv, out, softmax_lse0, softmax_lse1, cu_seqlens,
rng_state0, rng_state1)
ctx.dropout_p = dropout_p
ctx.max_seqlen0 = max_seqlen0
ctx.max_seqlen1 = max_seqlen1
ctx.batch_size0 = batch_size0
ctx.softmax_scale = softmax_scale
ctx.causal = causal
if not return_softmax:
return out
else:
max_seqlen_q = max(softmax_lse0.shape[2], softmax_lse1.shape[2])
max_seqlen_k = max(S_dmask0.shape[3], S_dmask1.shape[3])
softmax_lse = torch.cat([F.pad(softmax_lse0, (0, max_seqlen_q - softmax_lse0.shape[2])),
F.pad(softmax_lse1, (0, max_seqlen_q - softmax_lse1.shape[2]))],
dim=0)
return out, softmax_lse, S_dmask0, S_dmask1
@staticmethod
def backward(ctx, dout, *args):
qkv, out, softmax_lse0, softmax_lse1, cu_seqlens, rng_state0, rng_state1 = ctx.saved_tensors
batch_size0 = ctx.batch_size0
if rng_state0 is not None:
cur_rng_state = torch.cuda.get_rng_state()
torch.cuda.set_rng_state(rng_state0)
if rng_state1 is not None:
generator1 = torch.Generator(device='cuda')
generator1.set_state(rng_state1)
else:
generator1 = None
dqkv = torch.empty_like(qkv)
_flash_attn_backward(
dout, qkv[:, 0], qkv[:, 1], qkv[:, 2], out, softmax_lse0,
dqkv[:, 0], dqkv[:, 1], dqkv[:, 2], cu_seqlens[:batch_size0 + 1],
cu_seqlens[:batch_size0 + 1], ctx.max_seqlen0, ctx.max_seqlen0, ctx.dropout_p,
ctx.softmax_scale, ctx.causal
)
s = torch.cuda.Stream()
with torch.cuda.stream(s):
_flash_attn_backward(
dout, qkv[:, 0], qkv[:, 1], qkv[:, 2], out, softmax_lse1,
dqkv[:, 0], dqkv[:, 1], dqkv[:, 2], cu_seqlens[batch_size0:],
cu_seqlens[batch_size0:], ctx.max_seqlen1, ctx.max_seqlen1, ctx.dropout_p,
ctx.softmax_scale, ctx.causal, generator=generator1
)
torch.cuda.current_stream().wait_stream(s)
if rng_state0 is not None:
torch.cuda.set_rng_state(cur_rng_state)
return dqkv, None, None, None, None, None, None, None, None
def flash_attn_unpadded_qkvpacked_func(qkv, cu_seqlens, max_seqlen, dropout_p, softmax_scale=None, def flash_attn_unpadded_qkvpacked_func(qkv, cu_seqlens, max_seqlen, dropout_p, softmax_scale=None,
causal=False, return_attn_probs=False): causal=False, return_attn_probs=False):
"""dropout_p should be set to 0.0 during evaluation """dropout_p should be set to 0.0 during evaluation
...@@ -243,6 +325,43 @@ def flash_attn_unpadded_func(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, ...@@ -243,6 +325,43 @@ def flash_attn_unpadded_func(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q,
dropout_p, softmax_scale, causal, return_attn_probs) dropout_p, softmax_scale, causal, return_attn_probs)
def flash_attn_unpadded_qkvpacked_split_func(
qkv, cu_seqlens, max_seqlen0, max_seqlen1, batch_size0, dropout_p, softmax_scale=None,
causal=False, return_attn_probs=False):
"""
Split attention into 2 kernels running on 2 separate streams for performance reason:
e.g., if the batch has some sequences of length <= 128 and some > 128, it might be faster to
have one kernel dealing with seqlen <= 128 and one kernel for seqlen > 128.
dropout_p should be set to 0.0 during evaluation.
Arguments:
qkv: (total, 3, nheads, headdim), where total = total number of tokens in the batch.
cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
of the sequences in the batch, used to index into qkv.
max_seqlen0: int. Maximum sequence length in 1st part of the batch.
max_seqlen1: int. Maximum sequence length in 2nd part of the batch.
batch_size0: int. Number of sequences in the 1st part of the batch.
dropout_p: float. Dropout probability.
softmax_scale: float. The scaling of QK^T before applying softmax.
Default to 1 / sqrt(headdim).
causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
return_attn_probs: bool. Whether to return the attention probabilities. This option is for
testing only. The returned probabilities are not guaranteed to be correct
(they might not have the right scaling).
Return:
out: (total, nheads, headdim).
softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The
logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
normalization factor).
S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
The output of softmax (possibly with different scaling). It also encodes the dropout
pattern (negative means that location was dropped, nonnegative means it was kept).
"""
return FlashAttnQKVPackedSplitFunc.apply(qkv, cu_seqlens, max_seqlen0, max_seqlen1, batch_size0,
dropout_p, softmax_scale, causal, return_attn_probs)
def flash_attn_func(qkv, cu_seqlens, dropout_p, max_s, softmax_scale=None, causal=False, def flash_attn_func(qkv, cu_seqlens, dropout_p, max_s, softmax_scale=None, causal=False,
return_attn_probs=False): return_attn_probs=False):
"""For backward-compatibility only, will remove soon. """For backward-compatibility only, will remove soon.
......
setup.py
View file @ 35d589fa
...@@ -125,10 +125,11 @@ ext_modules.append( ...@@ -125,10 +125,11 @@ ext_modules.append(
"csrc/flash_attn/src/fmha_block_dgrad_fp16_kernel_loop.sm80.cu", "csrc/flash_attn/src/fmha_block_dgrad_fp16_kernel_loop.sm80.cu",
], ],
extra_compile_args={ extra_compile_args={
"cxx": ["-O3"] + generator_flag, "cxx": ["-O3", "-std=c++17"] + generator_flag,
"nvcc": append_nvcc_threads( "nvcc": append_nvcc_threads(
[ [
"-O3", "-O3",
"-std=c++17",
"-U__CUDA_NO_HALF_OPERATORS__", "-U__CUDA_NO_HALF_OPERATORS__",
"-U__CUDA_NO_HALF_CONVERSIONS__", "-U__CUDA_NO_HALF_CONVERSIONS__",
"--expt-relaxed-constexpr", "--expt-relaxed-constexpr",
......
tests/test_flash_attn.py
View file @ 35d589fa
...@@ -8,20 +8,28 @@ import pytest ...@@ -8,20 +8,28 @@ import pytest
from einops import rearrange, repeat from einops import rearrange, repeat
from flash_attn.flash_attn_interface import flash_attn_func, flash_attn_unpadded_qkvpacked_func, _get_block_size, flash_attn_unpadded_kvpacked_func, flash_attn_unpadded_func from flash_attn.flash_attn_interface import flash_attn_func, flash_attn_unpadded_qkvpacked_func, _get_block_size, flash_attn_unpadded_kvpacked_func, flash_attn_unpadded_func
from flash_attn.flash_attn_interface import flash_attn_unpadded_qkvpacked_split_func
from flash_attn.bert_padding import unpad_input, pad_input, index_first_axis from flash_attn.bert_padding import unpad_input, pad_input, index_first_axis
is_sm75 = torch.cuda.get_device_capability('cuda') == (7, 5) is_sm75 = torch.cuda.get_device_capability('cuda') == (7, 5)
is_sm80 = torch.cuda.get_device_capability('cuda') == (8, 0)
def generate_random_padding_mask(max_seqlen, batch_size, device, mode='random'): def generate_random_padding_mask(max_seqlen, batch_size, device, mode='random'):
assert mode in ['full', 'random', 'third'] assert mode in ['full', 'random', 'third', 'split']
if mode == 'full': if mode == 'full':
lengths = torch.full((batch_size, 1), max_seqlen, device=device, dtype=torch.int32) lengths = torch.full((batch_size, 1), max_seqlen, device=device, dtype=torch.int32)
elif mode == 'random': elif mode == 'random':
lengths = torch.randint(max(1, max_seqlen - 20), max_seqlen, (batch_size, 1), device=device) lengths = torch.randint(max(1, max_seqlen - 20), max_seqlen + 1, (batch_size, 1), device=device)
elif mode == 'third': elif mode == 'third':
lengths = torch.randint(max_seqlen // 3, max_seqlen, (batch_size, 1), device=device) lengths = torch.randint(max_seqlen // 3, max_seqlen + 1, (batch_size, 1), device=device)
elif mode == 'split':
lengths0 = torch.randint(min(128, max_seqlen), max_seqlen + 1,
(batch_size // 4 * 3, 1), device=device)
lengths1 = torch.randint(min(max(1, max_seqlen - 20), 128), min(max_seqlen, 128) + 1,
(batch_size - batch_size // 4 * 3, 1), device=device)
lengths = torch.cat([lengths0, lengths1], dim=0)
padding_mask = repeat(torch.arange(max_seqlen, device=device), 's -> b s', b=batch_size) < lengths padding_mask = repeat(torch.arange(max_seqlen, device=device), 's -> b s', b=batch_size) < lengths
return padding_mask return padding_mask
...@@ -331,6 +339,7 @@ def get_dropout_fraction(dropout_mask, query_padding_mask=None, key_padding_mask ...@@ -331,6 +339,7 @@ def get_dropout_fraction(dropout_mask, query_padding_mask=None, key_padding_mask
@pytest.mark.parametrize('dtype', ([torch.float16] if is_sm75 else [torch.float16, torch.bfloat16])) @pytest.mark.parametrize('dtype', ([torch.float16] if is_sm75 else [torch.float16, torch.bfloat16]))
# @pytest.mark.parametrize('dtype', [torch.float16]) # @pytest.mark.parametrize('dtype', [torch.float16])
@pytest.mark.parametrize('causal', [False, True]) @pytest.mark.parametrize('causal', [False, True])
# @pytest.mark.parametrize('causal', [False])
@pytest.mark.parametrize('d', [128, 64, 32, 16]) @pytest.mark.parametrize('d', [128, 64, 32, 16])
# @pytest.mark.parametrize('d', [64]) # @pytest.mark.parametrize('d', [64])
@pytest.mark.parametrize('seqlen', [97, 128, 200, 256, 257, 384, 512, 768, 1024, 1025, 2048]) @pytest.mark.parametrize('seqlen', [97, 128, 200, 256, 257, 384, 512, 768, 1024, 1025, 2048])
...@@ -385,7 +394,7 @@ def test_flash_attn_unpadded_qkvpacked(seqlen, d, dropout_p, causal, dtype): ...@@ -385,7 +394,7 @@ def test_flash_attn_unpadded_qkvpacked(seqlen, d, dropout_p, causal, dtype):
print(f'Attention max diff: {(attn - attn_ref).abs().max().item()}') print(f'Attention max diff: {(attn - attn_ref).abs().max().item()}')
print(f'Attention Pytorch max diff: {(attn_pt - attn_ref).abs().max().item()}') print(f'Attention Pytorch max diff: {(attn_pt - attn_ref).abs().max().item()}')
if not (is_sm75 and d == 128): if is_sm80 or d < 128: # Only run backward for d=128 on A100
g = torch.randn_like(output) g = torch.randn_like(output)
dqkv_unpad, = torch.autograd.grad(output, qkv_unpad, g) dqkv_unpad, = torch.autograd.grad(output, qkv_unpad, g)
dqkv = dqkv_pad_fn(dqkv_unpad) dqkv = dqkv_pad_fn(dqkv_unpad)
...@@ -411,7 +420,7 @@ def test_flash_attn_unpadded_qkvpacked(seqlen, d, dropout_p, causal, dtype): ...@@ -411,7 +420,7 @@ def test_flash_attn_unpadded_qkvpacked(seqlen, d, dropout_p, causal, dtype):
else: else:
assert 0.99 <= dropout_fraction / dropout_p <= 1.01 assert 0.99 <= dropout_fraction / dropout_p <= 1.01
if not (is_sm75 and d == 128): if is_sm80 or d < 128: # Only run backward for d=128 on A100
assert (dqkv - dqkv_ref).abs().max().item() <= 2 * (dqkv_pt - dqkv_ref).abs().max().item() assert (dqkv - dqkv_ref).abs().max().item() <= 2 * (dqkv_pt - dqkv_ref).abs().max().item()
# assert torch.allclose(dqkv, dqkv_ref, rtol=rtol, atol=atol) # assert torch.allclose(dqkv, dqkv_ref, rtol=rtol, atol=atol)
...@@ -476,7 +485,7 @@ def test_flash_attn_unpadded_kvpacked(seqlen, d, dropout_p, causal, dtype): ...@@ -476,7 +485,7 @@ def test_flash_attn_unpadded_kvpacked(seqlen, d, dropout_p, causal, dtype):
print(f'Attention max diff: {(attn - attn_ref).abs().max().item()}') print(f'Attention max diff: {(attn - attn_ref).abs().max().item()}')
print(f'Attention Pytorch max diff: {(attn_pt - attn_ref).abs().max().item()}') print(f'Attention Pytorch max diff: {(attn_pt - attn_ref).abs().max().item()}')
if not (is_sm75 and d == 128): if is_sm80 or d < 128: # Only run backward for d=128 on A100
g = torch.randn_like(output) g = torch.randn_like(output)
dq_unpad, dkv_unpad, = torch.autograd.grad(output, (q_unpad, kv_unpad), g) dq_unpad, dkv_unpad, = torch.autograd.grad(output, (q_unpad, kv_unpad), g)
dq = dq_pad_fn(dq_unpad) dq = dq_pad_fn(dq_unpad)
...@@ -501,7 +510,7 @@ def test_flash_attn_unpadded_kvpacked(seqlen, d, dropout_p, causal, dtype): ...@@ -501,7 +510,7 @@ def test_flash_attn_unpadded_kvpacked(seqlen, d, dropout_p, causal, dtype):
else: else:
assert 0.99 <= dropout_fraction / dropout_p <= 1.01 assert 0.99 <= dropout_fraction / dropout_p <= 1.01
if not (is_sm75 and d == 128): if is_sm80 or d < 128: # Only run backward for d=128 on A100
assert (dq - dq_ref).abs().max().item() <= 2 * (dq_pt - dq_ref).abs().max().item() assert (dq - dq_ref).abs().max().item() <= 2 * (dq_pt - dq_ref).abs().max().item()
assert (dkv - dkv_ref).abs().max().item() <= 2 * (dkv_pt - dkv_ref).abs().max().item() assert (dkv - dkv_ref).abs().max().item() <= 2 * (dkv_pt - dkv_ref).abs().max().item()
# assert torch.allclose(dq, dq_ref, rtol=rtol, atol=atol) # assert torch.allclose(dq, dq_ref, rtol=rtol, atol=atol)
...@@ -568,7 +577,7 @@ def test_flash_attn_unpadded(seqlen, d, dropout_p, causal, dtype): ...@@ -568,7 +577,7 @@ def test_flash_attn_unpadded(seqlen, d, dropout_p, causal, dtype):
print(f'Attention max diff: {(attn - attn_ref).abs().max().item()}') print(f'Attention max diff: {(attn - attn_ref).abs().max().item()}')
print(f'Attention Pytorch max diff: {(attn_pt - attn_ref).abs().max().item()}') print(f'Attention Pytorch max diff: {(attn_pt - attn_ref).abs().max().item()}')
if not (is_sm75 and d == 128): if is_sm80 or d < 128: # Only run backward for d=128 on A100
g = torch.randn_like(output) g = torch.randn_like(output)
dq_unpad, dk_unpad, dv_unpad, = torch.autograd.grad(output, (q_unpad, k_unpad, v_unpad), g) dq_unpad, dk_unpad, dv_unpad, = torch.autograd.grad(output, (q_unpad, k_unpad, v_unpad), g)
dq = dq_pad_fn(dq_unpad) dq = dq_pad_fn(dq_unpad)
...@@ -594,7 +603,7 @@ def test_flash_attn_unpadded(seqlen, d, dropout_p, causal, dtype): ...@@ -594,7 +603,7 @@ def test_flash_attn_unpadded(seqlen, d, dropout_p, causal, dtype):
else: else:
assert 0.99 <= dropout_fraction / dropout_p <= 1.01 assert 0.99 <= dropout_fraction / dropout_p <= 1.01
if not (is_sm75 and d == 128): if is_sm80 or d < 128: # Only run backward for d=128 on A100
assert (dq - dq_ref).abs().max().item() <= 2 * (dq_pt - dq_ref).abs().max().item() assert (dq - dq_ref).abs().max().item() <= 2 * (dq_pt - dq_ref).abs().max().item()
assert (dk - dk_ref).abs().max().item() <= 2 * (dk_pt - dk_ref).abs().max().item() assert (dk - dk_ref).abs().max().item() <= 2 * (dk_pt - dk_ref).abs().max().item()
assert (dv - dv_ref).abs().max().item() <= 2 * (dv_pt - dv_ref).abs().max().item() assert (dv - dv_ref).abs().max().item() <= 2 * (dv_pt - dv_ref).abs().max().item()
...@@ -603,6 +612,104 @@ def test_flash_attn_unpadded(seqlen, d, dropout_p, causal, dtype): ...@@ -603,6 +612,104 @@ def test_flash_attn_unpadded(seqlen, d, dropout_p, causal, dtype):
# assert torch.allclose(dv, dv_ref, rtol=rtol, atol=atol) # assert torch.allclose(dv, dv_ref, rtol=rtol, atol=atol)
@pytest.mark.parametrize('dtype', ([torch.float16] if is_sm75 else [torch.float16, torch.bfloat16]))
# @pytest.mark.parametrize('dtype', [torch.float16])
@pytest.mark.parametrize('causal', [False, True])
# @pytest.mark.parametrize('causal', [False])
@pytest.mark.parametrize('d', [128, 64, 32, 16])
# @pytest.mark.parametrize('d', [64])
@pytest.mark.parametrize('seqlen', [512])
@pytest.mark.parametrize('dropout_p', [0.0, 0.17])
# @pytest.mark.parametrize('dropout_p', [0.0])
def test_flash_attn_unpadded_qkvpacked_split(seqlen, d, dropout_p, causal, dtype):
if seqlen >= 2048 and torch.cuda.get_device_properties('cuda').total_memory <= 16 * 2**30:
pytest.skip() # Reference implementation OOM
device = 'cuda'
# if dtype == torch.float16:
# rtol, atol = (1e-3, 3e-4) if not causal else (1e-3, 1e-3)
# else: # torch.bfloat16
# rtol, atol = (3e-3, 3e-3) if not causal else (1e-3, 1e-3)
# set seed
torch.random.manual_seed(0)
batch_size = 32
nheads = 4
x = torch.randn(batch_size, seqlen, nheads * d, device=device, dtype=dtype, requires_grad=True)
Wqkv = torch.nn.Linear(nheads * d, 3 * nheads * d, device=device, dtype=dtype)
key_padding_mask = generate_random_padding_mask(seqlen, batch_size, device, mode='split')
batch_size0 = batch_size // 4 * 3 # this must match what's in generate_random_padding_mask
# key_padding_mask = generate_random_padding_mask(seqlen, batch_size, device, mode='full')
qkv_unpad, cu_seqlens, max_seqlen0, qkv, output_pad_fn, dqkv_pad_fn = generate_qkv(
x, Wqkv, nheads, key_padding_mask, key_padding_mask, qkvpacked=True
)
max_seqlen1 = 128
output_unpad, sm_lse, S_dmask0, S_dmask1 = flash_attn_unpadded_qkvpacked_split_func(
qkv_unpad, cu_seqlens, max_seqlen0, max_seqlen1, batch_size0, dropout_p,
return_attn_probs=True, causal=causal
)
output = output_pad_fn(output_unpad)
S_dmask0_converted = convert_flash_attn_S_to_softmax(
S_dmask0, key_padding_mask[:batch_size0], key_padding_mask[:batch_size0], d, dropout_p > 0.0, causal=causal
)
S_dmask1_converted = convert_flash_attn_S_to_softmax(
S_dmask1, key_padding_mask[batch_size0:, :max_seqlen1], key_padding_mask[batch_size0:, :max_seqlen1], d, dropout_p > 0.0, causal=causal
)
padding = (S_dmask0_converted.shape[-1] - S_dmask1_converted.shape[-1],
S_dmask0_converted.shape[-2] - S_dmask1_converted.shape[-2])
S_dmask_converted = torch.cat([S_dmask0_converted,
F.pad(S_dmask1_converted, (0, padding[0], 0, padding[1]))], dim=0)
dropout_mask = S_dmask_converted >= 0
attn_unnorm = S_dmask_converted.abs()
attn = normalize_flash_attn_S(attn_unnorm, qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2],
key_padding_mask, key_padding_mask, dropout_p > 0.0, causal=causal)
dropout_fraction = get_dropout_fraction(dropout_mask, key_padding_mask, key_padding_mask,
causal=causal).item()
output_ref, attn_ref = attention_qkvpacked_ref(qkv, key_padding_mask, dropout_p, dropout_mask,
causal=causal)
output_pt, attn_pt = attention_qkvpacked_ref(qkv, key_padding_mask, dropout_p, dropout_mask,
causal=causal, upcast=False, reorder_ops=True)
print(f'Actual dropout fraction: {dropout_fraction}')
print(f'Output max diff: {(output - output_ref).abs().max().item()}')
print(f'Output mean diff: {(output - output_ref).abs().mean().item()}')
print(f'Pytorch max diff: {(output_pt - output_ref).abs().max().item()}')
print(f'Pytorch mean diff: {(output_pt - output_ref).abs().mean().item()}')
print(f'Attention max diff: {(attn - attn_ref).abs().max().item()}')
print(f'Attention Pytorch max diff: {(attn_pt - attn_ref).abs().max().item()}')
if is_sm80 or d < 128: # Only run backward for d=128 on A100
g = torch.randn_like(output)
dqkv_unpad, = torch.autograd.grad(output, qkv_unpad, g)
dqkv = dqkv_pad_fn(dqkv_unpad)
dqkv_ref, = torch.autograd.grad(output_ref, qkv, g)
dqkv_pt, = torch.autograd.grad(output_pt, qkv, g)
print(f'dQ max diff: {(dqkv[:, :, 0] - dqkv_ref[:, :, 0]).abs().max().item()}')
print(f'dK max diff: {(dqkv[:, :, 1] - dqkv_ref[:, :, 1]).abs().max().item()}')
print(f'dV max diff: {(dqkv[:, :, 2] - dqkv_ref[:, :, 2]).abs().max().item()}')
print(f'dQKV mean diff: {(dqkv - dqkv_ref).abs().mean().item()}')
print(f'dQ Pytorch max diff: {(dqkv_pt[:, :, 0] - dqkv_ref[:, :, 0]).abs().max().item()}')
print(f'dK Pytorch max diff: {(dqkv_pt[:, :, 1] - dqkv_ref[:, :, 1]).abs().max().item()}')
print(f'dV Pytorch max diff: {(dqkv_pt[:, :, 2] - dqkv_ref[:, :, 2]).abs().max().item()}')
print(f'dQKV Pytorch mean diff: {(dqkv_pt - dqkv_ref).abs().mean().item()}')
# Check that FlashAttention's numerical error is at most twice the numerical error
# of a Pytorch implementation.
assert (output - output_ref).abs().max().item() <= 2 * (output_pt - output_ref).abs().max().item()
# assert torch.allclose(output, output_ref, rtol=rtol, atol=atol)
assert (attn - attn_ref).abs().max().item() <= 2 * (attn_pt - attn_ref).abs().max().item()
# assert torch.allclose(attn, attn_ref, rtol=rtol, atol=atol)
if dropout_p == 0.0:
assert dropout_mask.all()
else:
assert 0.99 <= dropout_fraction / dropout_p <= 1.01
if is_sm80 or d < 128: # Only run backward for d=128 on A100
assert (dqkv - dqkv_ref).abs().max().item() <= 2 * (dqkv_pt - dqkv_ref).abs().max().item()
# assert torch.allclose(dqkv, dqkv_ref, rtol=rtol, atol=atol)
@pytest.mark.parametrize('dtype', ([torch.float16] if is_sm75 else [torch.float16, torch.bfloat16])) @pytest.mark.parametrize('dtype', ([torch.float16] if is_sm75 else [torch.float16, torch.bfloat16]))
# @pytest.mark.parametrize('dtype', [torch.float16]) # @pytest.mark.parametrize('dtype', [torch.float16])
@pytest.mark.parametrize('causal', [False, True]) @pytest.mark.parametrize('causal', [False, True])
...@@ -640,7 +747,7 @@ def test_flash_attn_race_condition(seqlen, d, dropout_p, causal, dtype): ...@@ -640,7 +747,7 @@ def test_flash_attn_race_condition(seqlen, d, dropout_p, causal, dtype):
S_dmask_0, query_padding_mask, key_padding_mask, d, dropout_p > 0.0, causal=causal S_dmask_0, query_padding_mask, key_padding_mask, d, dropout_p > 0.0, causal=causal
) )
if not (is_sm75 and d == 128): if is_sm80 or d < 128: # Only run backward for d=128 on A100
g = torch.randn_like(output_unpad_0) g = torch.randn_like(output_unpad_0)
dq_unpad_0, dk_unpad_0, dv_unpad_0, = torch.autograd.grad(output_unpad_0, dq_unpad_0, dk_unpad_0, dv_unpad_0, = torch.autograd.grad(output_unpad_0,
(q_unpad, k_unpad, v_unpad), g) (q_unpad, k_unpad, v_unpad), g)
...@@ -659,9 +766,84 @@ def test_flash_attn_race_condition(seqlen, d, dropout_p, causal, dtype): ...@@ -659,9 +766,84 @@ def test_flash_attn_race_condition(seqlen, d, dropout_p, causal, dtype):
# assert torch.equal(sm_lse, sm_lse_0) # assert torch.equal(sm_lse, sm_lse_0)
assert torch.equal(S_dmask_converted, S_dmask_converted_0) assert torch.equal(S_dmask_converted, S_dmask_converted_0)
if not (is_sm75 and d == 128): if is_sm80 or d < 128: # Only run backward for d=128 on A100
dq_unpad, dk_unpad, dv_unpad, = torch.autograd.grad(output_unpad, dq_unpad, dk_unpad, dv_unpad, = torch.autograd.grad(output_unpad,
(q_unpad, k_unpad, v_unpad), g) (q_unpad, k_unpad, v_unpad), g)
assert torch.equal(dq_unpad, dq_unpad_0) assert torch.equal(dq_unpad, dq_unpad_0)
assert torch.equal(dk_unpad, dk_unpad_0) assert torch.equal(dk_unpad, dk_unpad_0)
assert torch.equal(dv_unpad, dv_unpad_0) assert torch.equal(dv_unpad, dv_unpad_0)
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason='requires multiple GPUs')
def test_flash_attn_multigpu():
seqlen = 256
d = 64
dropout_p = 0.0
causal = False
dtype = torch.float16
device = 'cuda:1'
torch.random.manual_seed(0)
batch_size = 32
nheads = 4
x = torch.randn(batch_size, seqlen, nheads * d, device=device, dtype=dtype, requires_grad=True)
Wqkv = torch.nn.Linear(nheads * d, 3 * nheads * d, device=device, dtype=dtype)
key_padding_mask = generate_random_padding_mask(seqlen, batch_size, device, mode='random')
# key_padding_mask = generate_random_padding_mask(seqlen, batch_size, device, mode='full')
qkv_unpad, cu_seqlens, max_seqlen, qkv, output_pad_fn, dqkv_pad_fn = generate_qkv(
x, Wqkv, nheads, key_padding_mask, key_padding_mask, qkvpacked=True
)
output_unpad, sm_lse, S_dmask = flash_attn_unpadded_qkvpacked_func(
qkv_unpad, cu_seqlens, max_seqlen, dropout_p, return_attn_probs=True, causal=causal
)
output = output_pad_fn(output_unpad)
S_dmask_converted = convert_flash_attn_S_to_softmax(
S_dmask, key_padding_mask, key_padding_mask, d, dropout_p > 0.0, causal=causal
)
dropout_mask = S_dmask_converted >= 0
attn_unnorm = S_dmask_converted.abs()
attn = normalize_flash_attn_S(attn_unnorm, qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2],
key_padding_mask, key_padding_mask, dropout_p > 0.0, causal=causal)
dropout_fraction = get_dropout_fraction(dropout_mask, key_padding_mask, key_padding_mask,
causal=causal).item()
output_ref, attn_ref = attention_qkvpacked_ref(qkv, key_padding_mask, dropout_p, dropout_mask,
causal=causal)
output_pt, attn_pt = attention_qkvpacked_ref(qkv, key_padding_mask, dropout_p, dropout_mask,
causal=causal, upcast=False, reorder_ops=True)
print(f'Actual dropout fraction: {dropout_fraction}')
print(f'Output max diff: {(output - output_ref).abs().max().item()}')
print(f'Output mean diff: {(output - output_ref).abs().mean().item()}')
print(f'Pytorch max diff: {(output_pt - output_ref).abs().max().item()}')
print(f'Pytorch mean diff: {(output_pt - output_ref).abs().mean().item()}')
print(f'Attention max diff: {(attn - attn_ref).abs().max().item()}')
print(f'Attention Pytorch max diff: {(attn_pt - attn_ref).abs().max().item()}')
g = torch.randn_like(output)
dqkv_unpad, = torch.autograd.grad(output, qkv_unpad, g)
dqkv = dqkv_pad_fn(dqkv_unpad)
dqkv_ref, = torch.autograd.grad(output_ref, qkv, g)
dqkv_pt, = torch.autograd.grad(output_pt, qkv, g)
print(f'dQ max diff: {(dqkv[:, :, 0] - dqkv_ref[:, :, 0]).abs().max().item()}')
print(f'dK max diff: {(dqkv[:, :, 1] - dqkv_ref[:, :, 1]).abs().max().item()}')
print(f'dV max diff: {(dqkv[:, :, 2] - dqkv_ref[:, :, 2]).abs().max().item()}')
print(f'dQKV mean diff: {(dqkv - dqkv_ref).abs().mean().item()}')
print(f'dQ Pytorch max diff: {(dqkv_pt[:, :, 0] - dqkv_ref[:, :, 0]).abs().max().item()}')
print(f'dK Pytorch max diff: {(dqkv_pt[:, :, 1] - dqkv_ref[:, :, 1]).abs().max().item()}')
print(f'dV Pytorch max diff: {(dqkv_pt[:, :, 2] - dqkv_ref[:, :, 2]).abs().max().item()}')
print(f'dQKV Pytorch mean diff: {(dqkv_pt - dqkv_ref).abs().mean().item()}')
# Check that FlashAttention's numerical error is at most twice the numerical error
# of a Pytorch implementation.
assert (output - output_ref).abs().max().item() <= 2 * (output_pt - output_ref).abs().max().item()
# assert torch.allclose(output, output_ref, rtol=rtol, atol=atol)
assert (attn - attn_ref).abs().max().item() <= 2 * (attn_pt - attn_ref).abs().max().item()
# assert torch.allclose(attn, attn_ref, rtol=rtol, atol=atol)
if dropout_p == 0.0:
assert dropout_mask.all()
else:
assert 0.99 <= dropout_fraction / dropout_p <= 1.01
assert (dqkv - dqkv_ref).abs().max().item() <= 2 * (dqkv_pt - dqkv_ref).abs().max().item()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment