Commit e1f7593c3 for llama.cpp
commit e1f7593c32e8dd782fc7addb908d0911260511ae
Author: shaofeiqi <shaoqi@qti.qualcomm.com>
Date: Wed Nov 11 16:18:24 2026 -0800
opencl: add decomposed FA for Adreno (#32046)
* opencl: add decomposed FA
Assisted-by: Claude Opus 5.5
Co-authored-by: Hongqiang Wang <wangh@qti.qualcomm.com>
Co-authored-by: Shaofei Qi <shaoqi@qti.qualcomm.com>
* opencl: build FA MQ kernel on first FA call
* so allows warmup to early build neccessary kernels
Assisted-by: Claude Opus 5.5
* opencl: handle inf for fully masked row
* opencl: refine condition for using fa decomp
* opencl: support dp4a for PV mm
Assisted-by: Codex
* opencl: add direct PV for FA decomp to avoid V transpose for small batch
Assisted-by: Codex
* opencl: enable direct pv by default
* opencl: expand X2E decomposed FA to smaller heads and batches
- Lower the default head dimension threshold from 256 to 128 on X2E.
- Allow DK64 with P8, no attention sinks, GQA >= 2 and n_kv >= max(1024, 2*n_q).
- Allow P8 with V transpose for 2-31 queries when direct PV is unavailable, DK >= 128 and n_kv >= 1024.
Assisted-by: Codex
* opencl: adjust format
---------
Co-authored-by: Li He <lih@qti.qualcomm.com>
Co-authored-by: Hongqiang Wang <wangh@qti.qualcomm.com>
diff --git a/ggml/src/ggml-opencl/CMakeLists.txt b/ggml/src/ggml-opencl/CMakeLists.txt
index e2a82d5bc..0570fb835 100644
--- a/ggml/src/ggml-opencl/CMakeLists.txt
+++ b/ggml/src/ggml-opencl/CMakeLists.txt
@@ -322,11 +322,14 @@ set(GGML_OPENCL_KERNELS
repeat
mul_mat_f16_f32
mul_mm_f16_f32_kq_kqv
+ mul_mm_f16_f32_kq_p8
+ mul_mm_q8_f32_kqv
conv2d
conv2d_f16_f32
conv2_dw
flash_attn_repack
flash_attn_pre_f16
+ flash_attn_pv_direct
flash_attn_f32_f16
flash_attn_f32_q8_0
flash_attn_f32_q4_0
diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp
index c25a13dc6..53ad1cc40 100644
--- a/ggml/src/ggml-opencl/ggml-opencl.cpp
+++ b/ggml/src/ggml-opencl/ggml-opencl.cpp
@@ -79,6 +79,11 @@ bool ggml_cl_compute_forward(ggml_backend_t backend, struct ggml_tensor * tensor
static bool ggml_cl_is_q4_0_soa(const ggml_tensor * tensor);
static bool ggml_cl_is_q8_0_soa(const ggml_tensor * tensor);
static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst);
+static void ggml_cl_cpy(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst);
+static void ggml_cl_fa_decompose_soft_max(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst, cl_mem sums);
+#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
+static void ggml_cl_mul_mat_kq_kqv_adreno(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool is_kq);
+#endif
// See https://gmplib.org/~tege/divcnst-pldi94.pdf figure 4.1.
// Precompute mp (m' in the paper) and L such that division
@@ -571,6 +576,8 @@ struct ggml_opencl_fa_kernels {
std::map<std::pair<int, int>, cl_kernel> kv_pad_f16;
std::map<std::pair<int, int>, cl_kernel> mask_pad_f16;
std::map<std::pair<int, int>, cl_kernel> blk_f16;
+ // decomposed prefill: V^T, built with the prepass
+ std::map<std::pair<int, int>, cl_kernel> v_transpose_f16;
// generic prefill tile dims (f16 / f32 paths)
std::map<std::pair<int, int>, int> bm;
std::map<std::pair<int, int>, int> bn;
@@ -760,6 +767,18 @@ struct ggml_backend_opencl_context {
ggml_cl_buffer prealloc_src1;
ggml_cl_buffer prealloc_top_k;
+ // scratch for the decomposed flash attention prefill, sized per query chunk
+ ggml_cl_buffer prealloc_fa_vt;
+ ggml_cl_buffer prealloc_fa_kq;
+ ggml_cl_buffer prealloc_fa_kqv;
+ ggml_cl_buffer prealloc_fa_sums;
+ ggml_cl_buffer prealloc_fa_vtq;
+ ggml_cl_buffer prealloc_fa_vtd;
+ ggml_cl_buffer prealloc_fa_pq;
+ ggml_cl_buffer prealloc_fa_pd;
+ ggml_cl_buffer prealloc_fa_bmax;
+ ggml_cl_buffer prealloc_fa_bsum;
+
#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
ggml_cl_buffer prealloc_adreno_xmem_const;
bool adreno_xmem_gemm_enabled = false;
@@ -871,6 +890,13 @@ struct ggml_backend_opencl_context {
cl_kernel kernel_diag_f32;
cl_kernel kernel_soft_max, kernel_soft_max_4;
cl_kernel kernel_soft_max_f16, kernel_soft_max_4_f16;
+ cl_kernel kernel_fa_decompose_soft_max_4_f16, kernel_fa_decompose_scale_rows_f32;
+ cl_kernel kernel_mul_mm_f16_f32_kq_p8 = nullptr;
+ cl_kernel kernel_mul_mm_q8_f32_kqv = nullptr;
+ cl_kernel kernel_fa_p8_fixup = nullptr;
+ cl_kernel kernel_fa_kqv_direct_f16 = nullptr;
+ cl_kernel kernel_fa_kqv_direct_reduce = nullptr;
+ cl_kernel kernel_fa_v_transpose_q8 = nullptr;
ggml_opencl_fa_kernels fa;
#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
ggml_cl_adreno_xmem_attn_state adreno_xmem_attn;
@@ -1473,6 +1499,11 @@ struct ggml_backend_opencl_context {
#endif // GGML_OPENCL_USE_ADRENO_KERNELS
~ggml_backend_opencl_context() {
+ for (cl_kernel kernel : {kernel_mul_mm_f16_f32_kq_p8, kernel_mul_mm_q8_f32_kqv, kernel_fa_p8_fixup, kernel_fa_v_transpose_q8, kernel_fa_kqv_direct_f16, kernel_fa_kqv_direct_reduce}) {
+ if (kernel != nullptr) {
+ CL_CHECK(clReleaseKernel(kernel));
+ }
+ }
if (kernel_dsv4_hc_pre) {
CL_CHECK(clReleaseKernel(kernel_dsv4_hc_pre));
}
@@ -3972,6 +4003,26 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
GGML_LOG_CONT(".");
}
+#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
+ // mul_mm_f16_f32_kq_p8
+ if (backend_ctx->gpu_family == ADRENO && backend_ctx->has_integer_dot) {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "mul_mm_f16_f32_kq_p8.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("mul_mm_f16_f32_kq_p8.cl");
+#endif
+ cl_program prog = build_program_from_source_ex(backend_ctx->context, backend_ctx->device, kernel_src.c_str(),
+ compile_opts, /*fatal=*/false);
+ if (prog != nullptr) {
+ CL_CHECK((backend_ctx->kernel_mul_mm_f16_f32_kq_p8 = clCreateKernel(prog, "mul_mm_f16_f32_kq_p8", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ }
+ GGML_LOG_CONT(".");
+ }
+#endif
+
// mul
{
#ifdef GGML_OPENCL_EMBED_KERNELS
@@ -4178,8 +4229,50 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
build_program_from_source(backend_ctx, kernel_src.c_str(), compile_opts);
CL_CHECK((backend_ctx->kernel_soft_max_4_f16 = clCreateKernel(backend_ctx->program_softmax_4_f16, "kernel_soft_max_4_f16", &err), err));
+ CL_CHECK((backend_ctx->kernel_fa_decompose_soft_max_4_f16 = clCreateKernel(backend_ctx->program_softmax_4_f16, "kernel_fa_decompose_soft_max_4_f16", &err), err));
+ CL_CHECK((backend_ctx->kernel_fa_decompose_scale_rows_f32 = clCreateKernel(backend_ctx->program_softmax_4_f16, "kernel_fa_decompose_scale_rows_f32", &err), err));
+ CL_CHECK((backend_ctx->kernel_fa_p8_fixup = clCreateKernel(backend_ctx->program_softmax_4_f16, "kernel_fa_p8_fixup", &err), err));
+ GGML_LOG_CONT(".");
+ }
+
+#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
+ if (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E ||
+ backend_ctx->adreno_gen == ADRENO_GPU_GEN::A8X ||
+ backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E) {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "flash_attn_pv_direct.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("flash_attn_pv_direct.cl");
+#endif
+ cl_program prog = build_program_from_source_ex(backend_ctx->context, backend_ctx->device, kernel_src.c_str(),
+ compile_opts, /*fatal=*/false);
+ if (prog != nullptr) {
+ CL_CHECK((backend_ctx->kernel_fa_kqv_direct_f16 = clCreateKernel(prog, "kernel_fa_kqv_direct_f16", &err), err));
+ CL_CHECK((backend_ctx->kernel_fa_kqv_direct_reduce = clCreateKernel(prog, "kernel_fa_kqv_direct_reduce", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ }
+ GGML_LOG_CONT(".");
+ }
+
+ if (backend_ctx->gpu_family == ADRENO && backend_ctx->has_integer_dot) {
+#ifdef GGML_OPENCL_EMBED_KERNELS
+ const std::string kernel_src {
+ #include "mul_mm_q8_f32_kqv.cl.h"
+ };
+#else
+ const std::string kernel_src = read_file("mul_mm_q8_f32_kqv.cl");
+#endif
+ cl_program prog = build_program_from_source_ex(backend_ctx->context, backend_ctx->device, kernel_src.c_str(),
+ compile_opts, /*fatal=*/false);
+ if (prog != nullptr) {
+ CL_CHECK((backend_ctx->kernel_mul_mm_q8_f32_kqv = clCreateKernel(prog, "kernel_mul_mm_q8_f32_kqv", &err), err));
+ CL_CHECK(clReleaseProgram(prog));
+ }
GGML_LOG_CONT(".");
}
+#endif
// div
{
@@ -7825,6 +7918,17 @@ static void ggml_opencl_ensure_fa_pre_kernels(ggml_backend_opencl_context * back
backend_ctx->fa.kv_pad_f16[{dk, dv}] = k_kv_pad_f16;
backend_ctx->fa.mask_pad_f16[{dk, dv}] = k_mask_pad_f16;
backend_ctx->fa.blk_f16[{dk, dv}] = k_blk_f16;
+ // optional, the decomposed prefill checks for it and falls back to the fused tile
+ cl_kernel k_v_transpose_f16 = clCreateKernel(prog_pre_f16, "flash_attn_v_transpose_f16", &err);
+ if (err == CL_SUCCESS) {
+ backend_ctx->fa.v_transpose_f16[{dk, dv}] = k_v_transpose_f16;
+ }
+ if (backend_ctx->kernel_fa_v_transpose_q8 == nullptr && backend_ctx->has_integer_dot) {
+ backend_ctx->kernel_fa_v_transpose_q8 = clCreateKernel(prog_pre_f16, "flash_attn_v_transpose_q8", &err);
+ if (err != CL_SUCCESS) {
+ backend_ctx->kernel_fa_v_transpose_q8 = nullptr;
+ }
+ }
clReleaseProgram(prog_pre_f16);
}
@@ -22699,6 +22803,610 @@ static void ggml_cl_flash_attn_prefill_bin(ggml_backend_t backend, const ggml_te
}
#endif // GGML_OPENCL_USE_ADRENO_KERNELS
+static void ggml_cl_fa_kq_p8_dispatch(
+ ggml_backend_t backend,
+ const ggml_tensor * q, const ggml_tensor * k,
+ const ggml_tensor * mask, const ggml_tensor * sinks,
+ const ggml_tensor * dst,
+ int64_t q0, int nqc, float scale, float max_bias, int p_pitch
+) {
+ auto * backend_ctx = (ggml_backend_opencl_context *) backend->context;
+ auto * extra_k = (ggml_tensor_extra_cl *) (k->view_src ? k->view_src->extra : k->extra);
+ auto * extra_q = (ggml_tensor_extra_cl *) q->extra;
+ auto * extra_m = (ggml_tensor_extra_cl *) mask->extra;
+
+ const int M = k->ne[1];
+ const int K = k->ne[0];
+ const int N = nqc;
+ const int D_A = k->ne[2];
+ const int D_B = q->ne[2];
+
+ cl_int err;
+
+ cl_buffer_region region = {extra_k->offset + k->view_offs, k->nb[1] * k->ne[1]};
+ cl_mem k_sub = clCreateSubBuffer(extra_k->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &err);
+ CL_CHECK(err);
+
+ cl_image_format fmt = {CL_RGBA, CL_FLOAT};
+ cl_image_desc desc = {};
+ desc.image_type = CL_MEM_OBJECT_IMAGE1D_BUFFER;
+ desc.image_width = region.size / 16;
+ desc.buffer = k_sub;
+ cl_mem k_img = clCreateImage(backend_ctx->context, CL_MEM_READ_ONLY, &fmt, &desc, nullptr, &err);
+ CL_CHECK(err);
+
+ region.origin = extra_q->offset + q->view_offs + q0 * q->nb[1];
+ region.size = (size_t) K * N * D_B * sizeof(float);
+ cl_mem q_sub = clCreateSubBuffer(extra_q->data_device, 0, CL_BUFFER_CREATE_TYPE_REGION, ®ion, &err);
+ CL_CHECK(err);
+
+ const cl_ulong off_mask = extra_m->offset + mask->view_offs + q0 * mask->nb[1];
+ const cl_ulong mask_nb1 = mask->nb[1];
+ const int n_head_log2 = 1u << (unsigned) floorf(log2f((float) D_B));
+ const float m0 = powf(2.0f, -max_bias / n_head_log2);
+ const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2);
+
+ cl_uint a = 0;
+
+ cl_kernel kernel = backend_ctx->kernel_mul_mm_f16_f32_kq_p8;
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &k_img));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &q_sub));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_pq.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_bmax.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_bsum.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &extra_m->data_device));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &off_mask));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &mask_nb1));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(float), &scale));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(float), &max_bias));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(float), &m0));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(float), &m1));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &n_head_log2));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &M));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &K));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &N));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &D_A));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &D_B));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &p_pitch));
+
+ const size_t n_tiles = (N + 31) / 32;
+ size_t gws[3] = {64, n_tiles, (size_t)(M / 64) * D_B};
+
+ const bool two_tiles = n_tiles % 2 == 0 && backend_ctx->get_kernel_workgroup_size(kernel) >= 128;
+ size_t lws[3] = {64, two_tiles ? (size_t) 2 : (size_t) 1, 1};
+
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws, lws, dst);
+ CL_CHECK(clReleaseMemObject(k_img));
+ CL_CHECK(clReleaseMemObject(k_sub));
+ CL_CHECK(clReleaseMemObject(q_sub));
+
+ cl_mem sinks_buf = extra_q->data_device;
+ cl_ulong sinks_off = 0;
+ const int has_sinks = sinks != nullptr;
+
+ if (sinks) {
+ auto * extra_s = (ggml_tensor_extra_cl *) sinks->extra;
+ sinks_buf = extra_s->data_device;
+ sinks_off = extra_s->offset + sinks->view_offs;
+ }
+
+ const int nblk = M / 32;
+ const int row_blk = p_pitch / 32;
+
+ kernel = backend_ctx->kernel_fa_p8_fixup;
+
+ a = 0;
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_bmax.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_bsum.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_pd.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_sums.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &sinks_buf));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &sinks_off));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &has_sinks));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &nblk));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &row_blk));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &N));
+
+ size_t gws_fixup[3] = {(size_t) 64 * N, (size_t) D_B, 1};
+ size_t lws_fixup[3] = {64, 1, 1};
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws_fixup, lws_fixup, dst);
+}
+
+static void ggml_cl_fa_decompose_scratch_tensor(
+ ggml_tensor & t, ggml_tensor_extra_cl & extra, cl_mem buffer, ggml_type type,
+ int64_t ne0, int64_t ne1, int64_t ne2
+) {
+ memset(&t, 0, sizeof(t));
+ t.type = type;
+ t.ne[0] = ne0; t.ne[1] = ne1; t.ne[2] = ne2; t.ne[3] = 1;
+ t.nb[0] = ggml_type_size(type);
+ t.nb[1] = t.nb[0]*ne0;
+ t.nb[2] = t.nb[1]*ne1;
+ t.nb[3] = t.nb[2]*ne2;
+
+ extra.reset();
+ extra.data_device = buffer;
+ t.extra = &extra;
+}
+
+static void ggml_cl_fa_decompose_write_chunk(
+ ggml_backend_t backend, const ggml_tensor & kqv, int64_t q0, ggml_tensor * dst
+) {
+ auto * backend_ctx = (ggml_backend_opencl_context *) backend->context;
+ const int64_t dv = kqv.ne[0];
+ const int64_t nqc = kqv.ne[1];
+ const int64_t n_head = kqv.ne[2];
+
+ // apply the deferred 1/sum to the KQV rows
+ {
+ cl_kernel kernel = backend_ctx->kernel_fa_decompose_scale_rows_f32;
+ const cl_ulong offset_kqv = 0;
+ const cl_ulong offset_sums = 0;
+ const int dv_i = (int) dv;
+ const int nqc_i = (int) nqc;
+
+ CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &backend_ctx->prealloc_fa_kqv.buffer));
+ CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset_kqv));
+ CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &backend_ctx->prealloc_fa_sums.buffer));
+ CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &offset_sums));
+ CL_CHECK(clSetKernelArg(kernel, 4, sizeof(int), &dv_i));
+ CL_CHECK(clSetKernelArg(kernel, 5, sizeof(int), &nqc_i));
+
+ size_t gws[3] = { (size_t) 64 * nqc, (size_t) n_head, 1 };
+ size_t lws[3] = { 64, 1, 1 };
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws, lws, dst);
+ }
+
+ // permute into dst: [dv, nqc, n_head] -> dst[dv, n_head, q0 .. q0+nqc)
+ {
+ ggml_tensor kqv_perm = kqv;
+ kqv_perm.ne[1] = n_head;
+ kqv_perm.ne[2] = nqc;
+ kqv_perm.nb[1] = kqv.nb[2];
+ kqv_perm.nb[2] = kqv.nb[1];
+
+ ggml_tensor dst_chunk = *dst;
+ dst_chunk.ne[2] = nqc;
+ dst_chunk.view_offs = dst->view_offs + (size_t) q0 * dst->nb[2];
+
+ ggml_cl_cpy(backend, &kqv_perm, &dst_chunk, nullptr);
+ }
+}
+
+static void ggml_cl_flash_attn_decompose(
+ ggml_backend_t backend,
+ const ggml_tensor * q, const ggml_tensor * k, ggml_tensor * dst,
+ int64_t n_q_chunk, float scale, float max_bias
+) {
+ auto * backend_ctx = (ggml_backend_opencl_context *) backend->context;
+ const ggml_tensor * v = dst->src[2];
+ const ggml_tensor * mask = dst->src[3];
+ const ggml_tensor * sinks = dst->src[4];
+ const int64_t n_q = q->ne[1];
+ const int64_t n_kv = k->ne[1];
+ const int64_t dv = v->ne[0];
+ const int64_t n_head = q->ne[2];
+ const int64_t n_head_kv = k->ne[2];
+
+ const std::pair<int, int> dk_dv = {(int) q->ne[0], (int) dv};
+
+ // The image GEMMs work in 64 x 32 tiles and can touch past the exact extent,
+ // so size each scratch by the padded shape plus one extra head.
+ auto pad_up = [](int64_t x, int64_t a) { return ((x + a - 1) / a) * a; };
+ const int64_t kq_m = pad_up(n_kv, 64);
+ const int64_t kqv_m = pad_up(dv, 64);
+ const int64_t tile_n = pad_up(n_q_chunk, 32);
+
+ backend_ctx->prealloc_fa_vt.allocate(backend_ctx->context, (size_t)(n_kv * kqv_m * (n_head_kv + 1)) * sizeof(ggml_fp16_t));
+ backend_ctx->prealloc_fa_kq.allocate(backend_ctx->context, (size_t)(kq_m * tile_n * (n_head + 1)) * sizeof(float));
+ backend_ctx->prealloc_fa_kqv.allocate(backend_ctx->context, (size_t)(kqv_m * tile_n * (n_head + 1)) * sizeof(float));
+ backend_ctx->prealloc_fa_sums.allocate(backend_ctx->context, (size_t)(n_q_chunk * n_head) * sizeof(float));
+
+ // V^T once per call: [dv, n_kv, n_head_kv] -> [n_kv, dv, n_head_kv]
+ {
+ ggml_tensor_extra_cl * extra_v = (ggml_tensor_extra_cl *)v->extra;
+ cl_kernel kernel = backend_ctx->fa.v_transpose_f16.at(dk_dv);
+
+ const cl_ulong offset_v = extra_v->offset + v->view_offs;
+ const cl_ulong v_nb1 = v->nb[1];
+ const cl_ulong v_nb2 = v->nb[2];
+ const int n_kv_i = (int) n_kv;
+ const int dv_i = (int) dv;
+ const int n_head_kv_i = (int) n_head_kv;
+
+ cl_uint a = 0;
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &extra_v->data_device));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &offset_v));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_vt.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &n_kv_i));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &dv_i));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &n_head_kv_i));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &v_nb1));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &v_nb2));
+
+ size_t gws[3] = { (size_t)((dv + 31) / 32) * 32, (size_t)((n_kv + 31) / 32) * 8, (size_t) n_head_kv };
+ size_t lws[3] = { 32, 8, 1 };
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws, lws, dst);
+ }
+
+ ggml_tensor_extra_cl extra_vt;
+ ggml_tensor_extra_cl extra_kq;
+ ggml_tensor_extra_cl extra_kqv;
+ ggml_tensor vt;
+ ggml_tensor kq;
+ ggml_tensor kqv;
+ ggml_cl_fa_decompose_scratch_tensor(vt, extra_vt, backend_ctx->prealloc_fa_vt.buffer, GGML_TYPE_F16, n_kv, dv, n_head_kv);
+
+ for (int64_t q0 = 0; q0 < n_q; q0 += n_q_chunk) {
+ const int64_t nqc = MIN(n_q_chunk, n_q - q0);
+
+ // KQ: [dk, n_kv, n_head_kv] x [dk, nqc, n_head] -> [n_kv, nqc, n_head]
+ ggml_cl_fa_decompose_scratch_tensor(kq, extra_kq, backend_ctx->prealloc_fa_kq.buffer, GGML_TYPE_F32, n_kv, nqc, n_head);
+
+ {
+ ggml_tensor q_chunk = *q;
+ q_chunk.ne[1] = nqc;
+ q_chunk.view_offs = q->view_offs + (size_t) q0 * q->nb[1];
+ ggml_cl_mul_mat(backend, k, &q_chunk, &kq);
+ }
+
+ // softmax in place, the row sums go to prealloc_fa_sums
+ {
+ memcpy((float *) kq.op_params + 0, &scale, sizeof(float));
+ memcpy((float *) kq.op_params + 1, &max_bias, sizeof(float));
+
+ ggml_tensor mask_chunk = *mask;
+ mask_chunk.ne[1] = nqc;
+ mask_chunk.view_offs = mask->view_offs + (size_t) q0 * mask->nb[1];
+
+ ggml_cl_fa_decompose_soft_max(backend, &kq, &mask_chunk, sinks, &kq, backend_ctx->prealloc_fa_sums.buffer);
+ }
+
+ // KQV: [n_kv, dv, n_head_kv] x [n_kv, nqc, n_head] -> [dv, nqc, n_head]
+ // V^T is contiguous, so mul_mat would not pick the KQV image kernel; call it here.
+ ggml_cl_fa_decompose_scratch_tensor(kqv, extra_kqv, backend_ctx->prealloc_fa_kqv.buffer, GGML_TYPE_F32, dv, nqc, n_head);
+#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
+ if (backend_ctx->kernel_mul_mm_f16_f32_kqv != nullptr &&
+ (dv % 64) == 0 && (n_kv % 16) == 0 && nqc >= 32 &&
+ ((size_t)(vt.nb[2] * n_head_kv / 4) / 4 <= backend_ctx->image_max_buffer_size) &&
+ ((size_t)(dv * nqc * n_head) <= backend_ctx->image_max_buffer_size)) {
+ ggml_cl_mul_mat_kq_kqv_adreno(backend, &vt, &kq, &kqv, /*is_kq =*/ false);
+ } else
+#endif
+ {
+ ggml_cl_mul_mat(backend, &vt, &kq, &kqv);
+ }
+
+ ggml_cl_fa_decompose_write_chunk(backend, kqv, q0, dst);
+ }
+}
+
+static void ggml_cl_fa_pv_direct(ggml_backend_t backend, ggml_tensor * dst, int kv_pitch) {
+ auto * backend_ctx = (ggml_backend_opencl_context *) backend->context;
+ const ggml_tensor * q = dst->src[0];
+ const ggml_tensor * v = dst->src[2];
+ auto * extra_v = (ggml_tensor_extra_cl *) v->extra;
+
+ const int dv = v->ne[0];
+ const int n_q = q->ne[1];
+ const int n_head = q->ne[2];
+ const int n_head_kv = v->ne[2];
+ const int gqa = n_head / n_head_kv;
+ const int nblk = v->ne[1] / 32;
+ const int d_blocks = dv / 128;
+ // Split the KV range to give a small query batch about 256 workgroups.
+ const int want = MAX(1, 256 / (d_blocks * n_head_kv));
+ const int blocks_per_split = (nblk + want - 1) / want;
+ const int n_split = (nblk + blocks_per_split - 1) / blocks_per_split;
+ const int n_elem = n_head * n_q * dv;
+ backend_ctx->prealloc_fa_kq.allocate(backend_ctx->context, (size_t) n_split * n_elem * sizeof(float));
+
+ const cl_ulong off_v = extra_v->offset + v->view_offs;
+ const cl_ulong v_nb1 = v->nb[1];
+ const cl_ulong v_nb2 = v->nb[2];
+ {
+ cl_kernel kernel = backend_ctx->kernel_fa_kqv_direct_f16;
+ cl_uint a = 0;
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &extra_v->data_device));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &off_v));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &v_nb1));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &v_nb2));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_pq.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_pd.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_kq.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &dv));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &n_q));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &gqa));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &n_head));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &nblk));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &blocks_per_split));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &kv_pitch));
+ size_t gws[3] = {(size_t) d_blocks * 64, (size_t) n_head_kv, (size_t) n_split};
+ size_t lws[3] = {64, 1, 1};
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws, lws, dst);
+ }
+ {
+ cl_kernel kernel = backend_ctx->kernel_fa_kqv_direct_reduce;
+ cl_uint a = 0;
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_kq.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_kqv.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &n_elem));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &n_split));
+ size_t gws[3] = {((size_t) n_elem + 63) / 64 * 64, 1, 1};
+ size_t lws[3] = {64, 1, 1};
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws, lws, dst);
+ }
+}
+
+static void ggml_cl_flash_attn_decompose_kq_p8(
+ ggml_backend_t backend,
+ const ggml_tensor * q, const ggml_tensor * k,
+ ggml_tensor * dst,
+ int64_t n_q_chunk, float scale, float max_bias, bool kqv_direct
+) {
+ auto * backend_ctx = (ggml_backend_opencl_context *) backend->context;
+ const ggml_tensor * v = dst->src[2];
+ const ggml_tensor * mask = dst->src[3];
+ const ggml_tensor * sinks = dst->src[4];
+ const int64_t n_q = q->ne[1];
+ const int64_t n_kv = k->ne[1];
+ const int64_t dv = v->ne[0];
+ const int64_t n_head = q->ne[2];
+ const int64_t n_head_kv = k->ne[2];
+
+ auto pad_up = [](int64_t x, int64_t a) { return ((x + a - 1) / a) * a; };
+ const int64_t kqv_m = pad_up(dv, 64);
+ const int64_t tile_n = pad_up(n_q_chunk, 32);
+
+ // Pad the int8 row pitch by one cache line, as in the P8 KQV kernel.
+ const int kv_pitch = (int) n_kv + 64;
+ const size_t nblk = (size_t) kv_pitch / 32;
+
+ backend_ctx->prealloc_fa_pq.allocate(backend_ctx->context, (size_t) kv_pitch * n_q_chunk * n_head);
+ backend_ctx->prealloc_fa_pd.allocate(backend_ctx->context, nblk * n_q_chunk * n_head * sizeof(float));
+ backend_ctx->prealloc_fa_bmax.allocate(backend_ctx->context, nblk * n_q_chunk * n_head * sizeof(float));
+ backend_ctx->prealloc_fa_bsum.allocate(backend_ctx->context, nblk * n_q_chunk * n_head * sizeof(float));
+ backend_ctx->prealloc_fa_kqv.allocate(backend_ctx->context, (size_t)(kqv_m * tile_n * (n_head + 1)) * sizeof(float));
+ backend_ctx->prealloc_fa_sums.allocate(backend_ctx->context, (size_t)(n_q_chunk * n_head) * sizeof(float));
+
+ // Direct PV reads the original V cache and needs no transposed copy.
+ if (!kqv_direct) {
+ backend_ctx->prealloc_fa_vtq.allocate(backend_ctx->context, (size_t) kv_pitch * dv * n_head_kv);
+ backend_ctx->prealloc_fa_vtd.allocate(backend_ctx->context, nblk * dv * n_head_kv * sizeof(cl_half));
+
+ ggml_tensor_extra_cl * extra_v = (ggml_tensor_extra_cl *)v->extra;
+ cl_kernel kernel = backend_ctx->kernel_fa_v_transpose_q8;
+
+ const cl_ulong offset_v = extra_v->offset + v->view_offs;
+ const cl_ulong v_nb1 = v->nb[1];
+ const cl_ulong v_nb2 = v->nb[2];
+ const cl_ulong v_nb3 = v->nb[3];
+ const int n_kv_i = (int) n_kv;
+ const int dv_i = (int) dv;
+ const int n_head_kv_i = (int) n_head_kv;
+
+ cl_uint a = 0;
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &extra_v->data_device));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &offset_v));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_vtq.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_vtd.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &n_kv_i));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &dv_i));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &n_head_kv_i));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &v_nb1));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &v_nb2));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &v_nb3));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &kv_pitch));
+
+ size_t gws[3] = { (size_t)((dv + 31) / 32) * 32, (size_t)((n_kv + 31) / 32) * 8, (size_t) n_head_kv };
+ size_t lws[3] = { 32, 8, 1 };
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws, lws, dst);
+ }
+
+ ggml_tensor_extra_cl extra_kqv;
+ ggml_tensor kqv;
+
+ for (int64_t q0 = 0; q0 < n_q; q0 += n_q_chunk) {
+ const int64_t nqc = MIN(n_q_chunk, n_q - q0);
+
+ ggml_cl_fa_kq_p8_dispatch(backend, q, k, mask, sinks, dst, q0, nqc, scale, max_bias, kv_pitch);
+
+ ggml_cl_fa_decompose_scratch_tensor(kqv, extra_kqv, backend_ctx->prealloc_fa_kqv.buffer, GGML_TYPE_F32, dv, nqc, n_head);
+
+ if (kqv_direct) {
+ GGML_ASSERT(q0 == 0 && nqc == n_q);
+ ggml_cl_fa_pv_direct(backend, dst, kv_pitch);
+ } else {
+ const cl_ulong zero = 0;
+ const int dv_i = dv;
+ const int nkv_i = n_kv;
+ const int nqc_i = nqc;
+ const int nh_i = n_head;
+ const int nhkv_i = n_head_kv;
+
+ cl_kernel kernel = backend_ctx->kernel_mul_mm_q8_f32_kqv;
+
+ cl_uint a = 0;
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_vtq.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &zero));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_vtd.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &zero));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_pq.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &zero));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_pd.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &zero));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_mem), &backend_ctx->prealloc_fa_kqv.buffer));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(cl_ulong), &zero));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &dv_i));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &nkv_i));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &nqc_i));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &nh_i));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &nhkv_i));
+ CL_CHECK(clSetKernelArg(kernel, a++, sizeof(int), &kv_pitch));
+ size_t gws[3] = {(size_t)(dv / 64) * 64, (size_t)(n_head / n_head_kv) * ((nqc + 31) / 32), (size_t) n_head_kv};
+ size_t lws[3] = {64, 1, 1};
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, gws, lws, dst);
+ }
+
+ ggml_cl_fa_decompose_write_chunk(backend, kqv, q0, dst);
+ }
+}
+
+// Decomposed flash attention prefill.
+//
+// On Adreno the fused prefill tile is slow at large head sizes, while the
+// KQ / KQV image GEMMs are fast. -fa 0 uses those GEMMs, but it keeps KQ for
+// the whole ubatch in the graph. Here the FLASH_ATTN_EXT node stays, and the
+// queries are split into chunks. For each chunk: KQ -> softmax -> KQV into
+// backend scratch, so the scratch size depends on the chunk and not on the
+// ubatch, and the KV cache keeps the flash attention layout.
+//
+// Returns false when the shape is not supported, then the fused tile runs.
+// GGML_OPENCL_FA_DECOMPOSE=0 disables it, =1 enables it for any Adreno GPU and
+// head size. GGML_OPENCL_FA_DECOMPOSE_MB sets the KQ scratch budget (default 256).
+// GGML_OPENCL_FA_KQ_P8=0 disables U8 probabilities and int8 KQV.
+// GGML_OPENCL_FA_KQV_DIRECT=0 disables direct PV for small query batches.
+static bool ggml_cl_flash_attn_decompose_dispatch(
+ ggml_backend_t backend,
+ const ggml_tensor * q, const ggml_tensor * k, ggml_tensor * dst
+) {
+ ggml_backend_opencl_context * backend_ctx = (ggml_backend_opencl_context *)backend->context;
+
+ const ggml_tensor * v = dst->src[2];
+ const ggml_tensor * mask = dst->src[3];
+
+ if (backend_ctx->gpu_family != ADRENO) {
+ return false;
+ }
+
+ static const int env = []{
+ const char * e = getenv("GGML_OPENCL_FA_DECOMPOSE");
+ return (e && e[0]) ? (e[0] != '0' ? 1 : 0) : -1;
+ }();
+
+ if (env == 0) {
+ return false;
+ }
+
+ const int64_t n_q = q->ne[1];
+ const int64_t n_kv = k->ne[1];
+ const int64_t dk = q->ne[0];
+ const int64_t dv = v->ne[0];
+ const int64_t n_head = q->ne[2];
+ const int64_t n_head_kv = k->ne[2];
+
+ const bool measured_gen = backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E ||
+ backend_ctx->adreno_gen == ADRENO_GPU_GEN::A8X ||
+ backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E;
+
+ if (n_q < 2) {
+ return false;
+ }
+ if (q->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32 ||
+ k->type != GGML_TYPE_F16 || v->type != GGML_TYPE_F16) {
+ return false;
+ }
+ if (q->ne[3] != 1 || k->ne[3] != 1 || v->ne[3] != 1 || dst->ne[3] != 1) {
+ return false;
+ }
+ if (n_head_kv <= 0 || n_head < n_head_kv || n_head % n_head_kv != 0) {
+ return false;
+ }
+ if (k->nb[0] != ggml_type_size(k->type) || v->nb[0] != ggml_type_size(v->type)) {
+ return false;
+ }
+ // K and Q must be packed [dk, head, row], as the KV cache and the permuted Q
+ // are. mul_mat sends the KQ to the image KQ kernel only for this layout;
+ // other K views can reach the KQV image kernel there and give wrong results.
+ if (k->nb[2] != (size_t) dk * sizeof(ggml_fp16_t) || k->nb[1] != (size_t) dk * n_head_kv * sizeof(ggml_fp16_t) ||
+ q->nb[0] != sizeof(float) || q->nb[2] != (size_t) dk * sizeof(float) || q->nb[1] != (size_t) dk * n_head * sizeof(float)) {
+ return false;
+ }
+
+ // soft_max applies scale, mask, ALiBi and sinks, but not the softcap.
+ float scale;
+ float max_bias;
+ float logit_softcap;
+ memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float));
+ memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float));
+ memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float));
+ if (logit_softcap != 0.0f) {
+ return false;
+ }
+
+ // The deferred softmax kernel needs an f16 mask and rows that divide by 4.
+ if (mask == nullptr || mask->type != GGML_TYPE_F16 || mask->ne[0] < n_kv || (n_kv % 4) != 0) {
+ return false;
+ }
+
+ const std::pair<int, int> dk_dv = {(int) dk, (int) dv};
+ if (backend_ctx->fa.v_transpose_f16.count(dk_dv) == 0) {
+ return false;
+ }
+
+ static const char * p8_env = getenv("GGML_OPENCL_FA_KQ_P8");
+ const bool kq_p8 = (!p8_env || p8_env[0] != '0') && measured_gen &&
+ backend_ctx->kernel_mul_mm_f16_f32_kq_p8 != nullptr &&
+ backend_ctx->kernel_mul_mm_q8_f32_kqv != nullptr &&
+ backend_ctx->kernel_fa_v_transpose_q8 != nullptr &&
+ n_kv <= INT_MAX - 64 && k->nb[1] <= INT_MAX &&
+ (n_kv % 64) == 0 && (dk % 16) == 0 && (dv % 64) == 0 &&
+ mask->ne[2] == 1 && mask->ne[3] == 1 && mask->nb[0] == sizeof(ggml_fp16_t) &&
+ (mask->nb[1] % sizeof(ggml_fp16_t)) == 0 &&
+ k->nb[1] * n_kv / 16 <= backend_ctx->image_max_buffer_size;
+
+ // Keep DK64 with attention sinks opt-in until its accuracy is validated.
+ const bool dk64_ok = kq_p8 && dst->src[4] == nullptr && n_head / n_head_kv >= 2 &&
+ n_kv >= 1024 && n_q <= n_kv / 2;
+ const int min_dk = backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E ? (dk64_ok ? 64 : 128) : 256;
+ if (env != 1 && !(measured_gen && dk >= min_dk)) {
+ return false;
+ }
+
+ static const char * direct_env = getenv("GGML_OPENCL_FA_KQV_DIRECT");
+ const bool kqv_direct = (!direct_env || direct_env[0] != '0') && kq_p8 &&
+ backend_ctx->kernel_fa_kqv_direct_f16 != nullptr && backend_ctx->kernel_fa_kqv_direct_reduce != nullptr &&
+ v->ne[1] == n_kv && v->ne[2] == n_head_kv &&
+ n_q <= 32 / (n_head / n_head_kv) && dv >= 128 && dv <= INT_MAX && (dv % 128) == 0 &&
+ n_head <= INT_MAX / (n_q * dv) && (v->nb[1] % 4) == 0 && (v->nb[2] % 4) == 0;
+ const bool small_p8 = backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E &&
+ kq_p8 && dk >= 128 && n_kv >= 1024;
+ // Small batches without direct PV need P8 and a long KV cache on X2E.
+ if (n_q < 32 && !kqv_direct && !small_p8) {
+ return false;
+ }
+
+ // Chunk size: a byte budget for KQ, at least 8 tiles of 32 queries so the
+ // GEMMs get usable shapes, and below the image1d_buffer limit of the KQ dst.
+ static const size_t budget_bytes = []{
+ const char * e = getenv("GGML_OPENCL_FA_DECOMPOSE_MB");
+ const int mb = (e && e[0]) ? atoi(e) : 256;
+ return (size_t)(mb > 0 ? mb : 256) * 1024 * 1024;
+ }();
+
+ const size_t kq_row_bytes = (size_t) n_kv * n_head * sizeof(float);
+ int64_t n_q_chunk = MIN((int64_t)(budget_bytes / kq_row_bytes), n_q);
+ n_q_chunk = MAX(n_q_chunk, MIN((int64_t) 8 * 32, n_q));
+
+ if (!kq_p8) {
+ n_q_chunk = MIN(n_q_chunk, (int64_t)(backend_ctx->image_max_buffer_size / (n_kv * n_head)));
+ }
+ if (n_q_chunk < 32 && n_q_chunk < n_q) {
+ return false;
+ }
+
+ // even out the chunks, then align to the 32-wide n-tile
+ const int64_t n_chunks = (n_q + n_q_chunk - 1) / n_q_chunk;
+ n_q_chunk = MIN(n_q_chunk, ((n_q + n_chunks - 1) / n_chunks + 31) & ~(int64_t) 31);
+
+ if (kq_p8) {
+ ggml_cl_flash_attn_decompose_kq_p8(backend, q, k, dst, n_q_chunk, scale, max_bias, kqv_direct);
+ } else {
+ ggml_cl_flash_attn_decompose(backend, q, k, dst, n_q_chunk, scale, max_bias);
+ }
+ return true;
+}
+
static void ggml_cl_flash_attn(ggml_backend_t backend, const ggml_tensor * q, const ggml_tensor * k, ggml_tensor * dst) {
const ggml_tensor * v = dst->src[2];
const ggml_tensor * mask = dst->src[3];
@@ -22808,15 +23516,23 @@ static void ggml_cl_flash_attn(ggml_backend_t backend, const ggml_tensor * q, co
((d_head_q == 256 && (mq_gqa == 2 || mq_gqa == 6)) ||
(d_head_q == 512 && (mq_gqa == 2 || mq_gqa == 16)) ||
(d_head_q == 128 && mq_gqa == 16 && backend_ctx->has_subgroup_shuffle));
- const bool mq_decode_shape = backend_ctx->gpu_family == ADRENO && is_mixed && n_q == 1 &&
- d_head_q == d_head_v && n_head_kv > 0 && n_head % n_head_kv == 0 &&
- ((((d_head_q == 64 && mq_gqa == 8) || (d_head_q == 128 && mq_gqa == 4)) &&
- backend_ctx->has_subgroup_shuffle) ||
- ((d_head_q == 256 || d_head_q == 512) && (mq_gqa == 4 || mq_gqa == 8)) ||
- mq_decode_x2e);
- if (mq_decode_shape && n_kv >= 32) {
+ const bool mq_decode_kernel = backend_ctx->gpu_family == ADRENO && is_mixed &&
+ d_head_q == d_head_v && n_head_kv > 0 && n_head % n_head_kv == 0 &&
+ ((((d_head_q == 64 && mq_gqa == 8) || (d_head_q == 128 && mq_gqa == 4)) &&
+ backend_ctx->has_subgroup_shuffle) ||
+ ((d_head_q == 256 || d_head_q == 512) && (mq_gqa == 4 || mq_gqa == 8)) ||
+ mq_decode_x2e);
+ const bool mq_decode_shape = mq_decode_kernel && n_q == 1;
+ // compile on the first call for this shape (the warmup is a prefill), not on the first decode
+ if (mq_decode_kernel) {
ggml_opencl_ensure_fa_f32_f16_mq_decode(backend_ctx, d_head_q, d_head_v, mq_gqa);
}
+
+ // after all lazy compiles, so the decode kernels are built during warmup as on the fused path
+ if (n_q > 1 && ggml_cl_flash_attn_decompose_dispatch(backend, q, k, dst)) {
+ return;
+ }
+
const bool use_native_q8_0_q1 = is_q8_0 && n_q == 1 &&
backend_ctx->fa.f32_q8_0_q1.count(dk_dv) > 0;
// Native q8_0 prefill — reads q8_0 directly, wg_size = cfg->bm.
@@ -37409,6 +38125,102 @@ static void ggml_cl_soft_max(ggml_backend_t backend, const ggml_tensor * src0, c
backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
}
+// Softmax for the decomposed flash attention: like ggml_cl_soft_max, but dst keeps
+// exp(x - max) and the row sums go to `sums`, one float per (row, head). src2 is the
+// optional sinks tensor.
+static void ggml_cl_fa_decompose_soft_max(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, ggml_tensor * dst, cl_mem sums) {
+ GGML_ASSERT(src0 && src0->extra);
+ GGML_ASSERT(src1 && src1->extra && src1->type == GGML_TYPE_F16);
+ GGML_ASSERT(dst && dst->extra);
+ GGML_ASSERT(sums != nullptr);
+ GGML_ASSERT(src0->ne[0] % 4 == 0);
+
+ if (src2) {
+ GGML_ASSERT(src2->extra);
+ }
+
+ ggml_backend_opencl_context *backend_ctx = (ggml_backend_opencl_context *)backend->context;
+
+ ggml_tensor_extra_cl * extra0 = (ggml_tensor_extra_cl *)src0->extra;
+ ggml_tensor_extra_cl * extra1 = (ggml_tensor_extra_cl *)src1->extra;
+ ggml_tensor_extra_cl * extra2 = src2 ? (ggml_tensor_extra_cl *)src2->extra : nullptr;
+ ggml_tensor_extra_cl * extrad = (ggml_tensor_extra_cl *)dst->extra;
+
+ cl_ulong offset0 = extra0->offset + src0->view_offs;
+ cl_ulong offset1 = extra1->offset + src1->view_offs;
+ cl_ulong offset2 = extra2 ? extra2->offset + src2->view_offs : offset0;
+ cl_ulong offsetd = extrad->offset + dst->view_offs;
+ cl_ulong offsets = 0;
+
+ const int ne00 = src0->ne[0];
+ const int ne01 = src0->ne[1];
+ const int ne02 = src0->ne[2];
+ const int ne03 = src0->ne[3];
+
+ const cl_ulong nb01 = src0->nb[1];
+ const cl_ulong nb02 = src0->nb[2];
+ const cl_ulong nb03 = src0->nb[3];
+
+ const int ne12 = src1->ne[2];
+ const int ne13 = src1->ne[3];
+
+ const cl_ulong nb11 = src1->nb[1];
+ const cl_ulong nb12 = src1->nb[2];
+ const cl_ulong nb13 = src1->nb[3];
+
+ const cl_ulong nb1 = dst->nb[1];
+ const cl_ulong nb2 = dst->nb[2];
+ const cl_ulong nb3 = dst->nb[3];
+
+ float scale, max_bias;
+ memcpy(&scale, dst->op_params + 0, sizeof(float));
+ memcpy(&max_bias, dst->op_params + 1, sizeof(float));
+
+ const int n_head = src0->ne[2];
+ const int n_head_log2 = 1u << (uint32_t) floorf(log2f((float) n_head));
+
+ const float m0 = powf(2.0f, -(max_bias ) / n_head_log2);
+ const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2);
+
+ // one wave per row, same as ggml_cl_soft_max on Adreno
+ const int nth = 64;
+
+ cl_kernel kernel = backend_ctx->kernel_fa_decompose_soft_max_4_f16;
+
+ CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra0->data_device));
+ CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset0));
+ CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &extra1->data_device));
+ CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &offset1));
+ CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_mem), extra2 ? &extra2->data_device : &extra0->data_device));
+ CL_CHECK(clSetKernelArg(kernel, 5, sizeof(cl_ulong), &offset2));
+ CL_CHECK(clSetKernelArg(kernel, 6, sizeof(cl_mem), &extrad->data_device));
+ CL_CHECK(clSetKernelArg(kernel, 7, sizeof(cl_ulong), &offsetd));
+ CL_CHECK(clSetKernelArg(kernel, 8, sizeof(cl_mem), &sums));
+ CL_CHECK(clSetKernelArg(kernel, 9, sizeof(cl_ulong), &offsets));
+ CL_CHECK(clSetKernelArg(kernel, 10, sizeof(int), &ne00));
+ CL_CHECK(clSetKernelArg(kernel, 11, sizeof(cl_ulong), &nb01));
+ CL_CHECK(clSetKernelArg(kernel, 12, sizeof(cl_ulong), &nb02));
+ CL_CHECK(clSetKernelArg(kernel, 13, sizeof(cl_ulong), &nb03));
+ CL_CHECK(clSetKernelArg(kernel, 14, sizeof(int), &ne12));
+ CL_CHECK(clSetKernelArg(kernel, 15, sizeof(int), &ne13));
+ CL_CHECK(clSetKernelArg(kernel, 16, sizeof(cl_ulong), &nb11));
+ CL_CHECK(clSetKernelArg(kernel, 17, sizeof(cl_ulong), &nb12));
+ CL_CHECK(clSetKernelArg(kernel, 18, sizeof(cl_ulong), &nb13));
+ CL_CHECK(clSetKernelArg(kernel, 19, sizeof(cl_ulong), &nb1));
+ CL_CHECK(clSetKernelArg(kernel, 20, sizeof(cl_ulong), &nb2));
+ CL_CHECK(clSetKernelArg(kernel, 21, sizeof(cl_ulong), &nb3));
+ CL_CHECK(clSetKernelArg(kernel, 22, sizeof(float), &scale));
+ CL_CHECK(clSetKernelArg(kernel, 23, sizeof(float), &max_bias));
+ CL_CHECK(clSetKernelArg(kernel, 24, sizeof(float), &m0));
+ CL_CHECK(clSetKernelArg(kernel, 25, sizeof(float), &m1));
+ CL_CHECK(clSetKernelArg(kernel, 26, sizeof(int), &n_head_log2));
+
+ size_t global_work_size[] = {(size_t)ne01*nth, (size_t)ne02, (size_t)ne03};
+ size_t local_work_size[] = {(size_t)nth, 1, 1};
+
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
+}
+
static void ggml_cl_rope(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
GGML_ASSERT(src0);
GGML_ASSERT(src0->extra);
diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_pre_f16.cl b/ggml/src/ggml-opencl/kernels/flash_attn_pre_f16.cl
index 88ead4bcb..74aa67b08 100644
--- a/ggml/src/ggml-opencl/kernels/flash_attn_pre_f16.cl
+++ b/ggml/src/ggml-opencl/kernels/flash_attn_pre_f16.cl
@@ -154,3 +154,123 @@ __kernel void flash_attn_blk_f16(
blk[((ulong) mask_slice * (ulong) n_q_blocks + (ulong) q_block_idx) * (ulong) n_kv_blocks + (ulong) kv_block_idx] = res;
}
+
+// V^T for the decomposed flash attention: [dv, n_kv, n_head_kv] (strided) -> [n_kv, dv, n_head_kv]
+// (contiguous). This is the layout the Adreno KQV image GEMM reads. A local tile keeps both the
+// read and the write coalesced, and the +1 column avoids bank conflicts.
+#define FA_VT_TILE 32
+#define FA_VT_ROWS 8
+__kernel void flash_attn_v_transpose_f16(
+ const global void * v_void, ulong v_offset,
+ global void * vt_void,
+ const int n_kv,
+ const int dv,
+ const int n_head_kv,
+ const ulong v_nb1,
+ const ulong v_nb2
+) {
+ __local half tile[FA_VT_TILE][FA_VT_TILE + 1];
+
+ const int d0 = get_group_id(0) * FA_VT_TILE;
+ const int k0 = get_group_id(1) * FA_VT_TILE;
+ const int head_kv_idx = get_group_id(2);
+
+ const int lx = get_local_id(0);
+ const int ly = get_local_id(1);
+
+ const global char * v_head = (const global char *) v_void + v_offset + (ulong) head_kv_idx * v_nb2;
+
+ for (int r = 0; r < FA_VT_TILE; r += FA_VT_ROWS) {
+ const int d = d0 + lx;
+ const int kv = k0 + ly + r;
+ half val = (half) 0.0f;
+ if (d < dv && kv < n_kv) {
+ val = ((const global half *) (v_head + (ulong) kv * v_nb1))[d];
+ }
+ tile[ly + r][lx] = val;
+ }
+
+ barrier(CLK_LOCAL_MEM_FENCE);
+
+ global half * vt_head = (global half *) vt_void + (ulong) head_kv_idx * (ulong) n_kv * (ulong) dv;
+
+ for (int r = 0; r < FA_VT_TILE; r += FA_VT_ROWS) {
+ const int kv = k0 + lx;
+ const int d = d0 + ly + r;
+ if (d < dv && kv < n_kv) {
+ vt_head[(ulong) d * (ulong) n_kv + (ulong) kv] = tile[lx][ly + r];
+ }
+ }
+}
+
+// Transpose V and quantize each block of 32 KV rows for the int8 KQV GEMM.
+__kernel void flash_attn_v_transpose_q8(
+ const global void * v_void, ulong v_offset,
+ global void * vt_q_void, // char, [n_head_kv][dv][kv_pitch]
+ global void * vt_d_void, // half, [n_head_kv][dv][kv_pitch/32]
+ const int n_kv,
+ const int dv,
+ const int n_head_kv,
+ const ulong v_nb1,
+ const ulong v_nb2,
+ const ulong v_nb3,
+ const int kv_pitch // bytes per V^T row, >= n_kv, multiple of 32
+) {
+ __local half tile[FA_VT_TILE][FA_VT_TILE + 1];
+
+ const int d0 = get_group_id(0) * FA_VT_TILE;
+ const int k0 = get_group_id(1) * FA_VT_TILE;
+ const int head_kv_idx = get_group_id(2) % n_head_kv;
+ const int batch_idx = get_group_id(2) / n_head_kv;
+
+ const int lx = get_local_id(0);
+ const int ly = get_local_id(1);
+
+ const global char * v_head = (const global char *) v_void + v_offset +
+ (ulong) batch_idx * v_nb3 + (ulong) head_kv_idx * v_nb2;
+
+ for (int r = 0; r < FA_VT_TILE; r += FA_VT_ROWS) {
+ const int d = d0 + lx;
+ const int kv = k0 + ly + r;
+ half val = (half) 0.0f;
+ if (d < dv && kv < n_kv) {
+ val = ((const global half *) (v_head + (ulong) kv * v_nb1))[d];
+ }
+ tile[ly + r][lx] = val;
+ }
+
+ barrier(CLK_LOCAL_MEM_FENCE);
+
+ const int lid = ly * FA_VT_TILE + lx;
+ const int nblk = kv_pitch / FA_VT_TILE;
+
+ if (lid < FA_VT_TILE) {
+ const int d = d0 + lid;
+ if (d < dv) {
+ const ulong hbase = ((ulong) batch_idx * (ulong) n_head_kv + (ulong) head_kv_idx);
+
+ global char * qrow = (global char *) vt_q_void + hbase*(ulong)kv_pitch*(ulong)dv + (ulong) d * (ulong) kv_pitch;
+ global half * drow = (global half *) vt_d_void + hbase*(ulong)nblk*(ulong)dv + (ulong) d * (ulong) nblk;
+
+ float v[FA_VT_TILE];
+ float amax = 0.0f;
+
+ for (int i = 0; i < FA_VT_TILE; ++i) {
+ v[i] = (float) tile[i][lid];
+ amax = fmax(amax, fabs(v[i]));
+ }
+
+ const float dq = amax / 127.0f;
+ const float id = amax > 0.0f ? 127.0f / amax : 0.0f;
+
+ drow[k0 / FA_VT_TILE] = (half) dq;
+
+ for (int i = 0; i < FA_VT_TILE; ++i) {
+ const int kv = k0 + i;
+ if (kv < n_kv) {
+ qrow[kv] = (char) clamp((int) rint(v[i] * id), -127, 127);
+ }
+ }
+ }
+ }
+}
diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_pv_direct.cl b/ggml/src/ggml-opencl/kernels/flash_attn_pv_direct.cl
new file mode 100644
index 000000000..201fe4443
--- /dev/null
+++ b/ggml/src/ggml-opencl/kernels/flash_attn_pv_direct.cl
@@ -0,0 +1,106 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+#pragma OPENCL EXTENSION cl_qcom_reqd_sub_group_size : enable
+
+#define FA_KQVD_MAXC 32
+
+// Read F16 V in its cache layout. Share each V row across the queries and heads of one KV group.
+// P8 and its block scales are expanded in local memory; KV splits write separate partial sums.
+__attribute__((qcom_reqd_sub_group_size("half")))
+kernel void kernel_fa_kqv_direct_f16(
+ global const char * v,
+ ulong off_v,
+ ulong v_nb1,
+ ulong v_nb2,
+ global const uchar * pq,
+ global const float * pd,
+ global float * part, // [split][head][n_q][dv]
+ int dv,
+ int n_q,
+ int gqa,
+ int n_head,
+ int nblk,
+ int blk_per_split,
+ int p_pitch
+) {
+ const int lid = get_local_id(0);
+ const int d = get_group_id(0)*128 + lid*2;
+ const int head_kv = get_group_id(1);
+ const int split = get_group_id(2);
+ const int ncol = n_q*gqa;
+ const int b0 = split*blk_per_split;
+ const int b1 = min(b0 + blk_per_split, nblk);
+ const int nbp = p_pitch / 32;
+
+ __local float sh_p[32][FA_KQVD_MAXC];
+
+ float2 acc[FA_KQVD_MAXC];
+ #pragma unroll
+ for (int c = 0; c < FA_KQVD_MAXC; ++c) {
+ acc[c] = (float2)(0.0f);
+ }
+
+ global const char * vbase = v + off_v + (size_t)head_kv*v_nb2 + (size_t)d*sizeof(half);
+
+ for (int b = b0; b < b1; ++b) {
+ barrier(CLK_LOCAL_MEM_FENCE);
+ // Each lane expands 16 KV entries for one column.
+ const int c = lid % FA_KQVD_MAXC;
+ const int r0 = (lid / FA_KQVD_MAXC)*16;
+ float16 f = (float16)(0.0f);
+ if (c < ncol) {
+ const int n = c / gqa;
+ const int h = head_kv*gqa + c % gqa;
+ const size_t row = (size_t)h*n_q + n;
+ const float sc = pd[row*nbp + b];
+ const uchar16 u = vload16(0, pq + row*p_pitch + (size_t)b*32 + r0);
+ f = convert_float16(u)*sc;
+ }
+ sh_p[r0+ 0][c] = f.s0; sh_p[r0+ 1][c] = f.s1; sh_p[r0+ 2][c] = f.s2; sh_p[r0+ 3][c] = f.s3;
+ sh_p[r0+ 4][c] = f.s4; sh_p[r0+ 5][c] = f.s5; sh_p[r0+ 6][c] = f.s6; sh_p[r0+ 7][c] = f.s7;
+ sh_p[r0+ 8][c] = f.s8; sh_p[r0+ 9][c] = f.s9; sh_p[r0+10][c] = f.sa; sh_p[r0+11][c] = f.sb;
+ sh_p[r0+12][c] = f.sc; sh_p[r0+13][c] = f.sd; sh_p[r0+14][c] = f.se; sh_p[r0+15][c] = f.sf;
+ barrier(CLK_LOCAL_MEM_FENCE);
+
+ global const char * vr = vbase + (size_t)b*32*v_nb1;
+ for (int r = 0; r < 32; ++r) {
+ const float2 x = vload_half2(0, (global const half *)(vr + (size_t)r*v_nb1));
+ #pragma unroll
+ for (int c = 0; c < FA_KQVD_MAXC; c += 4) {
+ if (c < ncol) {
+ const float4 pv = vload4(0, &sh_p[r][c]);
+ acc[c+0] = mad((float2)(pv.x), x, acc[c+0]);
+ acc[c+1] = mad((float2)(pv.y), x, acc[c+1]);
+ acc[c+2] = mad((float2)(pv.z), x, acc[c+2]);
+ acc[c+3] = mad((float2)(pv.w), x, acc[c+3]);
+ }
+ }
+ }
+ }
+
+ #pragma unroll
+ for (int c = 0; c < FA_KQVD_MAXC; ++c) {
+ if (c < ncol) {
+ const int n = c / gqa;
+ const int h = head_kv*gqa + c % gqa;
+ vstore2(acc[c], 0, part + (((size_t)split*n_head + h)*n_q + n)*dv + d);
+ }
+ }
+}
+
+kernel void kernel_fa_kqv_direct_reduce(
+ global const float * part,
+ global float * dst,
+ int n_elem,
+ int n_split
+) {
+ const int i = get_global_id(0);
+ if (i >= n_elem) {
+ return;
+ }
+ float s = 0.0f;
+ for (int k = 0; k < n_split; ++k) {
+ s += part[(size_t)k*n_elem + i];
+ }
+ dst[i] = s;
+}
diff --git a/ggml/src/ggml-opencl/kernels/mul_mm_f16_f32_kq_p8.cl b/ggml/src/ggml-opencl/kernels/mul_mm_f16_f32_kq_p8.cl
new file mode 100644
index 000000000..2c0a56c85
--- /dev/null
+++ b/ggml/src/ggml-opencl/kernels/mul_mm_f16_f32_kq_p8.cl
@@ -0,0 +1,250 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+
+#define LM_FIRST_256B 0
+#define LM_SECOND_256B 64
+#define LM_THIRD_256B 128
+#define LM_FOURTH_256B 192
+
+
+inline float16 mm_load_a(
+ image1d_buffer_t matrix_A,
+ uint subMatrixAStartInElements,
+ int line_stride_matrix_A_in_bytes
+) {
+ __private float8 regA;
+ size_t sub_block_id_m = get_local_id(0);
+
+ uint a_texCoord = subMatrixAStartInElements/2 + (sub_block_id_m * line_stride_matrix_A_in_bytes/4);
+
+ regA.s0123 = read_imagef(matrix_A, a_texCoord/4);
+ regA.s4567 = read_imagef(matrix_A, (a_texCoord+4)/4);
+
+ return convert_float16(as_half16(regA));
+}
+
+inline float4 alu_32(
+ float16 regA,
+ __local float4 * matrix_B_vec
+) {
+
+ __private float4 rC = 0;
+ int i = get_sub_group_id() * 64;
+
+ rC += regA.s0 * matrix_B_vec[i];
+ rC += regA.s1 * matrix_B_vec[i + 16];
+ rC += regA.s4 * matrix_B_vec[i + 1];
+ rC += regA.s5 * matrix_B_vec[i + 17];
+ rC += regA.s8 * matrix_B_vec[i + 2];
+ rC += regA.s9 * matrix_B_vec[i + 18];
+ rC += regA.sc * matrix_B_vec[i + 3];
+ rC += regA.sd * matrix_B_vec[i + 19];
+
+ i += 32;
+
+ rC += regA.s2 * matrix_B_vec[i];
+ rC += regA.s3 * matrix_B_vec[i + 16];
+ rC += regA.s6 * matrix_B_vec[i + 1];
+ rC += regA.s7 * matrix_B_vec[i + 17];
+ rC += regA.sa * matrix_B_vec[i + 2];
+ rC += regA.sb * matrix_B_vec[i + 18];
+ rC += regA.se * matrix_B_vec[i + 3];
+ rC += regA.sf * matrix_B_vec[i + 19];
+
+ return rC;
+}
+
+inline float16 alu_16(
+ float16 regA,
+ __local float * matrix_B_local
+) {
+ float16 out;
+ __local float4 * matrix_B_vec = (__local float4*)matrix_B_local;
+
+ out.s0123 = alu_32(regA, matrix_B_vec);
+ out.s4567 = alu_32(regA, matrix_B_vec + 4);
+ out.s89ab = alu_32(regA, matrix_B_vec + 8);
+ out.scdef = alu_32(regA, matrix_B_vec + 12);
+
+ return out;
+}
+
+inline void mm_mad(
+ __local float * matrix_B_local,
+ float16 regA,
+ float8 regB,
+ uint b_localOffsetInWords,
+ float16 * regC0_ptr,
+ float16 * regC1_ptr
+) {
+ int offset = b_localOffsetInWords + get_sub_group_id() * 256;
+
+ matrix_B_local[offset + LM_FIRST_256B] = regB.s0;
+ matrix_B_local[offset + LM_SECOND_256B] = regB.s1;
+ matrix_B_local[offset + LM_THIRD_256B] = regB.s2;
+ matrix_B_local[offset + LM_FOURTH_256B] = regB.s3;
+
+ float16 add0 = alu_16(regA, matrix_B_local);
+ *regC0_ptr += add0;
+
+ matrix_B_local[offset + LM_FIRST_256B] = regB.s4;
+ matrix_B_local[offset + LM_SECOND_256B] = regB.s5;
+ matrix_B_local[offset + LM_THIRD_256B] = regB.s6;
+ matrix_B_local[offset + LM_FOURTH_256B] = regB.s7;
+
+ float16 add1 = alu_16(regA, matrix_B_local);
+ *regC1_ptr += add1;
+}
+
+#define TILESIZE_K 16
+#define TILESIZE_M 64
+#define TILESIZE_N 32
+// Fuse KQ and block softmax into U8 P, block maxima and sums. No F32 score matrix is stored.
+__kernel void mul_mm_f16_f32_kq_p8(
+ __read_only image1d_buffer_t matrix_A,
+ __global float * matrix_B,
+ __global uchar * qp, // u8 P, [D_B][N][p_pitch]
+ __global float * bmax, // per 32-row block max, [D_B][N][p_pitch/32]
+ __global float * bsum, // per block sum of exp(s - bmax), [D_B][N][p_pitch/32]
+ __global char * mask, // f16 [>=M][N] view (chunk rows), row stride mask_nb1 bytes
+ ulong offset_mask,
+ ulong mask_nb1,
+ float scale,
+ float max_bias,
+ float m0,
+ float m1,
+ int n_head_log2,
+ int M,
+ int K,
+ int N,
+ int D_A,
+ int D_B,
+ int p_pitch // bytes per P row, >= M, multiple of 32
+) {
+
+ // Match the n-tile-first grid in ggml_cl_fa_kq_p8_dispatch.
+ const uint m_blocks = (M+TILESIZE_M-1)/TILESIZE_M;
+ uint block_id_n = get_global_id(1);
+ uint block_id_m = get_global_id(2) % m_blocks;
+ uint block_id_d = get_global_id(2) / m_blocks;
+
+ __private float16 regA;
+ __private float8 regB;
+ __private float16 regC0 = (float16)(0.0f);
+ __private float16 regC1 = (float16)(0.0f);
+
+ const uint col = block_id_m * TILESIZE_M;
+ const uint row = block_id_n * TILESIZE_N;
+ const uint depth_A = block_id_d / (D_B/D_A);
+ const uint depth_B = block_id_d;
+
+ int line_stride_matrix_A_in_bytes = K * D_A * 2;
+ int line_stride_matrix_B_in_bytes = K * D_B * 4;
+
+ const uint strideAinElements = line_stride_matrix_A_in_bytes / 2;
+ const uint strideBinElements = line_stride_matrix_B_in_bytes / 4;
+
+ size_t sub_block_id_m = get_local_id(0);
+
+ uint b_localOffsetInWords = (sub_block_id_m/16)*16
+ + ((((sub_block_id_m)>>0)&1)<<2)
+ + ((((sub_block_id_m)>>1)&1)<<3)
+ + ((((sub_block_id_m)>>2)&1)<<0)
+ + ((((sub_block_id_m)>>3)&1)<<1);
+
+ uint2 b_globalOffsetInWords_xy = {((sub_block_id_m%4)*4), (sub_block_id_m>>2)};
+ uint subMatrixAStartInElements = col * strideAinElements + depth_A * K;
+ uint subMatrixBStartInElements = 0;
+
+ __local float matrix_B_local[1024];
+
+ // Clamp Q loads in the last tile; columns past N are not stored.
+ const uint c0 = min(row + b_globalOffsetInWords_xy.y, (uint)N - 1);
+ const uint c1 = min(row + b_globalOffsetInWords_xy.y + 16, (uint)N - 1);
+ const uint b_globalOffsetInWords00 = c0*strideBinElements + depth_B*K + b_globalOffsetInWords_xy.x;
+ const uint b_globalOffsetInWords16 = c1*strideBinElements + depth_B*K + b_globalOffsetInWords_xy.x;
+
+ for (uint step=0; step < K; step+=TILESIZE_K) {
+ regA = mm_load_a(matrix_A, subMatrixAStartInElements, line_stride_matrix_A_in_bytes);
+
+ uint b_coordInWords00 = subMatrixBStartInElements + b_globalOffsetInWords00;
+ uint b_coordInWords16 = subMatrixBStartInElements + b_globalOffsetInWords16;
+
+ regB.s0123 = vload4(b_coordInWords00/4, matrix_B);
+ regB.s4567 = vload4(b_coordInWords16/4, matrix_B);
+
+ mm_mad(matrix_B_local, regA, regB, b_localOffsetInWords, ®C0, ®C1);
+
+ subMatrixAStartInElements += TILESIZE_K;
+ subMatrixBStartInElements += TILESIZE_K;
+ }
+
+ // Transpose each score tile in local memory so one lane owns a 32-row softmax block.
+ __local float p8_lds[2048];
+ __local float * lds = p8_lds + get_sub_group_id() * 1024;
+ const uint lane = get_local_id(0);
+ __global const half * mrow = (__global const half *)(mask + offset_mask);
+ const uint mstride = (uint)(mask_nb1 / 2);
+ const int nblk = p_pitch / 32;
+
+ #pragma unroll
+ for (int half_ = 0; half_ < 2; ++half_) {
+ barrier(CLK_LOCAL_MEM_FENCE);
+ const float16 cc = half_ == 0 ? regC0 : regC1;
+ lds[ 0*64 + lane] = cc.s0; lds[ 1*64 + lane] = cc.s1; lds[ 2*64 + lane] = cc.s2; lds[ 3*64 + lane] = cc.s3;
+ lds[ 4*64 + lane] = cc.s4; lds[ 5*64 + lane] = cc.s5; lds[ 6*64 + lane] = cc.s6; lds[ 7*64 + lane] = cc.s7;
+ lds[ 8*64 + lane] = cc.s8; lds[ 9*64 + lane] = cc.s9; lds[10*64 + lane] = cc.sa; lds[11*64 + lane] = cc.sb;
+ lds[12*64 + lane] = cc.sc; lds[13*64 + lane] = cc.sd; lds[14*64 + lane] = cc.se; lds[15*64 + lane] = cc.sf;
+ barrier(CLK_LOCAL_MEM_FENCE);
+ if (lane < 32) {
+ const int j = lane >> 1;
+ const int b = lane & 1;
+ const uint n = row + half_*16 + j;
+ const uint h = depth_B;
+ if (n < (uint)N) {
+ float slope = 1.0f;
+ if (max_bias > 0.0f) {
+ const float base = (int)h < n_head_log2 ? m0 : m1;
+ const int ex = (int)h < n_head_log2 ? (int)h + 1 : 2*((int)h - n_head_log2) + 1;
+ slope = pow(base, ex);
+ }
+ const uint kvb = col + b*32; // first kv row of this block
+ const __local float * src = lds + j*64 + b*32;
+ float16 s0 = vload16(0, src);
+ float16 s1 = vload16(1, src);
+ const __global half * mp = mrow + (size_t)n*mstride + kvb;
+ const float16 m0v = convert_float16(vload16(0, mp));
+ const float16 m1v = convert_float16(vload16(1, mp));
+ s0 = s0*scale + slope*m0v;
+ s1 = s1*scale + slope*m1v;
+ const float16 mx16 = fmax(s0, s1);
+ const float8 mx8 = fmax(mx16.lo, mx16.hi);
+ const float4 mx4 = fmax(mx8.lo, mx8.hi);
+ const float2 mx2 = fmax(mx4.lo, mx4.hi);
+ const float amax = fmax(mx2.x, mx2.y);
+ const uint prow = h*(uint)N + n;
+ const uint blk = kvb >> 5;
+ float sum = 0.0f;
+ uchar16 q0 = (uchar16)(0);
+ uchar16 q1 = (uchar16)(0);
+ // Use a finite sentinel: -cl-finite-math-only can remove infinity checks.
+ if (amax > -1.0e30f) {
+ const float16 e0 = exp(s0 - amax);
+ const float16 e1 = exp(s1 - amax);
+ const float16 a16 = e0 + e1;
+ const float8 a8 = a16.lo + a16.hi;
+ const float4 a4 = a8.lo + a8.hi;
+ const float2 a2 = a4.lo + a4.hi;
+ sum = a2.x + a2.y;
+ q0 = convert_uchar16_sat_rte(e0*255.0f);
+ q1 = convert_uchar16_sat_rte(e1*255.0f);
+ }
+ __global uchar * qdst = qp + (size_t)prow*(size_t)p_pitch + (size_t)kvb;
+ vstore16(q0, 0, qdst);
+ vstore16(q1, 1, qdst);
+ bmax[(size_t)prow*nblk + blk] = amax;
+ bsum[(size_t)prow*nblk + blk] = sum;
+ }
+ }
+ }
+}
diff --git a/ggml/src/ggml-opencl/kernels/mul_mm_q8_f32_kqv.cl b/ggml/src/ggml-opencl/kernels/mul_mm_q8_f32_kqv.cl
new file mode 100644
index 000000000..08d80961c
--- /dev/null
+++ b/ggml/src/ggml-opencl/kernels/mul_mm_q8_f32_kqv.cl
@@ -0,0 +1,131 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+#pragma OPENCL EXTENSION cl_khr_subgroups : enable
+#ifdef cl_khr_integer_dot_product
+#pragma OPENCL EXTENSION cl_khr_integer_dot_product : enable
+#endif
+
+// Multiply U8 P by S8 V^T, with one scale per 32 KV entries.
+// Group query heads by KV head to share V^T in L2.
+// The packed dot takes the V word (registers) first: swapped operands miscompile on Adreno.
+
+#define KQV_TM 64 // d rows per workgroup, one per lane
+#ifndef KQV_TN
+#define KQV_TN 32 // queries per workgroup, all owned by every lane
+#endif
+#ifndef KQV_NB
+#define KQV_NB 8 // 32-blocks of P staged per barrier
+#endif
+#define KQV_WG 64
+
+__attribute__((qcom_wave_pair_mode(1)))
+kernel void kernel_mul_mm_q8_f32_kqv(
+ global const uint * vq, // V^T int8, [head_kv][dv][kv_pitch] as packed uints
+ ulong off_vq,
+ global const half * vd, // V^T scales, [head_kv][dv][kv_pitch/32]
+ ulong off_vd,
+ global const uint * pq, // P u8, [head][n_q][kv_pitch] as packed uints
+ ulong off_pq,
+ global const float * pd, // P scales f32, [head][n_q][kv_pitch/32]
+ ulong off_pd,
+ global float * dst, // [head][n_q][dv]
+ ulong off_dst,
+ int dv,
+ int n_kv,
+ int n_q,
+ int n_head,
+ int n_head_kv,
+ int kv_pitch // bytes per V^T row and per P row, >= n_kv, multiple of 32
+) {
+ vq = (global const uint *)((global const char *)vq + off_vq);
+ vd = (global const half *)((global const char *)vd + off_vd);
+ pq = (global const uint *)((global const char *)pq + off_pq);
+ pd = (global const float *)((global const char *)pd + off_pd);
+ dst = (global float *)((global const char *)dst + off_dst);
+
+ const int lid = get_local_id(0);
+ const int gsz = n_head / n_head_kv; // query heads per KV head
+ const int g1 = get_group_id(1);
+ const int head_kv = get_group_id(2);
+ const int head = head_kv*gsz + (g1 % gsz);
+ const int qn0 = (g1 / gsz) * KQV_TN; // first query of this tile
+ const int d = get_group_id(0)*KQV_TM + lid; // this lane's d row
+
+ const int nblk = n_kv / 32; // dispatch guarantees n_kv % 32 == 0
+ const int nu = kv_pitch / 4; // uints per row
+ const int nbp = kv_pitch / 32; // scales per row
+
+ __local uint sh_pq[KQV_TN][KQV_NB][8];
+ __local float sh_pd[KQV_TN][KQV_NB];
+
+ float acc[KQV_TN];
+ #pragma unroll
+ for (int t = 0; t < KQV_TN; ++t) {
+ acc[t] = 0.0f;
+ }
+
+ const int dl = min(d, dv - 1); // clamp so a tail lane loads in bounds
+ const size_t vbase = ((size_t)head_kv*dv + dl)*nu;
+ const size_t vdbas = ((size_t)head_kv*dv + dl)*nbp;
+ const size_t pbase = (size_t)head*n_q;
+
+ for (int bg = 0; bg < nblk; bg += KQV_NB) {
+ const int nb_here = min(KQV_NB, nblk - bg);
+
+ // stage P: each (query, block) is 8 uints = two 16-byte loads
+ for (int i = lid; i < KQV_TN*KQV_NB*2; i += KQV_WG) {
+ const int t = i / (KQV_NB*2);
+ const int bb = (i / 2) % KQV_NB;
+ const int h = (i & 1) * 4;
+ const int qi = qn0 + t;
+ uint4 v = (uint4)(0u);
+ if (qi < n_q && bb < nb_here) {
+ v = vload4(0, &pq[(pbase + qi)*nu + (size_t)(bg + bb)*8 + h]);
+ }
+ vstore4(v, 0, &sh_pq[t][bb][h]);
+ }
+ for (int i = lid; i < KQV_TN*KQV_NB; i += KQV_WG) {
+ const int t = i / KQV_NB;
+ const int bb = i % KQV_NB;
+ const int qi = qn0 + t;
+ sh_pd[t][bb] = (qi < n_q && bb < nb_here)
+ ? pd[(pbase + qi)*nbp + bg + bb] : 0.0f;
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+
+ for (int bb = 0; bb < nb_here; ++bb) {
+ const uint4 w0 = vload4(0, &vq[vbase + (size_t)(bg + bb)*8]);
+ const uint4 w1 = vload4(0, &vq[vbase + (size_t)(bg + bb)*8 + 4]);
+ const float dvs = (float)vd[vdbas + bg + bb];
+
+ #pragma unroll
+ for (int t = 0; t < KQV_TN; ++t) {
+ const uint4 a0 = vload4(0, &sh_pq[t][bb][0]);
+ const uint4 a1 = vload4(0, &sh_pq[t][bb][4]);
+ int raw = 0;
+ raw = dot_acc_sat_4x8packed_su_int(w0.s0, a0.s0, raw);
+ raw = dot_acc_sat_4x8packed_su_int(w0.s1, a0.s1, raw);
+ raw = dot_acc_sat_4x8packed_su_int(w0.s2, a0.s2, raw);
+ raw = dot_acc_sat_4x8packed_su_int(w0.s3, a0.s3, raw);
+ raw = dot_acc_sat_4x8packed_su_int(w1.s0, a1.s0, raw);
+ raw = dot_acc_sat_4x8packed_su_int(w1.s1, a1.s1, raw);
+ raw = dot_acc_sat_4x8packed_su_int(w1.s2, a1.s2, raw);
+ raw = dot_acc_sat_4x8packed_su_int(w1.s3, a1.s3, raw);
+ acc[t] += dvs * sh_pd[t][bb] * (float)raw;
+ }
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+ }
+
+ if (d >= dv) {
+ return;
+ }
+
+ // dst is [head][n_q][dv] with dv contiguous, so the 64 lanes of a d-block write contiguously
+ #pragma unroll
+ for (int t = 0; t < KQV_TN; ++t) {
+ const int qi = qn0 + t;
+ if (qi < n_q) {
+ dst[(pbase + qi)*dv + d] = acc[t];
+ }
+ }
+}
diff --git a/ggml/src/ggml-opencl/kernels/softmax_4_f16.cl b/ggml/src/ggml-opencl/kernels/softmax_4_f16.cl
index 571d16507..8dfc5519d 100644
--- a/ggml/src/ggml-opencl/kernels/softmax_4_f16.cl
+++ b/ggml/src/ggml-opencl/kernels/softmax_4_f16.cl
@@ -106,3 +106,176 @@ kernel void kernel_soft_max_4_f16(
pdst4[i00] /= sum;
}
}
+
+// Same as kernel_soft_max_4_f16, but leaves exp(x - max) in dst and writes the row sum to psums.
+// The decomposed flash attention applies 1/sum to the KQV result, which is much smaller than P.
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_fa_decompose_soft_max_4_f16(
+ global char * src0,
+ ulong offset0,
+ global char * src1,
+ ulong offset1,
+ global char * src2,
+ ulong offset2,
+ global char * dst,
+ ulong offsetd,
+ global char * psums,
+ ulong offset_sums,
+ int ne00,
+ ulong nb01,
+ ulong nb02,
+ ulong nb03,
+ int ne12,
+ int ne13,
+ ulong nb11,
+ ulong nb12,
+ ulong nb13,
+ ulong nb1,
+ ulong nb2,
+ ulong nb3,
+ float scale,
+ float max_bias,
+ float m0,
+ float m1,
+ int n_head_log2
+) {
+ src0 = src0 + offset0;
+ src1 = src1 + offset1;
+ src2 = src2 + offset2;
+ dst = dst + offsetd;
+
+ int i03 = get_group_id(2);
+ int i02 = get_group_id(1);
+ int i01 = get_group_id(0);
+
+ int i13 = i03%ne13;
+ int i12 = i02%ne12;
+ int i11 = i01;
+
+ global float4 * psrc4 = (global float4 *)(src0 + i01*nb01 + i02*nb02 + i03*nb03);
+ global half4 * pmask = src1 != src0 ? (global half4 *)(src1 + i11*nb11 + i12*nb12 + i13*nb13) : 0;
+ global float * psrc2 = src2 != src0 ? (global float *)(src2) : 0;
+ global float4 * pdst4 = (global float4 *)(dst + i01*nb1 + i02*nb2 + i03*nb3);
+
+ float slope = 1.0f;
+
+ // ALiBi
+ if (max_bias > 0.0f) {
+ int h = i02;
+
+ float base = h < n_head_log2 ? m0 : m1;
+ int exp = h < n_head_log2 ? h + 1 : 2*(h - n_head_log2) + 1;
+
+ slope = pow(base, exp);
+ }
+
+ // parallel max
+ float4 lmax4 = psrc2 ? psrc2[i02] : -INFINITY;
+ for (int i00 = get_local_id(0); i00 < ne00/4; i00 += get_local_size(0)) {
+ lmax4 = fmax(lmax4, psrc4[i00]*scale + slope*(pmask ? convert_float4(pmask[i00]) : 0.0f));
+ }
+ float lmax = fmax(fmax(lmax4.s0, lmax4.s1), fmax(lmax4.s2, lmax4.s3));
+
+ float max = sub_group_reduce_max(lmax);
+ // Avoid -inf - (-inf) for fully masked rows.
+ if (max == -INFINITY) {
+ max = 0.0f;
+ }
+
+ // parallel sum
+ float4 lsum4 = 0.0f;
+ for (int i00 = get_local_id(0); i00 < ne00/4; i00 += get_local_size(0)) {
+ const float4 exp_psrc4 = exp((psrc4[i00]*scale + slope*(pmask ? convert_float4(pmask[i00]) : 0.0f)) - max);
+ lsum4 += exp_psrc4;
+ pdst4[i00] = exp_psrc4;
+ }
+ float lsum = lsum4.s0 + lsum4.s1 + lsum4.s2 + lsum4.s3;
+
+ float sum = sub_group_reduce_add(lsum);
+
+ if (psrc2) {
+ sum += exp(psrc2[i02] - max);
+ }
+
+ if (get_local_id(0) == 0) {
+ global float * sums = (global float *)(psums + offset_sums);
+ sums[(i03*get_num_groups(1) + i02)*get_num_groups(0) + i01] = sum;
+ }
+}
+
+// out[:, n, h] *= 1/sums[h*ne1 + n], out is [dv, ne1, n_head] contiguous.
+kernel void kernel_fa_decompose_scale_rows_f32(
+ global char * dst,
+ ulong offsetd,
+ global char * psums,
+ ulong offset_sums,
+ int dv,
+ int ne1
+) {
+ global float * out = (global float *)(dst + offsetd);
+ global float * sums = (global float *)(psums + offset_sums);
+
+ const int n = get_group_id(0);
+ const int h = get_group_id(1);
+
+ const float s = sums[h*ne1 + n];
+ const float r = s > 0.0f ? 1.0f/s : 0.0f;
+
+ global float * row = out + ((size_t)h*ne1 + n)*dv;
+ for (int i = get_local_id(0); i < dv; i += get_local_size(0)) {
+ row[i] *= r;
+ }
+}
+
+// Merge the P8 block maxima and sums into F32 block scales and row sums.
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_fa_p8_fixup(
+ global float * bmax,
+ global float * bsum,
+ global float * dp,
+ global float * psums,
+ global char * sinks,
+ ulong offset_sinks,
+ int has_sinks,
+ int nblk,
+ int row_blk, // blocks per row of bmax/bsum/dp (the padded pitch), >= nblk
+ int N
+) {
+ const int n = get_group_id(0);
+ const int head = get_group_id(1);
+ const uint row = (uint)head*(uint)N + (uint)n;
+ global float * bm = bmax + (size_t)row*row_blk;
+ global float * bs = bsum + (size_t)row*row_blk;
+ global float * dr = dp + (size_t)row*row_blk;
+
+ float lmax = -INFINITY;
+ for (int b = get_local_id(0); b < nblk; b += get_local_size(0)) {
+ lmax = fmax(lmax, bm[b]);
+ }
+ float rowmax = sub_group_reduce_max(lmax);
+ float sink = 0.0f;
+ if (has_sinks) {
+ sink = ((global float *)(sinks + offset_sinks))[head];
+ rowmax = fmax(rowmax, sink);
+ }
+
+ float lsum = 0.0f;
+ for (int b = get_local_id(0); b < nblk; b += get_local_size(0)) {
+ const float m = bm[b];
+ // Use a finite sentinel: -cl-finite-math-only can remove infinity checks.
+ const float f = (m > -1.0e30f && rowmax > -1.0e30f) ? exp(m - rowmax) : 0.0f;
+ dr[b] = f / 255.0f;
+ lsum += bs[b]*f;
+ }
+ float sum = sub_group_reduce_add(lsum);
+ if (has_sinks && rowmax > -1.0e30f) {
+ sum += exp(sink - rowmax);
+ }
+ if (get_local_id(0) == 0) {
+ psums[row] = sum;
+ }
+}