opencl: general flash attention decode performance optimizations (#25366)

* opencl: vec flash-attention decode kernels for f16/q8_0/q4_0 KV

* opencl: improve non FA KQ mv kernels

* opencl: tweaks for multiquery FA

* opencl: some tweaks for FA q1 kernels

* opencl: FA with DK=DV=512 for gemma-4

* opencl: various fixes

* opencl: cleanup

* opencl: fix FA decode crash for DK=512 (gemma-4)

The DK=512 decode-only program does not create the f32_f16 prefill
kernel, so the compiled check in ensure_fa_variant never hit and
supports_op gave inconsistent answers for the same op. block_n is also
unset for DK=512 decode; guard it to avoid an out-of-range read at
dispatch.

* opencl: run DK=512 FA decode on CPU

DK=512 decode is bandwidth-bound and faster on the CPU than the GPU,
increasingly so with depth. Decline it in supports_op; prefill stays on the GPU.

* opencl: compile MQ_GQA=8 FA kernels in a minimal program

The full program compiled with -D MQ_GQA=8 runs the Adreno compiler out
of memory at DK>=256. Only the vec_mq kernels are used from this
program, so compile it with FA_MQ_ONLY, which excludes everything else.
Also include the program name in the compile error log.

* opencl: remove stray token in flash_attn_f32_f16.cl

A stray "." broke the f32_f16 program build.

* opencl: split f16-KV FA decode finer (FD_KV_PER_SPLIT_F16)

The 2048 default under-fills the GPU on single-query f16-KV decode;
use 512 for f16 KV to get more splits. Quantized KV keeps 2048.

---------

Co-authored-by: Li He <lih@qti.qualcomm.com>
This commit is contained in:
Hongqiang Wang
2026-07-06 19:57:52 -07:00
committed by GitHub
co-authored by Li He
parent ee445f93d8
commit 6f8895feec
8 changed files with 4485 additions and 170 deletions
+15 -1
View File
@@ -11,7 +11,12 @@
#define DK_VEC (DK/4)
#define DV_VEC (DV/4)
#define WG_SIZE (BLOCK_M)
#define Q1_WG_SIZE 64
// q1 reduces over a Q1_WG_SIZE-wide WG via work-group barriers; the launch WG
// must match. Defaults to the Adreno sg (64); host passes -D FA_SG=32 on Intel.
#ifndef FA_SG
#define FA_SG 64
#endif
#define Q1_WG_SIZE FA_SG
// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs,
// infinite operand can cause undefined behavior and miscompilation for exp.
@@ -114,6 +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.
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;