wu-arch/kernels
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

flash: Optimize spad use, fix rowsum

Browse Source
This commit is contained in:
Hansung Kim
2024-08-15 16:54:15 -07:00
parent e809d25305
commit 745aa098ed
1 changed files with 6 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
tests/regression/flash_attention/kernel.cpp
View File

@@ -82,10 +82,8 @@ inline void thread_block_flashattn(float *S,
thread_offset += NUM_THREADS; thread_offset += NUM_THREADS;
} }
// stage per-thread max value in smem // stage per-thread max value in smem
// FIXME: we could warp_id instead of row here, but we need another barrier
// at the end of the loop iteration to prevent write-after-read hazard
// FIXME: threadblock_id needs to be in here too // FIXME: threadblock_id needs to be in here too
float *warp_smem = sharedmem_scratchpad + (row * NUM_THREADS); float *warp_smem = sharedmem_scratchpad + (warp_id * NUM_THREADS);
warp_smem[tid_in_warp] = per_thread_max; warp_smem[tid_in_warp] = per_thread_max;
// sync writes to warp_smem // sync writes to warp_smem
@@ -115,22 +113,20 @@ inline void thread_block_flashattn(float *S,
// const uint32_t row_stride = // const uint32_t row_stride =
// (exp_elem_per_thread * threads_per_threadblock) / B_COL; // (exp_elem_per_thread * threads_per_threadblock) / B_COL;
thread_offset = first_thread_offset + tid_in_warp;
// broadcast rowmax to all threads in the warp // broadcast rowmax to all threads in the warp
const float row_max = sharedmem_row_max_sum[row]; const float row_max = sharedmem_row_max_sum[row];
thread_offset = first_thread_offset + tid_in_warp;
#pragma GCC unroll #pragma GCC unroll
for (int i = 0; i < per_row_iter; i++) { for (int i = 0; i < per_row_iter; i++) {
float val = S[thread_offset]; float val = S[thread_offset];
// FIXME: placeholder for proper exp // FIXME: placeholder for proper exp
val = val; val -= row_max;
// update S in-place to P // update S in-place to P
// S[thread_offset] = val; S[thread_offset] = val;
gmem_tmp1[thread_offset] = val; gmem_tmp1[thread_offset] = val;
gmem_tmp2[thread_offset] = val - row_max;
thread_offset += NUM_THREADS; thread_offset += NUM_THREADS;
} }
@@ -142,7 +138,7 @@ inline void thread_block_flashattn(float *S,
// //
// two-level tree reduction, similar to rowmax // two-level tree reduction, similar to rowmax
#if 0 thread_offset = first_thread_offset + tid_in_warp;
float per_thread_sum = 0.0f; float per_thread_sum = 0.0f;
#pragma GCC unroll #pragma GCC unroll
for (int i = 0; i < per_row_iter; i++) { for (int i = 0; i < per_row_iter; i++) {
@@ -151,7 +147,7 @@ inline void thread_block_flashattn(float *S,
} }
// stage per-thread sum value in smem // stage per-thread sum value in smem
// FIXME: threadblock_id needs to be in here too // FIXME: threadblock_id needs to be in here too
warp_smem = sharedmem_scratchpad + (row * NUM_THREADS); warp_smem = sharedmem_scratchpad + (warp_id * NUM_THREADS);
warp_smem[tid_in_warp] = per_thread_sum; warp_smem[tid_in_warp] = per_thread_sum;
// sync writes to warp_smem // sync writes to warp_smem
@@ -167,7 +163,6 @@ inline void thread_block_flashattn(float *S,
sharedmem_row_max_sum[row] = per_thread_sum; sharedmem_row_max_sum[row] = per_thread_sum;
gmem_tmp2[row] = per_thread_sum; gmem_tmp2[row] = per_thread_sum;
} }
#endif
threadblock_barrier(threadblock_id_in_cluster, threadblock_barrier(threadblock_id_in_cluster,
warps_per_threadblock_per_core); warps_per_threadblock_per_core);