opencl: flash attention for Adreno — fused kernels, quant KV, flash-decoding

Fused FA kernels for f16, native q8_0, and native q4_0 KV on Qualcomm
Adreno. Covers prefill, n_q=1 decode, and a flash-decoding (K-split)
path for long-context decode.

- N_SPLIT splits DK/DV across threads in a WG; QK partials reduced via
  sub_group_shuffle_xor, with a 3-phase local-memory fallback when
  subgroup shuffle isn't available.
- Flash-Decoding fires at n_kv >= 2048 on n_q <= 8 non-causal shapes,
  DK <= 128 (DK <= 64 for multi-q FD).
- Asymmetric KV (-ctk X -ctv Y, X != Y) via on-GPU dequant of the quant
  side to F32.
- Per-(dk,dv) lazy compilation. f16 q1 / q1_split share Q via __local
  to avoid private-array spill.
- One-pass online softmax (FA-2 style) in q8_0 and q4_0 q1 / q1_split
  kernels: maintains per-thread (m_i, l_i, o_acc) on a single sweep,
  cross-thread merge rescales by alpha=exp(m_i - m_final). Eliminates
  the second K read of the original two-pass implementation. Measured
  on gpt-oss-20b at d=16k decode (Adreno X2-90): q8_0 +16.6%, q4_0
  +17.5% end-to-end vs two-pass baseline.
- Dim-table tuning across DK in {40, 64, 80, 96, 112, 128, 192, 256},
  with a per-generation selector hook (X1, X2 stubs).

Tested on Snapdragon X Elite (Adreno X1-85) and Snapdragon X2 Elite
Extreme (Adreno X2-90), OpenCL 3.0.

Author: wanghqc

ggml/src/ggml-opencl/CMakeLists.txt

@@ -169,7 +169,10 @@ set(GGML_OPENCL_KERNELS
     mul_mm_f16_f32_kq_kqv
     conv2d
     conv2d_f16_f32
+    flash_attn_pre_f16
     flash_attn_f32_f16
+    flash_attn_f32_q8_0
+    flash_attn_f32_q4_0
     flash_attn_f16
     flash_attn_f32
 )

ggml/src/ggml-opencl/ggml-opencl.cpp

@@ -583,6 +583,196 @@ kernel void kernel_restore_block_q8_0(
     }
 }
 
+// AoS q8_0 dequant → f16. One thread per 32-elem block.
+kernel void kernel_dequant_q8_0_f16_aos(
+    global char * src,
+    global half * dst,
+    int n_blocks
+) {
+    int blk = get_global_id(0);
+    if (blk >= n_blocks) return;
+
+    global char * block = src + blk * (QK8_0 + 2);
+    float d = vload_half(0, (global half *)block);
+    global char * qs = block + 2;
+
+    global half * out = dst + blk * QK8_0;
+    for (int i = 0; i < QK8_0; ++i) {
+        out[i] = (half)(d * (float)qs[i]);
+    }
+}
+
+// View-aware AoS q8_0 → f32 dequant (f32/f32 FA path).
+kernel void kernel_dequant_q8_0_f32_view_aos(
+    global char * src,
+    ulong         src_offset,
+    ulong         src_nb1,
+    ulong         src_nb2,
+    ulong         src_nb3,
+    int           nblk0,
+    int           ne1,
+    int           ne2,
+    int           ne3,
+    global float * dst
+) {
+    int blk_i0 = get_global_id(0);
+    int i1     = get_global_id(1);
+    int batch  = get_global_id(2);
+
+    if (blk_i0 >= nblk0) return;
+    if (i1     >= ne1)   return;
+
+    int i2 = batch % ne2;
+    int i3 = batch / ne2;
+    if (i3 >= ne3) return;
+
+    global char * block = src + src_offset + (ulong)i3*src_nb3 + (ulong)i2*src_nb2 + (ulong)i1*src_nb1 + (ulong)blk_i0 * (2 + QK8_0);
+    float d = vload_half(0, (global half *)block);
+    global char * qs = block + 2;
+
+    ulong dst_row_base = ((ulong)i3 * ne2 * ne1 + (ulong)i2 * ne1 + (ulong)i1) * nblk0;
+    global float * out = dst + (dst_row_base + blk_i0) * QK8_0;
+
+    for (int i = 0; i < QK8_0; ++i) {
+        out[i] = d * (float)qs[i];
+    }
+}
+
+// View-aware AoS q8_0 → f16 dequant. Rows tight, batch strides may be gapped.
+kernel void kernel_dequant_q8_0_f16_view_aos(
+    global char * src,
+    ulong         src_offset,
+    ulong         src_nb1,
+    ulong         src_nb2,
+    ulong         src_nb3,
+    int           nblk0,
+    int           ne1,
+    int           ne2,
+    int           ne3,
+    global half * dst
+) {
+    int blk_i0 = get_global_id(0);
+    int i1     = get_global_id(1);
+    int batch  = get_global_id(2);
+
+    if (blk_i0 >= nblk0) return;
+    if (i1     >= ne1)   return;
+
+    int i2 = batch % ne2;
+    int i3 = batch / ne2;
+    if (i3 >= ne3) return;
+
+    global char * block = src + src_offset + (ulong)i3*src_nb3 + (ulong)i2*src_nb2 + (ulong)i1*src_nb1 + (ulong)blk_i0 * (2 + QK8_0);
+    float d = vload_half(0, (global half *)block);
+    global char * qs = block + 2;
+
+    ulong dst_row_base = ((ulong)i3 * ne2 * ne1 + (ulong)i2 * ne1 + (ulong)i1) * nblk0;
+    global half * out = dst + (dst_row_base + blk_i0) * QK8_0;
+
+    for (int i = 0; i < QK8_0; ++i) {
+        out[i] = (half)(d * (float)qs[i]);
+    }
+}
+
+// View-aware AoS q4_0 → f32 dequant (mirrors the q8_0 view variant).
+kernel void kernel_dequant_q4_0_f32_view_aos(
+    global char * src,
+    ulong         src_offset,
+    ulong         src_nb1,
+    ulong         src_nb2,
+    ulong         src_nb3,
+    int           nblk0,
+    int           ne1,
+    int           ne2,
+    int           ne3,
+    global float * dst
+) {
+    int blk_i0 = get_global_id(0);
+    int i1     = get_global_id(1);
+    int batch  = get_global_id(2);
+
+    if (blk_i0 >= nblk0) return;
+    if (i1     >= ne1)   return;
+
+    int i2 = batch % ne2;
+    int i3 = batch / ne2;
+    if (i3 >= ne3) return;
+
+    global char * block = src + src_offset + (ulong)i3*src_nb3 + (ulong)i2*src_nb2 + (ulong)i1*src_nb1 + (ulong)blk_i0 * (2 + QK4_0/2);
+    float d = vload_half(0, (global half *)block);
+    global uchar * qs = (global uchar *)(block + 2);
+
+    ulong dst_row_base = ((ulong)i3 * ne2 * ne1 + (ulong)i2 * ne1 + (ulong)i1) * nblk0;
+    global float * out = dst + (dst_row_base + blk_i0) * QK4_0;
+
+    for (int i = 0; i < QK4_0/2; ++i) {
+        uchar byte = qs[i];
+        int q0 = (int)(byte & 0x0F) - 8;
+        int q1 = (int)(byte >> 4)   - 8;
+        out[i]            = d * (float)q0;
+        out[i + QK4_0/2]  = d * (float)q1;
+    }
+}
+
+// View-aware AoS q4_0 → f16 dequant (mirrors the q8_0 view variant).
+kernel void kernel_dequant_q4_0_f16_view_aos(
+    global char * src,
+    ulong         src_offset,
+    ulong         src_nb1,
+    ulong         src_nb2,
+    ulong         src_nb3,
+    int           nblk0,
+    int           ne1,
+    int           ne2,
+    int           ne3,
+    global half * dst
+) {
+    int blk_i0 = get_global_id(0);
+    int i1     = get_global_id(1);
+    int batch  = get_global_id(2);
+
+    if (blk_i0 >= nblk0) return;
+    if (i1     >= ne1)   return;
+
+    int i2 = batch % ne2;
+    int i3 = batch / ne2;
+    if (i3 >= ne3) return;
+
+    global char * block = src + src_offset + (ulong)i3*src_nb3 + (ulong)i2*src_nb2 + (ulong)i1*src_nb1 + (ulong)blk_i0 * (2 + QK4_0/2);
+    float d = vload_half(0, (global half *)block);
+    global uchar * qs = (global uchar *)(block + 2);
+
+    ulong dst_row_base = ((ulong)i3 * ne2 * ne1 + (ulong)i2 * ne1 + (ulong)i1) * nblk0;
+    global half * out = dst + (dst_row_base + blk_i0) * QK4_0;
+
+    for (int i = 0; i < QK4_0/2; ++i) {
+        uchar byte = qs[i];
+        int q0 = (int)(byte & 0x0F) - 8;
+        int q1 = (int)(byte >> 4)   - 8;
+        out[i]          = (half)(d * (float)q0);
+        out[i + QK4_0/2] = (half)(d * (float)q1);
+    }
+}
+
+// SoA q8_0 dequant; layout matches kernel_convert_block_q8_0.
+kernel void kernel_dequant_q8_0_f16_soa(
+    global char * src_q,
+    global char * src_d,
+    global half * dst,
+    int n_blocks
+) {
+    int blk = get_global_id(0);
+    if (blk >= n_blocks) return;
+
+    float d = vload_half(0, (global half *)src_d + blk);
+    global char * qs = src_q + blk * QK8_0;
+
+    global half * out = dst + blk * QK8_0;
+    for (int i = 0; i < QK8_0; ++i) {
+        out[i] = (half)(d * (float)qs[i]);
+    }
+}
+
 kernel void kernel_restore_block_q8_0_trans(
     global uchar * src_q,
     global half  * src_d,

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

@@ -45,7 +45,7 @@ kernel void kernel_exp_f16(
     src0 = (global half*)((global char*)src0 + offset0);
     dst  = (global half*)((global char*)dst + offsetd);
 
-    dst[get_global_id(0)] = exp(src0[get_global_id(0)]);
+    dst[get_global_id(0)] = convert_half(exp(convert_float(src0[get_global_id(0)])));
 }
 
 kernel void kernel_exp_f16_4(
@@ -61,7 +61,7 @@ kernel void kernel_exp_f16_4(
     src0 = (global half4*)((global char*)src0 + offset0);
     dst  = (global half4*)((global char*)dst + offsetd);
 
-    dst[get_global_id(0)] = exp(src0[get_global_id(0)]);
+    dst[get_global_id(0)] = convert_half4(exp(convert_float4(src0[get_global_id(0)])));
 }
 
 kernel void kernel_exp_f32_nc(
@@ -120,6 +120,6 @@ kernel void kernel_exp_f16_nc(
         global const half * x = (global const half *)(src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
         global       half * y = (global       half *)(dst  + i3*nb3  + i2*nb2  + i1*nb1  + i0*nb0);
 
-        *y = exp(*x);
+        *y = convert_half(exp(convert_float(*x)));
     }
 }

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

@@ -37,7 +37,10 @@ kernel void kernel_expm1_f16(
     src0 = (global half*)((global char*)src0 + offset0);
     dst  = (global half*)((global char*)dst + offsetd);
 
-    dst[get_global_id(0)] = exp(src0[get_global_id(0)]) - 1.0h;
+    {
+        const float x = convert_float(src0[get_global_id(0)]);
+        dst[get_global_id(0)] = convert_half(exp(x) - 1.0f);
+    }
 }
 
 kernel void kernel_expm1_f16_4(
@@ -49,7 +52,10 @@ kernel void kernel_expm1_f16_4(
     src0 = (global half4*)((global char*)src0 + offset0);
     dst  = (global half4*)((global char*)dst + offsetd);
 
-    dst[get_global_id(0)] = exp(src0[get_global_id(0)]) - 1.0h;
+    {
+        const float4 x = convert_float4(src0[get_global_id(0)]);
+        dst[get_global_id(0)] = convert_half4(exp(x) - 1.0f);
+    }
 }
 
 kernel void kernel_expm1_f32_nc(
@@ -108,6 +114,7 @@ kernel void kernel_expm1_f16_nc(
         global const half * x = (global const half *)(src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00);
         global       half * y = (global       half *)(dst  + i3*nb3  + i2*nb2  + i1*nb1  + i0*nb0);
 
-        *y = exp(*x) - 1.0f;
+        const float fx = convert_float(*x);
+        *y = convert_half(exp(fx) - 1.0f);
     }
 }