Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source 
* 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
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl
+438 -1
View File
@@ -24,7 +24,11 @@
#define DK_VEC (DK/4)
#define DV_VEC (DV/4)
#define Q1_WG_SIZE 64
#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.
@@ -310,6 +314,201 @@ __kernel void flash_attn_f32_q8_0_q1(
}
}
#ifdef cl_intel_subgroups
#pragma OPENCL EXTENSION cl_intel_subgroups : enable
#else
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
#endif
#ifdef cl_qcom_reqd_sub_group_size
#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
#define REQD_SUBGROUP_SIZE_64 __attribute__((qcom_reqd_sub_group_size("half")))
#else
#define REQD_SUBGROUP_SIZE_64
#endif
#define VEC_NSG 4
#define VEC_WG_SIZE (Q1_WG_SIZE * VEC_NSG)
#define Q1V_DV_PER_THREAD ((DV_VEC + Q1_WG_SIZE - 1) / Q1_WG_SIZE)
inline float4 dequant_q8_0_lane(const global char * block_ptr, int lane) {
const float d = vload_half(0, (const global half *)block_ptr);
const global char * qs = block_ptr + 2 + lane * 4;
return d * (float4)((float)qs[0], (float)qs[1], (float)qs[2], (float)qs[3]);
}
REQD_SUBGROUP_SIZE_64
__kernel void flash_attn_f32_q8_0_q1_vec(
const global void * q_void, ulong q_offset,
const global void * k_void, ulong k_offset,
const global void * v_void, ulong v_offset,
global void * o_void, ulong o_offset,
const float scale,
const int n_q,
const int n_kv,
const int is_causal,
const int n_head,
const ulong q_nb1, const ulong q_nb2, const ulong q_nb3,
const ulong k_nb1, const ulong k_nb2, const ulong k_nb3,
const ulong v_nb1, const ulong v_nb2, const ulong v_nb3,
const ulong o_nb1, const ulong o_nb2, const ulong o_nb3,
const float max_bias,
const float m0,
const float m1,
const int n_head_log2,
const float logit_softcap,
const int n_head_kv,
const global void* mask_void,
const ulong mask_offset,
const ulong mask_nb1,
const ulong mask_nb2,
const ulong mask_nb3,
const int mask_ne2,
const int mask_ne3,
const global void* sinks_void,
const ulong sinks_offset
) {
const int tid = get_local_id(0);
const int sgid = tid / Q1_WG_SIZE;
const int tid_sg = tid % Q1_WG_SIZE;
const int head_batch_idx = get_global_id(1);
const int batch_idx = head_batch_idx / n_head;
const int head_idx = head_batch_idx % n_head;
const int gqa_ratio = n_head / n_head_kv;
const int head_kv_idx = head_idx / gqa_ratio;
const global char * q_base = (const global char *) q_void + q_offset;
const global char * k_base = (const global char *) k_void + k_offset;
const global char * v_base = (const global char *) v_void + v_offset;
global char * o_base = (global char *) o_void + o_offset;
const global char * mask_base = NULL;
if (mask_void != NULL) {
const int mask_head_idx = head_idx % mask_ne2;
const int mask_batch_idx = batch_idx % mask_ne3;
mask_base = (const global char *) mask_void + mask_offset +
mask_batch_idx * mask_nb3 + mask_head_idx * mask_nb2;
}
__local ACC_TYPE4 q_shared[DK_VEC];
{
const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2;
const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset);
for (int i = tid; i < DK_VEC; i += VEC_WG_SIZE) {
q_shared[i] = CONVERT_Q_ACC4(q_ptr[i]);
}
}
barrier(CLK_LOCAL_MEM_FENCE);
const float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
const global ACC_TYPE * sinks_ptr = NULL;
if (sinks_void != NULL) {
sinks_ptr = (const global ACC_TYPE *) ((const global char *) sinks_void + sinks_offset);
}
ACC_TYPE4 o_acc[Q1V_DV_PER_THREAD];
#pragma unroll
for (int i = 0; i < Q1V_DV_PER_THREAD; ++i) o_acc[i] = (ACC_TYPE4)(0.0f);
ACC_TYPE m_i = FA_M_INIT;
ACC_TYPE l_i = 0.0f;
const int kv_per_sg = (n_kv + VEC_NSG - 1) / VEC_NSG;
const int kv_start = sgid * kv_per_sg;
const int kv_end = min(n_kv, kv_start + kv_per_sg);
for (int k_idx = kv_start; k_idx < kv_end; ++k_idx) {
const global char * k_row = k_base + batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
const global char * v_row = v_base + batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
ACC_TYPE4 dot4 = (ACC_TYPE4)(0.0f);
for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
const int block_idx = qk / 8;
const int lane = qk % 8;
const float4 k_v = dequant_q8_0_lane(k_row + block_idx * Q8_0_BLOCK_SIZE, lane);
dot4 = mad(q_shared[qk], k_v, dot4);
}
ACC_TYPE dot_partial = dot4.s0 + dot4.s1 + dot4.s2 + dot4.s3;
ACC_TYPE score = sub_group_reduce_add(dot_partial) * scale;
if (mask_base != NULL) {
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base;
score += slope * (ACC_TYPE) mask_ptr[k_idx];
}
if (logit_softcap > 0.0f) {
score = logit_softcap * tanh(score / logit_softcap);
}
const ACC_TYPE m_new = max(m_i, score);
const ACC_TYPE scale_prev = native_exp(m_i - m_new);
const ACC_TYPE p = native_exp(score - m_new);
int idx = 0;
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
const int block_idx = dv / 8;
const int lane = dv % 8;
const float4 v_v = dequant_q8_0_lane(v_row + block_idx * Q8_0_BLOCK_SIZE, lane);
o_acc[idx] = mad(p, v_v, o_acc[idx] * scale_prev);
}
l_i = l_i * scale_prev + p;
m_i = m_new;
}
__local ACC_TYPE sg_m[VEC_NSG];
__local ACC_TYPE sg_l[VEC_NSG];
__local ACC_TYPE4 sg_o[VEC_NSG][DV_VEC];
if (tid_sg == 0) {
sg_m[sgid] = m_i;
sg_l[sgid] = l_i;
}
{
int idx = 0;
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
sg_o[sgid][dv] = o_acc[idx];
}
}
barrier(CLK_LOCAL_MEM_FENCE);
if (sgid == 0) {
ACC_TYPE m_final = sg_m[0];
#pragma unroll
for (int s = 1; s < VEC_NSG; ++s) {
m_final = max(m_final, sg_m[s]);
}
if (sinks_ptr != NULL) {
m_final = max(m_final, sinks_ptr[head_idx]);
}
ACC_TYPE l_final = 0.0f;
#pragma unroll
for (int s = 0; s < VEC_NSG; ++s) {
l_final += sg_l[s] * native_exp(sg_m[s] - m_final);
}
if (sinks_ptr != NULL) {
l_final += native_exp(sinks_ptr[head_idx] - m_final);
}
const ACC_TYPE l_inv = (l_final > 0.0f) ? (1.0f / l_final) : 0.0f;
const ulong o_row_offset = batch_idx * o_nb3 + head_idx * o_nb1;
global O_DATA_TYPE4 * o_row = (global O_DATA_TYPE4 *) (o_base + o_row_offset);
int idx = 0;
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
ACC_TYPE4 o_merged = (ACC_TYPE4)(0.0f);
#pragma unroll
for (int s = 0; s < VEC_NSG; ++s) {
const ACC_TYPE alpha = native_exp(sg_m[s] - m_final);
o_merged = mad((ACC_TYPE4)(alpha), sg_o[s][dv], o_merged);
}
o_row[dv] = CONVERT_O_DATA4(o_merged * l_inv);
}
}
}
// Flash-decoding split pass for q8_0 KV. Partial record: [m, l, O[DV]].
// Merge kernel from flash_attn_f32_f16.cl is type-agnostic and reused.
#define FA_PARTIAL_FLOATS (2 + DV)
@@ -533,6 +732,244 @@ __kernel void flash_attn_f32_q8_0_q1_split(
#define FA_V_STRATEGY 0
#endif
#ifndef MQ_GQA
#define MQ_GQA 4
#endif
#ifndef MQ_NSG_SPLIT
#define MQ_NSG_SPLIT 4
#endif
#define MQ_SPLIT_WG_SIZE_Q8 (Q1_WG_SIZE * MQ_NSG_SPLIT)
REQD_SUBGROUP_SIZE_64
__kernel void flash_attn_f32_q8_0_q1_vec_mq_split(
const global void * q_void, ulong q_offset,
const global void * k_void, ulong k_offset,
const global void * v_void, ulong v_offset,
const float scale,
const int n_q,
const int n_kv,
const int n_head,
const ulong q_nb1, const ulong q_nb2, const ulong q_nb3,
const ulong k_nb1, const ulong k_nb2, const ulong k_nb3,
const ulong v_nb1, const ulong v_nb2, const ulong v_nb3,
const float max_bias,
const float m0,
const float m1,
const int n_head_log2,
const float logit_softcap,
const int n_head_kv,
const global void * mask_void,
const ulong mask_offset,
const ulong mask_nb1,
const ulong mask_nb2,
const ulong mask_nb3,
const int mask_ne2,
const int mask_ne3,
global float * partial_void,
const int n_splits,
const int kv_per_split
) {
const int tid = get_local_id(0);
const int sgid = tid / Q1_WG_SIZE;
const int tid_sg = tid % Q1_WG_SIZE;
const int kvhead_batch_idx = get_global_id(1);
const int split_q_idx = get_global_id(2);
const int split_idx = split_q_idx % n_splits;
const int q_idx = split_q_idx / n_splits;
const int batch_idx = kvhead_batch_idx / n_head_kv;
const int head_kv_idx = kvhead_batch_idx % n_head_kv;
const int kv_start = split_idx * kv_per_split;
const int kv_end = min(kv_start + kv_per_split, n_kv);
const ulong record_stride = (ulong) FA_PARTIAL_FLOATS;
if (kv_start >= kv_end) {
// Empty split — write sentinel for each of the MQ_GQA Q-heads.
if (tid == 0) {
#pragma unroll
for (int h = 0; h < MQ_GQA; ++h) {
const int head_idx = head_kv_idx * MQ_GQA + h;
const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
* n_splits + split_idx);
global float * rec = partial_void + rec_idx * record_stride;
rec[0] = FA_M_INIT;
rec[1] = 0.0f;
}
}
return;
}
const global char * q_base = (const global char *) q_void + q_offset;
const global char * k_base = (const global char *) k_void + k_offset;
const global char * v_base = (const global char *) v_void + v_offset;
__local ACC_TYPE4 q_shared[MQ_GQA * DK_VEC];
for (int i = tid; i < MQ_GQA * DK_VEC; i += MQ_SPLIT_WG_SIZE_Q8) {
const int h = i / DK_VEC;
const int k = i % DK_VEC;
const int head_idx = head_kv_idx * MQ_GQA + h;
const ulong q_row_offset = batch_idx * q_nb3 + head_idx * q_nb2 + (ulong) q_idx * q_nb1;
const global Q_DATA_TYPE4 * q_ptr = (const global Q_DATA_TYPE4 *) (q_base + q_row_offset);
q_shared[h * DK_VEC + k] = CONVERT_Q_ACC4(q_ptr[k]);
}
barrier(CLK_LOCAL_MEM_FENCE);
float slope[MQ_GQA];
#pragma unroll
for (int h = 0; h < MQ_GQA; ++h) {
slope[h] = get_alibi_slope(max_bias, head_kv_idx * MQ_GQA + h, n_head_log2, m0, m1);
}
const global char * mask_base[MQ_GQA];
if (mask_void != NULL) {
const int mask_batch_idx = batch_idx % mask_ne3;
const global char * mask_base_b = (const global char *) mask_void + mask_offset +
mask_batch_idx * mask_nb3 +
(ulong) q_idx * mask_nb1;
#pragma unroll
for (int h = 0; h < MQ_GQA; ++h) {
const int head_idx = head_kv_idx * MQ_GQA + h;
const int mask_head_idx = head_idx % mask_ne2;
mask_base[h] = mask_base_b + mask_head_idx * mask_nb2;
}
} else {
#pragma unroll
for (int h = 0; h < MQ_GQA; ++h) mask_base[h] = NULL;
}
ACC_TYPE4 o_acc[MQ_GQA][Q1V_DV_PER_THREAD];
ACC_TYPE m_i[MQ_GQA];
ACC_TYPE l_i[MQ_GQA];
#pragma unroll
for (int h = 0; h < MQ_GQA; ++h) {
m_i[h] = FA_M_INIT;
l_i[h] = 0.0f;
#pragma unroll
for (int i = 0; i < Q1V_DV_PER_THREAD; ++i) o_acc[h][i] = (ACC_TYPE4)(0.0f);
}
const int kv_len = kv_end - kv_start;
const int kv_per_sg = (kv_len + MQ_NSG_SPLIT - 1) / MQ_NSG_SPLIT;
const int kv_lo = kv_start + sgid * kv_per_sg;
const int kv_hi = min(kv_end, kv_lo + kv_per_sg);
for (int k_idx = kv_lo; k_idx < kv_hi; ++k_idx) {
const global char * k_row = k_base + batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
const global char * v_row = v_base + batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
ACC_TYPE4 dot4[MQ_GQA];
#pragma unroll
for (int h = 0; h < MQ_GQA; ++h) dot4[h] = (ACC_TYPE4)(0.0f);
for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
const int block_idx = qk / 8;
const int lane = qk % 8;
const float4 k_v = dequant_q8_0_lane(k_row + block_idx * Q8_0_BLOCK_SIZE, lane);
#pragma unroll
for (int h = 0; h < MQ_GQA; ++h) {
dot4[h] = mad(q_shared[h * DK_VEC + qk], k_v, dot4[h]);
}
}
ACC_TYPE score[MQ_GQA];
#pragma unroll
for (int h = 0; h < MQ_GQA; ++h) {
const ACC_TYPE dot_partial = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
ACC_TYPE s = sub_group_reduce_add(dot_partial) * scale;
if (mask_base[h] != NULL) {
const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
s += slope[h] * (ACC_TYPE) mask_ptr[k_idx];
}
if (logit_softcap > 0.0f) {
s = logit_softcap * tanh(s / logit_softcap);
}
score[h] = s;
}
ACC_TYPE p_h[MQ_GQA];
ACC_TYPE sp_h[MQ_GQA];
#pragma unroll
for (int h = 0; h < MQ_GQA; ++h) {
const ACC_TYPE m_new = max(m_i[h], score[h]);
sp_h[h] = native_exp(m_i[h] - m_new);
p_h[h] = native_exp(score[h] - m_new);
l_i[h] = l_i[h] * sp_h[h] + p_h[h];
m_i[h] = m_new;
}
int idx = 0;
for (int dv = tid_sg; dv < DV_VEC; dv += Q1_WG_SIZE, ++idx) {
const int block_idx = dv / 8;
const int lane = dv % 8;
const float4 v_v = dequant_q8_0_lane(v_row + block_idx * Q8_0_BLOCK_SIZE, lane);
#pragma unroll
for (int h = 0; h < MQ_GQA; ++h) {
o_acc[h][idx] = mad(p_h[h], v_v, o_acc[h][idx] * sp_h[h]);
}
}
}
__local ACC_TYPE sg_m[MQ_GQA][MQ_NSG_SPLIT];
__local ACC_TYPE sg_l[MQ_GQA][MQ_NSG_SPLIT];
__local ACC_TYPE4 sg_o[MQ_NSG_SPLIT][DV_VEC];
if (tid_sg == 0) {
#pragma unroll
for (int h = 0; h < MQ_GQA; ++h) {
sg_m[h][sgid] = m_i[h];
sg_l[h][sgid] = l_i[h];
}
}
#pragma unroll
for (int h = 0; h < MQ_GQA; ++h) {
{
int idx = 0;
for (int dv_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE, ++idx) {
sg_o[sgid][dv_idx] = o_acc[h][idx];
}
}
barrier(CLK_LOCAL_MEM_FENCE);
if (sgid == 0) {
const int head_idx = head_kv_idx * MQ_GQA + h;
ACC_TYPE m_c = sg_m[h][0];
#pragma unroll
for (int s = 1; s < MQ_NSG_SPLIT; ++s) {
m_c = max(m_c, sg_m[h][s]);
}
ACC_TYPE l_c = 0.0f;
#pragma unroll
for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
l_c += sg_l[h][s] * native_exp(sg_m[h][s] - m_c);
}
const ulong rec_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
* n_splits + split_idx);
global float * rec = partial_void + rec_idx * record_stride;
global float4 * rec_o = (global float4 *) (rec + 2);
if (tid_sg == 0) {
rec[0] = (float) m_c;
rec[1] = (float) l_c;
}
for (int dv_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE) {
ACC_TYPE4 o_merged = (ACC_TYPE4)(0.0f);
#pragma unroll
for (int s = 0; s < MQ_NSG_SPLIT; ++s) {
const ACC_TYPE alpha = native_exp(sg_m[h][s] - m_c);
o_merged = mad((ACC_TYPE4)(alpha), sg_o[s][dv_idx], o_merged);
}
rec_o[dv_idx] = o_merged;
}
}
barrier(CLK_LOCAL_MEM_FENCE);
}
}
__kernel void flash_attn_f32_q8_0(
const global void * q_void, ulong q_offset,
const global void * k_void, ulong k_offset,