Skip to content

GitLab

Commit 9e92a1f2 authored by Tri Dao's avatar Tri Dao
Browse files

Attempt to use atomicCAS to replace atomicAdd(bfloat16)

parent 6731855b
Show whitespace changes
Inline Side-by-side
Showing with 13 additions and 0 deletions
csrc/flash_attn/src/fmha/gmem_tile.h
View file @ 9e92a1f2
...@@ -34,6 +34,18 @@ ...@@ -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 { ...@@ -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]);
} }
} }
} }
......
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