opencl: fix WAR race in the generic FA tile kernels when the WG spans subgroups (llama/26434)

This commit is contained in:
Hongqiang Wang 2026-08-18 13:57:24 -07:00 committed by Georgi Gerganov
parent 4be3101670
commit fa1e2bcaaf
2 changed files with 20 additions and 7 deletions

View File

@ -118,6 +118,17 @@ __kernel void flash_attn_f16(
__local DATA_TYPE4 l_v[BLOCK_N][DV_VEC];
for (int k_start = 0; k_start < n_kv; k_start += BLOCK_N) {
#if WG_SIZE > FA_SG
// WAR on l_k/l_v: a thread that finishes the compute below early either
// it skipped it (my_query_row >= n_q, the continue) or its subgroup simply
// ran ahead wraps around and reloads the tiles while another subgroup is
// still reading them. Any WG that is exactly one lockstep subgroup
// (WG_SIZE == FA_SG) cannot diverge and hides this; a WG spanning multiple
// subgroups (Intel sg=32, or BLOCK_M > 64 on Adreno) corrupts the result.
// All threads reach this each iteration (no-op on the first), so it does
// not diverge with the continue. Compiled out when WG == one subgroup.
barrier(CLK_LOCAL_MEM_FENCE);
#endif
for (int i = tid; i < BLOCK_N * DK_VEC; i += WG_SIZE) {
const int row = i / DK_VEC;
const int col = i % DK_VEC;

View File

@ -119,13 +119,15 @@ __kernel void flash_attn_f32(
__local DATA_TYPE4 l_v[BLOCK_N][DV_VEC];
for (int k_start = 0; k_start < n_kv; k_start += BLOCK_N) {
#if FA_SG < 64
// WAR on l_k/l_v: threads with my_query_row >= n_q skip the compute below
// (continue) and would race ahead to reload the tiles while active threads
// still read them. A single 64-wide Adreno subgroup (WG == sg) runs lockstep
// and hides this; a WG that spans multiple narrower subgroups (Intel sg=32)
// corrupts the result. All threads reach this each iteration (no-op on the
// first), so it does not diverge with the continue. Compiled out at sg=64.
#if WG_SIZE > FA_SG
// WAR on l_k/l_v: a thread that finishes the compute below early either
// it skipped it (my_query_row >= n_q, the continue) or its subgroup simply
// ran ahead wraps around and reloads the tiles while another subgroup is
// still reading them. Any WG that is exactly one lockstep subgroup
// (WG_SIZE == FA_SG) cannot diverge and hides this; a WG spanning multiple
// subgroups (Intel sg=32, or BLOCK_M > 64 on Adreno) corrupts the result.
// All threads reach this each iteration (no-op on the first), so it does
// not diverge with the continue. Compiled out when WG == one subgroup.
barrier(CLK_LOCAL_MEM_FENCE);
#endif
for (int i = tid; i < BLOCK_N * DK_VEC; i += WG_SIZE) {