squall/flash-attention
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Compilation] Change BOOL_SWITCH to fix Windows compilation

Browse Source 
Follow xFormers's DISTPATCH_BOOL. Haven't tested it on Windows.
This commit is contained in:
Tri Dao 2023-01-06 14:40:58 -08:00
parent a668890fcd
commit a1f49a2b92
9 changed files with 29 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
csrc/flash_attn/src/fmha_bwd_hdim128.cu
View File

@ -5,7 +5,7 @@
#include "fmha_bwd_launch_template.h"
void run_fmha_bwd_hdim128(FMHA_dgrad_params &params, cudaStream_t stream, const bool configure) {
FP16_SWITCH(params.is_bf16, ({
FP16_SWITCH(params.is_bf16, ([&] {
using Kernel_traits = FMHA_kernel_traits<128, 128, 16, 1, 8, 0x100u, elem_type>;
run_fmha_bwd_loop<Kernel_traits>(params, stream, configure);
}));

2
csrc/flash_attn/src/fmha_bwd_hdim32.cu
View File

@ -5,7 +5,7 @@
#include "fmha_bwd_launch_template.h"
void run_fmha_bwd_hdim32(FMHA_dgrad_params &params, cudaStream_t stream, const bool configure) {
FP16_SWITCH(params.is_bf16, ({
FP16_SWITCH(params.is_bf16, ([&] {
if (params.seqlen_k == 128) {
using Kernel_traits = FMHA_kernel_traits<128, 32, 16, 1, 8, 0x08u, elem_type>;
run_fmha_bwd_loop<Kernel_traits>(params, stream, configure);

2
csrc/flash_attn/src/fmha_bwd_hdim64.cu
View File

@ -5,7 +5,7 @@
#include "fmha_bwd_launch_template.h"
void run_fmha_bwd_hdim64(FMHA_dgrad_params &params, cudaStream_t stream, const bool configure) {
FP16_SWITCH(params.is_bf16, ({
FP16_SWITCH(params.is_bf16, ([&] {
auto dprops = at::cuda::getCurrentDeviceProperties();
if (params.seqlen_k == 128) {
using Kernel_traits = FMHA_kernel_traits<128, 64, 16, 1, 8, 0x08u, elem_type>;

2
csrc/flash_attn/src/fmha_bwd_launch_template.h
View File

@ -61,7 +61,7 @@ void run_fmha_bwd_loop(FMHA_dgrad_params &params, cudaStream_t stream, const boo
bool is_dropout = params.p_dropout < 1.f; // params.p_dropout is the probability of "keeping"
// Work-around for gcc 7. It doesn't like nested BOOL_SWITCH.
BOOL_SWITCH(is_dropout, IsDropoutConst, ({
BOOL_SWITCH(is_dropout, IsDropoutConst, ([&] {
auto kernel = params.is_causal
? &fmha_bwd_dq_dk_dv_loop_kernel<Kernel_traits, IsDropoutConst, true>
: &fmha_bwd_dq_dk_dv_loop_kernel<Kernel_traits, IsDropoutConst, false>;

2
csrc/flash_attn/src/fmha_fwd_hdim128.cu
View File

@ -5,7 +5,7 @@
#include "fmha_fwd_launch_template.h"
void run_fmha_fwd_hdim128(Launch_params<FMHA_fprop_params> &launch_params) {
FP16_SWITCH(launch_params.params.is_bf16, ({
FP16_SWITCH(launch_params.params.is_bf16, ([&] {
using Kernel_traits = FMHA_kernel_traits<128, 128, 16, 1, 4, 0x08u, elem_type>;
run_fmha_fwd_loop<Kernel_traits>(launch_params);
}));

2
csrc/flash_attn/src/fmha_fwd_hdim32.cu
View File

@ -5,7 +5,7 @@
#include "fmha_fwd_launch_template.h"
void run_fmha_fwd_hdim32(Launch_params<FMHA_fprop_params> &launch_params) {
FP16_SWITCH(launch_params.params.is_bf16, ({
FP16_SWITCH(launch_params.params.is_bf16, ([&] {
if (launch_params.params.seqlen_k == 128) {
using Kernel_traits = FMHA_kernel_traits<128, 32, 16, 1, 4, 0x08u, elem_type>;
run_fmha_fwd_loop<Kernel_traits>(launch_params);

2
csrc/flash_attn/src/fmha_fwd_hdim64.cu
View File

@ -5,7 +5,7 @@
#include "fmha_fwd_launch_template.h"
void run_fmha_fwd_hdim64(Launch_params<FMHA_fprop_params> &launch_params) {
FP16_SWITCH(launch_params.params.is_bf16, ({
FP16_SWITCH(launch_params.params.is_bf16, ([&] {
if (launch_params.params.seqlen_k == 128) {
using Kernel_traits = FMHA_kernel_traits<128, 64, 16, 1, 4, 0x08u, elem_type>;
run_fmha_fwd_loop<Kernel_traits>(launch_params);

2
csrc/flash_attn/src/fmha_fwd_launch_template.h
View File

@ -56,7 +56,7 @@ void run_fmha_fwd_loop(Launch_params<FMHA_fprop_params> &launch_params) {
// Work-around for gcc 7. It doesn't like nested BOOL_SWITCH.
// https://github.com/kokkos/kokkos-kernels/issues/349
// https://github.com/HazyResearch/flash-attention/issues/21
BOOL_SWITCH(launch_params.is_dropout, IsDropoutConst, ({
BOOL_SWITCH(launch_params.is_dropout, IsDropoutConst, ([&] {
auto kernel = launch_params.params.is_causal
? (launch_params.return_softmax
? &fmha_fwd_loop_kernel<Kernel_traits, IsDropoutConst, true, true>

37
csrc/flash_attn/src/static_switch.h
View File

@ -1,5 +1,6 @@
// 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
// and https://github.com/facebookresearch/xformers/blob/main/xformers/csrc/attention/cuda/fmha/gemm_kernel_utils.h#L8
#pragma once
@ -9,27 +10,31 @@
///
/// Usage:
/// ```
/// BOOL_SWITCH(flag, BoolConst, ({
/// BOOL_SWITCH(flag, BoolConst, ([&] {
/// some_function<BoolConst>(...);
/// }));
/// ```
/// We need "({" and "})" to make sure that the code is a single argument being passed to the macro.
#define BOOL_SWITCH(COND, CONST_NAME, CODE) \
if (COND) { \
constexpr bool CONST_NAME = true; \
CODE; \
} else { \
constexpr bool CONST_NAME = false; \
CODE; \
#define BOOL_SWITCH(COND, CONST_NAME, F) \
{ \
if (COND) { \
constexpr bool CONST_NAME = true; \
F(); \
} else { \
constexpr bool CONST_NAME = false; \
F(); \
} \
}
// modified from BOOL_SWITCH
// because MSVC cannot handle std::conditional with constexpr variable
#define FP16_SWITCH(COND, CODE) \
if (COND) { \
using elem_type = __nv_bfloat16; \
CODE; \
} else { \
using elem_type = __half; \
CODE; \
} \
#define FP16_SWITCH(COND, F) \
{ \
if (COND) { \
using elem_type = __nv_bfloat16; \
F(); \
} else { \
using elem_type = __half; \
F(); \
} \
}