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;
+ // Cluster-parallel decode
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_c8;
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_g8_c8;
+ // NSG_SPLIT=2 specializations (WG=128): the c8 kernel's register footprint
+ // caps its per-kernel WG at 128 on X2, below the stock 256/192 requirement.
+ // 2 subgroups × FA_CL_NCL streams still gives 16 in-flight rows per WG.
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_c8_ns2;
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_g8_c8_ns2;
+ // FA_CL_C=32 / MQ_GQA=8 / NSG_SPLIT=2 specialization for the DK=DV=256
+ // GQA=8 class (Qwen3.5/3.6-35B-A3B: 16 Q heads, 2 KV heads). o_acc =
+ // DV_VEC/32 × 8 = 128B/lane (in budget); the baseline fa1 path for this
+ // shape has NO MQ/FD at all and pays an 8× KV re-read per Q head.
+ std::map<std::pair<int, int>, cl_kernel> f32_f16_q1_vec_mq_split_g8_c32;
// 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
// 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;
+ // Cluster-parallel q8_0 decode
+ std::map<std::pair<int, int>, cl_kernel> f32_q8_0_q1_vec_mq_split_c8;
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
// 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;
+ // Cluster-parallel q4_0 decode
+ std::map<std::pair<int, int>, cl_kernel> f32_q4_0_q1_vec_mq_split_g8_c8;
+ std::map<std::pair<int, int>, cl_kernel> f32_q4_0_q1_vec_mq_split_c8;
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;
variant == FA_VARIANT_Q4_0_SPLIT;
if (is_split) {
opts += " -D N_SPLIT=" + std::to_string(cfg->n_split);
- if (backend_ctx->has_subgroup_shuffle) {
- opts += backend_ctx->has_qcom_subgroup_shuffle
- ? " -D cl_qcom_subgroup_shuffle=1"
- : " -D cl_khr_subgroup_shuffle=1";
- }
+ }
+ // Shuffle define for the split tile paths AND the cluster-parallel decode
+ // kernel (q1_vec_mq_split_c8) in the plain F32_F16 program. Without it the
+ // c8 kernel is compiled out (HAS_SUBGROUP_SHUFFLE guard) and dispatch
+ // falls back to the baseline mq_split.
+ if ((is_split || variant == FA_VARIANT_F32_F16) && backend_ctx->has_subgroup_shuffle) {
+ opts += backend_ctx->has_qcom_subgroup_shuffle
+ ? " -D cl_qcom_subgroup_shuffle=1"
+ : " -D cl_khr_subgroup_shuffle=1";
+ }
+ // X1E drops the explicit sub-group size pin on the c8 kernels, compiler
+ // routes the fp16-heavy kernel to a slow variant with explicit subgroup size
+ if (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E) {
+ opts += " -D FA_C8_NO_SG_PIN";
}
return opts;
}
opts += " -D FA_DECODE_ONLY -D FA_DECODE_MINIMAL";
}
+ // c8 cluster width (GGML_OPENCL_FA_CL_C overrides): value = GQA4 cluster
+ // width (kernel default 8); the g8 programs use 2x the value (default 16).
+ // Wider clusters halve per-lane o_acc at the cost of position streams per
+ // subgroup
+ static const int fa_cl_c_env = []{
+ const char * e = std::getenv("GGML_OPENCL_FA_CL_C");
+ const int x = (e && e[0]) ? atoi(e) : 0;
+ return (x == 8 || x == 16 || x == 32) ? x : 0; // 0 = per-gen default
+ }();
+ const int fa_cl_c_gqa4 = fa_cl_c_env ? fa_cl_c_env
+ : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E ? 16 : 0);
+ const std::string opts_cl_c_gqa4 = fa_cl_c_gqa4
+ ? " -D FA_CL_C=" + std::to_string(fa_cl_c_gqa4) : std::string();
+ const std::string fa_cl_c_g8_val = std::to_string(fa_cl_c_gqa4 ? fa_cl_c_gqa4 * 2 : 16);
+
const char * tag = nullptr;
switch (variant) {
case FA_VARIANT_F16: tag = "fa f16"; break;
default: break;
}
cl_program prog = build_program_from_source_ex(
- backend_ctx->context, backend_ctx->device, src.c_str(), opts,
+ backend_ctx->context, backend_ctx->device, src.c_str(), opts + opts_cl_c_gqa4,
/*fatal=*/false, tag, backend_ctx->queue);
if (!prog) { return false; }
clReleaseKernel(k_q1_vec_mq_split_k_img);
}
}
+ // Cluster-parallel decode variant
+ cl_kernel k_q1_vec_mq_split_c8 = clCreateKernel(prog, "flash_attn_f32_f16_q1_vec_mq_split_c8", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_q1_vec_mq_split_c8, 256,
+ "flash_attn_f32_f16_q1_vec_mq_split_c8", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_c8[{dk, dv}] = k_q1_vec_mq_split_c8;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_q1_vec_mq_split_c8, "flash_attn_f32_f16_q1_vec_mq_split_c8", dk, dv);
+ } else {
+ clReleaseKernel(k_q1_vec_mq_split_c8);
+ }
+ }
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 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";
+ // FA_CL_C for the g8 program: MQ_GQA=8 doubles the c8 kernel's
+ // per-lane o_acc, so widen the cluster to keep the register
+ // footprint inside the 192-thread WG cap (see fa_cl_c_gqa4 above
+ // for the per-gen default).
+ const std::string opts_g8 = opts + " -D MQ_GQA=8 -D MQ_NSG=3 -D MQ_NSG_SPLIT=3 -D FA_MQ_ONLY -D FA_CL_C=" + fa_cl_c_g8_val;
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);
clReleaseKernel(k_q1_vec_mq_split_g8_k_img);
}
}
+ // Cluster-parallel decode, MQ_GQA=8 / FA_CL_C=16 specialization
+ cl_kernel k_q1_vec_mq_split_g8_c8 = clCreateKernel(prog_g8, "flash_attn_f32_f16_q1_vec_mq_split_c8", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_q1_vec_mq_split_g8_c8, mq_g8_required_wg,
+ "flash_attn_f32_f16_q1_vec_mq_split_c8 (g8)", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8[{dk, dv}] = k_q1_vec_mq_split_g8_c8;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_q1_vec_mq_split_g8_c8, "flash_attn_f32_f16_q1_vec_mq_split_g8_c8", dk, dv);
+ } else {
+ clReleaseKernel(k_q1_vec_mq_split_g8_c8);
+ }
+ }
// 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);
}
clReleaseProgram(prog_g8);
}
+ // NSG_SPLIT=2 programs for the cluster-parallel kernel: its register
+ // footprint caps the per-kernel WG at 128 on X2 (< the stock 256/192
+ // requirement), so it can never register from the stock programs.
+ // With FA_CL_NCL position streams per subgroup, 2 subgroups still
+ // carry 16 in-flight rows per WG (baseline WG has 4). FA_MQ_ONLY
+ // keeps these compiles minimal; skipped when the stock program c8
+ // registered (some other device) or shuffles are absent.
+ if (!fa_decode_only && backend_ctx->has_subgroup_shuffle &&
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_c8.count({dk, dv}) == 0) {
+ const std::string opts_c8_ns2 = opts + " -D FA_MQ_ONLY -D MQ_GQA=4 -D MQ_NSG=2 -D MQ_NSG_SPLIT=2" + opts_cl_c_gqa4;
+ cl_program prog_c8 = build_program_from_source_ex(
+ backend_ctx->context, backend_ctx->device, src.c_str(), opts_c8_ns2,
+ /*fatal=*/false, "fa f32_f16 c8 NSG2", backend_ctx->queue);
+ if (prog_c8) {
+ cl_kernel k_c8 = clCreateKernel(prog_c8, "flash_attn_f32_f16_q1_vec_mq_split_c8", &err);
+ if (err == CL_SUCCESS) {
+ // WG = MQ_NSG(2) × Q1_WG_SIZE(=FA_SG): 128 Adreno (64), 64 Intel (32).
+ const size_t c8_ns2_wg = backend_ctx->gpu_family == INTEL ? 64 : 128;
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_c8, c8_ns2_wg,
+ "flash_attn_f32_f16_q1_vec_mq_split_c8 (ns2)", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_c8_ns2[{dk, dv}] = k_c8;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_c8, "flash_attn_f32_f16_q1_vec_mq_split_c8_ns2", dk, dv);
+ } else {
+ clReleaseKernel(k_c8);
+ }
+ }
+ clReleaseProgram(prog_c8);
+ }
+ }
+ // FA_CL_C=32 g8 program for the DK=DV=256 GQA=8
+ if (!fa_decode_only && backend_ctx->has_subgroup_shuffle &&
+ dk == 256 && dv == 256) {
+ const std::string opts_g8_c32 = opts + " -D FA_MQ_ONLY -D MQ_GQA=8 -D MQ_NSG=2 -D MQ_NSG_SPLIT=2 -D FA_CL_C=32";
+ cl_program prog_g8_c32 = build_program_from_source_ex(
+ backend_ctx->context, backend_ctx->device, src.c_str(), opts_g8_c32,
+ /*fatal=*/false, "fa f32_f16 c32 g8 d256 NSG2", backend_ctx->queue);
+ if (prog_g8_c32) {
+ cl_kernel k_g8_c32 = clCreateKernel(prog_g8_c32, "flash_attn_f32_f16_q1_vec_mq_split_c8", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_g8_c32, 128,
+ "flash_attn_f32_f16_q1_vec_mq_split_c8 (g8 c32 d256)", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c32[{dk, dv}] = k_g8_c32;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_g8_c32, "flash_attn_f32_f16_q1_vec_mq_split_g8_c32", dk, dv);
+ } else {
+ clReleaseKernel(k_g8_c32);
+ }
+ }
+ clReleaseProgram(prog_g8_c32);
+ }
+ }
+ if (!fa_decode_only && backend_ctx->has_subgroup_shuffle &&
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8.count({dk, dv}) == 0) {
+ const std::string opts_g8_c8_ns2 = opts + " -D FA_MQ_ONLY -D MQ_GQA=8 -D MQ_NSG=2 -D MQ_NSG_SPLIT=2 -D FA_CL_C=" + fa_cl_c_g8_val;
+ cl_program prog_g8_c8 = build_program_from_source_ex(
+ backend_ctx->context, backend_ctx->device, src.c_str(), opts_g8_c8_ns2,
+ /*fatal=*/false, "fa f32_f16 c8 g8 NSG2", backend_ctx->queue);
+ if (prog_g8_c8) {
+ cl_kernel k_g8_c8 = clCreateKernel(prog_g8_c8, "flash_attn_f32_f16_q1_vec_mq_split_c8", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_g8_c8, 128,
+ "flash_attn_f32_f16_q1_vec_mq_split_c8 (g8 ns2)", dk, dv)) {
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8_ns2[{dk, dv}] = k_g8_c8;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_g8_c8, "flash_attn_f32_f16_q1_vec_mq_split_g8_c8_ns2", dk, dv);
+ } else {
+ clReleaseKernel(k_g8_c8);
+ }
+ }
+ clReleaseProgram(prog_g8_c8);
+ }
+ }
break;
}
case FA_VARIANT_Q8_0:
}
clReleaseProgram(prog_mq_g8);
}
+ // GQA=4 cluster-parallel program (NSG_SPLIT=2 / WG=128)
+ if (backend_ctx->has_subgroup_shuffle) {
+ auto & m_c8_gqa4 = is_q8 ? backend_ctx->fa.f32_q8_0_q1_vec_mq_split_c8
+ : backend_ctx->fa.f32_q4_0_q1_vec_mq_split_c8;
+ const std::string name_c8_gqa4 = name_q1 + "_vec_mq_split_c8";
+ const std::string opts_c8_gqa4 = opts + " -D MQ_GQA=4 -D MQ_NSG=2 -D MQ_NSG_SPLIT=2" + opts_cl_c_gqa4;
+ cl_program prog_c8_gqa4 = build_program_from_source_ex(
+ backend_ctx->context, backend_ctx->device, src.c_str(), opts_c8_gqa4,
+ /*fatal=*/false, is_q8 ? "fa q8_0 c8 GQA4 NSG2" : "fa q4_0 c8 GQA4 NSG2",
+ backend_ctx->queue);
+ if (prog_c8_gqa4) {
+ cl_kernel k_c8_gqa4 = clCreateKernel(prog_c8_gqa4, name_c8_gqa4.c_str(), &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_c8_gqa4, 128,
+ name_c8_gqa4.c_str(), dk, dv)) {
+ m_c8_gqa4[{dk, dv}] = k_c8_gqa4;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_c8_gqa4, name_c8_gqa4.c_str(), dk, dv);
+ } else {
+ clReleaseKernel(k_c8_gqa4);
+ }
+ }
+ clReleaseProgram(prog_c8_gqa4);
+ }
+ }
+ // Cluster-parallel q4_0 decode kernel
+ if (!is_q8 && backend_ctx->has_subgroup_shuffle) {
+ const std::string opts_c8 = opts + " -D MQ_GQA=8 -D MQ_NSG=2 -D MQ_NSG_SPLIT=2";
+ cl_program prog_c8 = build_program_from_source_ex(
+ backend_ctx->context, backend_ctx->device, src.c_str(), opts_c8,
+ /*fatal=*/false, "fa q4_0 c8 NSG2", backend_ctx->queue);
+ if (prog_c8) {
+ cl_kernel k_c8 = clCreateKernel(prog_c8, "flash_attn_f32_q4_0_q1_vec_mq_split_c8", &err);
+ if (err == CL_SUCCESS) {
+ if (ggml_opencl_fa_kernel_fits_wg(backend_ctx, k_c8, 128,
+ "flash_attn_f32_q4_0_q1_vec_mq_split_c8 (g8 ns2)", dk, dv)) {
+ backend_ctx->fa.f32_q4_0_q1_vec_mq_split_g8_c8[{dk, dv}] = k_c8;
+ ggml_opencl_log_fa_kernel_spill(backend_ctx, k_c8, "flash_attn_f32_q4_0_q1_vec_mq_split_g8_c8", dk, dv);
+ } else {
+ clReleaseKernel(k_c8);
+ }
+ }
+ clReleaseProgram(prog_c8);
+ }
+ }
break;
}
case FA_VARIANT_F32_F16_SPLIT: {
const bool nq_in_vec_range = (n_q >= 1) && (n_q <= N_MAX_VEC_NQ);
const bool nq1_only = (n_q == 1);
+
+ // Cluster-parallel decode default on for Adreno X2E/X1E
+ static const int c8_env_state = []{
+ const char * e = getenv("GGML_OPENCL_FA_C8");
+ if (e == NULL || e[0] == '\0') { return -1; }
+ return (e[0] != '0') ? 1 : 0;
+ }();
+ const bool c8_default_on = backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E ||
+ backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E;
+ const bool c8_f16_on = (c8_env_state >= 0) ? (c8_env_state == 1) : c8_default_on;
+ // Quant-KV (q4_0/q8_0) GQA4 c8: default-on X2E + X1E
+ const bool c8_quant_on = (c8_env_state >= 0) ? (c8_env_state == 1) : c8_default_on;
if (mq_enabled && mq_kv_ok && nq_in_vec_range && !is_causal &&
backend_ctx->gpu_family != INTEL &&
!use_local_tile &&
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) {
+ // Cluster-parallel decode
+ const bool c8_env = d_head_q == 128 && d_head_v == 128 && c8_f16_on;
+ if (c8_env && backend_ctx->fa.f32_f16_q1_vec_mq_split_c8.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_c8.at(dk_dv);
+ use_fd_mq = true;
+ } else if (c8_env && backend_ctx->fa.f32_f16_q1_vec_mq_split_c8_ns2.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_c8_ns2.at(dk_dv);
+ use_fd_mq = true;
+ fd_mq_wg = 128;
+ } else 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;
fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split.at(dk_dv);
use_fd_mq = true;
}
+ // Cluster-parallel decode, DK=DV=256 GQA=8
+ } else if (nq1_only && is_mixed && gqa_ratio_dispatch == 8 &&
+ d_head_q == 256 && d_head_v == 256 &&
+ c8_env_state == 1 &&
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c32.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c32.at(dk_dv);
+ use_fd_mq = true;
+ fd_mq_wg = 128;
+ // Cluster-parallel decode for the g8
+ } else if (is_mixed && gqa_ratio_dispatch == 8 &&
+ d_head_q == 128 && d_head_v == 128 &&
+ c8_f16_on &&
+ (backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8.count(dk_dv) > 0 ||
+ backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8_ns2.count(dk_dv) > 0)) {
+ if (backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8.at(dk_dv);
+ fd_mq_wg = 192;
+ } else {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_g8_c8_ns2.at(dk_dv);
+ fd_mq_wg = 128;
+ }
+ 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 &&
} 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);
+ // Cluster-parallel q8_0 GQA4
+ if (c8_quant_on &&
+ backend_ctx->fa.f32_q8_0_q1_vec_mq_split_c8.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_q8_0_q1_vec_mq_split_c8.at(dk_dv);
+ fd_mq_wg = 128;
+ } else {
+ 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 &&
+ // Cluster-parallel q4_0
+ const bool q4_c8_on = c8_env_state == 1 &&
+ backend_ctx->fa.f32_q4_0_q1_vec_mq_split_g8_c8.count(dk_dv) > 0;
+ if (q4_c8_on && gqa_ratio_dispatch == 8 &&
+ d_head_q == 64 && d_head_v == 64) {
+ fd_k_split = backend_ctx->fa.f32_q4_0_q1_vec_mq_split_g8_c8.at(dk_dv);
+ use_fd_mq = true;
+ fd_mq_wg = 128;
+ } else 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);
} 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);
+ // Cluster-parallel q4_0 GQA4
+ if (c8_quant_on &&
+ backend_ctx->fa.f32_q4_0_q1_vec_mq_split_c8.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_q4_0_q1_vec_mq_split_c8.at(dk_dv);
+ fd_mq_wg = 128;
+ } else {
+ fd_k_split = backend_ctx->fa.f32_q4_0_q1_vec_mq_split.at(dk_dv);
+ }
use_fd_mq = true;
}
}
}
}
+ // Intel cluster-parallel decode FA
+ if (fd_k_split == NULL && backend_ctx->gpu_family == INTEL && n_q == 1 && !is_causal &&
+ is_mixed && gqa_ratio_dispatch == 4 && d_head_q == 128 && d_head_v == 128 &&
+ n_kv >= FD_MIN_N_KV &&
+ getenv("GGML_OPENCL_FA_C8") != NULL && getenv("GGML_OPENCL_FA_C8")[0] != '0' &&
+ backend_ctx->fa.f32_merge.count(dk_dv) > 0) {
+ if (backend_ctx->fa.f32_f16_q1_vec_mq_split_c8.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_c8.at(dk_dv);
+ use_fd_mq = true;
+ fd_mq_wg = 128;
+ } else if (backend_ctx->fa.f32_f16_q1_vec_mq_split_c8_ns2.count(dk_dv) > 0) {
+ fd_k_split = backend_ctx->fa.f32_f16_q1_vec_mq_split_c8_ns2.at(dk_dv);
+ use_fd_mq = true;
+ fd_mq_wg = 64;
+ }
+ }
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 &&
}
}
+// Cluster-parallel variant of _q1_vec_mq_split
+//
+// Tthe baseline keeps one 256B K row in flight per subgroup (32 lanes cooperate
+// on one position, serialized by the reduce+exp chain). This kernel
+// takes q1_split's memory-level parallelism at MQ's read-once traffic:
+// - the 64-lane subgroup is split into FA_CL_NCL clusters of FA_CL_C lanes;
+// - each cluster owns its own KV position stream (positions strided by
+// FA_CL_NCL) with private per-cluster online-softmax state, hence FA_CL_NCL
+// independent K rows in flight per subgroup, no cross-cluster serial chain;
+// - within a cluster, lanes split DK for the dot (cluster-reduce via
+// sub_group_shuffle_xor, steps < FA_CL_C stay inside the cluster) and
+// split DV for o_acc (each lane owns dv indices {lic + FA_CL_C*i} — the
+// same slice for every position, so accumulation is lane-local);
+// - merge stage 1 folds the FA_CL_NCL cluster partials with cross-cluster
+// shuffles (distances >= FA_CL_C); stage 2 is the baseline cross-subgroup
+// LDS merge (o published by cluster 0's lanes, layout-identical to the
+// baseline's sg_o).
+// The KV sweep runs a UNIFORM trip count (max over clusters) with a clamped
+// row address + FA_M_INIT score on the tail — keeps every shuffle convergent
+// (p = exp(FA_M_INIT - m) underflows to 0, so clamped-row reads are inert).
+// Register cost vs baseline: o_acc grows from DV_VEC/64 to DV_VEC/FA_CL_C
+// float4 per lane per head — FA_CL_C=8 / MQ_GQA=4 => 16 float4 (256B).
+
+#ifdef HAS_SUBGROUP_SHUFFLE // cluster reduce/merge needs shuffles; absent -> kernel dropped, dispatch falls back
+
+#ifndef FA_CL_C
+#define FA_CL_C 8
+#endif
+
+// The lane striping requires DK/DV to divide evenly across the cluster;
+// otherwise (e.g. DK=40 with FA_CL_C=16 -> zero-size arrays) compile the
+// kernel out — host soft-create falls back silently.
+#if (DK_VEC % FA_CL_C) == 0 && (DV_VEC % FA_CL_C) == 0
+#define FA_CL_NCL (Q1_WG_SIZE / FA_CL_C) // clusters (position streams) per subgroup
+#define FA_CL_DK (DK_VEC / FA_CL_C) // half4s of K per lane per row
+#define FA_CL_DV (DV_VEC / FA_CL_C) // float4s of o_acc per lane per head
+
+// explicit "half" sub-group attribute routes this fp16-heavy kernel to a slow
+// codegen path on the X1 compiler. X2 keeps the pin: its driver miscompile
+// without it.
+#ifdef FA_C8_NO_SG_PIN
+#define FA_C8_SG_ATTR
+#else
+// REQD_FA_SG pins the HW subgroup on Intel (intel_reqd_sub_group_size(FA_SG),
+// host passes -D FA_SG=32); empty on Adreno. REQD_SUBGROUP_SIZE_64 pins 64 on
+// Adreno; empty on Intel.
+#define FA_C8_SG_ATTR REQD_FA_SG REQD_SUBGROUP_SIZE_64
+#endif
+
+FA_C8_SG_ATTR
+__kernel void flash_attn_f32_f16_q1_vec_mq_split_c8(
+ 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 cl = tid_sg / FA_CL_C; // cluster id
+ const int lic = tid_sg % FA_CL_C; // lane in cluster
+ 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;
+
+ // 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;
+ }
+
+ // Per-CLUSTER online-softmax state (uniform across the cluster's lanes);
+ // o_acc holds this lane's DV slice {lic + FA_CL_C*i}.
+ ACC_TYPE4 o_acc[MQ_GQA][FA_CL_DV];
+ 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 < FA_CL_DV; ++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);
+
+ // Uniform trip count across the subgroup: every cluster runs n_iter
+ // iterations; tail positions clamp the row address and drop the score to
+ // FA_M_INIT so shuffles stay convergent and the contribution is exactly 0.
+ const int n_iter = (kv_hi - kv_lo + FA_CL_NCL - 1) / FA_CL_NCL;
+ const ulong kv_row_base = batch_idx * k_nb3 + head_kv_idx * k_nb2;
+ const ulong v_row_base = batch_idx * v_nb3 + head_kv_idx * v_nb2;
+
+ for (int it = 0; it < n_iter; ++it) {
+ const int k_idx = kv_lo + cl + it * FA_CL_NCL;
+ const int valid = k_idx < kv_hi;
+ const int k_safe = valid ? k_idx : (kv_hi - 1);
+
+ const global KV_DATA_TYPE4 * k_ptr = (const global KV_DATA_TYPE4 *) (k_base + kv_row_base + (ulong) k_safe * k_nb1);
+ const global KV_DATA_TYPE4 * v_ptr = (const global KV_DATA_TYPE4 *) (v_base + v_row_base + (ulong) k_safe * v_nb1);
+
+ // Dot: this lane covers DK elements {lic + FA_CL_C*i} of the cluster's row.
+ ACC_TYPE4 dot4[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) dot4[h] = (ACC_TYPE4)(0.0f);
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DK; ++i) {
+ const int kk = lic + FA_CL_C * i;
+ const ACC_TYPE4 k_vec = CONVERT_KV_ACC4(k_ptr[kk]);
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ dot4[h] = mad(q_shared[h * DK_VEC + kk], k_vec, dot4[h]);
+ }
+ }
+
+ // Cluster-reduce (xor steps < FA_CL_C stay inside the cluster) + score.
+ ACC_TYPE score[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ ACC_TYPE s = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
+ #pragma unroll
+ for (int step = 1; step < FA_CL_C; step <<= 1) {
+ s += sub_group_shuffle_xor(s, step);
+ }
+ s *= 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_safe];
+ }
+ if (logit_softcap > 0.0f) {
+ s = logit_softcap * tanh(s / logit_softcap);
+ }
+ score[h] = valid ? s : FA_M_INIT;
+ }
+
+ // Per-cluster online update — identical math to the baseline, but the
+ // serial chain is per cluster (depth n_iter, not kv_per_sg).
+ 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;
+ }
+
+ // V accumulate on this lane's DV slice (p = 0 on tail -> inert).
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DV; ++i) {
+ const ACC_TYPE4 v_vec = CONVERT_KV_ACC4(v_ptr[lic + FA_CL_C * i]);
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ o_acc[h][i] = mad(p_h[h], v_vec, o_acc[h][i] * sp_h[h]);
+ }
+ }
+ }
+
+ // Merge stage 1: fold the FA_CL_NCL cluster partials inside the subgroup.
+ // Lanes with equal lic across clusters hold the SAME dv slice, so a
+ // cross-cluster xor-reduce (distances FA_CL_C..Q1_WG_SIZE/2) sums o
+ // slice-wise; m/l fold the same way. All shuffles are subgroup-convergent.
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ ACC_TYPE m_c = m_i[h];
+ #pragma unroll
+ for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
+ m_c = max(m_c, sub_group_shuffle_xor(m_c, step));
+ }
+ const ACC_TYPE alpha = native_exp(m_i[h] - m_c);
+ ACC_TYPE l_c = l_i[h] * alpha;
+ #pragma unroll
+ for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
+ l_c += sub_group_shuffle_xor(l_c, step);
+ }
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DV; ++i) {
+ ACC_TYPE4 o = o_acc[h][i] * alpha;
+ #pragma unroll
+ for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
+ o.s0 += sub_group_shuffle_xor(o.s0, step);
+ o.s1 += sub_group_shuffle_xor(o.s1, step);
+ o.s2 += sub_group_shuffle_xor(o.s2, step);
+ o.s3 += sub_group_shuffle_xor(o.s3, step);
+ }
+ o_acc[h][i] = o;
+ }
+ m_i[h] = m_c;
+ l_i[h] = l_c;
+ }
+
+ // Merge stage 2: baseline cross-subgroup LDS merge. Cluster 0's lanes hold
+ // the subgroup's merged o (dv indices {lic + FA_CL_C*i}) — same sg_o layout
+ // and fold loop as q1_vec_mq_split.
+ __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) {
+ if (cl == 0) {
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DV; ++i) {
+ sg_o[sgid][lic + FA_CL_C * i] = o_acc[h][i];
+ }
+ }
+ 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 // DK_VEC/DV_VEC divisible by FA_CL_C
+#endif // HAS_SUBGROUP_SHUFFLE (q1_vec_mq_split_c8)
+
REQD_SUBGROUP_SIZE_64
__kernel void flash_attn_f32_f16_q1_vec_mq_split_k_img(
const global void * q_void, ulong q_offset,
}
}
+// flash_attn_f32_q4_0_q1_vec_mq_split_c8 — cluster-parallel variant of the MQ
+// split, port of flash_attn_f32_f16_q1_vec_mq_split_c8
+// Requires dp4a + subgroup shuffles
+
+#if defined(FA_HAVE_INT_DOT) && defined(HAS_SUBGROUP_SHUFFLE)
+
+#ifndef FA_CL_C
+#define FA_CL_C 8
+#endif
+
+// Lane striping requires DK/DV to divide across the cluster (see f16 c8).
+#if (DK_VEC % FA_CL_C) == 0 && (DV_VEC % FA_CL_C) == 0
+#define FA_CL_NCL (Q1_WG_SIZE / FA_CL_C) // clusters (position streams) per subgroup
+#define FA_CL_DKQ (DK_VEC / FA_CL_C) // K quartets per lane per row
+#define FA_CL_DVQ (DV_VEC / FA_CL_C) // V quartets (o_acc float4s) per lane per head
+
+#ifdef FA_C8_NO_SG_PIN
+#define FA_C8_SG_ATTR_Q4
+#else
+#define FA_C8_SG_ATTR_Q4 REQD_SUBGROUP_SIZE_64
+#endif
+
+FA_C8_SG_ATTR_Q4
+__kernel void flash_attn_f32_q4_0_q1_vec_mq_split_c8(
+ 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 cl = tid_sg / FA_CL_C; // cluster id
+ const int lic = tid_sg % FA_CL_C; // lane in cluster
+ 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;
+
+ // Stage MQ_GQA Q rows in __local as float4 (source for the quantize pass).
+ __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);
+
+ // Per-(h, block) int8-packed Q + (qd, q_sum), quantized once per WG.
+ __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);
+
+ 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;
+ }
+
+ // Per-CLUSTER online state; o_acc holds this lane's V quartets {lic + FA_CL_C*i}.
+ ACC_TYPE4 o_acc[MQ_GQA][FA_CL_DVQ];
+ 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 < FA_CL_DVQ; ++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);
+
+ // Uniform trip count; tail clamps the row address and drops the score to
+ // FA_M_INIT (p underflows to 0) so shuffles stay convergent.
+ const int n_iter = (kv_hi - kv_lo + FA_CL_NCL - 1) / FA_CL_NCL;
+ const ulong k_row_base = batch_idx * k_nb3 + head_kv_idx * k_nb2;
+ const ulong v_row_base = batch_idx * v_nb3 + head_kv_idx * v_nb2;
+
+ for (int it = 0; it < n_iter; ++it) {
+ const int k_idx = kv_lo + cl + it * FA_CL_NCL;
+ const int valid = k_idx < kv_hi;
+ const int k_safe = valid ? k_idx : (kv_hi - 1);
+
+ const global char * k_row = k_base + k_row_base + (ulong) k_safe * k_nb1;
+ const global char * v_row = v_base + v_row_base + (ulong) k_safe * v_nb1;
+
+ // dp4a K dot over this lane's quartets of the cluster's row.
+ ACC_TYPE lane_contrib[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) lane_contrib[h] = 0.0f;
+
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DKQ; ++i) {
+ const int qk = lic + FA_CL_C * i;
+ 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;
+ }
+ }
+
+ // Cluster-reduce + score.
+ ACC_TYPE score[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ ACC_TYPE s = lane_contrib[h];
+ #pragma unroll
+ for (int step = 1; step < FA_CL_C; step <<= 1) {
+ s += sub_group_shuffle_xor(s, step);
+ }
+ s *= 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_safe];
+ }
+ if (logit_softcap > 0.0f) {
+ s = logit_softcap * tanh(s / logit_softcap);
+ }
+ score[h] = valid ? s : FA_M_INIT;
+ }
+
+ // Per-cluster online update (serial chain depth n_iter, not kv_per_sg).
+ 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;
+ }
+
+ // V accumulate on this lane's quartets (p = 0 on tail -> inert).
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DVQ; ++i) {
+ const int dv = lic + FA_CL_C * i;
+ const float4 v_v = dequant_q4_0_lane(v_row + (dv / 8) * Q4_0_BLOCK_SIZE, dv % 8);
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ o_acc[h][i] = mad(p_h[h], v_v, o_acc[h][i] * sp_h[h]);
+ }
+ }
+ }
+
+ // Merge stage 1: fold cluster partials inside the subgroup via shuffles.
+ // Lanes with equal lic across clusters hold the SAME dv slice.
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ ACC_TYPE m_c = m_i[h];
+ #pragma unroll
+ for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
+ m_c = max(m_c, sub_group_shuffle_xor(m_c, step));
+ }
+ const ACC_TYPE alpha = native_exp(m_i[h] - m_c);
+ ACC_TYPE l_c = l_i[h] * alpha;
+ #pragma unroll
+ for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
+ l_c += sub_group_shuffle_xor(l_c, step);
+ }
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DVQ; ++i) {
+ ACC_TYPE4 o = o_acc[h][i] * alpha;
+ #pragma unroll
+ for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
+ o.s0 += sub_group_shuffle_xor(o.s0, step);
+ o.s1 += sub_group_shuffle_xor(o.s1, step);
+ o.s2 += sub_group_shuffle_xor(o.s2, step);
+ o.s3 += sub_group_shuffle_xor(o.s3, step);
+ }
+ o_acc[h][i] = o;
+ }
+ m_i[h] = m_c;
+ l_i[h] = l_c;
+ }
+
+ // Merge stage 2: baseline cross-subgroup LDS merge (o published by
+ // cluster 0's lanes; layout identical to the baseline sg_o).
+ __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) {
+ if (cl == 0) {
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DVQ; ++i) {
+ sg_o[sgid][lic + FA_CL_C * i] = o_acc[h][i];
+ }
+ }
+ 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 // DK_VEC/DV_VEC divisible by FA_CL_C
+#endif // FA_HAVE_INT_DOT && HAS_SUBGROUP_SHUFFLE (q1_vec_mq_split_c8)
+
__kernel void flash_attn_f32_q4_0(
const global void * q_void, ulong q_offset,
const global void * k_void, ulong k_offset,
}
}
+// flash_attn_f32_q8_0_q1_vec_mq_split_c8 — cluster-parallel variant of the MQ
+// split above, port of the f16/q4_0 c8 kernels
+
+#ifdef HAS_SUBGROUP_SHUFFLE
+
+#ifndef FA_CL_C
+#define FA_CL_C 8
+#endif
+
+// Lane striping requires DK/DV to divide across the cluster (see f16 c8).
+#if (DK_VEC % FA_CL_C) == 0 && (DV_VEC % FA_CL_C) == 0
+#define FA_CL_NCL (Q1_WG_SIZE / FA_CL_C) // clusters (position streams) per subgroup
+#define FA_CL_DKQ (DK_VEC / FA_CL_C) // K quartets per lane per row
+#define FA_CL_DVQ (DV_VEC / FA_CL_C) // V quartets (o_acc float4s) per lane per head
+
+#ifdef FA_C8_NO_SG_PIN
+#define FA_C8_SG_ATTR_Q8
+#else
+#define FA_C8_SG_ATTR_Q8 REQD_SUBGROUP_SIZE_64
+#endif
+
+FA_C8_SG_ATTR_Q8
+__kernel void flash_attn_f32_q8_0_q1_vec_mq_split_c8(
+ 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 cl = tid_sg / FA_CL_C; // cluster id
+ const int lic = tid_sg % FA_CL_C; // lane in cluster
+ 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;
+
+ // 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_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;
+ }
+
+ // Per-CLUSTER online state; o_acc holds this lane's V quartets {lic + FA_CL_C*i}.
+ ACC_TYPE4 o_acc[MQ_GQA][FA_CL_DVQ];
+ 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 < FA_CL_DVQ; ++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);
+
+ // Uniform trip count; tail clamps the row address and drops the score to
+ // FA_M_INIT (p underflows to 0) so shuffles stay convergent.
+ const int n_iter = (kv_hi - kv_lo + FA_CL_NCL - 1) / FA_CL_NCL;
+ const ulong k_row_base = batch_idx * k_nb3 + head_kv_idx * k_nb2;
+ const ulong v_row_base = batch_idx * v_nb3 + head_kv_idx * v_nb2;
+
+ for (int it = 0; it < n_iter; ++it) {
+ const int k_idx = kv_lo + cl + it * FA_CL_NCL;
+ const int valid = k_idx < kv_hi;
+ const int k_safe = valid ? k_idx : (kv_hi - 1);
+
+ const global char * k_row = k_base + k_row_base + (ulong) k_safe * k_nb1;
+ const global char * v_row = v_base + v_row_base + (ulong) k_safe * v_nb1;
+
+ // Float-dequant K dot over this lane's quartets of the cluster's row.
+ ACC_TYPE4 dot4[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) dot4[h] = (ACC_TYPE4)(0.0f);
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DKQ; ++i) {
+ const int qk = lic + FA_CL_C * i;
+ const float4 k_v = dequant_q8_0_lane(k_row + (qk / 8) * Q8_0_BLOCK_SIZE, qk % 8);
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ dot4[h] = mad(q_shared[h * DK_VEC + qk], k_v, dot4[h]);
+ }
+ }
+
+ // Cluster-reduce (xor steps < FA_CL_C stay inside the cluster) + score.
+ ACC_TYPE score[MQ_GQA];
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ ACC_TYPE s = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3;
+ #pragma unroll
+ for (int step = 1; step < FA_CL_C; step <<= 1) {
+ s += sub_group_shuffle_xor(s, step);
+ }
+ s *= 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_safe];
+ }
+ if (logit_softcap > 0.0f) {
+ s = logit_softcap * tanh(s / logit_softcap);
+ }
+ score[h] = valid ? s : FA_M_INIT;
+ }
+
+ // Per-cluster online update (serial chain depth n_iter, not kv_per_sg).
+ 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;
+ }
+
+ // V accumulate on this lane's quartets (p = 0 on tail -> inert).
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DVQ; ++i) {
+ const int dv = lic + FA_CL_C * i;
+ const float4 v_v = dequant_q8_0_lane(v_row + (dv / 8) * Q8_0_BLOCK_SIZE, dv % 8);
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ o_acc[h][i] = mad(p_h[h], v_v, o_acc[h][i] * sp_h[h]);
+ }
+ }
+ }
+
+ // Merge stage 1: fold cluster partials inside the subgroup via shuffles.
+ #pragma unroll
+ for (int h = 0; h < MQ_GQA; ++h) {
+ ACC_TYPE m_c = m_i[h];
+ #pragma unroll
+ for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
+ m_c = max(m_c, sub_group_shuffle_xor(m_c, step));
+ }
+ const ACC_TYPE alpha = native_exp(m_i[h] - m_c);
+ ACC_TYPE l_c = l_i[h] * alpha;
+ #pragma unroll
+ for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
+ l_c += sub_group_shuffle_xor(l_c, step);
+ }
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DVQ; ++i) {
+ ACC_TYPE4 o = o_acc[h][i] * alpha;
+ #pragma unroll
+ for (int step = FA_CL_C; step < Q1_WG_SIZE; step <<= 1) {
+ o.s0 += sub_group_shuffle_xor(o.s0, step);
+ o.s1 += sub_group_shuffle_xor(o.s1, step);
+ o.s2 += sub_group_shuffle_xor(o.s2, step);
+ o.s3 += sub_group_shuffle_xor(o.s3, step);
+ }
+ o_acc[h][i] = o;
+ }
+ m_i[h] = m_c;
+ l_i[h] = l_c;
+ }
+
+ // Merge stage 2: baseline cross-subgroup LDS merge (o published by
+ // cluster 0's lanes; layout identical to the baseline sg_o).
+ __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) {
+ if (cl == 0) {
+ #pragma unroll
+ for (int i = 0; i < FA_CL_DVQ; ++i) {
+ sg_o[sgid][lic + FA_CL_C * i] = o_acc[h][i];
+ }
+ }
+ 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 // DK_VEC/DV_VEC divisible by FA_CL_C
+#endif // HAS_SUBGROUP_SHUFFLE (q1_vec_mq_split_c8)
+
__kernel void flash_attn_f32_q8_0(
const global void * q_void, ulong q_offset,
const global void * k_void, ulong k_offset,
overview / pkg / ggml / sources / llama.cpp / commitdiff