leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 27 Packages Projects Releases Wiki Activity

xFormer updates to fMHA FW (#773)

Browse Source 
* xFormer updates to fMHA FW

* Convert format to BMHK for '41_fused_multi_head_attention_fixed_seqlen'

* Add missing files

* Remove xFormers specific code

* Update fused_multihead_attention_fixed_seqlen.cu

* rebase and solve conflicts

* remove white space

---------

Co-authored-by: danthe3rd <danthe3rd>
Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
This commit is contained in:
dan_the_3rd
2023-02-08 23:00:10 -05:00
committed by GitHub
co-authored by danthe3rd <danthe3rd> Haicheng Wu
parent 5ff5209ed5
commit 2e10404d26
10 changed files with 645 additions and 361 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/41_fused_multi_head_attention/kernel_forward.h
+9 -108
View File
@@ -29,15 +29,6 @@
*
**************************************************************************************************/
#pragma once
#ifdef HAS_PYTORCH
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/library.h>
#endif
#include <cmath>
#include <vector>
@@ -137,6 +128,9 @@ struct AttentionKernel {
output_accum_ptr; // [num_queries, num_heads, head_dim_value]
lse_scalar_t* logsumexp_ptr; // [num_heads, num_queries] - can be null
// Scale
accum_t scale;
// Dimensions/strides
int32_t head_dim;
int32_t head_dim_value;
@@ -154,18 +148,15 @@ struct AttentionKernel {
int32_t q_strideH;
int32_t k_strideH;
int32_t v_strideH;
int32_t o_strideH;
int64_t q_strideB;
int64_t k_strideB;
int64_t v_strideB;
int64_t o_strideB;
int32_t num_batches;
int32_t num_heads;
CUTLASS_HOST_DEVICE int32_t o_strideM() const {
return head_dim_value;
return head_dim_value * num_heads;
}
// Moves pointers to what we should process
// Returns "false" if there is no work to do
CUTLASS_DEVICE bool advance_to_block() {
@@ -195,9 +186,9 @@ struct AttentionKernel {
query_ptr += batch_id * q_strideB;
key_ptr += batch_id * k_strideB;
value_ptr += batch_id * v_strideB;
output_ptr += batch_id * o_strideB;
output_ptr += int64_t(batch_id * num_queries) * o_strideM();
if (output_accum_ptr != nullptr) {
output_accum_ptr += batch_id * o_strideB;
output_accum_ptr += int64_t(batch_id * num_queries) * o_strideM();
}
q_start = 0;
k_start = 0;
@@ -208,11 +199,11 @@ struct AttentionKernel {
key_ptr += k_start * k_strideM + head_id * k_strideH;
value_ptr += k_start * v_strideM + head_id * v_strideH;
output_ptr += int64_t(q_start + query_start) * o_strideM() +
head_id * o_strideH;
head_id * head_dim_value;
if (output_accum_ptr != nullptr) {
output_accum_ptr += int64_t(q_start + query_start) * o_strideM() +
head_id * o_strideH;
head_id * head_dim_value;
} else {
// Accumulate directly in the destination buffer (eg for f32)
output_accum_ptr = (accum_t*)output_ptr;
@@ -652,7 +643,7 @@ struct AttentionKernel {
warp_id(),
p.num_keys - iter_key_start,
iteratorC_tile_offset,
1.0f / cutlass::fast_sqrt(float(p.head_dim)));
p.scale);
}));
}));
@@ -858,93 +849,3 @@ __global__ void __launch_bounds__(AK::kNumThreads, AK::kMinBlocksPerSm)
}
AK::attention_kernel(p);
}
template <typename AK>
__global__ void __launch_bounds__(AK::kNumThreads, AK::kMinBlocksPerSm)
attention_kernel_batched(typename AK::Params params);
#define _ATTENTION_KERNEL_FORWARD_BEGIN(...) \
template <> \
__global__ void __launch_bounds__( \
__VA_ARGS__::kNumThreads, __VA_ARGS__::kMinBlocksPerSm) \
attention_kernel_batched<__VA_ARGS__>(typename __VA_ARGS__::Params p) { \
using Kernel = __VA_ARGS__;
#define _ATTENTION_KERNEL_FORWARD_END() }
#ifdef __CUDA_ARCH__
#define __CUDA_ARCH_OR_ZERO__ __CUDA_ARCH__
#else
#define __CUDA_ARCH_OR_ZERO__ 0
#endif
#define INSTANTIATE_ATTENTION_KERNEL_FORWARD( \
ARCH, \
SCALAR_T, \
IS_ALIGNED, \
QUERIES_PER_BLOCK, \
KEYS_PER_BLOCK, \
SINGLE_VALUE_ITER) \
_ATTENTION_KERNEL_FORWARD_BEGIN(AttentionKernel< \
SCALAR_T, \
cutlass::arch::Sm##ARCH, \
IS_ALIGNED, \
QUERIES_PER_BLOCK, \
KEYS_PER_BLOCK, \
SINGLE_VALUE_ITER>) \
if (!p.advance_to_block()) { \
return; \
} \
Kernel::attention_kernel(p); \
_ATTENTION_KERNEL_FORWARD_END();
#define INSTANTIATE_ATTENTION_KERNEL_FORWARD_DISABLED( \
ARCH, \
SCALAR_T, \
IS_ALIGNED, \
QUERIES_PER_BLOCK, \
KEYS_PER_BLOCK, \
SINGLE_VALUE_ITER) \
_ATTENTION_KERNEL_FORWARD_BEGIN(AttentionKernel< \
SCALAR_T, \
cutlass::arch::Sm##ARCH, \
IS_ALIGNED, \
QUERIES_PER_BLOCK, \
KEYS_PER_BLOCK, \
SINGLE_VALUE_ITER>) \
printf( \
"FATAL: this function is for sm%d, but was built for sm%d\n", \
int(ARCH), \
int(__CUDA_ARCH_OR_ZERO__)); \
_ATTENTION_KERNEL_FORWARD_END();
// All kernels are disabled by default
#define INSTANTIATE_ATTENTION_KERNEL_FORWARD_SM50(...) \
INSTANTIATE_ATTENTION_KERNEL_FORWARD_DISABLED(50, __VA_ARGS__)
#define INSTANTIATE_ATTENTION_KERNEL_FORWARD_SM70(...) \
INSTANTIATE_ATTENTION_KERNEL_FORWARD_DISABLED(70, __VA_ARGS__)
#define INSTANTIATE_ATTENTION_KERNEL_FORWARD_SM75(...) \
INSTANTIATE_ATTENTION_KERNEL_FORWARD_DISABLED(75, __VA_ARGS__)
#define INSTANTIATE_ATTENTION_KERNEL_FORWARD_SM80(...) \
INSTANTIATE_ATTENTION_KERNEL_FORWARD_DISABLED(80, __VA_ARGS__)
// Enable the right one based on __CUDA_ARCH__
#ifndef __CUDA_ARCH__
#elif __CUDA_ARCH__ < 500
#error "Need cuda arch at least 5.0"
#elif __CUDA_ARCH__ < 700
#undef INSTANTIATE_ATTENTION_KERNEL_FORWARD_SM50
#define INSTANTIATE_ATTENTION_KERNEL_FORWARD_SM50(...) \
INSTANTIATE_ATTENTION_KERNEL_FORWARD(50, __VA_ARGS__)
#elif __CUDA_ARCH__ < 750
#undef INSTANTIATE_ATTENTION_KERNEL_FORWARD_SM70
#define INSTANTIATE_ATTENTION_KERNEL_FORWARD_SM70(...) \
INSTANTIATE_ATTENTION_KERNEL_FORWARD(70, __VA_ARGS__)
#elif __CUDA_ARCH__ < 800
#undef INSTANTIATE_ATTENTION_KERNEL_FORWARD_SM75
#define INSTANTIATE_ATTENTION_KERNEL_FORWARD_SM75(...) \
INSTANTIATE_ATTENTION_KERNEL_FORWARD(75, __VA_ARGS__)
#elif __CUDA_ARCH__ >= 800
#undef INSTANTIATE_ATTENTION_KERNEL_FORWARD_SM80
#define INSTANTIATE_ATTENTION_KERNEL_FORWARD_SM80(...) \
INSTANTIATE_ATTENTION_KERNEL_FORWARD(80, __VA_ARGS__)
#endif