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:
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user