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

Update fMHA kernels (#992)

Browse Source 
* Update fMHA kernels

Upstream recent changes to fMHA that we did in xFormers.
Previous version in CUTLASS: facebookresearch/xformers@b6be33a
Updating to: facebookresearch/xformers@55a4798

* minor changes

* make var work

---------

Co-authored-by: danthe3rd <danthe3rd>
Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
This commit is contained in:
dan_the_3rd
2023-07-12 22:30:46 -04:00
committed by GitHub
co-authored by danthe3rd <danthe3rd> Haicheng Wu
parent f679663224
commit 146d314057
16 changed files with 1419 additions and 920 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/41_fused_multi_head_attention/gemm_kernel_utils.h
+11 -2
View File
@@ -228,8 +228,17 @@ struct call_conditional<false, TA, TB> {
// The cheapest way to do it is just to broadcast it from lane 0
////////////////////////////////////////////////////////////////////////////////
CUTLASS_DEVICE int32_t warp_uniform(int32_t value) {
return (int32_t)__shfl_sync(0xffffffff, (unsigned)value, 0);
template <typename T>
CUTLASS_DEVICE T warp_uniform(T value) {
struct {
union {
T value;
uint32_t asInt;
};
} p;
p.value = value;
p.asInt = __shfl_sync(0xffffffff, (unsigned)p.asInt, 0);
return p.value;
}
template <typename T>