{192, 128, 16, 16, 1, 0},
{192, 192, 16, 16, 1, 0},
{256, 256, 16, 16, 16, 0},
+ {512, 512, 8, 16, 64, 0},
};
struct ggml_opencl_fa_dim_table {
// f32 Q / f16 KV (mixed)
std::map<std::pair<int, int>, cl_kernel> f32_f16;
std::map<std::pair<int, int>, cl_kernel> f32_f16_split; // N_SPLIT>1 variant
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_split_k_img; // DK=512 prefill split, K via image1d_buffer_t
std::map<std::pair<int, int>, cl_kernel> f32_f16_q1;
std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_split; // flash-decoding K-split
+ // vec decode
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec;
+ // kv-head-coalesced vec decode
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq;
+ // kv-head-coalesced + flash-decoding split
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split;
+ // MQ_GQA=8 specializations
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_g8;
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_g8;
+ // k-image variant of MQ_G8 vec_mq_split
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_g8_k_img;
+ // k-image variant of MQ_GQA=4 vec_mq_split
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_k_img;
+ // alternative decode
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_local_tile;
+ // hybrid local-tile + MQ + FD-split kernel for DK=DV=128 only
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_local_mq_split;
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_local_mq_split_g8;
std::map<std::pair<int, int>, int> f32_f16_bm;
std::map<std::pair<int, int>, int> f32_f16_bn;
std::map<std::pair<int, int>, int> f32_f16_wg_size;
std::map<std::pair<int, int>, int> f32_f16_split_nkv_threshold;
// f32 Q / native q8_0 KV
std::map<std::pair<int, int>, cl_kernel> f32_q8_0_q1; // decode
+ std::map<std::pair<int, int>, cl_kernel> f32_q8_0_q1_vec; // DV-split + multi-subgroup decode
std::map<std::pair<int, int>, cl_kernel> f32_q8_0_q1_split; // flash-decoding pass 1
+ // KV-head-coalesced + flash-decoding split for q8_0 KV
+ std::map<std::pair<int, int>, cl_kernel> f32_q8_0_q1_vec_mq_split;
+ std::map<std::pair<int, int>, cl_kernel> f32_q8_0_q1_vec_mq_split_g8;
std::map<std::pair<int, int>, cl_kernel> f32_q8_0; // prefill (baseline)
std::map<std::pair<int, int>, cl_kernel> f32_q8_0_split; // N_SPLIT>1 variant
std::map<std::pair<int, int>, int> f32_q8_0_split_wg_size; // wg_size = bm*n_split
std::map<std::pair<int, int>, int> f32_q8_0_split_bm; // per-split BLOCK_M
// f32 Q / native q4_0 KV
std::map<std::pair<int, int>, cl_kernel> f32_q4_0_q1;
+ std::map<std::pair<int, int>, cl_kernel> f32_q4_0_q1_vec; // DV-split + multi-subgroup decode
std::map<std::pair<int, int>, cl_kernel> f32_q4_0_q1_split;
+ // kv-head-coalesced + flash-decoding split for q4_0 kv (dp4a K dot)
+ std::map<std::pair<int, int>, cl_kernel> f32_q4_0_q1_vec_mq_split;
+ std::map<std::pair<int, int>, cl_kernel> f32_q4_0_q1_vec_mq_split_g8;
std::map<std::pair<int, int>, cl_kernel> f32_q4_0;
std::map<std::pair<int, int>, cl_kernel> f32_q4_0_split;
std::map<std::pair<int, int>, int> f32_q4_0_split_wg_size;
ggml_cl_buffer prealloc_scales_trans;
ggml_cl_buffer prealloc_act_trans;
+ // pool of persistent image1d_buffer views over kv-cache layers, keyed by
+ // (parent buffer, offset within parent)
+ // used by the img-variant KQ/KQV dispatch paths to avoid per-call
+ // clCreateSubBuffer + clCreateImage + pending-release-queue on long-context decode
+ struct ImagePoolKey {
+ uintptr_t buf;
+ uint64_t offset;
+ bool operator<(const ImagePoolKey & o) const {
+ if (buf != o.buf) return buf < o.buf;
+ return offset < o.offset;
+ }
+ };
+ struct ImagePoolEntry {
+ cl_mem sub_buffer = nullptr;
+ cl_mem image = nullptr;
+ size_t k_bytes = 0;
+ cl_channel_type channel_data_type = CL_FLOAT;
+ };
+ std::map<ImagePoolKey, ImagePoolEntry> kq_img_pool;
+ std::map<ImagePoolKey, ImagePoolEntry> kqv_img_pool;
+
+ // pool for the on-device f16 buffer for kv-cache with non-FA quantized-K (q8_0/q4_0)
+ std::map<ImagePoolKey, ImagePoolEntry> dequant_f16_pool;
+
// prealloc buffers for src0 and src1
ggml_cl_buffer prealloc_src0;
ggml_cl_buffer prealloc_src1;
cl_kernel kernel_mul_mat_f16_f32_l4_dr;
cl_kernel kernel_mul_mat_f16_f32_l4_dr_ls;
cl_kernel kernel_mul_mat_f16_f32_l4_dr_lq;
+ cl_kernel kernel_mul_mat_f16_f32_l4_x8 = nullptr;
+ cl_kernel kernel_mul_mat_f16_f32_l4_x8_pair = nullptr;
+ cl_kernel kernel_mul_mat_f16_f32_l4_x8_gqa4 = nullptr;
+ cl_kernel kernel_mul_mat_f16_f32_l4_x8_gqa4_img = nullptr;
+ cl_kernel kernel_mul_mat_f16_f32_l4_x8_gqa_r4_img = nullptr;
+ cl_kernel kernel_mul_mat_f16_f32_l4_x8_gqa_r2_dk256_img = nullptr;
+ cl_kernel kernel_mul_mat_f16_f32_l4_y8 = nullptr;
+ cl_kernel kernel_mul_mat_f16_f32_l4_y8_gqa = nullptr;
+ cl_kernel kernel_mul_mat_f16_f32_l4_y8_gqa_img = nullptr;
cl_kernel kernel_mul_mat_f16_f32_tiled;
cl_kernel kernel_adreno_xmem_pack_src_f32;
cl_kernel kernel_adreno_xmem_prepack_weight_f16;
write_profiling_info();
profiling_results.clear();
#endif
+ // release pooled image1d_buffer views over KV cache layers.
+ for (auto & kv : kq_img_pool) {
+ if (kv.second.image) { CL_CHECK(clReleaseMemObject(kv.second.image)); }
+ if (kv.second.sub_buffer) { CL_CHECK(clReleaseMemObject(kv.second.sub_buffer)); }
+ }
+ kq_img_pool.clear();
+ for (auto & kv : kqv_img_pool) {
+ if (kv.second.image) { CL_CHECK(clReleaseMemObject(kv.second.image)); }
+ if (kv.second.sub_buffer) { CL_CHECK(clReleaseMemObject(kv.second.sub_buffer)); }
+ }
+ kqv_img_pool.clear();
+ for (auto & kv : dequant_f16_pool) {
+ if (kv.second.image) { CL_CHECK(clReleaseMemObject(kv.second.image)); }
+ }
+ dequant_f16_pool.clear();
}
}
};
// fatal=false returns NULL on compile failure instead of aborting; used for
// optional FA variants that may exhaust the Adreno compiler at large DK.
-static cl_program build_program_from_source_ex(cl_context ctx, cl_device_id dev, const char* program_buffer, const std::string &compile_opts, bool fatal, const char *tag = nullptr) {
- if (tag) {
- GGML_LOG_INFO("ggml_opencl: compiling %s\n", tag);
- }
-
+// when the compiler returns CL_OUT_OF_HOST_MEMORY/CL_OUT_OF_RESOURCES (seen with DK>=256/512)
+// for FA programs, do clFinish the queue to free up resources, then rebuild (up to 3x)
+// if retry_queue is provided
+static cl_program build_program_from_source_ex(cl_context ctx, cl_device_id dev, const char* program_buffer, const std::string &compile_opts, bool fatal, const char *tag = nullptr, cl_command_queue retry_queue = nullptr) {
+ if (tag) { GGML_LOG_INFO("ggml_opencl: compiling %s\n", tag); }
cl_program p;
char *program_log;
size_t program_size;
program_size = strlen(program_buffer);
- p = clCreateProgramWithSource(ctx, 1, (const char**)&program_buffer, &program_size, &err);
- if(err < 0) {
- GGML_LOG_ERROR("OpenCL error creating program");
- if (fatal) {
- exit(1);
+ const int max_attempts = retry_queue ? 3 : 1;
+ for (int attempt = 0; attempt < max_attempts; ++attempt) {
+ p = clCreateProgramWithSource(ctx, 1, (const char**)&program_buffer, &program_size, &err);
+ if(err < 0) {
+ GGML_LOG_ERROR("OpenCL error creating program");
+ if (fatal) exit(1);
+ return NULL;
+ }
+
+ err = clBuildProgram(p, 0, NULL, compile_opts.c_str(), NULL, NULL);
+ if (err == CL_SUCCESS) {
+ return p;
+ }
+
+ const bool transient = (err == CL_OUT_OF_HOST_MEMORY || err == CL_OUT_OF_RESOURCES);
+ if (retry_queue && transient && attempt + 1 < max_attempts) {
+ clReleaseProgram(p);
+ GGML_LOG_WARN("ggml_opencl: transient compile failure (err=%d)%s%s — clFinish + retry (%d/%d)\n",
+ err, tag ? " building " : "", tag ? tag : "", attempt + 2, max_attempts);
+ clFinish(retry_queue); // drain in-flight ops holding driver host-heap
+ continue;
}
- return nullptr;
- }
- err = clBuildProgram(p, 0, NULL, compile_opts.c_str(), NULL, NULL);
- if(err < 0) {
clGetProgramBuildInfo(p, dev, CL_PROGRAM_BUILD_LOG, 0, NULL, &log_size);
program_log = (char*) malloc(log_size + 1);
program_log[log_size] = '\0';
clGetProgramBuildInfo(p, dev, CL_PROGRAM_BUILD_LOG, log_size + 1, program_log, NULL);
- GGML_LOG_ERROR("ggml_opencl: kernel compile error (err=%d):\n\n%s\n", err, program_log);
+ GGML_LOG_ERROR("ggml_opencl: kernel compile error (err=%d)%s%s:\n\n%s\n", err, tag ? " building " : "", tag ? tag : "", program_log);
free(program_log);
clReleaseProgram(p);
if (fatal) {
}
return nullptr;
}
-
- return p;
+ return NULL;
}
static cl_program build_program_from_source(cl_context ctx, cl_device_id dev, const char* program_buffer, const std::string &compile_opts) {
CL_CHECK((backend_ctx->kernel_mul_mat_f16_f32_l4_dr_ls = clCreateKernel(backend_ctx->program_mul_mv_f16_f32_l4, "kernel_mul_mat_f16_f32_l4_dr_ls", &err), err));
CL_CHECK((backend_ctx->kernel_mul_mat_f16_f32_l4_dr_lq = clCreateKernel(backend_ctx->program_mul_mv_f16_f32_l4, "kernel_mul_mat_f16_f32_l4_dr_lq", &err), err));
}
+
+ cl_int err_x8 = CL_SUCCESS;
+ backend_ctx->kernel_mul_mat_f16_f32_l4_x8 =
+ clCreateKernel(backend_ctx->program_mul_mv_f16_f32_l4, "kernel_mul_mat_f16_f32_l4_x8", &err_x8);
+ if (err_x8 != CL_SUCCESS) { backend_ctx->kernel_mul_mat_f16_f32_l4_x8 = nullptr; }
+
+ cl_int err_x8p = CL_SUCCESS;
+ backend_ctx->kernel_mul_mat_f16_f32_l4_x8_pair =
+ clCreateKernel(backend_ctx->program_mul_mv_f16_f32_l4, "kernel_mul_mat_f16_f32_l4_x8_pair", &err_x8p);
+ if (err_x8p != CL_SUCCESS) { backend_ctx->kernel_mul_mat_f16_f32_l4_x8_pair = nullptr; }
+
+ cl_int err_x8g = CL_SUCCESS;
+ backend_ctx->kernel_mul_mat_f16_f32_l4_x8_gqa4 =
+ clCreateKernel(backend_ctx->program_mul_mv_f16_f32_l4, "kernel_mul_mat_f16_f32_l4_x8_gqa4", &err_x8g);
+ if (err_x8g != CL_SUCCESS) { backend_ctx->kernel_mul_mat_f16_f32_l4_x8_gqa4 = nullptr; }
+
+ cl_int err_x8gi = CL_SUCCESS;
+ backend_ctx->kernel_mul_mat_f16_f32_l4_x8_gqa4_img =
+ clCreateKernel(backend_ctx->program_mul_mv_f16_f32_l4, "kernel_mul_mat_f16_f32_l4_x8_gqa4_img", &err_x8gi);
+ if (err_x8gi != CL_SUCCESS) { backend_ctx->kernel_mul_mat_f16_f32_l4_x8_gqa4_img = nullptr; }
+
+ cl_int err_x8gi_r4 = CL_SUCCESS;
+ backend_ctx->kernel_mul_mat_f16_f32_l4_x8_gqa_r4_img =
+ clCreateKernel(backend_ctx->program_mul_mv_f16_f32_l4, "kernel_mul_mat_f16_f32_l4_x8_gqa_r4_img", &err_x8gi_r4);
+ if (err_x8gi_r4 != CL_SUCCESS) { backend_ctx->kernel_mul_mat_f16_f32_l4_x8_gqa_r4_img = nullptr; }
+
+ cl_int err_r2dk256 = CL_SUCCESS;
+ backend_ctx->kernel_mul_mat_f16_f32_l4_x8_gqa_r2_dk256_img =
+ clCreateKernel(backend_ctx->program_mul_mv_f16_f32_l4, "kernel_mul_mat_f16_f32_l4_x8_gqa_r2_dk256_img", &err_r2dk256);
+ if (err_r2dk256 != CL_SUCCESS) { backend_ctx->kernel_mul_mat_f16_f32_l4_x8_gqa_r2_dk256_img = nullptr; }
+
+ cl_int err_y8 = CL_SUCCESS;
+ backend_ctx->kernel_mul_mat_f16_f32_l4_y8 =
+ clCreateKernel(backend_ctx->program_mul_mv_f16_f32_l4, "kernel_mul_mat_f16_f32_l4_y8", &err_y8);
+ if (err_y8 != CL_SUCCESS) { backend_ctx->kernel_mul_mat_f16_f32_l4_y8 = nullptr; }
+
+ cl_int err_y8g = CL_SUCCESS;
+ backend_ctx->kernel_mul_mat_f16_f32_l4_y8_gqa =
+ clCreateKernel(backend_ctx->program_mul_mv_f16_f32_l4, "kernel_mul_mat_f16_f32_l4_y8_gqa", &err_y8g);
+ if (err_y8g != CL_SUCCESS) { backend_ctx->kernel_mul_mat_f16_f32_l4_y8_gqa = nullptr; }
+
+ cl_int err_y8gi = CL_SUCCESS;
+ backend_ctx->kernel_mul_mat_f16_f32_l4_y8_gqa_img =
+ clCreateKernel(backend_ctx->program_mul_mv_f16_f32_l4, "kernel_mul_mat_f16_f32_l4_y8_gqa_img", &err_y8gi);
+ if (err_y8gi != CL_SUCCESS) { backend_ctx->kernel_mul_mat_f16_f32_l4_y8_gqa_img = nullptr; }
GGML_LOG_CONT(".");
}
" -D BLOCK_M=" + std::to_string(cfg->bm) +
" -D BLOCK_N=" + std::to_string(cfg->bn);
+ if (backend_ctx->gpu_family == INTEL) {
+ opts += " -D FA_SG=32";
+ }
+
const bool is_split = variant == FA_VARIANT_F32_F16_SPLIT ||
variant == FA_VARIANT_Q8_0_SPLIT ||
variant == FA_VARIANT_Q4_0_SPLIT;
return opts;
}
+// only register when the kernel's required dispatch workgroup size is within
+// the limit of the device's maximum workgroup size
+static bool ggml_opencl_fa_kernel_fits_wg(ggml_backend_opencl_context * backend_ctx,
+ cl_kernel kernel, size_t required_wg,
+ const char * name, int dk, int dv) {
+ if (kernel == NULL) { return false; }
+ const size_t dev_max = backend_ctx->max_workgroup_size;
+ if (dev_max < required_wg) {
+ GGML_LOG_INFO("ggml_opencl: %s DK=%d DV=%d requires WG %zu > device max %zu; skipping registration (will fall back)\n",
+ name, dk, dv, required_wg, dev_max);
+ return false;
+ }
+ size_t kwg = 0;
+ cl_int err = clGetKernelWorkGroupInfo(kernel, backend_ctx->device,
+ CL_KERNEL_WORK_GROUP_SIZE,
+ sizeof(kwg), &kwg, NULL);
+ if (err != CL_SUCCESS) {
+ GGML_LOG_INFO("ggml_opencl: clGetKernelWorkGroupInfo failed for %s DK=%d DV=%d (err=%d); skipping registration\n",
+ name, dk, dv, err);
+ return false;
+ }
+ if (kwg < required_wg) {
+ GGML_LOG_INFO("ggml_opencl: %s DK=%d DV=%d per-kernel max %zu < required %zu; skipping registration (will fall back)\n",
+ name, dk, dv, kwg, required_wg);
+ return false;
+ }
+ return true;
+}
+
// Log private memory for an FA kernel. Enable via `GGML_OPENCL_FA_LOG_SPILL=1`.
// On Adreno non-zero private_mem means spilling to global memory due to resource
// constraint and usually causes performance degradation.
static void ggml_opencl_ensure_fa_pre_kernels(ggml_backend_opencl_context * backend_ctx, int dk, int dv) {
const std::pair<int, int> dk_dv = {dk, dv};
- if (backend_ctx->fa.kv_pad_f16.count(dk_dv) > 0) {
- return;
- }
const ggml_opencl_fa_dim * cfg = nullptr;
for (const auto & d : g_opencl_fa_dims) {
GGML_ABORT("ggml_opencl: no flash_attn config for DK=%d DV=%d", dk, dv);
}
- GGML_LOG_INFO("ggml_opencl: lazy-compiling flash_attn prepass for DK=%d DV=%d\n", dk, dv);
+ // BM-tile metadata is consumed by the prefill dispatch (n_q_blocks / wg
+ // sizing) regardless of whether the prepass kernels are needed for this
+ // n_kv — set it unconditionally
+ backend_ctx->fa.f32_f16_bm[{dk, dv}] = cfg->bm;
+ backend_ctx->fa.f32_f16_bn[{dk, dv}] = cfg->bn;
+ backend_ctx->fa.f32_f16_wg_size[{dk, dv}] = cfg->bm;
+ backend_ctx->fa.bm[{dk, dv}] = cfg->bm;
+ backend_ctx->fa.bn[{dk, dv}] = cfg->bn;
+ if (backend_ctx->fa.kv_pad_f16.count(dk_dv) > 0) { return; }
+
+ GGML_LOG_INFO("ggml_opencl: lazy-compiling flash_attn prepass for DK=%d DV=%d\n", dk, dv);
cl_int err;
const std::string src = ggml_opencl_fa_kernel_src(FA_VARIANT_PRE);
const std::string opts = ggml_opencl_fa_compile_opts(backend_ctx, cfg, FA_VARIANT_PRE);
-
- cl_program prog = build_program_from_source(backend_ctx->context, backend_ctx->device, src.c_str(), opts);
-
- cl_kernel k_kv_pad_f16, k_mask_pad_f16, k_blk_f16;
- CL_CHECK((k_kv_pad_f16 = clCreateKernel(prog, "flash_attn_kv_pad_f16", &err), err));
- CL_CHECK((k_mask_pad_f16 = clCreateKernel(prog, "flash_attn_mask_pad_f16", &err), err));
- CL_CHECK((k_blk_f16 = clCreateKernel(prog, "flash_attn_blk_f16", &err), err));
+ // retry when kernel compile fails
+ cl_program prog_pre_f16 = build_program_from_source_ex(
+ backend_ctx->context, backend_ctx->device, src.c_str(), opts,
+ /*fatal=*/false, "fa prepass f16", backend_ctx->queue);
+ if (!prog_pre_f16) { return; }
+ cl_kernel k_kv_pad_f16 = clCreateKernel(prog_pre_f16, "flash_attn_kv_pad_f16", &err);
+ if (err != CL_SUCCESS) { clReleaseProgram(prog_pre_f16); return; }
+ cl_kernel k_mask_pad_f16 = clCreateKernel(prog_pre_f16, "flash_attn_mask_pad_f16", &err);
+ if (err != CL_SUCCESS) { clReleaseKernel(k_kv_pad_f16); clReleaseProgram(prog_pre_f16); return; }
+ cl_kernel k_blk_f16 = clCreateKernel(prog_pre_f16, "flash_attn_blk_f16", &err);
+ if (err != CL_SUCCESS) { clReleaseKernel(k_kv_pad_f16); clReleaseKernel(k_mask_pad_f16); clReleaseProgram(prog_pre_f16); return; }
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;
- CL_CHECK(clReleaseProgram(prog));
+ clReleaseProgram(prog_pre_f16);
+}
- backend_ctx->fa.f32_f16_bm[{dk, dv}] = cfg->bm;
- backend_ctx->fa.f32_f16_bn[{dk, dv}] = cfg->bn;
- backend_ctx->fa.f32_f16_wg_size[{dk, dv}] = cfg->bm;
- backend_ctx->fa.bm[{dk, dv}] = cfg->bm;
- backend_ctx->fa.bn[{dk, dv}] = cfg->bn;
+// DK=512 prefill BM-tile
+static bool ggml_opencl_ensure_fa_f32_f16_prefill_512(ggml_backend_opencl_context * backend_ctx, bool split) {
+ const int dk = 512, dv = 512;
+ const std::pair<int, int> dk_dv = {dk, dv};
+ auto & target = split ? backend_ctx->fa.f32_f16_split : backend_ctx->fa.f32_f16;
+ if (target.count(dk_dv) > 0) { return true; }
+
+ static bool failed[2] = { false, false };
+ if (failed[split ? 1 : 0]) { return false; }
+
+ const ggml_opencl_fa_dim * cfg = nullptr;
+ for (const auto & d : g_opencl_fa_dims) {
+ if (d.dk == dk && d.dv == dv) { cfg = &d; break; }
+ }
+ if (cfg == nullptr) { failed[split ? 1 : 0] = true; return false; }
+ if (split && cfg->n_split <= 1) { failed[1] = true; return false; }
+
+ const ggml_opencl_fa_variant variant = split ? FA_VARIANT_F32_F16_SPLIT : FA_VARIANT_F32_F16;
+ std::string opts = ggml_opencl_fa_compile_opts(backend_ctx, cfg, variant) + " -D FA_PREFILL_ONLY";
+ cl_program prog = build_program_from_source_ex(
+ backend_ctx->context, backend_ctx->device,
+ ggml_opencl_fa_kernel_src(FA_VARIANT_F32_F16).c_str(), opts,
+ /*fatal=*/false, split ? "fa f32_f16 prefill512 split" : "fa f32_f16 prefill512",
+ backend_ctx->queue);
+ if (!prog) { failed[split ? 1 : 0] = true; return false; }
+
+ cl_int err;
+ cl_kernel k = clCreateKernel(prog, "flash_attn_f32_f16", &err);
+ if (err != CL_SUCCESS) { clReleaseProgram(prog); failed[split ? 1 : 0] = true; return false; }
+ target[dk_dv] = k;
+ if (split) {
+ backend_ctx->fa.f32_f16_split_wg_size[dk_dv] = cfg->bm * cfg->n_split;
+ backend_ctx->fa.f32_f16_split_nkv_threshold[dk_dv] = cfg->nkv_split_threshold;
+ }
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k,
+ split ? "flash_attn_f32_f16 (prefill512 split)" : "flash_attn_f32_f16 (prefill512)", dk, dv);
+ clReleaseProgram(prog);
+
+ // determine whether to use the K-image variant of the split tile
+ static const char * pkimg_build_env = getenv("GGML_OPENCL_FA_PREFILL_K_IMG");
+ const bool pkimg_build = (pkimg_build_env != NULL) && (pkimg_build_env[0] != '0');
+ if (split && pkimg_build && backend_ctx->fa.f32_f16_split_k_img.count(dk_dv) == 0) {
+ std::string opts_img = ggml_opencl_fa_compile_opts(backend_ctx, cfg, variant) +
+ " -D FA_PREFILL_ONLY -D FA_K_IMG -D FA_TILE_NAME=flash_attn_f32_f16_k_img";
+ cl_program prog_img = build_program_from_source_ex(
+ backend_ctx->context, backend_ctx->device,
+ ggml_opencl_fa_kernel_src(FA_VARIANT_F32_F16).c_str(), opts_img,
+ /*fatal=*/false, "fa f32_f16 prefill512 split k_img", backend_ctx->queue);
+ if (prog_img) {
+ cl_int err_img;
+ cl_kernel k_img = clCreateKernel(prog_img, "flash_attn_f32_f16_k_img", &err_img);
+ if (err_img == CL_SUCCESS) {
+ backend_ctx->fa.f32_f16_split_k_img[dk_dv] = k_img;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_img,
+ "flash_attn_f32_f16 (prefill512 split k_img)", dk, dv);
+ }
+ clReleaseProgram(prog_img);
+ }
+ }
+ return true;
}
// Compile one (variant, dk, dv); memoised. false = compiler rejected.
break;
}
case FA_VARIANT_F32_F16: {
- if (backend_ctx->fa.f32_f16.count(dk_dv)) {
+ // The DK=512 decode-only program does not create the f32_f16
+ // prefill kernel; check the q1 kernel instead so that repeated
+ // calls return a consistent result.
+ const bool decode_only = (dk == 512);
+ if (decode_only ? (backend_ctx->fa.f32_f16_q1.count(dk_dv) > 0)
+ : (backend_ctx->fa.f32_f16.count(dk_dv) > 0)) {
return true;
}
break;
}
const std::string src = ggml_opencl_fa_kernel_src(variant);
+ if (src.empty()) { return false; }
+ std::string opts = ggml_opencl_fa_compile_opts(backend_ctx, cfg, variant);
- if (src.empty()) {
- return false;
+ // bypass kernels for DK=512
+ const bool fa_decode_only = (variant == FA_VARIANT_F32_F16 && dk == 512);
+ if (fa_decode_only) {
+ opts += " -D FA_DECODE_ONLY -D FA_DECODE_MINIMAL";
}
- const std::string opts = ggml_opencl_fa_compile_opts(backend_ctx, cfg, variant);
const char * tag = nullptr;
switch (variant) {
default: break;
}
cl_program prog = build_program_from_source_ex(
- backend_ctx->context, backend_ctx->device, src.c_str(), opts, /*fatal=*/false, tag);
-
- if (!prog) {
- return false;
- }
+ backend_ctx->context, backend_ctx->device, src.c_str(), opts,
+ /*fatal=*/false, tag, backend_ctx->queue);
+ if (!prog) { return false; }
cl_int err;
switch (variant) {
break;
}
case FA_VARIANT_F32_F16: {
- cl_kernel k, kq1;
- CL_CHECK((k = clCreateKernel(prog, "flash_attn_f32_f16", &err), err));
+ cl_kernel kq1;
+ // BM-tile prefill kernel is excluded from the decode-only (DK=512)
+ if (!fa_decode_only) {
+ cl_kernel k;
+ CL_CHECK((k = clCreateKernel(prog, "flash_attn_f32_f16", &err), err));
+ backend_ctx->fa.f32_f16[{dk, dv}] = k;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k, "flash_attn_f32_f16", dk, dv);
+ }
CL_CHECK((kq1 = clCreateKernel(prog, "flash_attn_f32_f16_q1", &err), err));
- backend_ctx->fa.f32_f16[{dk, dv}] = k;
backend_ctx->fa.f32_f16_q1[{dk, dv}] = kq1;
- ggml_opencl_log_fa_kernel_spill(backend_ctx, k, "flash_attn_f32_f16", dk, dv);
ggml_opencl_log_fa_kernel_spill(backend_ctx, kq1, "flash_attn_f32_f16_q1", dk, dv);
cl_kernel k_split = clCreateKernel(prog, "flash_attn_f32_f16_q1_split", &err);
if (err == CL_SUCCESS) {
backend_ctx->fa.f32_f16_q1_split[{dk, dv}] = k_split;
ggml_opencl_log_fa_kernel_spill(backend_ctx, k_split, "flash_attn_f32_f16_q1_split", dk, dv);
}
+ // q1_vec decode kernel (DV-split + subgroup reduce)
+ cl_kernel k_q1_vec = clCreateKernel(prog, "flash_attn_f32_f16_q1_vec", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_q1_vec, 256,
+ "flash_attn_f32_f16_q1_vec", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_vec[{dk, dv}] = k_q1_vec;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_q1_vec, "flash_attn_f32_f16_q1_vec", dk, dv);
+ } else {
+ clReleaseKernel(k_q1_vec);
+ }
+ }
+ // KV-head-coalesced vec for high-GQA small models
+ cl_kernel k_q1_vec_mq = clCreateKernel(prog, "flash_attn_f32_f16_q1_vec_mq", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_q1_vec_mq, 256,
+ "flash_attn_f32_f16_q1_vec_mq", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_vec_mq[{dk, dv}] = k_q1_vec_mq;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_q1_vec_mq, "flash_attn_f32_f16_q1_vec_mq", dk, dv);
+ } else {
+ clReleaseKernel(k_q1_vec_mq);
+ }
+ }
+ // KV-head-coalesced + flash-decoding split, reuses merge kernel
+ cl_kernel k_q1_vec_mq_split = clCreateKernel(prog, "flash_attn_f32_f16_q1_vec_mq_split", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_q1_vec_mq_split, 256,
+ "flash_attn_f32_f16_q1_vec_mq_split", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_vec_mq_split[{dk, dv}] = k_q1_vec_mq_split;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_q1_vec_mq_split, "flash_attn_f32_f16_q1_vec_mq_split", dk, dv);
+ } else {
+ clReleaseKernel(k_q1_vec_mq_split);
+ }
+ }
+ // K-image variant of MQ_GQA=4 split
+ cl_kernel k_q1_vec_mq_split_k_img = clCreateKernel(prog, "flash_attn_f32_f16_q1_vec_mq_split_k_img", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_q1_vec_mq_split_k_img, 256,
+ "flash_attn_f32_f16_q1_vec_mq_split_k_img", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_k_img[{dk, dv}] = k_q1_vec_mq_split_k_img;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_q1_vec_mq_split_k_img, "flash_attn_f32_f16_q1_vec_mq_split_k_img", dk, dv);
+ } else {
+ clReleaseKernel(k_q1_vec_mq_split_k_img);
+ }
+ }
cl_kernel k_merge = clCreateKernel(prog, "flash_attn_f32_merge", &err);
if (err == CL_SUCCESS) {
backend_ctx->fa.f32_merge[{dk, dv}] = k_merge;
}
+ // local-tile decode variant
+ if (dk == 128 && dv == 128) {
+ cl_kernel k_lt = clCreateKernel(prog, "flash_attn_f32_f16_q1_local_tile", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_lt, 128,
+ "flash_attn_f32_f16_q1_local_tile", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_local_tile[{dk, dv}] = k_lt;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_lt, "flash_attn_f32_f16_q1_local_tile", dk, dv);
+ } else {
+ clReleaseKernel(k_lt);
+ }
+ }
+ // hybrid local-tile + MQ + FD-split
+ cl_kernel k_lmq = clCreateKernel(prog, "flash_attn_f32_f16_q1_local_mq_split", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_lmq, 64,
+ "flash_attn_f32_f16_q1_local_mq_split", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_local_mq_split[{dk, dv}] = k_lmq;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_lmq, "flash_attn_f32_f16_q1_local_mq_split", dk, dv);
+ } else {
+ clReleaseKernel(k_lmq);
+ }
+ }
+ }
+
+ // second compile of the same source with -DMQ_GQA=8.
+ // FA_MQ_ONLY keeps only the vec_mq kernels so that the program
+ // compiles within the Adreno compiler's memory budget at DK>=256.
+ const std::string opts_g8 = opts + " -D MQ_GQA=8 -D MQ_NSG=3 -D MQ_NSG_SPLIT=3 -D FA_MQ_ONLY";
+ cl_program prog_g8 = fa_decode_only ? nullptr : build_program_from_source_ex(
+ backend_ctx->context, backend_ctx->device, src.c_str(), opts_g8,
+ /*fatal=*/false, "fa f32_f16 MQ_GQA=8", backend_ctx->queue);
+ if (prog_g8) {
+ const size_t mq_g8_required_wg = 192; // Q1_WG_SIZE(64) * MQ_NSG_SPLIT(3)
+ cl_kernel k_q1_vec_mq_g8 = clCreateKernel(prog_g8, "flash_attn_f32_f16_q1_vec_mq", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_q1_vec_mq_g8, mq_g8_required_wg,
+ "flash_attn_f32_f16_q1_vec_mq (g8)", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_vec_mq_g8[{dk, dv}] = k_q1_vec_mq_g8;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_q1_vec_mq_g8, "flash_attn_f32_f16_q1_vec_mq_g8", dk, dv);
+ } else {
+ clReleaseKernel(k_q1_vec_mq_g8);
+ }
+ }
+ cl_kernel k_q1_vec_mq_split_g8 = clCreateKernel(prog_g8, "flash_attn_f32_f16_q1_vec_mq_split", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_q1_vec_mq_split_g8, mq_g8_required_wg,
+ "flash_attn_f32_f16_q1_vec_mq_split (g8)", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_g8[{dk, dv}] = k_q1_vec_mq_split_g8;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_q1_vec_mq_split_g8, "flash_attn_f32_f16_q1_vec_mq_split_g8", dk, dv);
+ } else {
+ clReleaseKernel(k_q1_vec_mq_split_g8);
+ }
+ }
+ // K-image variant
+ cl_kernel k_q1_vec_mq_split_g8_k_img = clCreateKernel(prog_g8, "flash_attn_f32_f16_q1_vec_mq_split_k_img", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_q1_vec_mq_split_g8_k_img, mq_g8_required_wg,
+ "flash_attn_f32_f16_q1_vec_mq_split_k_img (g8)", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_k_img[{dk, dv}] = k_q1_vec_mq_split_g8_k_img;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_q1_vec_mq_split_g8_k_img, "flash_attn_f32_f16_q1_vec_mq_split_g8_k_img", dk, dv);
+ } else {
+ clReleaseKernel(k_q1_vec_mq_split_g8_k_img);
+ }
+ }
+ // hybrid local-tile + MQ_GQA=8
+ if (dk == 128 && dv == 128) {
+ cl_kernel k_lmq_g8 = clCreateKernel(prog_g8, "flash_attn_f32_f16_q1_local_mq_split", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_lmq_g8, 64,
+ "flash_attn_f32_f16_q1_local_mq_split (g8)", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_local_mq_split_g8[{dk, dv}] = k_lmq_g8;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_lmq_g8, "flash_attn_f32_f16_q1_local_mq_split_g8", dk, dv);
+ } else {
+ clReleaseKernel(k_lmq_g8);
+ }
+ }
+ }
+ clReleaseProgram(prog_g8);
+ }
break;
}
case FA_VARIANT_Q8_0:
m_q1_split[{dk, dv}] = k_split;
ggml_opencl_log_fa_kernel_spill(backend_ctx, k_split, name_q1_split.c_str(), dk, dv);
}
+
+ // DV-split decode variant (q1_vec)
+ auto & m_q1_vec = is_q8 ? backend_ctx->fa.f32_q8_0_q1_vec : backend_ctx->fa.f32_q4_0_q1_vec;
+ const std::string name_q1_vec = name_q1 + "_vec";
+ cl_kernel k_q1_vec = clCreateKernel(prog, name_q1_vec.c_str(), &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_q1_vec, 256,
+ name_q1_vec.c_str(), dk, dv)) {
+ m_q1_vec[{dk, dv}] = k_q1_vec;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_q1_vec, name_q1_vec.c_str(), dk, dv);
+ } else {
+ clReleaseKernel(k_q1_vec);
+ }
+ }
+
+ // KV-head-coalesced + flash-decoding split
+ auto & m_mq_split = is_q8 ? backend_ctx->fa.f32_q8_0_q1_vec_mq_split
+ : backend_ctx->fa.f32_q4_0_q1_vec_mq_split;
+ const std::string name_mq_split = name_q1 + "_vec_mq_split";
+ cl_kernel k_q1_vec_mq_split = clCreateKernel(prog, name_mq_split.c_str(), &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_q1_vec_mq_split, 256,
+ name_mq_split.c_str(), dk, dv)) {
+ m_mq_split[{dk, dv}] = k_q1_vec_mq_split;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_q1_vec_mq_split, name_mq_split.c_str(), dk, dv);
+ } else {
+ clReleaseKernel(k_q1_vec_mq_split);
+ }
+ }
if (!backend_ctx->fa.f32_merge.count({dk, dv})) {
cl_kernel k_merge = clCreateKernel(prog, "flash_attn_f32_merge", &err);
if (err == CL_SUCCESS) {
backend_ctx->fa.f32_merge[{dk, dv}] = k_merge;
}
}
+ // Second compile with MQ_GQA=8, MQ_NSG=3, MQ_NSG_SPLIT=3
+ auto & m_mq_split_g8 = is_q8 ? backend_ctx->fa.f32_q8_0_q1_vec_mq_split_g8
+ : backend_ctx->fa.f32_q4_0_q1_vec_mq_split_g8;
+ const std::string opts_mq_g8 = opts + " -D MQ_GQA=8 -D MQ_NSG=3 -D MQ_NSG_SPLIT=3";
+ cl_program prog_mq_g8 = build_program_from_source_ex(
+ backend_ctx->context, backend_ctx->device, src.c_str(), opts_mq_g8,
+ /*fatal=*/false, is_q8 ? "fa q8_0 MQ_GQA=8" : "fa q4_0 MQ_GQA=8",
+ backend_ctx->queue);
+ if (prog_mq_g8) {
+ const size_t mq_g8_required_wg = 192;
+ cl_kernel k_g8 = clCreateKernel(prog_mq_g8, name_mq_split.c_str(), &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_g8, mq_g8_required_wg,
+ name_mq_split.c_str(), dk, dv)) {
+ m_mq_split_g8[{dk, dv}] = k_g8;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_g8, name_mq_split.c_str(), dk, dv);
+ } else {
+ clReleaseKernel(k_g8);
+ }
+ }
+ clReleaseProgram(prog_mq_g8);
+ }
break;
}
case FA_VARIANT_F32_F16_SPLIT: {
const std::string tag = std::string("fa ") + (is_q8_0 ? "q8_0" : "q4_0") +
" split DK=" + std::to_string(dk);
cl_program prog = build_program_from_source_ex(
- backend_ctx->context, backend_ctx->device, src.c_str(), opts, /*fatal=*/false, tag.c_str());
-
- if (!prog) {
- return false;
- }
-
+ backend_ctx->context, backend_ctx->device, src.c_str(), opts,
+ /*fatal=*/false, tag.c_str(), backend_ctx->queue);
+ if (!prog) { return false; }
cl_int err;
cl_kernel k;
if (is_q8_0) {
{ 40, 40}, { 64, 64}, { 80, 80}, { 96, 96},
{112, 112}, {128, 128}, {192, 128},
{192, 192}, {256, 256},
+ {512, 512},
};
bool dims_supported = false;
k->type != v->type &&
is_kv_type_ok(k->type) && is_kv_type_ok(v->type);
- return is_f32_f32 || is_f16_f16 || is_f32_f16 || is_f32_q8_0 || is_f32_q4_0 || is_f32_asym;
+ const bool kv_combo_ok = is_f32_f32 || is_f16_f16 || is_f32_f16 ||
+ is_f32_q8_0 || is_f32_q4_0 || is_f32_asym;
+ if (!kv_combo_ok) {
+ return false;
+ }
+
+ if (dk == 512) {
+ if (backend_ctx->gpu_family == INTEL) {
+ return false;
+ }
+ if (!is_f32_f16) {
+ return false;
+ }
+ if (q->ne[1] == 1) {
+ // DK=512 decode is bandwidth-bound and slower on the GPU
+ // than on the CPU; decline it here so it runs on the CPU.
+ // Prefill (n_q > 1) stays on the GPU.
+ return false;
+ } else {
+ // prefill, BM-tile in its own FA_PREFILL_ONLY program
+ if (!ggml_opencl_ensure_fa_f32_f16_prefill_512(backend_ctx, /*split=*/false)) {
+ return false;
+ }
+ }
+ }
+ return true;
}
default:
return false;
GGML_ASSERT(dst_off == total_bytes);
}
+// forward decl: used by the FA decode dispatch (K-image variant) below.
+static cl_mem ggml_cl_img_pool_get_or_create(
+ ggml_backend_opencl_context * backend_ctx,
+ std::map<ggml_backend_opencl_context::ImagePoolKey,
+ ggml_backend_opencl_context::ImagePoolEntry> & pool,
+ cl_mem data_device,
+ cl_ulong offset0,
+ size_t required_bytes,
+ cl_channel_type channel_data_type);
+
// Rebuild AoS q8_0/q4_0 bytes from a SoA tensor into a temp buffer.
// Returns false if the tensor is not SoA-quantised (already AoS).
static bool ggml_cl_flash_attn_reconstruct_aos(
// (d_head <= FD_MAX_DK_MULTI) and capped at FD_MAX_N_Q_MULTI queries.
static constexpr int FD_MIN_N_KV = 2048;
static constexpr int FD_KV_PER_SPLIT = 2048;
+// f16 KV decode wants more splits than the 2048 default; quantized KV keeps 2048.
+static constexpr int FD_KV_PER_SPLIT_F16 = 512;
static constexpr int FD_MIN_SPLITS = 2;
static constexpr int FD_MAX_SPLITS = 16;
static constexpr int FD_MAX_DK = 128;
static constexpr int FD_MAX_DK_MULTI = 64;
static constexpr int FD_MAX_N_Q_MULTI = 8;
+// MQ FD split-groups have few subgroups (MQ_NSG_SPLIT), so use a smaller
+// kv_per_split to keep the softmax recurrence short; non-MQ keeps FD_KV_PER_SPLIT.
+static constexpr int FD_MQ_KV_PER_SPLIT = 256;
+static constexpr int FD_MQ_MAX_SPLITS = 128;
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 int n_head_kv = k->ne[2];
const int n_batch = q->ne[3];
- // Per-variant lazy compile for this (dk, dv).
- ggml_opencl_ensure_fa_pre_kernels(backend_ctx, d_head_q, d_head_v);
+ // DK=512 (Gemma-4 global layers) runs decode-only (q1 / q1_split) on
+ // Adreno - it never uses the BM-tile path, and the prepass + split-tile
+ // programs OOM the compiler at DK=512; supports_op only admits
+ // n_q==1 here and prefill goes to CPU
+ const bool fa_decode_only_512 = (d_head_q == 512);
+
+ // per-variant lazy compile for this (dk, dv)
+ // DK=512 decode (n_q==1) needs no prepass
+ // DK=512 prefill (n_q>1) does, so compile it only when needed
+ if (!fa_decode_only_512 || n_q > 1) {
+ ggml_opencl_ensure_fa_pre_kernels(backend_ctx, d_head_q, d_head_v);
+ }
cl_kernel kernel = NULL;
+ bool use_prefill_k_img = false; // K is image1d_buffer_t for DK=512 prefill
const bool is_f16 = q->type == GGML_TYPE_F16;
const bool is_mixed = q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16;
ggml_opencl_ensure_fa_variant(backend_ctx, d_head_q, d_head_v, FA_VARIANT_F16);
} else if (is_mixed) {
ggml_opencl_ensure_fa_variant(backend_ctx, d_head_q, d_head_v, FA_VARIANT_F32_F16);
- ggml_opencl_ensure_fa_variant(backend_ctx, d_head_q, d_head_v, FA_VARIANT_F32_F16_SPLIT);
+ if (fa_decode_only_512) {
+ // DK=512: the BM-tile prefill kernels are specifically compiled from
+ // FA_PREFILL_ONLY
+ if (n_q > 1) {
+ ggml_opencl_ensure_fa_f32_f16_prefill_512(backend_ctx, /*split=*/false);
+ ggml_opencl_ensure_fa_f32_f16_prefill_512(backend_ctx, /*split=*/true);
+ }
+ } else {
+ ggml_opencl_ensure_fa_variant(backend_ctx, d_head_q, d_head_v, FA_VARIANT_F32_F16_SPLIT);
+ }
} else if (is_q8_0) {
ggml_opencl_ensure_fa_variant(backend_ctx, d_head_q, d_head_v, FA_VARIANT_Q8_0);
if (d_head_q == 96 && d_head_v == 96) {
const int block_m = n_q > 1
? (is_mixed ? backend_ctx->fa.f32_f16_bm.at(dk_dv) : backend_ctx->fa.bm.at(dk_dv))
: 0;
- const int block_n = is_mixed
- ? backend_ctx->fa.f32_f16_bn.at(dk_dv)
- : backend_ctx->fa.bn.at(dk_dv);
+ // block_n is only used by the n_q > 1 prefill path; its map is not
+ // populated for DK=512 decode, so do not read it for decode.
+ const int block_n = (n_q > 1)
+ ? (is_mixed ? backend_ctx->fa.f32_f16_bn.at(dk_dv)
+ : backend_ctx->fa.bn.at(dk_dv))
+ : 0;
// Pick split variant only when n_kv crosses the per-(dk,dv) threshold.
+ // the N_SPLIT>1 prefill tile reduces DK partials via subgroup shuffle,
+ // on Intel it uses the non-split BM tile and does not depend on subgroup size
const bool use_split_kernel = (n_q > 1 && is_mixed &&
+ backend_ctx->gpu_family != INTEL &&
backend_ctx->fa.f32_f16_split.count(dk_dv) > 0 &&
n_kv >= backend_ctx->fa.f32_f16_split_nkv_threshold.at(dk_dv));
- const bool use_split_q8_0 = (use_native_q8_0 &&
+ const bool use_split_q8_0 = (use_native_q8_0 && backend_ctx->gpu_family != INTEL &&
backend_ctx->fa.f32_q8_0_split.count(dk_dv) > 0 &&
n_kv >= backend_ctx->fa.f32_q8_0_split_nkv_threshold.at(dk_dv));
- const bool use_split_q4_0 = (use_native_q4_0 &&
+ const bool use_split_q4_0 = (use_native_q4_0 && backend_ctx->gpu_family != INTEL &&
backend_ctx->fa.f32_q4_0_split.count(dk_dv) > 0 &&
n_kv >= backend_ctx->fa.f32_q4_0_split_nkv_threshold.at(dk_dv));
const int wg_size_fa = (n_q > 1 && is_mixed)
max_bias = params[1];
logit_softcap = params[2];
+ bool use_q1_vec = false;
+ bool use_q1_vec_mq = false;
+ bool use_local_tile = false;
+ // KV-head-coalesced gate: gqa_ratio == compile-time MQ_GQA
+ // restricts to DK=DV=256 for now due to local memory size
+ const int gqa_ratio_dispatch = n_head_kv > 0 ? (n_head / n_head_kv) : 0;
if (n_q == 1) {
if (use_native_q8_0_q1) {
- kernel = backend_ctx->fa.f32_q8_0_q1.at(dk_dv);
+ if (d_head_v >= 256 &&
+ backend_ctx->fa.f32_q8_0_q1_vec.count(dk_dv) > 0) {
+ kernel = backend_ctx->fa.f32_q8_0_q1_vec.at(dk_dv);
+ use_q1_vec = true;
+ } else {
+ kernel = backend_ctx->fa.f32_q8_0_q1.at(dk_dv);
+ }
} else if (use_native_q4_0_q1) {
- kernel = backend_ctx->fa.f32_q4_0_q1.at(dk_dv);
+ // q4_0 vec kernel uses per-lane dp4a (cl_khr_integer_dot_product)
+ const char * q4vec_env = getenv("GGML_OPENCL_FA_Q4_VEC");
+ const bool q4vec_off = (q4vec_env != NULL) && (q4vec_env[0] == '0');
+ if (!q4vec_off && d_head_v >= 256 &&
+ backend_ctx->fa.f32_q4_0_q1_vec.count(dk_dv) > 0) {
+ kernel = backend_ctx->fa.f32_q4_0_q1_vec.at(dk_dv);
+ use_q1_vec = true;
+ } else {
+ kernel = backend_ctx->fa.f32_q4_0_q1.at(dk_dv);
+ }
} else if (is_mixed) {
- kernel = backend_ctx->fa.f32_f16_q1.at(dk_dv);
+ static const char * lt_env = getenv("GGML_OPENCL_FA_LOCAL_TILE");
+ static const bool lt_on = (lt_env != NULL) && (lt_env[0] != '0');
+ if (lt_on && d_head_q == 128 && d_head_v == 128 &&
+ backend_ctx->fa.f32_f16_q1_local_tile.count(dk_dv) > 0) {
+ kernel = backend_ctx->fa.f32_f16_q1_local_tile.at(dk_dv);
+ use_local_tile = true;
+ } else {
+ static const char * f16_vec_dk128_env = getenv("GGML_OPENCL_FA_F16_VEC_DK128");
+ static const bool f16_vec_dk128_off = (f16_vec_dk128_env != NULL) && (f16_vec_dk128_env[0] == '0');
+ const int dv_gate = f16_vec_dk128_off ? 256 : 128;
+ if (d_head_v >= dv_gate &&
+ backend_ctx->fa.f32_f16_q1_vec.count(dk_dv) > 0) {
+ kernel = backend_ctx->fa.f32_f16_q1_vec.at(dk_dv);
+ use_q1_vec = true;
+ } else {
+ kernel = backend_ctx->fa.f32_f16_q1.at(dk_dv);
+ }
+ }
} else if (is_f16) {
kernel = backend_ctx->fa.f16_q1.at(dk_dv);
} else {
? backend_ctx->fa.f32_q4_0_split.at(dk_dv)
: backend_ctx->fa.f32_q4_0.at(dk_dv);
} else if (is_mixed) {
- kernel = use_split_kernel
- ? backend_ctx->fa.f32_f16_split.at(dk_dv)
- : backend_ctx->fa.f32_f16.at(dk_dv);
+ if (use_split_kernel) {
+ // DK=512 prefill: opt-in texture-cache K reads (image1d_buffer_t).
+ static const char * pkimg_env = getenv("GGML_OPENCL_FA_PREFILL_K_IMG");
+ const bool pkimg_on = (pkimg_env != NULL) && (pkimg_env[0] != '0');
+ if (d_head_q == 512 && pkimg_on &&
+ backend_ctx->fa.f32_f16_split_k_img.count(dk_dv) > 0) {
+ kernel = backend_ctx->fa.f32_f16_split_k_img.at(dk_dv);
+ use_prefill_k_img = true;
+ } else {
+ kernel = backend_ctx->fa.f32_f16_split.at(dk_dv);
+ }
+ } else {
+ kernel = backend_ctx->fa.f32_f16.at(dk_dv);
+ }
} else if (is_f16) {
kernel = backend_ctx->fa.f16.at(dk_dv);
} else {
kernel = backend_ctx->fa.f32.at(dk_dv);
}
}
+
+ // Intel goes to the basic q1 kernel
+ if (backend_ctx->gpu_family == INTEL && n_q == 1) {
+ use_q1_vec = use_q1_vec_mq = use_local_tile = false;
+ if (is_mixed && backend_ctx->fa.f32_f16_q1.count(dk_dv)) { kernel = backend_ctx->fa.f32_f16_q1.at(dk_dv); }
+ else if (is_f16 && backend_ctx->fa.f16_q1.count(dk_dv)) { kernel = backend_ctx->fa.f16_q1.at(dk_dv); }
+ else if (is_q8_0 && backend_ctx->fa.f32_q8_0_q1.count(dk_dv)) { kernel = backend_ctx->fa.f32_q8_0_q1.at(dk_dv); }
+ else if (is_q4_0 && backend_ctx->fa.f32_q4_0_q1.count(dk_dv)) { kernel = backend_ctx->fa.f32_q4_0_q1.at(dk_dv); }
+ else if (backend_ctx->fa.f32_q1.count(dk_dv)) { kernel = backend_ctx->fa.f32_q1.at(dk_dv); }
+ }
GGML_ASSERT(kernel != NULL);
ggml_cl_flash_attn_temp_buffer temp_k;
const int is_causal = (mask == NULL && n_q > 1 && n_q == n_kv);
const int fd_max_n_q = (d_head_q <= FD_MAX_DK_MULTI) ? FD_MAX_N_Q_MULTI : 1;
cl_kernel fd_k_split = NULL;
- if (n_q >= 1 && n_q <= fd_max_n_q && n_kv >= FD_MIN_N_KV && !is_causal &&
+ bool use_fd_mq = false;
+ size_t fd_mq_wg = 256; // MQ_GQA=4 kernel: Q1_WG_SIZE(64) * MQ_NSG_SPLIT(4)
+ bool use_fa_k_img = false; // K bound as image1d_buffer_t instead of (buf, offset)
+
+ {
+ const char * mq_env = getenv("GGML_OPENCL_FA_MQ");
+ const bool mq_enabled = (mq_env == NULL) ? true : (mq_env[0] != '0');
+ const bool mq_kv_ok = is_mixed || is_q8_0 || is_q4_0;
+
+ const char * lmq_env = getenv("GGML_OPENCL_FA_LOCAL_MQ_SPLIT");
+ const bool lmq_on = (lmq_env != NULL) && (lmq_env[0] != '0');
+
+ static const char * vec_nq_env = getenv("GGML_OPENCL_FA_VEC_NQ");
+ static const int N_MAX_VEC_NQ = (vec_nq_env != NULL && vec_nq_env[0] != '\0')
+ ? atoi(vec_nq_env) : 1;
+
+ const bool nq_in_vec_range = (n_q >= 1) && (n_q <= N_MAX_VEC_NQ);
+ const bool nq1_only = (n_q == 1);
+ if (mq_enabled && mq_kv_ok && nq_in_vec_range && !is_causal &&
+ backend_ctx->gpu_family != INTEL &&
+ !use_local_tile &&
+ n_kv >= FD_MIN_N_KV &&
+ backend_ctx->fa.f32_merge.count(dk_dv) > 0) {
+ if (nq1_only && lmq_on && is_mixed && d_head_q == 128 && d_head_v == 128 &&
+ gqa_ratio_dispatch == 8 &&
+ backend_ctx->fa.f32_f16_q1_local_mq_split_g8.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_local_mq_split_g8.at(dk_dv);
+ use_fd_mq = true;
+ fd_mq_wg = 64;
+ } else if (nq1_only && lmq_on && is_mixed && d_head_q == 128 && d_head_v == 128 &&
+ gqa_ratio_dispatch == 4 &&
+ backend_ctx->fa.f32_f16_q1_local_mq_split.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_local_mq_split.at(dk_dv);
+ use_fd_mq = true;
+ fd_mq_wg = 64;
+ } else if (nq1_only && is_mixed && gqa_ratio_dispatch == 4 &&
+ ((d_head_q == 256 && d_head_v == 256) ||
+ (d_head_q == 128 && d_head_v == 128)) &&
+ backend_ctx->fa.f32_f16_q1_vec_mq_split.count(dk_dv) > 0) {
+ const bool k_img_on = d_head_q == 128 && d_head_v == 128 &&
+ getenv("GGML_OPENCL_FA_K_IMG") != NULL &&
+ getenv("GGML_OPENCL_FA_K_IMG")[0] != '0' &&
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_k_img.count(dk_dv) > 0;
+ if (k_img_on) {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_k_img.at(dk_dv);
+ use_fd_mq = true;
+ use_fa_k_img = true;
+ } else {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split.at(dk_dv);
+ use_fd_mq = true;
+ }
+ } else if (is_mixed && gqa_ratio_dispatch == 8 &&
+ d_head_q == 128 && d_head_v == 128 &&
+ getenv("GGML_OPENCL_FA_K_IMG") != NULL &&
+ getenv("GGML_OPENCL_FA_K_IMG")[0] != '0' &&
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_k_img.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_k_img.at(dk_dv);
+ use_fd_mq = true;
+ fd_mq_wg = 192;
+ use_fa_k_img = true;
+ } else if (is_mixed && gqa_ratio_dispatch == 8 &&
+ d_head_q == 128 && d_head_v == 128 &&
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_g8.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_g8.at(dk_dv);
+ use_fd_mq = true;
+ fd_mq_wg = 192;
+ } else if (nq1_only && is_q8_0 && gqa_ratio_dispatch == 8 &&
+ d_head_q == 128 && d_head_v == 128 &&
+ backend_ctx->fa.f32_q8_0_q1_vec_mq_split_g8.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_q8_0_q1_vec_mq_split_g8.at(dk_dv);
+ use_fd_mq = true;
+ fd_mq_wg = 192;
+ } else if (nq1_only && is_q8_0 && gqa_ratio_dispatch == 4 &&
+ d_head_q == 128 && d_head_v == 128 &&
+ backend_ctx->fa.f32_q8_0_q1_vec_mq_split.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_q8_0_q1_vec_mq_split.at(dk_dv);
+ use_fd_mq = true;
+ } else if (nq1_only && is_q4_0) {
+ const char * q4_mq_env = getenv("GGML_OPENCL_FA_Q4_MQ");
+ const bool q4_mq_on = (q4_mq_env != NULL) && (q4_mq_env[0] != '0');
+ if (q4_mq_on && gqa_ratio_dispatch == 8 &&
+ d_head_q == 128 && d_head_v == 128 &&
+ backend_ctx->fa.f32_q4_0_q1_vec_mq_split_g8.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_q4_0_q1_vec_mq_split_g8.at(dk_dv);
+ use_fd_mq = true;
+ fd_mq_wg = 192;
+ } else if (gqa_ratio_dispatch == 4 &&
+ d_head_q == 128 && d_head_v == 128 &&
+ backend_ctx->fa.f32_q4_0_q1_vec_mq_split.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_q4_0_q1_vec_mq_split.at(dk_dv);
+ use_fd_mq = true;
+ }
+ }
+ }
+ }
+ if (fd_k_split == NULL &&
+ n_q >= 1 && n_q <= fd_max_n_q && n_kv >= FD_MIN_N_KV && !is_causal &&
d_head_q <= FD_MAX_DK &&
backend_ctx->fa.f32_merge.count(dk_dv) > 0) {
if (is_mixed && backend_ctx->fa.f32_f16_q1_split.count(dk_dv) > 0) {
const bool use_fd = (fd_k_split != NULL);
const int n_q_blocks = n_q > 1 ? (n_q + block_m - 1) / block_m : 0;
- const int n_kv_blocks = n_kv > 0 ? (n_kv + block_n - 1) / block_n : 0;
+ const int n_kv_blocks = (n_kv > 0 && block_n > 0) ? (n_kv + block_n - 1) / block_n : 0;
// KV pad + blk prepass are pure overhead when FD will fire — skip them.
const bool use_mixed_prepass = is_mixed && n_q > 1 && !use_fd;
- const bool use_kv_pad = use_mixed_prepass && (n_kv % block_n != 0);
+ // make sure prepass kernels are compiled
+ const bool have_kv_pad = backend_ctx->fa.kv_pad_f16.count(dk_dv) > 0;
+ const bool have_blk = backend_ctx->fa.blk_f16.count(dk_dv) > 0;
+ const bool use_kv_pad = use_mixed_prepass && (n_kv % block_n != 0) && have_kv_pad;
// blk prepass: per-KV-tile mask class (0=masked, 1=mixed, 2=unmasked).
// Consumed identically by f32_f16, q8_0 and q4_0 prefill kernels.
const bool use_quant_prepass = (use_native_q8_0 || use_native_q4_0) && !use_fd;
- const bool use_blk_mask = (use_mixed_prepass || use_quant_prepass) && mask_buffer != NULL;
+ const bool use_blk_mask = (use_mixed_prepass || use_quant_prepass) && mask_buffer != NULL && have_blk;
if (use_kv_pad) {
cl_int err;
const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2_f);
if (use_fd) {
- int n_splits = (n_kv + FD_KV_PER_SPLIT - 1) / FD_KV_PER_SPLIT;
- if (n_splits < FD_MIN_SPLITS) {
- n_splits = FD_MIN_SPLITS;
- }
- if (n_splits > FD_MAX_SPLITS) {
- n_splits = FD_MAX_SPLITS;
- }
+ static const int fd_env_kv_per_split = []{
+ const char * e = getenv("GGML_OPENCL_FD_KV_PER_SPLIT");
+ return (e && e[0]) ? atoi(e) : 0;
+ }();
+ static const int fd_env_max_splits = []{
+ const char * e = getenv("GGML_OPENCL_FD_MAX_SPLITS");
+ return (e && e[0]) ? atoi(e) : 0;
+ }();
+
+ int fd_kv_per_split = use_fd_mq ? FD_MQ_KV_PER_SPLIT
+ : (is_mixed ? FD_KV_PER_SPLIT_F16 : FD_KV_PER_SPLIT);
+ int fd_max_splits = use_fd_mq ? FD_MQ_MAX_SPLITS : FD_MAX_SPLITS;
+ if (fd_env_kv_per_split > 0) { fd_kv_per_split = fd_env_kv_per_split; }
+ if (fd_env_max_splits > 0) { fd_max_splits = fd_env_max_splits; }
+ int n_splits = (n_kv + fd_kv_per_split - 1) / fd_kv_per_split;
+ if (n_splits < FD_MIN_SPLITS) { n_splits = FD_MIN_SPLITS; }
+ if (n_splits > fd_max_splits) { n_splits = fd_max_splits; }
const int kv_per_split = (n_kv + n_splits - 1) / n_splits;
const int fa_partial_floats = 2 + d_head_v;
int argi = 0;
CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_mem), &extra_q->data_device));
CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &offset_q));
- CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_mem), &k_data_device));
- CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &offset_k));
+ if (use_fa_k_img) {
+ const size_t nb00_bytes = sizeof(uint16_t);
+ const size_t k_bytes_span =
+ (size_t)(n_kv > 0 ? n_kv - 1 : 0) * (size_t)k_nb1 +
+ (size_t)(n_head_kv > 0 ? n_head_kv - 1 : 0) * (size_t)k_nb2 +
+ (size_t)(n_batch > 0 ? n_batch - 1 : 0) * (size_t)k_nb3 +
+ (size_t)d_head_q * nb00_bytes;
+ const size_t k_bytes = (k_bytes_span + 7) & ~(size_t)7;
+ const size_t k_pixels = k_bytes >> 3;
+ cl_mem k_img = nullptr;
+ if (k_pixels > 0 && k_pixels <= backend_ctx->image_max_buffer_size) {
+ k_img = ggml_cl_img_pool_get_or_create(
+ backend_ctx, backend_ctx->kq_img_pool,
+ k_data_device, offset_k, k_bytes, CL_HALF_FLOAT);
+ }
+
+ // if image creation fails, fallback to buffer based kernels
+ if (k_img == nullptr) {
+ if (gqa_ratio_dispatch == 4 &&
+ backend_ctx->fa.f32_f16_q1_vec_mq_split.count(dk_dv) > 0) {
+ k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split.at(dk_dv);
+ } else {
+ k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_g8.at(dk_dv);
+ }
+ use_fa_k_img = false;
+ CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_mem), &k_data_device));
+ CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &offset_k));
+ } else {
+ CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_mem), &k_img));
+ }
+ } else {
+ CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_mem), &k_data_device));
+ CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &offset_k));
+ }
CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_mem), &v_data_device));
CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(cl_ulong), &offset_v));
CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(float), &scale));
CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(int), &n_splits));
CL_CHECK(clSetKernelArg(k_split, argi++, sizeof(int), &kv_per_split));
- const size_t fd_wg = 64; // matches Q1_WG_SIZE in the kernel
+ // MQ split kernel uses MQ_NSG_SPLIT subgroups and one WG per (kv_head, batch, split)
+ // matches Q1_WG_SIZE * NSG (MQ_GQA=4 -> 256; MQ_GQA=8 -> 192)
+ const size_t fd_wg = use_fd_mq ? fd_mq_wg : 64;
+ const size_t fd_head_dim = use_fd_mq
+ ? (size_t)(n_head_kv * n_batch)
+ : (size_t)(n_head * n_batch);
size_t fd_lws[3] = { fd_wg, 1, 1 };
// gid(2) packs q_idx * n_splits + split_idx.
- size_t fd_gws[3] = { fd_wg, (size_t)(n_head * n_batch), (size_t)(n_splits * n_q) };
+ size_t fd_gws[3] = { fd_wg, fd_head_dim, (size_t)(n_splits * n_q) };
backend_ctx->enqueue_ndrange_kernel(k_split, 3, fd_gws, fd_lws, dst);
cl_kernel k_merge = backend_ctx->fa.f32_merge.at(dk_dv);
return;
}
+ cl_mem prefill_k_img = nullptr;
+ if (use_prefill_k_img) {
+ const size_t nb00_bytes = sizeof(uint16_t);
+ const size_t k_bytes_span =
+ (size_t)(n_kv > 0 ? n_kv - 1 : 0) * (size_t)k_nb1 +
+ (size_t)(n_head_kv > 0 ? n_head_kv - 1 : 0) * (size_t)k_nb2 +
+ (size_t)(n_batch > 0 ? n_batch - 1 : 0) * (size_t)k_nb3 +
+ (size_t)d_head_q * nb00_bytes;
+ const size_t k_bytes = (k_bytes_span + 7) & ~(size_t)7;
+ const size_t k_pixels = k_bytes >> 3;
+ if (k_pixels > 0 && k_pixels <= backend_ctx->image_max_buffer_size) {
+ prefill_k_img = ggml_cl_img_pool_get_or_create(
+ backend_ctx, backend_ctx->kq_img_pool,
+ k_data_device, offset_k, k_bytes, CL_HALF_FLOAT);
+ }
+ if (prefill_k_img == nullptr) {
+ kernel = backend_ctx->fa.f32_f16_split.at(dk_dv);
+ use_prefill_k_img = false;
+ }
+ }
+
CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra_q->data_device));
CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_ulong), &offset_q));
- CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &k_data_device));
+ CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), use_prefill_k_img ? &prefill_k_img : &k_data_device));
CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_ulong), &offset_k));
CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_mem), &v_data_device));
CL_CHECK(clSetKernelArg(kernel, 5, sizeof(cl_ulong), &offset_v));
}
if (n_q == 1) {
- const size_t wg_size = 64;
- size_t local_work_size[] = { wg_size, 1 };
- size_t global_work_size[] = { wg_size, (size_t)(n_head * n_batch) };
- backend_ctx->enqueue_ndrange_kernel(kernel, 2, global_work_size, local_work_size, dst);
+ if (use_local_tile) {
+ const size_t lt_wg = 128;
+ size_t local_work_size[] = { lt_wg, 1, 1 };
+ size_t global_work_size[] = { lt_wg, (size_t) n_head, (size_t) n_batch };
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
+ } else {
+ // q1_vec dispatches with NSG subgroups
+ const size_t q1_wg = backend_ctx->gpu_family == INTEL ? 32 : 64;
+ const size_t wg_size = use_q1_vec ? 256 : q1_wg;
+ const size_t head_dim_global = use_q1_vec_mq
+ ? (size_t)(n_head_kv * n_batch)
+ : (size_t)(n_head * n_batch);
+ size_t local_work_size[] = { wg_size, 1 };
+ size_t global_work_size[] = { wg_size, head_dim_global };
+ backend_ctx->enqueue_ndrange_kernel(kernel, 2, global_work_size, local_work_size, dst);
+ }
} else if (use_native_q8_0 || use_native_q4_0) {
// Native quant prefill. The split variant may override BLOCK_M
// (e.g. DK=96 quant uses BM=16).
cl_ulong src_nb2;
cl_ulong src_nb3;
+ uintptr_t pool_key_buf = 0;
+ cl_ulong pool_key_off = (cl_ulong) tensor->view_offs;
+
const bool is_soa = tensor->type == GGML_TYPE_Q4_0
? ggml_cl_is_q4_0_soa(tensor)
: ggml_cl_is_q8_0_soa(tensor);
cl_mem aos = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, parent_nbytes, NULL, &err);
CL_CHECK(err);
- cl_kernel kernel;
- if (tensor->type == GGML_TYPE_Q8_0) {
+ // large q4_0/q8_0 WEIGHTS are stored transposed and small weights
+ // (and the AoS KV-cache, handled in the else branch above) are not.
+ // choose a proper restore kernel based on this.
+ bool restored = false;
+#ifdef GGML_OPENCL_USE_ADRENO_KERNELS
+ const int p_ne00 = (int) parent->ne[0];
+ const int p_ne01 = (int) parent->ne[1];
+ if (tensor->type == GGML_TYPE_Q8_0 && enable_adreno_trans_weight(backend_ctx, parent)) {
auto * extra = (ggml_tensor_extra_cl_q8_0 *) soa_src->extra;
- kernel = backend_ctx->kernel_restore_block_q8_0;
+ pool_key_buf = (uintptr_t) extra->q;
+ cl_kernel kernel = backend_ctx->kernel_restore_block_q8_0_trans;
CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra->q));
CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &extra->d));
CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &aos));
- } else {
+ CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_int), &p_ne00));
+ CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_int), &p_ne01));
+ size_t gws[] = { (size_t)(((p_ne01 + 63) / 64) * 64), 1, 1 };
+ size_t lws[] = { 64, 1, 1 };
+ CL_CHECK(clEnqueueNDRangeKernel(backend_ctx->queue, kernel, 3, NULL, gws, lws, 0, NULL, NULL));
+ restored = true;
+ } else if (tensor->type == GGML_TYPE_Q4_0 &&
+ use_adreno_kernels(backend_ctx, parent) &&
+ !use_adreno_moe_kernels(backend_ctx, parent)) {
auto * extra = (ggml_tensor_extra_cl_q4_0 *) soa_src->extra;
- kernel = backend_ctx->kernel_restore_block_q4_0;
- CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra->q));
- CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &extra->d));
- CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &aos));
+ pool_key_buf = (uintptr_t) extra->q;
+ const size_t size_q = (size_t) ggml_nelements(parent) / blck_size * (blck_size / 2);
+ const size_t size_d = (size_t) ggml_nelements(parent) / blck_size * sizeof(ggml_fp16_t);
+ cl_int err2 = CL_SUCCESS;
+ cl_mem buf_tq = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, size_q, NULL, &err2); CL_CHECK(err2);
+ cl_mem buf_td = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, size_d, NULL, &err2); CL_CHECK(err2);
+ transpose_2d_as_16b(backend_ctx, extra->q, buf_tq, size_q, p_ne01, p_ne00 / 4);
+ transpose_2d_as_16b(backend_ctx, extra->d, buf_td, size_d, p_ne01, p_ne00 / 32);
+ cl_uchar mask_0F = 0x0F, mask_F0 = 0xF0;
+ cl_kernel kernel = backend_ctx->kernel_restore_block_q4_0_noshuffle;
+ CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &buf_tq));
+ CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &buf_td));
+ CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &aos));
+ CL_CHECK(clSetKernelArg(kernel, 3, sizeof(cl_uchar), &mask_0F));
+ CL_CHECK(clSetKernelArg(kernel, 4, sizeof(cl_uchar), &mask_F0));
+ const size_t n_blk = parent_nbytes / block_bytes;
+ size_t gws[] = { n_blk, 1, 1 };
+ size_t lws[] = { 1, 1, 1 };
+ CL_CHECK(clEnqueueNDRangeKernel(backend_ctx->queue, kernel, 3, NULL, gws, lws, 0, NULL, NULL));
+
+ CL_CHECK(clReleaseMemObject(buf_tq));
+ CL_CHECK(clReleaseMemObject(buf_td));
+ restored = true;
}
+#endif // GGML_OPENCL_USE_ADRENO_KERNELS
+
+ if (!restored) {
+ cl_kernel kernel;
+ if (tensor->type == GGML_TYPE_Q8_0) {
+ auto * extra = (ggml_tensor_extra_cl_q8_0 *) soa_src->extra;
+ kernel = backend_ctx->kernel_restore_block_q8_0;
+ CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra->q));
+ CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &extra->d));
+ CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &aos));
+ pool_key_buf = (uintptr_t) extra->q;
+ } else {
+ auto * extra = (ggml_tensor_extra_cl_q4_0 *) soa_src->extra;
+ kernel = backend_ctx->kernel_restore_block_q4_0;
+ CL_CHECK(clSetKernelArg(kernel, 0, sizeof(cl_mem), &extra->q));
+ CL_CHECK(clSetKernelArg(kernel, 1, sizeof(cl_mem), &extra->d));
+ CL_CHECK(clSetKernelArg(kernel, 2, sizeof(cl_mem), &aos));
+ pool_key_buf = (uintptr_t) extra->q;
+ }
- const size_t n_blocks = parent_nbytes / block_bytes;
- size_t gws_rec[] = { n_blocks, 1, 1 };
- size_t lws_rec[] = { 1, 1, 1 };
- CL_CHECK(clEnqueueNDRangeKernel(backend_ctx->queue, kernel, 3, NULL, gws_rec, lws_rec, 0, NULL, NULL));
+ const size_t n_blocks = parent_nbytes / block_bytes;
+ size_t gws_rec[] = { n_blocks, 1, 1 };
+ size_t lws_rec[] = { 1, 1, 1 };
+ CL_CHECK(clEnqueueNDRangeKernel(backend_ctx->queue, kernel, 3, NULL, gws_rec, lws_rec, 0, NULL, NULL));
+ }
(void) parent_row_blocks;
(void) parent_row_bytes;
src_nb1 = tensor->nb[1];
src_nb2 = tensor->nb[2];
src_nb3 = tensor->nb[3];
+ pool_key_buf = (uintptr_t) extra->data_device;
+ pool_key_off = (cl_ulong) src_offset;
}
const cl_int nblk0 = (cl_int) (tensor->ne[0] / ggml_blck_size(tensor->type));
const size_t out_bytes = (size_t) ggml_nelements(tensor) * sizeof(ggml_fp16_t);
- cl_int err;
- cl_mem out = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, out_bytes, NULL, &err);
- CL_CHECK(err);
+ // reuse a pooled f16 buffer for this KV-cache view across decode steps instead of
+ // allocating new one per attention op
+ cl_mem out = nullptr;
+ {
+ auto & pool = backend_ctx->dequant_f16_pool;
+ ggml_backend_opencl_context::ImagePoolKey key{pool_key_buf, (uint64_t) pool_key_off};
+ auto it = pool.find(key);
+ if (it != pool.end() && it->second.k_bytes >= out_bytes && it->second.image) {
+ out = it->second.image;
+ } else {
+ if (it != pool.end()) {
+ if (it->second.image) { CL_CHECK(clReleaseMemObject(it->second.image)); }
+ pool.erase(it);
+ }
+ cl_int err = CL_SUCCESS;
+ out = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, out_bytes, NULL, &err);
+ CL_CHECK(err);
+ ggml_backend_opencl_context::ImagePoolEntry entry;
+ entry.image = out;
+ entry.k_bytes = out_bytes;
+ pool[key] = entry;
+ }
+ }
cl_kernel dq_kernel = tensor->type == GGML_TYPE_Q8_0
? backend_ctx->kernel_dequant_q8_0_f16_view_aos
return out;
}
+// look up or create a pooled image1d_buffer over a KV-cache view.
+static cl_mem ggml_cl_img_pool_get_or_create(
+ ggml_backend_opencl_context * backend_ctx,
+ std::map<ggml_backend_opencl_context::ImagePoolKey,
+ ggml_backend_opencl_context::ImagePoolEntry> & pool,
+ cl_mem data_device,
+ cl_ulong offset0,
+ size_t required_bytes,
+ cl_channel_type channel_data_type
+) {
+ ggml_backend_opencl_context::ImagePoolKey key{(uintptr_t)data_device, (uint64_t)offset0};
+ auto it = pool.find(key);
+ if (it != pool.end()
+ && it->second.k_bytes >= required_bytes
+ && it->second.channel_data_type == channel_data_type
+ && it->second.image != nullptr) {
+ return it->second.image;
+ }
+
+ // need to create or recreate and release any stale entry first.
+ if (it != pool.end()) {
+ if (it->second.image) { CL_CHECK(clReleaseMemObject(it->second.image)); }
+ if (it->second.sub_buffer) {CL_CHECK(clReleaseMemObject(it->second.sub_buffer)); }
+ pool.erase(it);
+ }
+
+ cl_int status = CL_SUCCESS;
+ cl_buffer_region region = {};
+ region.origin = (size_t)offset0;
+ region.size = required_bytes;
+ cl_mem sub = clCreateSubBuffer(data_device, 0,
+ CL_BUFFER_CREATE_TYPE_REGION, ®ion, &status);
+ if (status != CL_SUCCESS) {
+ return nullptr;
+ }
+
+ const size_t pixel_size = (channel_data_type == CL_HALF_FLOAT) ? 8 : 16;
+ cl_image_format fmt = {CL_RGBA, channel_data_type};
+ cl_image_desc desc = {};
+ desc.image_type = CL_MEM_OBJECT_IMAGE1D_BUFFER;
+ desc.image_width = required_bytes / pixel_size;
+ desc.buffer = sub;
+ cl_mem img = clCreateImage(backend_ctx->context, CL_MEM_READ_ONLY,
+ &fmt, &desc, NULL, &status);
+ if (status != CL_SUCCESS) {
+ CL_CHECK(clReleaseMemObject(sub));
+ return nullptr;
+ }
+
+ ggml_backend_opencl_context::ImagePoolEntry entry;
+ entry.sub_buffer = sub;
+ entry.image = img;
+ entry.k_bytes = required_bytes;
+ entry.channel_data_type = channel_data_type;
+ pool[key] = entry;
+ return img;
+}
+
static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
GGML_ASSERT(src0);
GGML_ASSERT(src0->extra);
ggml_backend_opencl_context *backend_ctx = (ggml_backend_opencl_context *)backend->context;
- // Non-contig quant src0: on-device dequant to f16 then native f16 MUL_MAT.
- if ((src0t == GGML_TYPE_Q4_0 || src0t == GGML_TYPE_Q8_0) && !ggml_is_contiguous(src0)) {
+ // quant kv without FA
+ // used for non-contiguous src0 (the usual head-major permuted K view when n_head_kv>1)
+ // AND for the contiguous case that occurs when n_head_kv==1 (e.g. Gemma-4 E2B)
+ if ((src0t == GGML_TYPE_Q4_0 || src0t == GGML_TYPE_Q8_0) &&
+ (!ggml_is_contiguous(src0) || src1->ne[2] > src0->ne[2])) {
cl_mem f16_buf = ggml_cl_mul_mat_dequant_quant_to_f16(backend_ctx, src0, nullptr);
ggml_tensor fake_src0 = *src0;
fake_src0.nb[3] = fake_src0.nb[2] * src0->ne[2];
ggml_cl_mul_mat(backend, &fake_src0, src1, dst);
-
- // Safe to release now: OpenCL retains the memobj while queued
- // kernels that reference it are still in flight.
- CL_CHECK(clReleaseMemObject(f16_buf));
return;
}
return;
}
}
+
+ static const char * mm_kq_gqa_img_env = getenv("GGML_OPENCL_MM_KQ_GQA_IMG");
+ static const bool mm_kq_gqa_img_on = (mm_kq_gqa_img_env == nullptr || mm_kq_gqa_img_env[0] != '0');
+ static const char * mm_kq_gqa_r4_img_env = getenv("GGML_OPENCL_MM_KQ_GQA_R4_IMG");
+ static const bool mm_kq_gqa_r4_img_on = (mm_kq_gqa_r4_img_env == nullptr || mm_kq_gqa_r4_img_env[0] != '0');
+ const bool img_r4_gate =
+ mm_kq_gqa_r4_img_on &&
+ backend_ctx->kernel_mul_mat_f16_f32_l4_x8_gqa_r4_img != nullptr &&
+ ne11 == 1 && ne01 >= 64 && (ne01 % 16) == 0 && ne00 == 128 &&
+ (ne12 % ne02) == 0 && (ne12 / ne02) == 4 && (ne13 / ne03) == 1;
+ if (mm_kq_gqa_img_on &&
+ backend_ctx->kernel_mul_mat_f16_f32_l4_x8_gqa4_img != nullptr &&
+ ne11 == 1 && ne01 >= 64 && (ne01 % 16) == 0 && ne00 == 128 &&
+ (ne12 % ne02) == 0 && (ne12 / ne02) == 8 && (ne13 / ne03) == 1) {
+ const size_t nb00_bytes = sizeof(uint16_t);
+ const size_t k_bytes_span =
+ (size_t)(ne01 > 0 ? ne01 - 1 : 0) * (size_t)nb01 +
+ (size_t)(ne02 > 0 ? ne02 - 1 : 0) * (size_t)nb02 +
+ (size_t)(ne03 > 0 ? ne03 - 1 : 0) * (size_t)nb03 +
+ (size_t)ne00 * nb00_bytes;
+
+ const size_t k_bytes = (k_bytes_span + 15) & ~(size_t)15;
+ const size_t k_pixels = k_bytes >> 4;
+ if (k_pixels > 0 && k_pixels <= backend_ctx->image_max_buffer_size) {
+ cl_kernel kernel = backend_ctx->kernel_mul_mat_f16_f32_l4_x8_gqa4_img;
+ cl_mem K_img = ggml_cl_img_pool_get_or_create(
+ backend_ctx, backend_ctx->kq_img_pool,
+ extra0->data_device, offset0, k_bytes, CL_FLOAT);
+ if (K_img != nullptr) {
+ cl_uint k_arg = 0;
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &K_img));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &extra1->data_device));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &offset1));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &extrad->data_device));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &offsetd));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne00));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne01));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne02));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb01));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb02));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb03));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne10));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne11));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne12));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb10));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb11));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb12));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb13));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne0));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne1));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &r2));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &r3));
+
+ const int nth0_d = 64;
+ const int64_t n_wg_x = ne01 / 16;
+ size_t global_work_size[] = {(size_t)n_wg_x * nth0_d, (size_t)1, (size_t)ne02 * ne13};
+ size_t local_work_size[] = {(size_t)nth0_d, (size_t)1, 1};
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
+ return;
+ }
+ }
+ }
+
+ // r2=4 specialization
+ if (img_r4_gate) {
+ const size_t nb00_bytes = sizeof(uint16_t);
+ const size_t k_bytes_span =
+ (size_t)(ne01 > 0 ? ne01 - 1 : 0) * (size_t)nb01 +
+ (size_t)(ne02 > 0 ? ne02 - 1 : 0) * (size_t)nb02 +
+ (size_t)(ne03 > 0 ? ne03 - 1 : 0) * (size_t)nb03 +
+ (size_t)ne00 * nb00_bytes;
+ const size_t k_bytes = (k_bytes_span + 15) & ~(size_t)15;
+ const size_t k_pixels = k_bytes >> 4;
+ if (k_pixels > 0 && k_pixels <= backend_ctx->image_max_buffer_size) {
+ cl_kernel kernel = backend_ctx->kernel_mul_mat_f16_f32_l4_x8_gqa_r4_img;
+ cl_mem K_img = ggml_cl_img_pool_get_or_create(
+ backend_ctx, backend_ctx->kq_img_pool,
+ extra0->data_device, offset0, k_bytes, CL_FLOAT);
+ if (K_img != nullptr) {
+ cl_uint k_arg = 0;
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &K_img));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &extra1->data_device));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &offset1));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &extrad->data_device));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &offsetd));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne00));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne01));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne02));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb01));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb02));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb03));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne10));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne11));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne12));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb10));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb11));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb12));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb13));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne0));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne1));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &r2));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &r3));
+
+ const int nth0_d = 64;
+ const int64_t n_wg_x = ne01 / 16;
+ size_t global_work_size[] = {(size_t)n_wg_x * nth0_d, (size_t)1, (size_t)ne02 * ne13};
+ size_t local_work_size[] = {(size_t)nth0_d, (size_t)1, 1};
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
+ return;
+ }
+ }
+ }
+
+ // DK=256, r2=2 specialization
+ static const char * mm_kq_r2_dk256_env = getenv("GGML_OPENCL_MM_KQ_GQA_R2_DK256_IMG");
+ static const bool mm_kq_r2_dk256_on = (mm_kq_r2_dk256_env != nullptr && mm_kq_r2_dk256_env[0] != '0');
+ if (mm_kq_r2_dk256_on &&
+ backend_ctx->kernel_mul_mat_f16_f32_l4_x8_gqa_r2_dk256_img != nullptr &&
+ ne11 == 1 && ne01 >= 64 && (ne01 % 16) == 0 && ne00 == 256 &&
+ (ne12 % ne02) == 0 && (ne12 / ne02) == 2 && (ne13 / ne03) == 1) {
+ const size_t nb00_bytes = sizeof(uint16_t);
+ const size_t k_bytes_span =
+ (size_t)(ne01 > 0 ? ne01 - 1 : 0) * (size_t)nb01 +
+ (size_t)(ne02 > 0 ? ne02 - 1 : 0) * (size_t)nb02 +
+ (size_t)(ne03 > 0 ? ne03 - 1 : 0) * (size_t)nb03 +
+ (size_t)ne00 * nb00_bytes;
+ const size_t k_bytes = (k_bytes_span + 15) & ~(size_t)15;
+ const size_t k_pixels = k_bytes >> 4;
+ if (k_pixels > 0 && k_pixels <= backend_ctx->image_max_buffer_size) {
+ cl_kernel kernel = backend_ctx->kernel_mul_mat_f16_f32_l4_x8_gqa_r2_dk256_img;
+ cl_mem K_img = ggml_cl_img_pool_get_or_create(
+ backend_ctx, backend_ctx->kq_img_pool,
+ extra0->data_device, offset0, k_bytes, CL_FLOAT);
+ if (K_img != nullptr) {
+ cl_uint k_arg = 0;
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &K_img));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &extra1->data_device));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &offset1));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &extrad->data_device));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &offsetd));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne00));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne01));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne02));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb01));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb02));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb03));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne10));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne11));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne12));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb10));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb11));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb12));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb13));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne0));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne1));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &r2));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &r3));
+
+ const int nth0_d = 64;
+ const int64_t n_wg_x = ne01 / 16;
+ size_t global_work_size[] = {(size_t)n_wg_x * nth0_d, (size_t)1, (size_t)ne02 * ne13};
+ size_t local_work_size[] = {(size_t)nth0_d, (size_t)1, 1};
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
+ return;
+ }
+ }
+ }
+
+ // GQA-coalesced KQV for decode using image1d_buffer_t
+ static const char * mm_kqv_gqa_img_env = getenv("GGML_OPENCL_MM_KQV_GQA_IMG");
+ static const bool mm_kqv_gqa_img_on = (mm_kqv_gqa_img_env != nullptr && mm_kqv_gqa_img_env[0] != '0');
+ if (mm_kqv_gqa_img_on &&
+ backend_ctx->kernel_mul_mat_f16_f32_l4_y8_gqa_img != nullptr &&
+ ne11 == 1 && ne01 == 128 &&
+ (ne12 % ne02) == 0 && (ne12 / ne02) == 8 && (ne13 / ne03) == 1) {
+ const size_t nb00_bytes = sizeof(uint16_t);
+ const size_t v_bytes_span =
+ (size_t)(ne01 > 0 ? ne01 - 1 : 0) * (size_t)nb01 +
+ (size_t)(ne02 > 0 ? ne02 - 1 : 0) * (size_t)nb02 +
+ (size_t)(ne03 > 0 ? ne03 - 1 : 0) * (size_t)nb03 +
+ (size_t)ne00 * nb00_bytes;
+ const size_t v_bytes = (v_bytes_span + 7) & ~(size_t)7;
+ const size_t v_pixels = v_bytes >> 3;
+ if (v_pixels > 0 && v_pixels <= backend_ctx->image_max_buffer_size) {
+ cl_kernel kernel = backend_ctx->kernel_mul_mat_f16_f32_l4_y8_gqa_img;
+ cl_mem V_img = ggml_cl_img_pool_get_or_create(
+ backend_ctx, backend_ctx->kqv_img_pool,
+ extra0->data_device, offset0, v_bytes, CL_HALF_FLOAT);
+ if (V_img != nullptr) {
+ cl_uint k_arg = 0;
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &V_img));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &extra1->data_device));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &offset1));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_mem), &extrad->data_device));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &offsetd));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne00));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne01));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne02));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb01));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb02));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb03));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne10));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne11));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne12));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb10));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb11));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb12));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(cl_ulong), &nb13));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne0));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &ne1));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &r2));
+ CL_CHECK(clSetKernelArg(kernel, k_arg++, sizeof(int), &r3));
+
+ const int nth0_d = 64;
+ const int64_t n_wg_x = ne01 / 8;
+ size_t global_work_size[] = {(size_t)n_wg_x * nth0_d, (size_t)1, (size_t)ne02 * ne13};
+ size_t local_work_size[] = {(size_t)nth0_d, (size_t)1, 1};
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
+ return;
+ }
+ }
+ }
}
if (ne01 && ne1 && use_adreno_kernels(backend_ctx, src0)) {
kernel = backend_ctx->kernel_mul_mat_f16_f32_l4_dr_ls;
nrows = 1;
} else if (ne00 >= 128 && ne01 >= 8 && ne00%4 == 0) {
- if (ne11 == 1) {
+ // multi-output decode variants when Q is a single row
+ static const char * mm_force_l4_env = getenv("GGML_OPENCL_MM_F16_FORCE_L4");
+ static const bool mm_force_l4_on = (mm_force_l4_env != nullptr && mm_force_l4_env[0] != '0');
+ const bool can_multi_out = !mm_force_l4_on && ne11 == 1 && ne01 >= 64 && ne01 % 8 == 0;
+ // paired-K-row variant that doubles per-wave-cycle
+ static const char * mm_kq_pair_env = getenv("GGML_OPENCL_MM_KQ_PAIR");
+ static const bool mm_kq_pair_on = (mm_kq_pair_env != nullptr && mm_kq_pair_env[0] != '0');
+ // GQA-coalesced variant that reads each K-row once and
+ // emits gqa_ratio outputs
+ static const char * mm_kq_gqa_env = getenv("GGML_OPENCL_MM_KQ_GQA");
+ static const bool mm_kq_gqa_on = (mm_kq_gqa_env != nullptr && mm_kq_gqa_env[0] != '0');
+ // GQA-coalesced KQV variant (DK=128/r2=8/r3=1) that reads
+ // each V slab once per K-head and emits all r2 Q-heads
+ static const char * mm_kqv_gqa_env = getenv("GGML_OPENCL_MM_KQV_GQA");
+ static const bool mm_kqv_gqa_on = (mm_kqv_gqa_env != nullptr && mm_kqv_gqa_env[0] != '0');
+ if (can_multi_out && (ne01 % 16) == 0 && ne00 == 128 && r2 == 8 && r3 == 1 && mm_kq_gqa_on &&
+ backend_ctx->kernel_mul_mat_f16_f32_l4_x8_gqa4 != nullptr) {
+ kernel = backend_ctx->kernel_mul_mat_f16_f32_l4_x8_gqa4;
+ nrows = 1;
+ } else if (can_multi_out && ne00 <= 256 && mm_kq_pair_on &&
+ backend_ctx->kernel_mul_mat_f16_f32_l4_x8_pair != nullptr) {
+ kernel = backend_ctx->kernel_mul_mat_f16_f32_l4_x8_pair;
+ nrows = 1;
+ } else if (can_multi_out && ne00 <= 256 &&
+ backend_ctx->kernel_mul_mat_f16_f32_l4_x8 != nullptr) {
+ kernel = backend_ctx->kernel_mul_mat_f16_f32_l4_x8;
+ nrows = 1;
+ } else if (can_multi_out && ne01 == 128 && r2 == 8 && r3 == 1 && mm_kqv_gqa_on &&
+ backend_ctx->kernel_mul_mat_f16_f32_l4_y8_gqa != nullptr) {
+ kernel = backend_ctx->kernel_mul_mat_f16_f32_l4_y8_gqa;
+ nrows = 1;
+ } else if (can_multi_out &&
+ backend_ctx->kernel_mul_mat_f16_f32_l4_y8 != nullptr) {
+ kernel = backend_ctx->kernel_mul_mat_f16_f32_l4_y8;
+ nrows = 1;
+ } else if (ne11 == 1) {
+ // Decode shapes that don't satisfy the x8/y8 row
+ // constraints (ne01 < 64 or ne01 % 8 != 0) fall back to
+ // upstream's 4-output _dr kernel.
kernel = backend_ctx->kernel_mul_mat_f16_f32_l4_dr;
nrows = 1; // not used by this kernel
} else {
kernel = backend_ctx->kernel_mul_mat_f16_f32_l4;
- nrows = ne11;
+ nrows = ne11;
}
} else {
kernel = backend_ctx->kernel_mul_mat_f16_f32;
size_t global_work_size[] = {(size_t)(ne01+ndst*nth1-1)/(ndst*nth1)*nth0, (size_t)ne11*nth1, (size_t)ne12*ne13};
size_t local_work_size[] = {(size_t)nth0, (size_t)nth1, 1};
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
+ } else if (kernel == backend_ctx->kernel_mul_mat_f16_f32_l4_x8 ||
+ kernel == backend_ctx->kernel_mul_mat_f16_f32_l4_x8_pair ||
+ kernel == backend_ctx->kernel_mul_mat_f16_f32_l4_y8) {
+ // multi-output decode variants: each WG processes 8 outputs along ne01, ne11 == 1
+ const int64_t n_wg_x = ne01 / 8;
+ size_t global_work_size[] = {(size_t)n_wg_x*nth0, (size_t)nth1, (size_t)ne12*ne13};
+ size_t local_work_size[] = {(size_t)nth0, (size_t)nth1, 1};
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
+ } else if (kernel == backend_ctx->kernel_mul_mat_f16_f32_l4_x8_gqa4) {
+ // GQA-coalesced KQ: one WG per K-head emits N_K_ROWS_GQA=16 K-rows * r2 Q-heads
+ const int64_t n_wg_x = ne01 / 16;
+ size_t global_work_size[] = {(size_t)n_wg_x*nth0, (size_t)nth1, (size_t)ne02*ne13};
+ size_t local_work_size[] = {(size_t)nth0, (size_t)nth1, 1};
+ backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
+ } else if (kernel == backend_ctx->kernel_mul_mat_f16_f32_l4_y8_gqa) {
+ // GQA-coalesced KQV: one WG per K-head emits 8 DV-rows * r2 Q-heads
+ const int64_t n_wg_x = ne01 / 8;
+ size_t global_work_size[] = {(size_t)n_wg_x*nth0, (size_t)nth1, (size_t)ne02*ne13};
+ size_t local_work_size[] = {(size_t)nth0, (size_t)nth1, 1};
backend_ctx->enqueue_ndrange_kernel(kernel, 3, global_work_size, local_work_size, dst);
} else {
if (kernel == backend_ctx->kernel_mul_mat_f16_f32_l4_dr) {
#define DK_VEC (DK/4)
#define DV_VEC (DV/4)
#define WG_SIZE (BLOCK_M)
-#define Q1_WG_SIZE 64
+// q1 reduces over a Q1_WG_SIZE-wide WG via work-group barriers; the launch WG
+// must match. Defaults to the Adreno sg (64); host passes -D FA_SG=32 on Intel.
+#ifndef FA_SG
+#define FA_SG 64
+#endif
+#define Q1_WG_SIZE FA_SG
// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs,
// infinite operand can cause undefined behavior and miscompilation for exp.
#define DK_VEC (DK/4)
#define DV_VEC (DV/4)
#define WG_SIZE (BLOCK_M)
-#define Q1_WG_SIZE 64
+// q1 reduces over a Q1_WG_SIZE-wide WG via work-group barriers; the launch WG
+// must match. Defaults to the Adreno sg (64); host passes -D FA_SG=32 on Intel.
+#ifndef FA_SG
+#define FA_SG 64
+#endif
+#define Q1_WG_SIZE FA_SG
// The kernels are built with -cl-finite-math-only. On some older Adreno GPUs,
// infinite operand can cause undefined behavior and miscompilation for exp.
__local DATA_TYPE4 l_v[BLOCK_N][DV_VEC];
for (int k_start = 0; k_start < n_kv; k_start += BLOCK_N) {
+#if FA_SG < 64
+ // WAR on l_k/l_v: threads with my_query_row >= n_q skip the compute below
+ // (continue) and would race ahead to reload the tiles while active threads
+ // still read them. A single 64-wide Adreno subgroup (WG == sg) runs lockstep
+ // and hides this; a WG that spans multiple narrower subgroups (Intel sg=32)
+ // corrupts the result. All threads reach this each iteration (no-op on the
+ // first), so it does not diverge with the continue. Compiled out at sg=64.
+ barrier(CLK_LOCAL_MEM_FENCE);
+#endif
for (int i = tid; i < BLOCK_N * DK_VEC; i += WG_SIZE) {
const int row = i / DK_VEC;
const int col = i % DK_VEC;
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+#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
+
+// subgroup size for q1 kernels
+#ifndef FA_SG
+#define FA_SG 64
+#endif
+#ifdef cl_intel_required_subgroup_size
+#pragma OPENCL EXTENSION cl_intel_required_subgroup_size : enable
+#define REQD_FA_SG __attribute__((intel_reqd_sub_group_size(FA_SG)))
+#else
+#define REQD_FA_SG
+#endif
+
#ifdef cl_khr_subgroup_shuffle
#pragma OPENCL EXTENSION cl_khr_subgroup_shuffle : enable
#define HAS_SUBGROUP_SHUFFLE 1
#define DK_VEC (DK/4)
#define DV_VEC (DV/4)
-#define Q1_WG_SIZE 64
+
+#ifndef FA_PARTIAL_FLOATS
+#define FA_PARTIAL_FLOATS (2 + DV)
+#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.
return pow(base, exph);
}
-__kernel void flash_attn_f32_f16(
+
+// Adreno compiler crashes when attempting to compile the entire program for DK=512,
+// FA_DECODE_ONLY allows bypass the encoding kernel.
+#if !defined(FA_DECODE_ONLY) && !defined(FA_MQ_ONLY)
+#ifndef FA_TILE_NAME
+#define FA_TILE_NAME flash_attn_f32_f16
+#endif
+__kernel void FA_TILE_NAME(
const global void * q_void, ulong q_offset,
+#ifdef FA_K_IMG
+ __read_only image1d_buffer_t k_img, ulong k_offset_unused,
+#else
const global void * k_void, ulong k_offset,
+#endif
const global void * v_void, ulong v_offset,
global void * o_void, ulong o_offset,
const float scale,
const int mask_batch_idx = mask_void != NULL ? batch_idx % mask_ne3 : 0;
const global char* q_base = (const global char*)q_void + q_offset;
+#ifndef FA_K_IMG
const global char* k_base = (const global char*)k_void + k_offset;
+#endif
const global char* v_base = (const global char*)v_void + v_offset;
global char* o_base = (global char*)o_void + o_offset;
const ulong k_tile_nb3 = use_kv_pad ? (ulong) n_head_kv * k_tile_nb2 : k_nb3;
const ulong v_tile_nb2 = use_kv_pad ? (ulong) BLOCK_N * v_nb1 : v_nb2;
const ulong v_tile_nb3 = use_kv_pad ? (ulong) n_head_kv * v_tile_nb2 : v_nb3;
+#ifdef FA_K_IMG
+ // K via texture cache for the bulk (aligned) tiles; the ragged last
+ // tile (use_kv_pad) still reads the f32-strided pad buffer from global.
+ const global char* k_tile_base = use_kv_pad ? (const global char*) k_pad_void : (const global char*) 0;
+ const int k_pitch_px_row = (int)(k_nb1 >> 3);
+ const int k_pitch_px_head = (int)(k_nb2 >> 3);
+ const int k_pitch_px_batch = (int)(k_nb3 >> 3);
+#else
const global char* k_tile_base = use_kv_pad ? (const global char*) k_pad_void : k_base;
+#endif
const global char* v_tile_base = use_kv_pad ? (const global char*) v_pad_void : v_base;
for (int i = tid; i < BLOCK_N * DK_VEC; i += WG_SIZE) {
const int col = i % DK_VEC;
const int k_row_idx = k_tile_start + row;
if (use_kv_pad || k_row_idx < n_kv) {
+#ifdef FA_K_IMG
+ if (use_kv_pad) {
+ const ulong k_row_offset = batch_idx * k_tile_nb3 + head_kv_idx * k_tile_nb2 + k_row_idx * k_nb1;
+ l_k[row][col] = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col];
+ } else {
+ const int k_row_px = batch_idx * k_pitch_px_batch + head_kv_idx * k_pitch_px_head + k_row_idx * k_pitch_px_row;
+ l_k[row][col] = read_imageh(k_img, k_row_px + col);
+ }
+#else
const ulong k_row_offset = batch_idx * k_tile_nb3 + head_kv_idx * k_tile_nb2 + k_row_idx * k_nb1;
l_k[row][col] = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col];
+#endif
} else {
l_k[row][col] = (KV_DATA_TYPE4)(0.0h);
}
}
#endif
}
+#endif // !FA_DECODE_ONLY
+// allow bypassing decode kernels to avoid compiler crash for DK=512 on Adreno GPUs
+#ifndef FA_PREFILL_ONLY
+#ifndef FA_MQ_ONLY // q1 excluded from the MQ-only (g8) program
+REQD_FA_SG
__kernel void flash_attn_f32_f16_q1(
const global void * q_void, ulong q_offset,
const global void * k_void, ulong k_offset,
const int head_kv_idx = head_idx / gqa_ratio;
const global char* q_base = (const global char*)q_void + q_offset;
+#ifndef FA_K_IMG
const global char* k_base = (const global char*)k_void + k_offset;
+#endif
const global char* v_base = (const global char*)v_void + v_offset;
global char* o_base = (global char*)o_void + o_offset;
for (int i = tid; i < DK_VEC; i += Q1_WG_SIZE) {
q_shared[i] = CONVERT_Q_ACC4(q_ptr[i]);
}
- barrier(CLK_LOCAL_MEM_FENCE);
+ sub_group_barrier(CLK_LOCAL_MEM_FENCE);
float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
m_i = max(m_i, score);
}
- __local ACC_TYPE local_m[Q1_WG_SIZE];
- local_m[tid] = m_i;
- barrier(CLK_LOCAL_MEM_FENCE);
- FA_UNROLL
- for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) {
- if (tid < s) local_m[tid] = max(local_m[tid], local_m[tid + s]);
- barrier(CLK_LOCAL_MEM_FENCE);
- }
- const ACC_TYPE m_final = local_m[0];
+ const ACC_TYPE m_final = sub_group_reduce_max(m_i);
ACC_TYPE4 o_acc[DV_VEC];
FA_UNROLL
}
}
- __local ACC_TYPE local_l[Q1_WG_SIZE];
__local ACC_TYPE4 local_o_comp[Q1_WG_SIZE];
- local_l[tid] = l_i;
- barrier(CLK_LOCAL_MEM_FENCE);
- FA_UNROLL
- for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) {
- if (tid < s) local_l[tid] += local_l[tid + s];
- barrier(CLK_LOCAL_MEM_FENCE);
- }
+ const ACC_TYPE l_red = sub_group_reduce_add(l_i);
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);
- ACC_TYPE l_final = local_l[0];
+ ACC_TYPE l_final = l_red;
if (sinks_ptr != NULL) {
l_final += exp(sinks_ptr[head_idx] - m_final);
const ACC_TYPE l_inv = 1.0f / l_final;
for (int i = 0; i < DV_VEC; i++) {
local_o_comp[tid] = o_acc[i];
- barrier(CLK_LOCAL_MEM_FENCE);
+ sub_group_barrier(CLK_LOCAL_MEM_FENCE);
FA_UNROLL
for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) {
if (tid < s) local_o_comp[tid] += local_o_comp[tid + s];
- barrier(CLK_LOCAL_MEM_FENCE);
+ sub_group_barrier(CLK_LOCAL_MEM_FENCE);
}
if (tid == 0) {
o_row[i] = CONVERT_O_DATA4(local_o_comp[0] * l_inv);
}
}
-// Flash-decoding split pass. gid(2) = q_idx * n_splits + split_idx.
-// Partial record per split: [m, l, O[DV]]. Merge kernel applies sink + norm.
-#define FA_PARTIAL_FLOATS (2 + DV)
+#endif // !FA_MQ_ONLY (q1)
+// decode variant for large DV (e.g. Gemma-4 DK=DV=512 global layers).
+#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)
-__kernel void flash_attn_f32_f16_q1_split(
+// allow bypassing the kernel to avoid compiler crash for DK=512 on Adreno GPUs
+#if !defined(FA_DECODE_MINIMAL) && !defined(FA_MQ_ONLY)
+REQD_SUBGROUP_SIZE_64
+__kernel void flash_attn_f32_f16_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 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 global void* sinks_void,
+ const ulong sinks_offset
) {
- const int tid = get_local_id(0);
- const int head_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 = 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 int tid = get_local_id(0);
+ const int sgid = tid / Q1_WG_SIZE; // subgroup index (0..VEC_NSG-1)
+ const int tid_sg = tid % Q1_WG_SIZE; // lane within subgroup
+ const int head_batch_idx = get_global_id(1);
- 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;
- const ulong record_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
- * n_splits + split_idx);
- global float * rec = partial_void + record_idx * record_stride;
- global float4 * rec_o = (global float4 *) (rec + 2);
+ const int batch_idx = head_batch_idx / n_head;
+ const int head_idx = head_batch_idx % n_head;
- if (kv_start >= kv_end) {
- // Empty split: leave sentinel partial for merge.
- if (tid == 0) {
- rec[0] = FA_M_INIT;
- rec[1] = 0.0f;
- }
- return;
- }
+ 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 +
- (ulong) q_idx * mask_nb1;
+ mask_batch_idx * mask_nb3 + mask_head_idx * mask_nb2;
}
- // Share Q via local memory (n_q=1 per split -> uniform across WG).
+ // Q is uniform across the WG — stage in __local once. All WG threads load.
__local ACC_TYPE4 q_shared[DK_VEC];
- 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);
- for (int i = tid; i < DK_VEC; i += Q1_WG_SIZE) {
- q_shared[i] = CONVERT_Q_ACC4(q_ptr[i]);
+ {
+ 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);
- // Pass 1a — split-local max.
+ const global ACC_TYPE * sinks_ptr = NULL;
+ if (sinks_void != NULL) {
+ sinks_ptr = (const global ACC_TYPE *) ((const global char *) sinks_void + sinks_offset);
+ }
+
+ // per-thread DV slice within its subgroup
+ // DV=512 -> 2x float4 = 32 bytes; DV=256 -> 1x float4 - no spill
+ 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);
+
+ // each subgroup independently runs the FA-2 online softmax over its slice of n_kv.
+ // sinks are not folded into per-subgroup m_i — they're added once in
+ // the cross-subgroup merge to avoid double-counting.
ACC_TYPE m_i = FA_M_INIT;
- for (int k_idx = kv_start + tid; k_idx < kv_end; k_idx += Q1_WG_SIZE) {
- const ulong k_row_offset = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
- const global KV_DATA_TYPE4 * k_ptr = (const global KV_DATA_TYPE4 *) (k_base + k_row_offset);
- ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f);
- #pragma unroll
- for (int k = 0; k < DK_VEC; ++k) {
- dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc);
+ 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 ulong k_row_off = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
+ const ulong v_row_off = batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
+ const global KV_DATA_TYPE4 * k_ptr = (const global KV_DATA_TYPE4 *) (k_base + k_row_off);
+ const global KV_DATA_TYPE4 * v_ptr = (const global KV_DATA_TYPE4 *) (v_base + v_row_off);
+
+ // Q*K^T: each thread accumulates its DK slice; subgroup-reduce the partial.
+ ACC_TYPE4 dot4 = (ACC_TYPE4)(0.0f);
+ for (int k = tid_sg; k < DK_VEC; k += Q1_WG_SIZE) {
+ dot4 = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot4);
}
- ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale;
+ 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);
+ 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);
}
- m_i = max(m_i, score);
- }
- __local ACC_TYPE local_m[Q1_WG_SIZE];
- local_m[tid] = m_i;
- barrier(CLK_LOCAL_MEM_FENCE);
- #pragma unroll
- for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) {
- if (tid < s) local_m[tid] = max(local_m[tid], local_m[tid + s]);
- barrier(CLK_LOCAL_MEM_FENCE);
+ // FA-2 online update. All threads in the subgroup see the same score,
+ // so m_i and l_i evolve identically across lanes within the subgroup.
+ 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_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE, ++idx) {
+ o_acc[idx] = mad(p, CONVERT_KV_ACC4(v_ptr[dv_idx]), o_acc[idx] * scale_prev);
+ }
+ l_i = l_i * scale_prev + p;
+ m_i = m_new;
}
- const ACC_TYPE m_c = local_m[0];
- // Pass 1b — softmax-weighted V accumulate.
- ACC_TYPE4 o_acc[DV_VEC];
- #pragma unroll
- for (int i = 0; i < DV_VEC; ++i) o_acc[i] = (ACC_TYPE4)(0.0f);
- ACC_TYPE l_i = 0.0f;
+ // Cross-subgroup merge via __local. Each subgroup publishes (m_i, l_i)
+ // and its o_acc slice; subgroup 0 then folds them into the final norm
+ // and writes the row.
+ __local ACC_TYPE sg_m[VEC_NSG];
+ __local ACC_TYPE sg_l[VEC_NSG];
+ __local ACC_TYPE4 sg_o[VEC_NSG][DV_VEC];
- for (int k_idx = kv_start + tid; k_idx < kv_end; k_idx += Q1_WG_SIZE) {
- const ulong k_row_offset = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
- const ulong v_row_offset = batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
- const global KV_DATA_TYPE4 * k_ptr = (const global KV_DATA_TYPE4 *) (k_base + k_row_offset);
- const global KV_DATA_TYPE4 * v_ptr = (const global KV_DATA_TYPE4 *) (v_base + v_row_offset);
- ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f);
- #pragma unroll
- for (int k = 0; k < DK_VEC; ++k) {
- dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc);
+ if (tid_sg == 0) {
+ sg_m[sgid] = m_i;
+ sg_l[sgid] = l_i;
+ }
+ {
+ 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[idx];
}
- ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * 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];
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+
+ if (sgid == 0) {
+ // m_final = max over all subgroups' m_i, plus the sink (if any).
+ 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 (logit_softcap > 0.0f) {
- score = logit_softcap * tanh(score / logit_softcap);
+ if (sinks_ptr != NULL) {
+ m_final = max(m_final, sinks_ptr[head_idx]);
}
- const ACC_TYPE p = exp(score - m_c);
- l_i += p;
+
+ ACC_TYPE l_final = 0.0f;
#pragma unroll
- for (int i = 0; i < DV_VEC; ++i) {
- o_acc[i] = mad(p, CONVERT_KV_ACC4(v_ptr[i]), o_acc[i]);
+ 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);
+
+ // Each thread in subgroup 0 writes its DV slice, folding all subgroups'
+ // contributions with the rescale factor.
+ int idx = 0;
+ for (int dv_idx = tid_sg; dv_idx < DV_VEC; dv_idx += 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_idx], o_merged);
+ }
+ o_row[dv_idx] = CONVERT_O_DATA4(o_merged * l_inv);
}
}
+}
- __local ACC_TYPE local_l[Q1_WG_SIZE];
- __local ACC_TYPE4 local_o[Q1_WG_SIZE];
- local_l[tid] = l_i;
- barrier(CLK_LOCAL_MEM_FENCE);
- #pragma unroll
- for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) {
- if (tid < s) local_l[tid] += local_l[tid + s];
- barrier(CLK_LOCAL_MEM_FENCE);
- }
- const ACC_TYPE l_c = local_l[0];
+#endif // !FA_DECODE_MINIMAL
- if (tid == 0) {
- rec[0] = (float) m_c;
- rec[1] = (float) l_c;
+#ifndef FA_DECODE_ONLY
+
+// flash_attn_f32_f16_q1_local_tile
+// one WG per (q_idx, q_head)
+
+#define LT_KC 32
+#define LT_WG 128
+
+#ifndef FA_MQ_ONLY // q1_local_tile excluded from the MQ-only (g8) program
+REQD_SUBGROUP_SIZE_64
+__kernel void flash_attn_f32_f16_q1_local_tile(
+ 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 q_idx = get_global_id(0) / LT_WG;
+ const int head_idx = get_global_id(1);
+ const int batch_idx = get_global_id(2);
+ const int tid = get_local_id(0);
+
+ const int gqa_ratio = n_head_kv > 0 ? (n_head / n_head_kv) : 1;
+ const int head_kv_idx = head_idx / gqa_ratio;
+
+ const float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
+
+ __local half k_tile[LT_KC * DK]; // 32*128*2 = 8 KB at DK=128
+ __local half v_tile[LT_KC * DV]; // 8 KB
+ __local float red[LT_WG]; // 512 B reduction scratch
+ __local float score_shared; // broadcast score (each K-step)
+
+ // Each thread owns one float of Q at index `tid` (assumes LT_WG == DK).
+ const global char * q_row_base = (const global char *) q_void + q_offset +
+ batch_idx * q_nb3 + head_idx * q_nb2 + (ulong) q_idx * q_nb1;
+ float q_val = ((const global float *) q_row_base)[tid];
+
+ 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 +
+ (ulong) q_idx * mask_nb1;
}
- for (int i = 0; i < DV_VEC; ++i) {
- local_o[tid] = o_acc[i];
+
+ float o_val = 0.0f;
+ float m_i = FA_M_INIT;
+ float l_i = 0.0f;
+
+ for (int kb = 0; kb < n_kv; kb += LT_KC) {
+ const int tile_len = min(LT_KC, n_kv - kb);
+
+ // Stage K and V tiles into __local.
+ for (int i = tid; i < tile_len * DK; i += LT_WG) {
+ const int j = i / DK;
+ const int d = i % DK;
+ const int kv_idx = kb + j;
+ const global char * k_row = (const global char *) k_void + k_offset +
+ batch_idx * k_nb3 + head_kv_idx * k_nb2 +
+ (ulong) kv_idx * k_nb1;
+ const global char * v_row = (const global char *) v_void + v_offset +
+ batch_idx * v_nb3 + head_kv_idx * v_nb2 +
+ (ulong) kv_idx * v_nb1;
+ k_tile[j * DK + d] = ((const global half *) k_row)[d];
+ v_tile[j * DV + d] = ((const global half *) v_row)[d];
+ }
barrier(CLK_LOCAL_MEM_FENCE);
- #pragma unroll
- for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) {
- if (tid < s) local_o[tid] += local_o[tid + s];
+
+ for (int j = 0; j < tile_len; ++j) {
+ const int kv_idx = kb + j;
+
+ // Q·K dot via __local tree-reduce.
+ red[tid] = q_val * convert_float(k_tile[j * DK + tid]);
+ barrier(CLK_LOCAL_MEM_FENCE);
+ for (int stride = LT_WG >> 1; stride > 0; stride >>= 1) {
+ if (tid < stride) {
+ red[tid] += red[tid + stride];
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+ }
+
+ if (tid == 0) {
+ float s = red[0] * scale;
+ if (mask_base != NULL) {
+ const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base;
+ s += slope * (float) mask_ptr[kv_idx];
+ }
+ if (logit_softcap > 0.0f) {
+ s = logit_softcap * tanh(s / logit_softcap);
+ }
+ score_shared = s;
+ }
barrier(CLK_LOCAL_MEM_FENCE);
+
+ const float s = score_shared;
+ const float m_new = fmax(m_i, s);
+ const float alpha = native_exp(m_i - m_new);
+ const float beta = native_exp(s - m_new);
+
+ o_val = o_val * alpha + beta * convert_float(v_tile[j * DV + tid]);
+ l_i = l_i * alpha + beta;
+ m_i = m_new;
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+ }
+
+ // Fold attention sinks into the running (m, l, o), if present.
+ if (sinks_void != NULL) {
+ const global float * sinks_ptr =
+ (const global float *) ((const global char *) sinks_void + sinks_offset);
+ const float m_sink = sinks_ptr[head_idx];
+ const float m_new = fmax(m_i, m_sink);
+ const float alpha = native_exp(m_i - m_new);
+ const float beta = native_exp(m_sink - m_new);
+ o_val = o_val * alpha;
+ l_i = l_i * alpha + beta;
+ m_i = m_new;
+ }
+
+ const float l_inv = (l_i > 0.0f) ? (1.0f / l_i) : 0.0f;
+ global float * o_row = (global float *) ((global char *) o_void + o_offset +
+ batch_idx * o_nb3 + head_idx * o_nb1 +
+ (ulong) q_idx * o_nb2);
+ o_row[tid] = o_val * l_inv;
+}
+
+// flash_attn_f32_f16_q1_local_mq_split
+
+#define LMQ_WG 64
+#define LMQ_KC 32
+#define LMQ_DPL 2 // DK / LMQ_WG at DK=128
+
+#endif // !FA_MQ_ONLY (q1_local_tile)
+#ifndef MQ_GQA
+#define MQ_GQA 4
+#endif
+
+#ifndef FA_PARTIAL_FLOATS
+#define FA_PARTIAL_FLOATS (2 + DV)
+#endif
+
+#ifndef FA_MQ_ONLY // q1_local_mq_split excluded from the MQ-only (g8) program
+REQD_SUBGROUP_SIZE_64
+__kernel void flash_attn_f32_f16_q1_local_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); // 0..LMQ_WG-1
+ 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 Q-head so merge treats it as 0.
+ 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;
+
+ // Stage MQ_GQA Q rows in __local (MQ_GQA × DK floats).
+ __local float q_shared[MQ_GQA * DK];
+ for (int i = tid; i < MQ_GQA * DK; i += LMQ_WG) {
+ const int h = i / DK;
+ const int d = i % DK;
+ const int head_idx = head_kv_idx * MQ_GQA + h;
+ const ulong q_row_off = batch_idx * q_nb3 + head_idx * q_nb2 + (ulong) q_idx * q_nb1;
+ const global float * q_ptr = (const global float *) (q_base + q_row_off);
+ q_shared[h * DK + d] = q_ptr[d];
+ }
+
+ // K/V tile staging buffers (16 KB combined at DK=DV=128 KC=32).
+ __local half k_tile[LMQ_KC * DK];
+ __local half v_tile[LMQ_KC * DV];
+
+ // Per-h state held in private registers.
+ float o_acc[MQ_GQA][LMQ_DPL];
+ float m_i[MQ_GQA];
+ float l_i[MQ_GQA];
+ float slope[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ m_i[h] = FA_M_INIT;
+ l_i[h] = 0.0f;
+ slope[h] = get_alibi_slope(max_bias, head_kv_idx * MQ_GQA + h, n_head_log2, m0, m1);
+ #pragma unroll
+ for (int p = 0; p < LMQ_DPL; ++p) o_acc[h][p] = 0.0f;
+ }
+
+ // Per-h mask pointers.
+ 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;
+ }
+
+ barrier(CLK_LOCAL_MEM_FENCE); // Ensure Q staged before first dot.
+
+ for (int kb = kv_start; kb < kv_end; kb += LMQ_KC) {
+ const int tile_len = min((int) LMQ_KC, kv_end - kb);
+
+ // Cooperative load K + V tile.
+ for (int i = tid; i < tile_len * DK; i += LMQ_WG) {
+ const int j = i / DK;
+ const int d = i % DK;
+ const int kv_idx = kb + j;
+ const global char * k_row = k_base + batch_idx * k_nb3 + head_kv_idx * k_nb2 + (ulong) kv_idx * k_nb1;
+ const global char * v_row = v_base + batch_idx * v_nb3 + head_kv_idx * v_nb2 + (ulong) kv_idx * v_nb1;
+ k_tile[j * DK + d] = ((const global half *) k_row)[d];
+ v_tile[j * DV + d] = ((const global half *) v_row)[d];
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+
+ // Process each cache row in the tile.
+ for (int j = 0; j < tile_len; ++j) {
+ const int kv_idx = kb + j;
+
+ // Dot product per h: lane owns LMQ_DPL D-elements at (tid*LMQ_DPL..).
+ float score[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ float contrib = 0.0f;
+ #pragma unroll
+ for (int p = 0; p < LMQ_DPL; ++p) {
+ const int d = tid * LMQ_DPL + p;
+ contrib += q_shared[h * DK + d] * (float) k_tile[j * DK + d];
+ }
+ float s = sub_group_reduce_add(contrib) * scale;
+ if (mask_base[h] != NULL) {
+ const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h];
+ s += slope[h] * (float) mask_ptr[kv_idx];
+ }
+ if (logit_softcap > 0.0f) {
+ s = logit_softcap * tanh(s / logit_softcap);
+ }
+ score[h] = s;
+ }
+
+ // Online softmax update + V accumulation per h.
+ float p_h[MQ_GQA];
+ float sp_h[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ const float m_new = fmax(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;
+ }
+
+ #pragma unroll
+ for (int p = 0; p < LMQ_DPL; ++p) {
+ const int d = tid * LMQ_DPL + p;
+ const float v_val = (float) v_tile[j * DV + d];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ o_acc[h][p] = o_acc[h][p] * sp_h[h] + p_h[h] * v_val;
+ }
+ }
+ }
+ barrier(CLK_LOCAL_MEM_FENCE); // Before next tile load overwrites k/v_tile.
+ }
+
+ // write partial records: one per (h, split)
+ #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;
+ global float * rec_o = rec + 2;
+
+ if (tid == 0) {
+ rec[0] = m_i[h];
+ rec[1] = l_i[h];
+ }
+ #pragma unroll
+ for (int p = 0; p < LMQ_DPL; ++p) {
+ const int d = tid * LMQ_DPL + p;
+ rec_o[d] = o_acc[h][p];
+ }
+ }
+}
+
+#endif // !FA_MQ_ONLY (q1_local_mq_split)
+#ifndef MQ_NSG
+#define MQ_NSG 4
+#endif
+#define MQ_WG_SIZE (Q1_WG_SIZE * MQ_NSG)
+
+REQD_SUBGROUP_SIZE_64
+__kernel void flash_attn_f32_f16_q1_vec_mq(
+ 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; // subgroup 0..MQ_NSG-1
+ const int tid_sg = tid % Q1_WG_SIZE; // lane 0..63
+ const int kvhead_batch_idx = get_global_id(1);
+
+ const int batch_idx = kvhead_batch_idx / n_head_kv;
+ const int head_kv_idx = kvhead_batch_idx % n_head_kv;
+
+ 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;
+
+ __local ACC_TYPE4 q_shared[MQ_GQA * DK_VEC];
+ for (int i = tid; i < MQ_GQA * DK_VEC; i += MQ_WG_SIZE) {
+ 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;
+ 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);
+
+ // per-h ALiBi slope
+ 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);
+ }
+
+ // per-h mask row pointer
+ 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;
+ #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;
+ }
+
+ const global ACC_TYPE * sinks_ptr = NULL;
+ if (sinks_void != NULL) {
+ sinks_ptr = (const global ACC_TYPE *) ((const global char *) sinks_void + sinks_offset);
+ }
+
+ // per-thread per-h DV slice.
+ 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);
+ }
+
+ // each subgroup independently sweeps its slice of n_kv.
+ const int kv_per_sg = (n_kv + MQ_NSG - 1) / MQ_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 ulong k_row_off = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
+ const ulong v_row_off = batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
+ const global KV_DATA_TYPE4 * k_ptr = (const global KV_DATA_TYPE4 *) (k_base + k_row_off);
+ const global KV_DATA_TYPE4 * v_ptr = (const global KV_DATA_TYPE4 *) (v_base + v_row_off);
+
+ // Q*K^T: load each K stride once, dot against all MQ_GQA Q rows.
+ ACC_TYPE4 dot4[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) dot4[h] = (ACC_TYPE4)(0.0f);
+ for (int k = tid_sg; k < DK_VEC; k += Q1_WG_SIZE) {
+ const ACC_TYPE4 k_vec = CONVERT_KV_ACC4(k_ptr[k]);
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ dot4[h] = mad(q_shared[h * DK_VEC + k], k_vec, 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;
+ }
+
+ // FA-2 online softmax update — V load amortized across MQ_GQA heads.
+ // p, scale_prev are computed per h; the V vector is loaded once
+ // per dv stride and reused MQ_GQA times.
+ 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_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE, ++idx) {
+ const ACC_TYPE4 v_vec = CONVERT_KV_ACC4(v_ptr[dv_idx]);
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ o_acc[h][idx] = mad(p_h[h], v_vec, o_acc[h][idx] * sp_h[h]);
+ }
+ }
+ }
+
+ // cross subgroup merge
+ __local ACC_TYPE sg_m[MQ_GQA][MQ_NSG];
+ __local ACC_TYPE sg_l[MQ_GQA][MQ_NSG];
+ __local ACC_TYPE4 sg_o[MQ_NSG][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) {
+ // each subgroup publishes its o_acc slice for head 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_final = sg_m[h][0];
+ #pragma unroll
+ for (int s = 1; s < MQ_NSG; ++s) {
+ m_final = max(m_final, sg_m[h][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 < MQ_NSG; ++s) {
+ l_final += sg_l[h][s] * native_exp(sg_m[h][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);
+ 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; ++s) {
+ const ACC_TYPE alpha = native_exp(sg_m[h][s] - m_final);
+ o_merged = mad((ACC_TYPE4)(alpha), sg_o[s][dv_idx], o_merged);
+ }
+ o_row[dv_idx] = CONVERT_O_DATA4(o_merged * l_inv);
+ }
+ }
+ // Barrier guards next h's overwrite of sg_o.
+ barrier(CLK_LOCAL_MEM_FENCE);
+ }
+}
+
+#ifndef MQ_NSG_SPLIT
+#define MQ_NSG_SPLIT 4
+#endif
+#define MQ_SPLIT_WG_SIZE (Q1_WG_SIZE * MQ_NSG_SPLIT)
+
+#ifndef FA_PARTIAL_FLOATS
+#define FA_PARTIAL_FLOATS (2 + DV)
+#endif
+
+REQD_SUBGROUP_SIZE_64
+__kernel void flash_attn_f32_f16_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) {
+ // write sentinel for each of the MQ_GQA Q-heads so the
+ // merge pass treats this slot as dropped
+ 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;
+
+ // stage MQ_GQA Q rows in __local once (uniform across WG)
+ __local ACC_TYPE4 q_shared[MQ_GQA * DK_VEC];
+ for (int i = tid; i < MQ_GQA * DK_VEC; i += MQ_SPLIT_WG_SIZE) {
+ 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);
+ }
+
+ // each subgroup independently sweeps its slice of the split's kv range.
+ 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 ulong k_row_off = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
+ const ulong v_row_off = batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
+ const global KV_DATA_TYPE4 * k_ptr = (const global KV_DATA_TYPE4 *) (k_base + k_row_off);
+ const global KV_DATA_TYPE4 * v_ptr = (const global KV_DATA_TYPE4 *) (v_base + v_row_off);
+
+ ACC_TYPE4 dot4[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) dot4[h] = (ACC_TYPE4)(0.0f);
+ for (int k = tid_sg; k < DK_VEC; k += Q1_WG_SIZE) {
+ const ACC_TYPE4 k_vec = CONVERT_KV_ACC4(k_ptr[k]);
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ dot4[h] = mad(q_shared[h * DK_VEC + k], k_vec, 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_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE, ++idx) {
+ const ACC_TYPE4 v_vec = CONVERT_KV_ACC4(v_ptr[dv_idx]);
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ o_acc[h][idx] = mad(p_h[h], v_vec, o_acc[h][idx] * sp_h[h]);
+ }
+ }
+ }
+
+ // per-h cross-subgroup merge
+ __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;
+
+ // fold per-subgroup (m, l) into split-level (m_c, l_c)
+ 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;
+ }
+ // each thread writes its DV slice of the merged O.
+ 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);
+ }
+}
+
+REQD_SUBGROUP_SIZE_64
+__kernel void flash_attn_f32_f16_q1_vec_mq_split_k_img(
+ const global void * q_void, ulong q_offset,
+ __read_only image1d_buffer_t k_img,
+ 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) {
+ 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 * 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) {
+ 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);
+ }
+
+ // K pitches in pixel units, pixel = 1 half4 = 8 B -> byte_stride >> 3.
+ const int pitch_px_row = (int)(k_nb1 >> 3);
+ const int pitch_px_head = (int)(k_nb2 >> 3);
+ const int pitch_px_batch = (int)(k_nb3 >> 3);
+
+ 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 int k_row_px = batch_idx * pitch_px_batch +
+ head_kv_idx * pitch_px_head +
+ k_idx * pitch_px_row;
+
+ const ulong v_row_off = batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
+ const global KV_DATA_TYPE4 * v_ptr = (const global KV_DATA_TYPE4 *) (v_base + v_row_off);
+
+ ACC_TYPE4 dot4[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) dot4[h] = (ACC_TYPE4)(0.0f);
+ for (int k = tid_sg; k < DK_VEC; k += Q1_WG_SIZE) {
+ const half4 k_h4 = read_imageh(k_img, k_row_px + k);
+ const ACC_TYPE4 k_vec = CONVERT_KV_ACC4(k_h4);
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ dot4[h] = mad(q_shared[h * DK_VEC + k], k_vec, 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_idx = tid_sg; dv_idx < DV_VEC; dv_idx += Q1_WG_SIZE, ++idx) {
+ const ACC_TYPE4 v_vec = CONVERT_KV_ACC4(v_ptr[dv_idx]);
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ o_acc[h][idx] = mad(p_h[h], v_vec, 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);
+ }
+}
+#endif // !FA_DECODE_ONLY
+
+#ifndef FA_MQ_ONLY // q1_split + merge excluded from the MQ-only (g8) program
+__kernel void flash_attn_f32_f16_q1_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 head_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 = 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 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;
+ const ulong record_idx = ((((ulong) batch_idx * n_head + head_idx) * n_q + q_idx)
+ * n_splits + split_idx);
+ global float * rec = partial_void + record_idx * record_stride;
+ global float4 * rec_o = (global float4 *) (rec + 2);
+
+ if (kv_start >= kv_end) {
+ // Empty split: leave sentinel partial for merge.
+ if (tid == 0) {
+ 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;
+
+ 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 +
+ (ulong) q_idx * mask_nb1;
+ }
+
+ // share Q via local memory (n_q=1 per split -> uniform across WG).
+ __local ACC_TYPE4 q_shared[DK_VEC];
+ 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);
+ for (int i = tid; i < DK_VEC; i += Q1_WG_SIZE) {
+ q_shared[i] = CONVERT_Q_ACC4(q_ptr[i]);
+ }
+ sub_group_barrier(CLK_LOCAL_MEM_FENCE);
+
+ const float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1);
+
+ // pass 1a — split-local max.
+ ACC_TYPE m_i = FA_M_INIT;
+ for (int k_idx = kv_start + tid; k_idx < kv_end; k_idx += Q1_WG_SIZE) {
+ const ulong k_row_offset = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
+ const global KV_DATA_TYPE4 * k_ptr = (const global KV_DATA_TYPE4 *) (k_base + k_row_offset);
+ ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f);
+ #pragma unroll
+ for (int k = 0; k < DK_VEC; ++k) {
+ dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc);
+ }
+ ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * 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);
+ }
+ m_i = max(m_i, score);
+ }
+
+ const ACC_TYPE m_c = sub_group_reduce_max(m_i);
+
+ // pass 1b — softmax-weighted V accumulate.
+ ACC_TYPE4 o_acc[DV_VEC];
+ #pragma unroll
+ for (int i = 0; i < DV_VEC; ++i) o_acc[i] = (ACC_TYPE4)(0.0f);
+ ACC_TYPE l_i = 0.0f;
+
+ for (int k_idx = kv_start + tid; k_idx < kv_end; k_idx += Q1_WG_SIZE) {
+ const ulong k_row_offset = batch_idx * k_nb3 + head_kv_idx * k_nb2 + k_idx * k_nb1;
+ const ulong v_row_offset = batch_idx * v_nb3 + head_kv_idx * v_nb2 + k_idx * v_nb1;
+ const global KV_DATA_TYPE4 * k_ptr = (const global KV_DATA_TYPE4 *) (k_base + k_row_offset);
+ const global KV_DATA_TYPE4 * v_ptr = (const global KV_DATA_TYPE4 *) (v_base + v_row_offset);
+ ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f);
+ #pragma unroll
+ for (int k = 0; k < DK_VEC; ++k) {
+ dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc);
+ }
+ ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * 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 p = exp(score - m_c);
+ l_i += p;
+ #pragma unroll
+ for (int i = 0; i < DV_VEC; ++i) {
+ o_acc[i] = mad(p, CONVERT_KV_ACC4(v_ptr[i]), o_acc[i]);
+ }
+ }
+
+ __local ACC_TYPE4 local_o[Q1_WG_SIZE];
+ const ACC_TYPE l_c = sub_group_reduce_add(l_i);
+
+ if (tid == 0) {
+ rec[0] = (float) m_c;
+ rec[1] = (float) l_c;
+ }
+ for (int i = 0; i < DV_VEC; ++i) {
+ local_o[tid] = o_acc[i];
+ sub_group_barrier(CLK_LOCAL_MEM_FENCE);
+ #pragma unroll
+ for (int s = Q1_WG_SIZE / 2; s > 0; s >>= 1) {
+ if (tid < s) local_o[tid] += local_o[tid + s];
+ sub_group_barrier(CLK_LOCAL_MEM_FENCE);
}
if (tid == 0) {
rec_o[i] = local_o[0];
}
}
-// FD Pass 2: merge per-split partials into final O. Empty splits drop via exp(-INF)=0.
+// FD Pass 2: merge per-split partials into final O
+// empty splits drop via exp(-INF)=0.
__kernel void flash_attn_f32_merge(
const global float * partial_void,
global void * o_void,
global O_DATA_TYPE4 * o_row = (global O_DATA_TYPE4 *) ((global char *) o_void + o_offset + o_row_offset);
o_row[lane] = CONVERT_O_DATA4(o);
}
+#endif // !FA_MQ_ONLY (q1_split + merge)
+#endif // !FA_PREFILL_ONLY (decode kernels)
#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.
}
}
+#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)
+
+// Dequant one float4 lane (0..7) from a q4_0 block.
+// Lanes 0..3 → low nibbles of qs[0..15], lanes 4..7 → high nibbles.
+inline float4 dequant_q4_0_lane(const global char * block_ptr, int lane) {
+ const float d = vload_half(0, (const global half *)block_ptr);
+ const global uchar * qs = (const global uchar *)(block_ptr + 2);
+ const int g = lane & 3;
+ const int shift = (lane < 4) ? 0 : 4;
+ return d * (float4)((float)((qs[g*4+0] >> shift) & 0x0F) - 8.0f,
+ (float)((qs[g*4+1] >> shift) & 0x0F) - 8.0f,
+ (float)((qs[g*4+2] >> shift) & 0x0F) - 8.0f,
+ (float)((qs[g*4+3] >> shift) & 0x0F) - 8.0f);
+}
+
+REQD_SUBGROUP_SIZE_64
+__kernel void flash_attn_f32_q4_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);
+
+#ifdef FA_HAVE_INT_DOT
+ // quantize Q to int8-packed uints + per-block (qd, q_sum) once per WG for dp4a
+ // one thread per Q block, remaining threads idle this step
+ __local uint q_packed_shared[DK_Q4_BLOCKS * 8];
+ __local float q_d_shared[DK_Q4_BLOCKS];
+ __local int q_sum_shared[DK_Q4_BLOCKS];
+ if (tid < DK_Q4_BLOCKS) {
+ ACC_TYPE4 q_block[8];
+ #pragma unroll
+ for (int i = 0; i < 8; ++i) q_block[i] = q_shared[tid * 8 + i];
+ uint packed[8];
+ q4_q_block_info info = quant_q_block_int8_packed_q4(q_block, packed);
+ #pragma unroll
+ for (int i = 0; i < 8; ++i) q_packed_shared[tid * 8 + i] = packed[i];
+ q_d_shared[tid] = info.qd;
+ q_sum_shared[tid] = info.q_sum;
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+#endif
+
+ 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;
+
+#ifdef FA_HAVE_INT_DOT
+ // per-lane dp4a: each lane packs 4 raw q4_0 nibbles into a uint,
+ // then dot_acc_sat_4x8packed_ss_int against the matching uint.
+ ACC_TYPE lane_contrib = 0.0f;
+ for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
+ const int block_idx = qk / 8;
+ const int lane_in_block = qk % 8;
+ const int g = lane_in_block & 3;
+ const int shift = (lane_in_block < 4) ? 0 : 4;
+ const global char * k_block = k_row + block_idx * Q4_0_BLOCK_SIZE;
+ const float kd = vload_half(0, (const global half *)k_block);
+ const global uchar * k_qs = (const global uchar *)(k_block + 2);
+ const uchar b0 = k_qs[g*4 + 0];
+ const uchar b1 = k_qs[g*4 + 1];
+ const uchar b2 = k_qs[g*4 + 2];
+ const uchar b3 = k_qs[g*4 + 3];
+ const uint k_packed = ((uint)((b0 >> shift) & 0x0F)) |
+ ((uint)((b1 >> shift) & 0x0F)) << 8 |
+ ((uint)((b2 >> shift) & 0x0F)) << 16 |
+ ((uint)((b3 >> shift) & 0x0F)) << 24;
+ const uint q_packed_lane = q_packed_shared[block_idx * 8 + lane_in_block];
+ const int raw_dot = dot_acc_sat_4x8packed_ss_int(q_packed_lane, k_packed, 0);
+ const float qd = q_d_shared[block_idx];
+ const float block_scale = qd * kd;
+ float contrib = (float)raw_dot * block_scale;
+ if (lane_in_block == 0) {
+ // block bias correction is per-block
+ const int q_sum_b = q_sum_shared[block_idx];
+ contrib -= 8.0f * block_scale * (float)q_sum_b;
+ }
+ lane_contrib += contrib;
+ }
+ ACC_TYPE score = sub_group_reduce_add(lane_contrib) * scale;
+#else
+ 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_q4_0_lane(k_row + block_idx * Q4_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;
+#endif
+
+ 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_q4_0_lane(v_row + block_idx * Q4_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 q4_0 KV. Merge kernel is type-agnostic and
// shared with the f16/q8_0 FA kernels.
#define FA_PARTIAL_FLOATS (2 + DV)
#define WG_SIZE BLOCK_M
#endif
+#ifndef MQ_GQA
+#define MQ_GQA 4
+#endif
+#ifndef MQ_NSG_SPLIT
+#define MQ_NSG_SPLIT 4
+#endif
+#define MQ_SPLIT_WG_SIZE_Q4 (Q1_WG_SIZE * MQ_NSG_SPLIT)
+
+REQD_SUBGROUP_SIZE_64
+__kernel void flash_attn_f32_q4_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) {
+ 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_Q4) {
+ 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);
+
+#ifdef FA_HAVE_INT_DOT
+ __local uint q_packed_shared[MQ_GQA * DK_Q4_BLOCKS * 8];
+ __local float q_d_shared[MQ_GQA * DK_Q4_BLOCKS];
+ __local int q_sum_shared[MQ_GQA * DK_Q4_BLOCKS];
+ {
+ const int active = MQ_GQA * DK_Q4_BLOCKS;
+ if (tid < active) {
+ const int h = tid / DK_Q4_BLOCKS;
+ const int block_id = tid % DK_Q4_BLOCKS;
+ ACC_TYPE4 q_block[8];
+ #pragma unroll
+ for (int i = 0; i < 8; ++i) q_block[i] = q_shared[h * DK_VEC + block_id * 8 + i];
+ uint packed[8];
+ q4_q_block_info info = quant_q_block_int8_packed_q4(q_block, packed);
+ #pragma unroll
+ for (int i = 0; i < 8; ++i) q_packed_shared[(h * DK_Q4_BLOCKS + block_id) * 8 + i] = packed[i];
+ q_d_shared[h * DK_Q4_BLOCKS + block_id] = info.qd;
+ q_sum_shared[h * DK_Q4_BLOCKS + block_id] = info.q_sum;
+ }
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+#endif
+
+ 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;
+
+#ifdef FA_HAVE_INT_DOT
+ ACC_TYPE lane_contrib[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) lane_contrib[h] = 0.0f;
+
+ for (int qk = tid_sg; qk < DK_VEC; qk += Q1_WG_SIZE) {
+ const int block_idx = qk / 8;
+ const int lane_in_block = qk % 8;
+ const int g = lane_in_block & 3;
+ const int shift = (lane_in_block < 4) ? 0 : 4;
+ const global char * k_block = k_row + block_idx * Q4_0_BLOCK_SIZE;
+ const float kd = vload_half(0, (const global half *)k_block);
+ const global uchar * k_qs = (const global uchar *)(k_block + 2);
+ const uchar b0 = k_qs[g*4 + 0];
+ const uchar b1 = k_qs[g*4 + 1];
+ const uchar b2 = k_qs[g*4 + 2];
+ const uchar b3 = k_qs[g*4 + 3];
+ const uint k_packed = ((uint)((b0 >> shift) & 0x0F)) |
+ ((uint)((b1 >> shift) & 0x0F)) << 8 |
+ ((uint)((b2 >> shift) & 0x0F)) << 16 |
+ ((uint)((b3 >> shift) & 0x0F)) << 24;
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ const uint q_packed_lane = q_packed_shared[(h * DK_Q4_BLOCKS + block_idx) * 8 + lane_in_block];
+ const int raw_dot = dot_acc_sat_4x8packed_ss_int(q_packed_lane, k_packed, 0);
+ const float qd = q_d_shared[h * DK_Q4_BLOCKS + block_idx];
+ const float block_scale = qd * kd;
+ float contrib = (float) raw_dot * block_scale;
+ if (lane_in_block == 0) {
+ const int q_sum_b = q_sum_shared[h * DK_Q4_BLOCKS + block_idx];
+ contrib -= 8.0f * block_scale * (float) q_sum_b;
+ }
+ lane_contrib[h] += contrib;
+ }
+ }
+
+ ACC_TYPE score[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ ACC_TYPE s = sub_group_reduce_add(lane_contrib[h]) * 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;
+ }
+#else
+ // fallback float-dequant K dot
+ 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_q4_0_lane(k_row + block_idx * Q4_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;
+ }
+#endif
+
+ 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_q4_0_lane(v_row + block_idx * Q4_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]);
+ }
+ }
+ }
+
+ // per-h cross-subgroup merge
+ __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_q4_0(
const global void * q_void, ulong q_offset,
const global void * k_void, ulong k_offset,
#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.
}
}
+#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)
#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,
#define REQD_SUBGROUP_SIZE_128 __attribute__((qcom_reqd_sub_group_size("full")))
#endif
+#ifdef cl_khr_subgroup_shuffle
+#pragma OPENCL EXTENSION cl_khr_subgroup_shuffle : enable
+#define HAS_SUBGROUP_SHUFFLE 1
+#elif defined(cl_qcom_subgroup_shuffle)
+#pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable
+#define HAS_SUBGROUP_SHUFFLE 1
+#endif
+
// Assumes row size (ne00) is a multiple of 4
#ifdef ADRENO_GPU
REQD_SUBGROUP_SIZE_64
}
}
#endif // ADRENO_GPU
+
+#define N_ROWS_PER_WG 8
+#define N_OUTS_PER_WG 8
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mat_f16_f32_l4_x8(
+ global char * src0,
+ ulong offset0,
+ global char * src1,
+ ulong offset1,
+ global float * dst,
+ ulong offsetd,
+ int ne00,
+ int ne01,
+ int ne02,
+ ulong nb00,
+ ulong nb01,
+ ulong nb02,
+ ulong nb03,
+ int ne10,
+ int ne11,
+ int ne12,
+ ulong nb10,
+ ulong nb11,
+ ulong nb12,
+ ulong nb13,
+ int ne0,
+ int ne1,
+ int r2,
+ int r3
+) {
+ src0 = (global char *)((global char *)src0 + offset0);
+ src1 = (global char *)((global char *)src1 + offset1);
+ dst = (global float*)((global char *)dst + offsetd);
+
+ const int sgs_lid = get_sub_group_local_id();
+ const int sgs_sz = get_max_sub_group_size();
+
+ const int r0_base = get_group_id(0) * N_ROWS_PER_WG;
+ const int im = get_group_id(2);
+
+ const int i12 = im % ne12;
+ const int i13 = im / ne12;
+
+ const ulong offset_src1 = (i12) * nb12 + (i13) * nb13;
+ global float4 * y4 = (global float4 *)(src1 + offset_src1);
+
+ __local float4 q_loc[64]; // ne00/4 max for sub_group_size 64
+ if (sgs_lid < ne00 / 4) {
+ q_loc[sgs_lid] = y4[sgs_lid];
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+
+ #pragma unroll
+ for (int dr = 0; dr < N_ROWS_PER_WG; ++dr) {
+ const int r0 = r0_base + dr;
+ if (r0 >= ne01) return;
+
+ const ulong offset_src0 = r0 * nb01 + (i12 / r2) * nb02 + (i13 / r3) * nb03;
+ global half4 * x4 = (global half4 *)(src0 + offset_src0);
+
+ float sumf = 0.0f;
+ for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
+ const half4 k4 = x4[i];
+ const float4 q = q_loc[i];
+ sumf += convert_float(k4.s0) * q.s0
+ + convert_float(k4.s1) * q.s1
+ + convert_float(k4.s2) * q.s2
+ + convert_float(k4.s3) * q.s3;
+ }
+
+ const float all_sum = sub_group_reduce_add(sumf);
+ if (sgs_lid == 0) {
+ dst[im * ne1 * ne0 + r0] = all_sum; // ne11 == 1, so r1==0
+ }
+ }
+}
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mat_f16_f32_l4_y8(
+ global char * src0,
+ ulong offset0,
+ global char * src1,
+ ulong offset1,
+ global float * dst,
+ ulong offsetd,
+ int ne00,
+ int ne01,
+ int ne02,
+ ulong nb00,
+ ulong nb01,
+ ulong nb02,
+ ulong nb03,
+ int ne10,
+ int ne11,
+ int ne12,
+ ulong nb10,
+ ulong nb11,
+ ulong nb12,
+ ulong nb13,
+ int ne0,
+ int ne1,
+ int r2,
+ int r3
+) {
+ src0 = (global char *)((global char *)src0 + offset0);
+ src1 = (global char *)((global char *)src1 + offset1);
+ dst = (global float*)((global char *)dst + offsetd);
+
+ const int sgs_lid = get_sub_group_local_id();
+ const int sgs_sz = get_max_sub_group_size();
+
+ const int r0_base = get_group_id(0) * N_OUTS_PER_WG;
+ const int im = get_group_id(2);
+
+ const int i12 = im % ne12;
+ const int i13 = im / ne12;
+
+ const ulong offset_src1 = (i12) * nb12 + (i13) * nb13;
+ global float4 * y4 = (global float4 *)(src1 + offset_src1);
+
+ global half4 * x4_o[N_OUTS_PER_WG];
+ #pragma unroll
+ for (int o = 0; o < N_OUTS_PER_WG; ++o) {
+ const int r0 = r0_base + o;
+ const int r0c = (r0 < ne01) ? r0 : 0;
+ const ulong off = r0c * nb01 + (i12 / r2) * nb02 + (i13 / r3) * nb03;
+ x4_o[o] = (global half4 *)(src0 + off);
+ }
+
+ float sum[N_OUTS_PER_WG] = { 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f };
+
+ for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
+ const float4 q4 = y4[i];
+ #pragma unroll
+ for (int o = 0; o < N_OUTS_PER_WG; ++o) {
+ const half4 v4 = x4_o[o][i];
+ sum[o] += convert_float(v4.s0) * q4.s0
+ + convert_float(v4.s1) * q4.s1
+ + convert_float(v4.s2) * q4.s2
+ + convert_float(v4.s3) * q4.s3;
+ }
+ }
+
+ #pragma unroll
+ for (int o = 0; o < N_OUTS_PER_WG; ++o) {
+ const int r0 = r0_base + o;
+ const float s = sub_group_reduce_add(sum[o]);
+ if (sgs_lid == 0 && r0 < ne01) {
+ dst[im * ne1 * ne0 + r0] = s;
+ }
+ }
+}
+
+#define N_OUTS_PAIR 8
+#define N_PAIRS_PAIR (N_OUTS_PAIR / 2)
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mat_f16_f32_l4_x8_pair(
+ global char * src0,
+ ulong offset0,
+ global char * src1,
+ ulong offset1,
+ global float * dst,
+ ulong offsetd,
+ int ne00,
+ int ne01,
+ int ne02,
+ ulong nb00,
+ ulong nb01,
+ ulong nb02,
+ ulong nb03,
+ int ne10,
+ int ne11,
+ int ne12,
+ ulong nb10,
+ ulong nb11,
+ ulong nb12,
+ ulong nb13,
+ int ne0,
+ int ne1,
+ int r2,
+ int r3
+) {
+ src0 = (global char *)((global char *)src0 + offset0);
+ src1 = (global char *)((global char *)src1 + offset1);
+ dst = (global float*)((global char *)dst + offsetd);
+
+ const int sgs_lid = get_sub_group_local_id();
+ const int half_id = sgs_lid >> 5; // 0 = lower half, 1 = upper half
+ const int lane_h = sgs_lid & 31; // lane 0..31 within half
+
+ const int r0_base = get_group_id(0) * N_OUTS_PAIR;
+ const int im = get_group_id(2);
+
+ const int i12 = im % ne12;
+ const int i13 = im / ne12;
+
+ const ulong offset_src1 = (i12) * nb12 + (i13) * nb13;
+ global float4 * y4 = (global float4 *)(src1 + offset_src1);
+
+ __local float4 q_loc[64]; // ne00/4 max for sub_group_size 64
+ if (sgs_lid < ne00 / 4) {
+ q_loc[sgs_lid] = y4[sgs_lid];
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+
+ const int dk_vec = ne00 / 4;
+
+ #pragma unroll
+ for (int p = 0; p < N_PAIRS_PAIR; ++p) {
+ const int r0 = r0_base + 2 * p + half_id;
+
+ const ulong offset_src0 = r0 * nb01 + (i12 / r2) * nb02 + (i13 / r3) * nb03;
+ global half4 * x4 = (global half4 *)(src0 + offset_src0);
+
+ float sumf = 0.0f;
+ for (int i = lane_h; i < dk_vec; i += 32) {
+ const half4 k4 = x4[i];
+ const float4 q = q_loc[i];
+ sumf += convert_float(k4.s0) * q.s0
+ + convert_float(k4.s1) * q.s1
+ + convert_float(k4.s2) * q.s2
+ + convert_float(k4.s3) * q.s3;
+ }
+
+ sumf += sub_group_shuffle_xor(sumf, 16);
+ sumf += sub_group_shuffle_xor(sumf, 8);
+ sumf += sub_group_shuffle_xor(sumf, 4);
+ sumf += sub_group_shuffle_xor(sumf, 2);
+ sumf += sub_group_shuffle_xor(sumf, 1);
+
+ if (lane_h == 0) {
+ dst[im * ne1 * ne0 + r0] = sumf;
+ }
+ }
+}
+
+#define N_K_ROWS_GQA 16
+#define GQA_RATIO_GQA 8
+#define LANES_PER_QH 8 // 64 / GQA_RATIO_GQA
+#define DK_VEC_GQA 32 // DK / 4 for DK=128
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mat_f16_f32_l4_x8_gqa4(
+ global char * src0,
+ ulong offset0,
+ global char * src1,
+ ulong offset1,
+ global float * dst,
+ ulong offsetd,
+ int ne00,
+ int ne01,
+ int ne02,
+ ulong nb00,
+ ulong nb01,
+ ulong nb02,
+ ulong nb03,
+ int ne10,
+ int ne11,
+ int ne12,
+ ulong nb10,
+ ulong nb11,
+ ulong nb12,
+ ulong nb13,
+ int ne0,
+ int ne1,
+ int r2,
+ int r3
+) {
+ src0 = (global char *)((global char *)src0 + offset0);
+ src1 = (global char *)((global char *)src1 + offset1);
+ dst = (global float*)((global char *)dst + offsetd);
+
+ const int sgs_lid = get_sub_group_local_id();
+ const int q_id = sgs_lid >> 3; // 0..7: which Q-head (8 per WG)
+ const int lane_q = sgs_lid & 7; // 0..7: lane within Q-head partition
+
+ const int r0_base = get_group_id(0) * N_K_ROWS_GQA;
+ const int im_kv = get_group_id(2);
+
+ const int i02 = im_kv % ne02; // K-head index (also K2 batch)
+ const int i03 = im_kv / ne02; // n13 batch index
+
+ const int q_head_lo = i02 * GQA_RATIO_GQA;
+
+ __local float4 q_loc[GQA_RATIO_GQA * DK_VEC_GQA]; // 4 × 32 = 128 float4
+ #pragma unroll
+ for (int qh = 0; qh < GQA_RATIO_GQA; ++qh) {
+ const int qh_idx = q_head_lo + qh;
+ global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
+
+ if (sgs_lid < DK_VEC_GQA) {
+ q_loc[qh * DK_VEC_GQA + sgs_lid] = y4[sgs_lid];
+ }
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+
+ // K base offset for this WG. All 8 K-rows × 4 Q-heads share this K-head.
+ const ulong offset_src0_base = (i02) * nb02 + (i03 / r3) * nb03;
+
+ #pragma unroll
+ for (int dr = 0; dr < N_K_ROWS_GQA; ++dr) {
+ const int r0 = r0_base + dr;
+
+ const ulong offset_src0 = r0 * nb01 + offset_src0_base;
+ global half4 * x4 = (global half4 *)(src0 + offset_src0);
+
+ float sumf = 0.0f;
+ #pragma unroll
+ for (int t = 0; t < 4; ++t) {
+ const int i = lane_q + t * LANES_PER_QH; // 8, 16, 24-step
+ const half4 k4 = x4[i];
+ const float4 q = q_loc[q_id * DK_VEC_GQA + i];
+ sumf += convert_float(k4.s0) * q.s0
+ + convert_float(k4.s1) * q.s1
+ + convert_float(k4.s2) * q.s2
+ + convert_float(k4.s3) * q.s3;
+ }
+
+ sumf += sub_group_shuffle_xor(sumf, 4);
+ sumf += sub_group_shuffle_xor(sumf, 2);
+ sumf += sub_group_shuffle_xor(sumf, 1);
+
+ if (lane_q == 0) {
+ const int im_out = i03 * ne12 + (q_head_lo + q_id);
+ dst[im_out * ne1 * ne0 + r0] = sumf;
+ }
+ }
+}
+
+#define N_DV_ROWS_Y8GQA 8
+#define GQA_RATIO_Y8GQA 8
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mat_f16_f32_l4_y8_gqa(
+ global char * src0,
+ ulong offset0,
+ global char * src1,
+ ulong offset1,
+ global float * dst,
+ ulong offsetd,
+ int ne00,
+ int ne01,
+ int ne02,
+ ulong nb00,
+ ulong nb01,
+ ulong nb02,
+ ulong nb03,
+ int ne10,
+ int ne11,
+ int ne12,
+ ulong nb10,
+ ulong nb11,
+ ulong nb12,
+ ulong nb13,
+ int ne0,
+ int ne1,
+ int r2,
+ int r3
+) {
+ src0 = (global char *)((global char *)src0 + offset0);
+ src1 = (global char *)((global char *)src1 + offset1);
+ dst = (global float*)((global char *)dst + offsetd);
+
+ const int sgs_lid = get_sub_group_local_id();
+ const int sgs_sz = get_max_sub_group_size();
+
+ const int r0_base = get_group_id(0) * N_DV_ROWS_Y8GQA;
+ const int im_kv = get_group_id(2);
+
+ const int i02 = im_kv % ne02; // K-head index
+ const int i03 = im_kv / ne02; // n13 batch index
+
+ // GQA Q-heads sharing this K-head.
+ const int q_head_lo = i02 * GQA_RATIO_Y8GQA;
+
+ global float4 * y4_q[GQA_RATIO_Y8GQA];
+ #pragma unroll
+ for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
+ const int qh_idx = q_head_lo + qh;
+ y4_q[qh] = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
+ }
+
+ global half4 * x4_o[N_DV_ROWS_Y8GQA];
+ #pragma unroll
+ for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
+ const int r0 = r0_base + o;
+ const int r0c = (r0 < ne01) ? r0 : 0;
+ const ulong off = r0c * nb01 + (i02) * nb02 + (i03 / r3) * nb03;
+ x4_o[o] = (global half4 *)(src0 + off);
+ }
+
+ float sum[N_DV_ROWS_Y8GQA][GQA_RATIO_Y8GQA] = { {0.0f} };
+
+ for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
+ // load 8 V values (one per DV row), same K-head, K-pos = i.
+ half4 v[N_DV_ROWS_Y8GQA];
+ #pragma unroll
+ for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
+ v[o] = x4_o[o][i];
+ }
+
+ // load 8 softmax values (one per Q-head).
+ float4 q[GQA_RATIO_Y8GQA];
+ #pragma unroll
+ for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
+ q[qh] = y4_q[qh][i];
+ }
+
+ #pragma unroll
+ for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
+ const float4 vf = (float4)(convert_float(v[o].s0),
+ convert_float(v[o].s1),
+ convert_float(v[o].s2),
+ convert_float(v[o].s3));
+ #pragma unroll
+ for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
+ sum[o][qh] += vf.s0 * q[qh].s0
+ + vf.s1 * q[qh].s1
+ + vf.s2 * q[qh].s2
+ + vf.s3 * q[qh].s3;
+ }
+ }
+ }
+
+ #pragma unroll
+ for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
+ const int r0 = r0_base + o;
+ #pragma unroll
+ for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
+ const float s = sub_group_reduce_add(sum[o][qh]);
+ if (sgs_lid == 0 && r0 < ne01) {
+ const int im_out = i03 * ne12 + (q_head_lo + qh);
+ dst[im_out * ne1 * ne0 + r0] = s;
+ }
+ }
+ }
+}
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mat_f16_f32_l4_x8_gqa4_img(
+ __read_only image1d_buffer_t src0_img,
+ global char * src1,
+ ulong offset1,
+ global float * dst,
+ ulong offsetd,
+ int ne00,
+ int ne01,
+ int ne02,
+ ulong nb01,
+ ulong nb02,
+ ulong nb03,
+ int ne10,
+ int ne11,
+ int ne12,
+ ulong nb10,
+ ulong nb11,
+ ulong nb12,
+ ulong nb13,
+ int ne0,
+ int ne1,
+ int r2,
+ int r3
+) {
+ src1 = (global char *)((global char *)src1 + offset1);
+ dst = (global float*)((global char *)dst + offsetd);
+
+ const int sgs_lid = get_sub_group_local_id();
+ const int q_id = sgs_lid >> 3; // 0..7: which Q-head (8 per WG)
+ const int lane_q = sgs_lid & 7; // 0..7: lane within Q-head partition
+
+ const int r0_base = get_group_id(0) * N_K_ROWS_GQA;
+ const int im_kv = get_group_id(2);
+
+ const int i02 = im_kv % ne02;
+ const int i03 = im_kv / ne02;
+
+ const int q_head_lo = i02 * GQA_RATIO_GQA;
+
+ __local float4 q_loc[GQA_RATIO_GQA * DK_VEC_GQA];
+ #pragma unroll
+ for (int qh = 0; qh < GQA_RATIO_GQA; ++qh) {
+ const int qh_idx = q_head_lo + qh;
+ global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
+ if (sgs_lid < DK_VEC_GQA) {
+ q_loc[qh * DK_VEC_GQA + sgs_lid] = y4[sgs_lid];
+ }
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+
+ const int pitch_px_row = (int)(nb01 >> 4);
+ const int pitch_px_head = (int)(nb02 >> 4);
+ const int pitch_px_n13 = (int)(nb03 >> 4);
+
+ const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
+
+ #pragma unroll
+ for (int dr = 0; dr < N_K_ROWS_GQA; ++dr) {
+ const int r0 = r0_base + dr;
+ const int row_px_base = r0 * pitch_px_row + head_px_base;
+
+ float sumf = 0.0f;
+ #pragma unroll
+ for (int t = 0; t < 2; ++t) {
+ const int p = lane_q + t * LANES_PER_QH; // pixel idx in row, 0..15
+ const half8 k8 = as_half8(read_imagef(src0_img, row_px_base + p));
+ const int i0 = 2 * p; // first half4 idx
+ const float4 qa = q_loc[q_id * DK_VEC_GQA + i0 ];
+ const float4 qb = q_loc[q_id * DK_VEC_GQA + i0 + 1];
+ sumf += convert_float(k8.s0) * qa.s0
+ + convert_float(k8.s1) * qa.s1
+ + convert_float(k8.s2) * qa.s2
+ + convert_float(k8.s3) * qa.s3
+ + convert_float(k8.s4) * qb.s0
+ + convert_float(k8.s5) * qb.s1
+ + convert_float(k8.s6) * qb.s2
+ + convert_float(k8.s7) * qb.s3;
+ }
+
+ sumf += sub_group_shuffle_xor(sumf, 4);
+ sumf += sub_group_shuffle_xor(sumf, 2);
+ sumf += sub_group_shuffle_xor(sumf, 1);
+
+ if (lane_q == 0) {
+ const int im_out = i03 * ne12 + (q_head_lo + q_id);
+ dst[im_out * ne1 * ne0 + r0] = sumf;
+ }
+ }
+}
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mat_f16_f32_l4_y8_gqa_img(
+ __read_only image1d_buffer_t src0_img,
+ global char * src1,
+ ulong offset1,
+ global float * dst,
+ ulong offsetd,
+ int ne00,
+ int ne01,
+ int ne02,
+ ulong nb01,
+ ulong nb02,
+ ulong nb03,
+ int ne10,
+ int ne11,
+ int ne12,
+ ulong nb10,
+ ulong nb11,
+ ulong nb12,
+ ulong nb13,
+ int ne0,
+ int ne1,
+ int r2,
+ int r3
+) {
+ src1 = (global char *)((global char *)src1 + offset1);
+ dst = (global float*)((global char *)dst + offsetd);
+
+ const int sgs_lid = get_sub_group_local_id();
+ const int sgs_sz = get_max_sub_group_size();
+
+ const int r0_base = get_group_id(0) * N_DV_ROWS_Y8GQA;
+ const int im_kv = get_group_id(2);
+
+ const int i02 = im_kv % ne02;
+ const int i03 = im_kv / ne02;
+
+ const int q_head_lo = i02 * GQA_RATIO_Y8GQA;
+
+ // Q (= softmax(KQ)) base pointers per Q-head
+ global float4 * y4_q[GQA_RATIO_Y8GQA];
+ #pragma unroll
+ for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
+ const int qh_idx = q_head_lo + qh;
+ y4_q[qh] = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
+ }
+
+ const int pitch_px_row = (int)(nb01 >> 3);
+ const int pitch_px_head = (int)(nb02 >> 3);
+ const int pitch_px_n13 = (int)(nb03 >> 3);
+
+ const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
+
+ // per-DV-row pixel base
+ int row_px_base[N_DV_ROWS_Y8GQA];
+ #pragma unroll
+ for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
+ const int r0 = r0_base + o;
+ const int r0c = (r0 < ne01) ? r0 : 0;
+ row_px_base[o] = r0c * pitch_px_row + head_px_base;
+ }
+
+ float sum[N_DV_ROWS_Y8GQA][GQA_RATIO_Y8GQA] = { {0.0f} };
+
+ for (int i = sgs_lid; i < ne00 / 4; i += sgs_sz) {
+ half4 v[N_DV_ROWS_Y8GQA];
+
+ #pragma unroll
+ for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
+ v[o] = read_imageh(src0_img, row_px_base[o] + i);
+ }
+
+ float4 q[GQA_RATIO_Y8GQA];
+ #pragma unroll
+ for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
+ q[qh] = y4_q[qh][i];
+ }
+ // 64 mads.
+ #pragma unroll
+ for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
+ const float4 vf = (float4)(convert_float(v[o].s0),
+ convert_float(v[o].s1),
+ convert_float(v[o].s2),
+ convert_float(v[o].s3));
+ #pragma unroll
+ for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
+ sum[o][qh] += vf.s0 * q[qh].s0
+ + vf.s1 * q[qh].s1
+ + vf.s2 * q[qh].s2
+ + vf.s3 * q[qh].s3;
+ }
+ }
+ }
+
+ #pragma unroll
+ for (int o = 0; o < N_DV_ROWS_Y8GQA; ++o) {
+ const int r0 = r0_base + o;
+ #pragma unroll
+ for (int qh = 0; qh < GQA_RATIO_Y8GQA; ++qh) {
+ const float s = sub_group_reduce_add(sum[o][qh]);
+ if (sgs_lid == 0 && r0 < ne01) {
+ const int im_out = i03 * ne12 + (q_head_lo + qh);
+ dst[im_out * ne1 * ne0 + r0] = s;
+ }
+ }
+ }
+}
+
+#define N_K_ROWS_GQA_R4 16
+#define GQA_RATIO_R4 4
+#define LANES_PER_QH_R4 16 // = 64 / GQA_RATIO_R4
+#define DK_VEC_R4 32 // DK / 4 for DK=128
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mat_f16_f32_l4_x8_gqa_r4_img(
+ __read_only image1d_buffer_t src0_img,
+ global char * src1,
+ ulong offset1,
+ global float * dst,
+ ulong offsetd,
+ int ne00,
+ int ne01,
+ int ne02,
+ ulong nb01,
+ ulong nb02,
+ ulong nb03,
+ int ne10,
+ int ne11,
+ int ne12,
+ ulong nb10,
+ ulong nb11,
+ ulong nb12,
+ ulong nb13,
+ int ne0,
+ int ne1,
+ int r2,
+ int r3
+) {
+ src1 = (global char *)((global char *)src1 + offset1);
+ dst = (global float*)((global char *)dst + offsetd);
+
+ const int sgs_lid = get_sub_group_local_id();
+ const int q_id = sgs_lid >> 4; // 0..3
+ const int lane_q = sgs_lid & 15; // 0..15
+
+ const int r0_base = get_group_id(0) * N_K_ROWS_GQA_R4;
+ const int im_kv = get_group_id(2);
+
+ const int i02 = im_kv % ne02;
+ const int i03 = im_kv / ne02;
+
+ const int q_head_lo = i02 * GQA_RATIO_R4;
+
+ __local float4 q_loc[GQA_RATIO_R4 * DK_VEC_R4];
+ #pragma unroll
+ for (int qh = 0; qh < GQA_RATIO_R4; ++qh) {
+ const int qh_idx = q_head_lo + qh;
+ global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
+ if (sgs_lid < DK_VEC_R4) {
+ q_loc[qh * DK_VEC_R4 + sgs_lid] = y4[sgs_lid];
+ }
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+
+ const int pitch_px_row = (int)(nb01 >> 4);
+ const int pitch_px_head = (int)(nb02 >> 4);
+ const int pitch_px_n13 = (int)(nb03 >> 4);
+
+ const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
+
+ #pragma unroll
+ for (int dr = 0; dr < N_K_ROWS_GQA_R4; ++dr) {
+ const int r0 = r0_base + dr;
+ const int row_px_base = r0 * pitch_px_row + head_px_base;
+
+ const int p = lane_q;
+ const half8 k8 = as_half8(read_imagef(src0_img, row_px_base + p));
+ const int i0 = 2 * p;
+ const float4 qa = q_loc[q_id * DK_VEC_R4 + i0 ];
+ const float4 qb = q_loc[q_id * DK_VEC_R4 + i0 + 1];
+
+ float sumf =
+ convert_float(k8.s0) * qa.s0
+ + convert_float(k8.s1) * qa.s1
+ + convert_float(k8.s2) * qa.s2
+ + convert_float(k8.s3) * qa.s3
+ + convert_float(k8.s4) * qb.s0
+ + convert_float(k8.s5) * qb.s1
+ + convert_float(k8.s6) * qb.s2
+ + convert_float(k8.s7) * qb.s3;
+
+ sumf += sub_group_shuffle_xor(sumf, 8);
+ sumf += sub_group_shuffle_xor(sumf, 4);
+ sumf += sub_group_shuffle_xor(sumf, 2);
+ sumf += sub_group_shuffle_xor(sumf, 1);
+
+ if (lane_q == 0) {
+ const int im_out = i03 * ne12 + (q_head_lo + q_id);
+ dst[im_out * ne1 * ne0 + r0] = sumf;
+ }
+ }
+}
+
+#define N_K_ROWS_GQA_R2_DK256 16
+#define GQA_RATIO_R2 2
+#define LANES_PER_QH_R2 32 // = 64 / GQA_RATIO_R2
+#define DK_VEC_DK256 64 // DK / 4 for DK=256
+
+#ifdef ADRENO_GPU
+REQD_SUBGROUP_SIZE_64
+#endif
+kernel void kernel_mul_mat_f16_f32_l4_x8_gqa_r2_dk256_img(
+ __read_only image1d_buffer_t src0_img,
+ global char * src1,
+ ulong offset1,
+ global float * dst,
+ ulong offsetd,
+ int ne00,
+ int ne01,
+ int ne02,
+ ulong nb01,
+ ulong nb02,
+ ulong nb03,
+ int ne10,
+ int ne11,
+ int ne12,
+ ulong nb10,
+ ulong nb11,
+ ulong nb12,
+ ulong nb13,
+ int ne0,
+ int ne1,
+ int r2,
+ int r3
+) {
+ src1 = (global char *)((global char *)src1 + offset1);
+ dst = (global float*)((global char *)dst + offsetd);
+
+ const int sgs_lid = get_sub_group_local_id();
+ const int q_id = sgs_lid >> 5; // 0..1
+ const int lane_q = sgs_lid & 31; // 0..31
+
+ const int r0_base = get_group_id(0) * N_K_ROWS_GQA_R2_DK256;
+ const int im_kv = get_group_id(2);
+
+ const int i02 = im_kv % ne02;
+ const int i03 = im_kv / ne02;
+
+ const int q_head_lo = i02 * GQA_RATIO_R2;
+
+ __local float4 q_loc[GQA_RATIO_R2 * DK_VEC_DK256];
+ #pragma unroll
+ for (int qh = 0; qh < GQA_RATIO_R2; ++qh) {
+ const int qh_idx = q_head_lo + qh;
+ global float4 * y4 = (global float4 *)(src1 + qh_idx * nb12 + i03 * nb13);
+ q_loc[qh * DK_VEC_DK256 + sgs_lid] = y4[sgs_lid];
+ }
+ barrier(CLK_LOCAL_MEM_FENCE);
+
+ const int pitch_px_row = (int)(nb01 >> 4);
+ const int pitch_px_head = (int)(nb02 >> 4);
+ const int pitch_px_n13 = (int)(nb03 >> 4);
+
+ const int head_px_base = i02 * pitch_px_head + (i03 / r3) * pitch_px_n13;
+
+ #pragma unroll
+ for (int dr = 0; dr < N_K_ROWS_GQA_R2_DK256; ++dr) {
+ const int r0 = r0_base + dr;
+ const int row_px_base = r0 * pitch_px_row + head_px_base;
+
+ const int p = lane_q;
+ const half8 k8 = as_half8(read_imagef(src0_img, row_px_base + p));
+ const int i0 = 2 * p;
+ const float4 qa = q_loc[q_id * DK_VEC_DK256 + i0 ];
+ const float4 qb = q_loc[q_id * DK_VEC_DK256 + i0 + 1];
+
+ float sumf =
+ convert_float(k8.s0) * qa.s0
+ + convert_float(k8.s1) * qa.s1
+ + convert_float(k8.s2) * qa.s2
+ + convert_float(k8.s3) * qa.s3
+ + convert_float(k8.s4) * qb.s0
+ + convert_float(k8.s5) * qb.s1
+ + convert_float(k8.s6) * qb.s2
+ + convert_float(k8.s7) * qb.s3;
+
+ sumf += sub_group_shuffle_xor(sumf, 16);
+ sumf += sub_group_shuffle_xor(sumf, 8);
+ sumf += sub_group_shuffle_xor(sumf, 4);
+ sumf += sub_group_shuffle_xor(sumf, 2);
+ sumf += sub_group_shuffle_xor(sumf, 1);
+
+ if (lane_q == 0) {
+ const int im_out = i03 * ne12 + (q_head_lo + q_id);
+ dst[im_out * ne1 * ne0 + r0] = sumf;
+ }
+ }
+}
overview / pkg / ggml / sources / llama.cpp / commitdiff