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

Attempt to use atomicCAS to replace atomicAdd(bfloat16)

Browse Source
This commit is contained in:
Tri Dao 2022-10-23 16:22:43 -07:00
parent 6731855b1f
commit 9e92a1f2d2
1 changed files with 13 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
csrc/flash_attn/src/fmha/gmem_tile.h
View File

@ -34,6 +34,18 @@
namespace fmha { namespace fmha {
// template <typename half2_t>
// inline __device__ void atomic_add_CAS(half2_t *address, const half2_t val) {
// uint32_t *address_as_ui = (uint32_t *)address;
// uint32_t old = *address_as_ui;
// uint32_t assumed;
// do {
// assumed = old;
// half2_t sum = __hadd2(val, reinterpret_cast<half2_t(&)>(old));
// old = atomicCAS(address_as_ui, assumed, reinterpret_cast<uint32_t(&)>(sum));
// } while (assumed != old);
// }
//////////////////////////////////////////////////////////////////////////////////////////////////// ////////////////////////////////////////////////////////////////////////////////////////////////////
template< template<
@ -146,6 +158,7 @@ struct Gmem_tile_qkv {
#pragma unroll #pragma unroll
for (int jj = 0; jj < 4; ++jj) { for (int jj = 0; jj < 4; ++jj) {
atomicAdd(ptr_ + jj, reinterpret_cast<const elem2_type(&)[4]>(data[ii])[jj]); atomicAdd(ptr_ + jj, reinterpret_cast<const elem2_type(&)[4]>(data[ii])[jj]);
// atomic_add_CAS(ptr_ + jj, reinterpret_cast<const elem2_type(&)[4]>(data[ii])[jj]);
} }
} }
} }