Skip to content

Commit dc64397

Browse files
olyasiriancris
authored andcommitted
ggml-opencl: add trailing barrier in f32 flash-attn tile loop
The f32 flash-attention kernel loads K/V tiles into local memory, barriers, reads them, then loops to overwrite the tiles for the next K/V block without a trailing barrier. Out-of-range lanes (the last partial BLOCK_M block) could also race ahead into the next tile load while active lanes were still reading. With n_kv > BLOCK_N (e.g. the bidirectional vision tower, n_kv=247) this corrupts the shared tiles. Add a trailing barrier(CLK_LOCAL_MEM_FENCE) at the end of the K/V block loop and guard the score computation with the query-row range check instead of an early continue.
1 parent c1dace7 commit dc64397

1 file changed

Lines changed: 14 additions & 4 deletions

File tree

ggml/src/ggml-opencl/kernels/flash_attn_f32.cl

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -122,10 +122,15 @@ __kernel void flash_attn_f32(
122122
}
123123
barrier(CLK_LOCAL_MEM_FENCE);
124124

125-
if (my_query_row >= n_q) {
126-
continue;
127-
}
128-
125+
// NOTE: do NOT `continue` for out-of-range query rows here. Every
126+
// work-item must reach the trailing barrier at the end of this loop,
127+
// otherwise the extra lanes (my_query_row >= n_q in the last partial
128+
// BLOCK_M block) race ahead to the next tile's load and overwrite
129+
// l_k/l_v while active lanes are still reading them. That shared-memory
130+
// race silently corrupts the K/V tiles for any sequence spanning more
131+
// than one BLOCK_N tile (e.g. the bidirectional Qwen3-VL vision tower,
132+
// n_kv=247), degrading the encode. Guard the score loop instead.
133+
if (my_query_row < n_q) {
129134
for (int j = 0; j < BLOCK_N; j += 2) {
130135
const int k_row0 = k_start + j;
131136
const int k_row1 = k_start + j + 1;
@@ -171,6 +176,11 @@ __kernel void flash_attn_f32(
171176
l_i = l_i * scale_prev + p0 + p1;
172177
m_i = m_new;
173178
}
179+
} // end if (my_query_row < n_q)
180+
181+
// Ensure every work-item has finished reading l_k/l_v before the next
182+
// iteration overwrites the shared K/V tiles.
183+
barrier(CLK_LOCAL_MEM_FENCE);
174184
}
175185

176186
if (my_query_row < n_q) {

0 commit comments

Comments
 (0)