From fa1e2bcaafcd33182f3f8d46b11683466ccf0781 Mon Sep 17 00:00:00 2001 From: Hongqiang Wang Date: Tue, 18 Aug 2026 13:57:24 -0700 Subject: [PATCH] opencl: fix WAR race in the generic FA tile kernels when the WG spans subgroups (llama/26434) --- ggml/src/ggml-opencl/kernels/flash_attn_f16.cl | 11 +++++++++++ ggml/src/ggml-opencl/kernels/flash_attn_f32.cl | 16 +++++++++------- 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f16.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f16.cl index fc58a22ec..f9797d346 100644 --- a/ggml/src/ggml-opencl/kernels/flash_attn_f16.cl +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f16.cl @@ -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; diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32.cl index 599877bdb..5911524e1 100644 --- a/ggml/src/ggml-opencl/kernels/flash_attn_f32.cl +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32.cl @@ -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) {