Commit a963f6595 for llama.cpp
commit a963f659578ba9437b53066ee89e3ff16c2dc84a
Author: sliu39 <russell.liu@intel.com>
Date: Thu Nov 12 18:19:17 2026 +0800
vulkan : Intel FA kernel for prefill. (#29357)
* vulkan : Intel customized FA shader for prefill usage
* vulkan : adding check of max_bias/logit_softcap for Intel FA kernel path
* vulkan : add m16n16k16/SIMD32 support for Xe FA prefill shader
* vulkan : Intel FA shader and host refinement, with accuracy issues fixes
* vulkan : Limit the SLM usage to 48 KiB for Intel prefill FA shader
* vulkan : Update Intel FA shaders to handle cases with large head dimension K, unaligned Q head stride and multi head mask cases
* vulkan : Intel FA shader and host refinement to align with existing design, extreme cases handling, shape check relaxing.
* vulkan : Intel FA shader and host update to support q as tensor view, unaligned batch stride q. Add env variable for enforcing XE FA path for other vendor GPUs with 16x16x16 CM1 and SIMD32 support
* vulkan : Inte FA shader and host update to fix the potential SLM over 48 KiB issue, enhance the buffer access robust and add asymmetric qk/pv ration support.
* vulkan : Inte FA shader and host update to add SLM 48 KiB size check, remove unnecessary extensions and finetune the unroll/dont_unroll loops
* vulkan : Intel FA shader and host update to add SLM size check function, use real device max SLM size in shader specialization constants.
* vulkan : Intel FA shader and host update to consolidate host helper function, remove redundant buffer sync, bypass unused pipeline and buffer, fix subgroup shmem conflict and add negative case handling.
* vulkan : Intel FA host update to correct get_xe_fa_tuning_params() function log message, and add workgroupMemoryExplicitLayout/workgroupMemoryExplicitLayout16BitAccess prerequisite feature check.
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h b/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h
index 0b3ebe9b8..e78f87b1b 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h
+++ b/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h
@@ -124,22 +124,30 @@ struct vk_flash_attn_push_constants {
static_assert(sizeof(vk_flash_attn_push_constants) <= 128, "sizeof(vk_flash_attn_push_constants) must be <= 128");
-struct vk_fa_xe_opt_push_constants {
- uint32_t kv_seq_len;
- uint32_t activation_length;
- uint32_t q_head;
- uint32_t kv_head;
- uint32_t qk_ratio;
- uint32_t qk_sub_groups;
+struct vk_xe_fa_push_constants {
+ uint32_t neq1;
+ uint32_t neq2;
+ uint32_t neq3;
+ uint32_t nek1;
+ uint32_t nek2;
+ uint32_t nek3;
+ uint32_t nev2;
+ uint32_t nev3;
+ uint32_t nem1;
+ uint32_t nem2;
+ uint32_t nem3;
+ uint32_t nbq1;
+ uint32_t nbq2;
+ uint32_t nbk1;
+ uint32_t nbk2;
+ uint32_t nbv1;
+ uint32_t nbv2;
+ uint32_t nbq3;
+ uint32_t nbk3;
+ uint32_t nbv3;
+ uint32_t nb3;
uint32_t flag;
- uint32_t nbkv_tok;
- uint32_t nbkv_head;
- uint32_t batch_stride_q;
- uint32_t batch_stride_k;
- uint32_t batch_stride_v;
- uint32_t batch_stride_m;
- uint32_t batch_stride_o;
- float softmax_scale;
+ float scale;
};
struct vk_op_push_constants {
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-types.h b/ggml/src/ggml-vulkan/ggml-vulkan-types.h
index bf2c02d4d..02ccb8ba0 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan-types.h
+++ b/ggml/src/ggml-vulkan/ggml-vulkan-types.h
@@ -433,6 +433,21 @@ struct vk_fa_pipeline_state {
}
};
+struct vk_xe_fa_pipeline_state {
+ uint32_t HSK, HSV;
+ uint32_t qk_ratio, pv_ratio;
+ uint32_t decode_n;
+ bool unaligned_seq_len;
+ bool unaligned_q;
+ bool partial_row_mask;
+ bool multi_head_mask;
+
+ bool operator<(const vk_xe_fa_pipeline_state &b) const {
+ return std::tie(HSK, HSV, qk_ratio, pv_ratio, decode_n, unaligned_seq_len, unaligned_q, partial_row_mask, multi_head_mask) <
+ std::tie(b.HSK, b.HSV, b.qk_ratio, b.pv_ratio, b.decode_n, b.unaligned_seq_len, b.unaligned_q, b.partial_row_mask, b.multi_head_mask);
+ }
+};
+
struct vk_conv2d_pipeline_state {
vk_conv2d_pipeline_state(uint32_t s0, uint32_t s1, uint32_t p0, uint32_t p1, uint32_t d0, uint32_t d1, uint32_t KW, uint32_t KH, uint32_t aligned)
: s0(s0), s1(s1), p0(p0), p1(p1), d0(d0), d1(d1), KW(KW), KH(KH), aligned(aligned) {}
@@ -763,6 +778,9 @@ struct vk_device_struct {
bool dot2_f16 {};
bool ocp_fp4 {};
+ bool workgroup_memory_explicit_layout {};
+ bool workgroup_memory_explicit_layout_16bit {};
+
bool pipeline_executable_properties_support {};
bool device_fault {};
@@ -1006,7 +1024,8 @@ struct vk_device_struct {
bool fa_sparse_compact_use_subgroups;
vk_pipeline pipeline_flash_attn_split_k_reduce;
- std::map<std::tuple<uint32_t, uint32_t, uint32_t, uint32_t>, std::pair<vk_pipeline, vk_pipeline>> pipeline_xe_fa_decode_dual_phases;
+ std::map<vk_xe_fa_pipeline_state, std::pair<vk_pipeline, vk_pipeline>> pipeline_xe_fa_decode_dual_phases;
+ std::map<vk_xe_fa_pipeline_state, vk_pipeline> pipeline_xe_fa_prefill;
vk_pipeline pipeline_count_experts;
// [2] is for whether to take n_experts from spec constant (0) or push constant (1)
@@ -1028,6 +1047,7 @@ struct vk_device_struct {
bool disable_host_visible_vidmem;
bool allow_sysmem_fallback;
bool disable_graph_optimize;
+ bool xe_fa_support {};
std::unique_ptr<vk_memory_logger> memory_logger;
@@ -1407,6 +1427,34 @@ struct vk_fa_tuning_params {
}
};
+struct vk_xe_fa_tuning_params {
+ size_t shmem_sz;
+ std::array<uint32_t, 2> workgroup_size;
+ std::array<uint32_t, 2> subgroup_size;
+ uint32_t tm;
+ uint32_t tn;
+ uint32_t tk;
+
+ struct {
+ uint32_t wg_row;
+ uint32_t gqa_fold;
+ uint32_t q_in_reg;
+ uint32_t p_per_iter;
+ uint32_t sg_col;
+ uint32_t acc_count;
+ uint32_t pipeline_flag;
+ } prefill_params;
+
+ struct {
+ uint32_t p_chunk_ph1;
+ uint32_t out_per_wg_ph1;
+ uint32_t split_p_per_iter_ph2;
+ uint32_t out_row_per_wg_ph2;
+ uint32_t out_col_per_wg_ph2;
+ uint32_t pipeline_flag_ph2;
+ } decode_params;
+};
+
struct GpuPipelineConfig {
// GPU architecture identifier.
// Example: vk_device_architecture::AMD_GCN
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
index e1de84191..041497d6a 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp
+++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
@@ -1434,6 +1434,186 @@ static std::vector<uint32_t> get_fa_spec_constants(const vk_fa_pipeline_state& s
};
}
+static vk_xe_fa_tuning_params get_xe_fa_tuning_params(const vk_device_struct * device, const vk_xe_fa_pipeline_state& pipeline_state, bool use_prefill) {
+ vk_xe_fa_tuning_params params {};
+ size_t shmem_sz = 0;
+ uint32_t sg_sz = 32u;
+ std::array<uint32_t, 3> matrix_mnk = { 16u, 16u, 16u };
+ bool intel_gpu = device->vendor_id == VK_VENDOR_ID_INTEL;
+
+ if (device->architecture == INTEL_XE1) {
+ sg_sz = 8u;
+ matrix_mnk = { 8u, 8u, 16u };
+ } else if (device->architecture == INTEL_XE2) {
+ sg_sz = 16u;
+ matrix_mnk = { 8u, 16u, 16u };
+ }
+
+ uint32_t tm = matrix_mnk[0];
+ uint32_t tn = matrix_mnk[1];
+ uint32_t tk = matrix_mnk[2];
+ uint32_t hsk = pipeline_state.HSK;
+ uint32_t hsv = pipeline_state.HSV;
+ uint32_t qk_ratio = pipeline_state.qk_ratio;
+ uint32_t pv_ratio = pipeline_state.pv_ratio;
+ uint32_t sg_per_gw;
+
+ params.tm = tm;
+ params.tn = tn;
+ params.tk = tk;
+ GGML_ASSERT(hsv % tm == 0);
+
+ if (use_prefill) {
+ sg_per_gw = sg_sz == 8 ? 32 : 16;
+ uint32_t wg_sz = sg_per_gw * sg_sz;
+ uint32_t output_div = hsk > 256 ? 8 : hsk > 160 ? 4 : hsk > 64 ? 2 : 1;
+ uint32_t p_per_iter = intel_gpu ? (hsk > 256 ? 16 : 32) : 16;
+ uint32_t q_in_reg = hsk > 512 ? hsk - 384 : std::min(128u, hsk);
+ uint32_t q_in_shmem = hsk - q_in_reg;
+ uint32_t wg_row = 256u / output_div;
+ uint32_t gqa_fold = (pipeline_state.multi_head_mask || qk_ratio != pv_ratio) ? 1 : 16;
+ uint32_t max_fold_rows = sg_sz <= 16 ? 16 : 8;
+
+ while (gqa_fold > 1 && (qk_ratio % gqa_fold != 0 || gqa_fold * output_div > max_fold_rows)) {
+ gqa_fold = gqa_fold / 2;
+ }
+
+ uint32_t mask_row = wg_row / gqa_fold;
+ uint32_t k_col = (hsk % 64 == 0 && hsk <= 512 && wg_sz <= 256) ? 8 : 4;
+ uint32_t k_loops = CEIL_DIV(hsk * p_per_iter, k_col * wg_sz);
+
+ bool k_in_shmem = hsk <= 576u;
+ bool v_in_shmem = hsk <= 128u && hsv <= 128u;
+ bool emf_in_shmem = wg_row / matrix_mnk[1] != sg_per_gw;
+ bool mask_in_shmem = mask_row <= 64;
+ size_t shmem_q_sz = q_in_shmem * wg_row * sizeof(ggml_fp16_t);
+ size_t shmem_oob_staging_sz = sg_per_gw * tm * tk * sizeof(ggml_fp16_t);
+ size_t shmem_mask_sz = mask_in_shmem ? (mask_row * p_per_iter * sizeof(ggml_fp16_t)) : 0;
+ size_t shmem_k_sz = k_in_shmem ? (k_loops * k_col * wg_sz * sizeof(ggml_fp16_t)) : shmem_oob_staging_sz;
+ size_t shmem_v_sz = v_in_shmem ? (hsv * p_per_iter * sizeof(ggml_fp16_t)) : 0;
+ size_t shmem_p_sz = wg_row * p_per_iter * sizeof(float);
+ size_t shmem_emf_sz = emf_in_shmem ? (wg_row * sizeof(float)) : 0;
+ size_t shmem_pv_sz = intel_gpu ? 0 : (sg_per_gw * tm * tn * sizeof(float));
+
+ // Reuse section sits at tail of shmem, reused with k + mask/kv oob staging/pv staging
+ size_t shmem_reuse_sz = shmem_k_sz + shmem_mask_sz;
+ // When v is not in shmem, its staging reuses the k + mask region; size the reuse section to the larger of the two
+ shmem_reuse_sz = v_in_shmem ? shmem_reuse_sz : std::max(shmem_reuse_sz, shmem_oob_staging_sz);
+ shmem_reuse_sz = std::max(shmem_pv_sz, shmem_reuse_sz);
+
+ // Q shmem size check for prologue staging is safe to bypass:
+ // 1. Prologue q staging shmem is released right after prologue, so it never overlaps other shmem sections
+ // 2. Extreme large hsk one-time q staging is capped within shader (max_q_reg_staging = shmem_size / 2 / Wr)
+ shmem_sz = shmem_q_sz + shmem_reuse_sz + shmem_v_sz + shmem_p_sz + shmem_emf_sz;
+
+ uint32_t sg_col = hsv / tm;
+ uint32_t pipeline_flag = 0;
+ sg_col = (sg_col + output_div - 1) / output_div;
+ sg_col = sg_col * tm;
+ if (pipeline_state.unaligned_seq_len) {
+ pipeline_flag |= 1 << 0;
+ }
+
+ if (pipeline_state.unaligned_q) {
+ pipeline_flag |= 1 << 1;
+ }
+
+ if (pipeline_state.partial_row_mask) {
+ pipeline_flag |= 1 << 2;
+ }
+
+ if (intel_gpu) {
+ pipeline_flag |= 1 << 3;
+ }
+
+ params.workgroup_size = { wg_sz };
+ params.subgroup_size = { sg_sz };
+ params.shmem_sz = shmem_sz;
+ params.prefill_params.pipeline_flag = pipeline_flag;
+ params.prefill_params.wg_row = wg_row;
+ params.prefill_params.gqa_fold = gqa_fold;
+ params.prefill_params.q_in_reg = q_in_reg;
+ params.prefill_params.p_per_iter = p_per_iter;
+ params.prefill_params.sg_col = sg_col;
+ params.prefill_params.acc_count = tm == 8 ? (sg_col % 16 == 0 ? 2 : 1) : 1;
+ } else {
+ auto upper_power_of_2 = [&](uint32_t in) {
+ GGML_ASSERT(in != 0);
+ if (in <= 1) return 1u;
+ uint32_t ret = in - 1;
+ ret |= ret >> 1;
+ ret |= ret >> 2;
+ ret |= ret >> 4;
+ ret |= ret >> 8;
+ ret |= ret >> 16;
+ return ret + 1;
+ };
+
+ size_t shmem_ph1_sz = 0;
+ size_t shmem_ph2_sz = 0;
+
+ // Phase 1 shmem size calculation
+ const uint32_t p_chunk_ph1 = 32;
+ // Phase 1 k shmem staging for extreme large hsk is capped within shader, bypass that check
+ uint32_t pmat_row = CEIL_DIV(qk_ratio, tn);
+ uint32_t aligned_q_len = upper_power_of_2(pipeline_state.decode_n);
+ uint32_t group_sz_ph1 = std::min(std::max(aligned_q_len * sg_sz, 64u), 256u);
+ uint32_t out_per_wg_ph1 = std::min(pipeline_state.decode_n, 256u / sg_sz);
+ size_t shmem_k_sz = group_sz_ph1 / sg_sz * tm * tk * sizeof(ggml_fp16_t);
+ size_t shmem_mask_sz = out_per_wg_ph1 * p_chunk_ph1 * sizeof(float);
+ size_t shmem_p_sz = p_chunk_ph1 * pmat_row * tn * sizeof(float);
+ // K shmem is reused with mask/p, calculate the max size
+ shmem_ph1_sz = std::max(shmem_k_sz, shmem_mask_sz + shmem_p_sz);
+
+ // Phase 2 shmem size calculation
+ uint32_t aligned_pv_ratio = upper_power_of_2(pv_ratio);
+ uint32_t split_p_per_iter_ph2 = (pv_ratio <= 32) ? 256 : 128;
+ uint32_t out_col_per_wg_ph2 = intel_gpu ? (pv_ratio > 16 ? 8 : 16) : 16;
+ uint32_t split_p_per_warp = 16;
+ uint32_t group_sz_ph2 = (split_p_per_iter_ph2 / split_p_per_warp) * sg_sz;
+ sg_per_gw = group_sz_ph2 / sg_sz;
+ uint32_t out_row_per_wg_ph2 = std::min(std::max(16u / aligned_pv_ratio, 1u), pipeline_state.decode_n);
+ uint32_t o_row = pv_ratio * out_row_per_wg_ph2;
+ uint32_t omat_row = CEIL_DIV(o_row, tn) * tn;
+ uint32_t global_max_loops = CEIL_DIV(o_row, sg_per_gw);
+ size_t shmem_global_max_sz = CEIL_DIV(global_max_loops * sg_per_gw, sg_sz) * sg_sz * sizeof(float);
+
+ uint32_t size_max = o_row * split_p_per_iter_ph2 / p_chunk_ph1;
+ size_t shmem_group_max_sz = CEIL_DIV(size_max, group_sz_ph2) * group_sz_ph2 * sizeof(float);
+
+ uint32_t p_loop_col = group_sz_ph2 * 4;
+ uint32_t p_loops = CEIL_DIV(o_row * split_p_per_iter_ph2, p_loop_col);
+ shmem_p_sz = std::max(p_loop_col * p_loops * sizeof(ggml_fp16_t), omat_row * split_p_per_iter_ph2 * sizeof(ggml_fp16_t));
+
+ size_t shmem_lf_sz = CEIL_DIV(sg_per_gw * p_loops, sg_sz) * sg_sz * sizeof(float);
+ size_t shmem_oob_staging_sz = sg_per_gw * tm * tk * sizeof(ggml_fp16_t);
+ size_t shmem_o_sz = sg_per_gw * out_col_per_wg_ph2 * omat_row * sizeof(float);
+ size_t shmem_reuse_sz = 2 * shmem_group_max_sz + std::max(shmem_oob_staging_sz, 2 * shmem_p_sz);
+
+ shmem_ph2_sz = shmem_lf_sz + shmem_global_max_sz;
+ // Shmem of p/group max is reused with output accumulation, calculate the max size
+ shmem_ph2_sz += std::max(shmem_reuse_sz, shmem_o_sz);
+ shmem_sz = std::max(shmem_ph1_sz, shmem_ph2_sz);
+
+ params.workgroup_size = { group_sz_ph1, group_sz_ph2 };
+ params.subgroup_size = { sg_sz, sg_sz };
+ params.shmem_sz = shmem_sz;
+ params.decode_params.p_chunk_ph1 = p_chunk_ph1;
+ params.decode_params.out_per_wg_ph1 = out_per_wg_ph1;
+ params.decode_params.split_p_per_iter_ph2 = split_p_per_iter_ph2;
+ params.decode_params.out_row_per_wg_ph2 = out_row_per_wg_ph2;
+ params.decode_params.out_col_per_wg_ph2 = out_col_per_wg_ph2;
+ params.decode_params.pipeline_flag_ph2 = pipeline_state.unaligned_seq_len ? (1u << 0) : 0;
+ }
+
+ if (shmem_sz > device->properties.limits.maxComputeSharedMemorySize) {
+ VK_LOG_DEBUG("get_xe_fa_tuning_params(HSK=" << hsk << ", HSV=" << hsv << ", qk_ratio=" << qk_ratio << ", pv_ratio=" << pv_ratio << "), "
+ "shmem_sz = " << shmem_sz << ", exceeding max shmem size = " << device->properties.limits.maxComputeSharedMemorySize);
+ }
+
+ return params;
+}
+
static bool ggml_vk_matmul_shmem_support(const vk_device& device, const std::vector<uint32_t>& warptile, bool mul_mat_id, ggml_type src0_type) {
uint32_t lut_size = 0;
@@ -3202,43 +3382,36 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_flash_attn_split_k_reduce, "fa_split_k_reduce", fa_split_k_reduce_len, fa_split_k_reduce_data, "main", 3, sizeof(vk_op_flash_attn_split_k_reduce_push_constants), {1, device->subgroup_size, 1}, {device->subgroup_size}, 1, true);
#if defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT)
- if (device->vendor_id == VK_VENDOR_ID_INTEL && (device->architecture == INTEL_XE2 || (device->architecture == INTEL_XE1 && device->coopmat_support && device->uma))) {
- auto upper_power_of_2 = [&](uint32_t in) {
- GGML_ASSERT(in != 0);
- if (in <= 1) return 1u;
- uint32_t ret = in - 1;
- ret |= ret >> 1;
- ret |= ret >> 2;
- ret |= ret >> 4;
- ret |= ret >> 8;
- ret |= ret >> 16;
- return ret + 1;
- };
-
- uint32_t xe_native_sub_group_size = 16;
- if (device->architecture == INTEL_XE1) {
- xe_native_sub_group_size = 8;
- }
+ if (device->xe_fa_support) {
+ uint32_t shmem_sz = (uint32_t)device->properties.limits.maxComputeSharedMemorySize;
+ vk_xe_fa_tuning_params xe_fa_params {};
for (auto& it : device->pipeline_xe_fa_decode_dual_phases) {
- const uint32_t split_p_chunk = 32;
- auto HdQk = it.first;
+ auto pipeline_state = it.first;
auto& pipelines = it.second;
- uint32_t head_dim_qk = std::get<0>(HdQk);
- uint32_t head_dim_pv = std::get<1>(HdQk);
- uint32_t gqa_ratio = std::get<2>(HdQk);
- uint32_t q_len = std::get<3>(HdQk);
- const uint32_t out_dim_per_wg = gqa_ratio > 16 ? 8 : 16;
- uint32_t aligned_q_len = upper_power_of_2(q_len);
- uint32_t group_sz_ph1 = std::min(std::max(aligned_q_len * xe_native_sub_group_size, 64u), 256u);
- uint32_t out_per_wg_ph1 = std::min(q_len, 256u / xe_native_sub_group_size);
- uint32_t aligned_gqa_ratio = upper_power_of_2(gqa_ratio);
- uint32_t split_p_per_iter_ph2 = 256;
- uint32_t split_p_per_warp = 16;
- uint32_t group_sz_ph2 = (split_p_per_iter_ph2 / split_p_per_warp) * xe_native_sub_group_size;
- uint32_t out_per_wg_ph2 = std::min(std::max(16u / aligned_gqa_ratio, 1u), q_len);
- ggml_vk_create_pipeline(device, pipelines.first, "xe_fa_decode_ph1", fa_decode_ph1_cm1_len, fa_decode_ph1_cm1_data, "main", 5, sizeof(vk_fa_xe_opt_push_constants), { 1, 32, 1 }, { group_sz_ph1, gqa_ratio, head_dim_qk, xe_native_sub_group_size, split_p_chunk, out_per_wg_ph1 }, 1, false, true, xe_native_sub_group_size);
- ggml_vk_create_pipeline(device, pipelines.second, "xe_fa_decode_ph2", fa_decode_ph2_cm1_len, fa_decode_ph2_cm1_data, "main", 5, sizeof(vk_fa_xe_opt_push_constants), { 1, 1, 1 }, { group_sz_ph2, gqa_ratio, head_dim_pv, out_per_wg_ph2, xe_native_sub_group_size, split_p_per_iter_ph2, split_p_chunk, out_dim_per_wg }, 1, false, true, xe_native_sub_group_size);
+ xe_fa_params = get_xe_fa_tuning_params(device.get(), pipeline_state, false);
+ std::vector<uint32_t> sc_ph1 = { xe_fa_params.workgroup_size[0], pipeline_state.qk_ratio, pipeline_state.HSK,
+ xe_fa_params.subgroup_size[0], xe_fa_params.decode_params.p_chunk_ph1,
+ xe_fa_params.decode_params.out_per_wg_ph1, xe_fa_params.decode_params.p_chunk_ph1 / xe_fa_params.tm,
+ xe_fa_params.tm, xe_fa_params.tn, shmem_sz };
+ std::vector<uint32_t> sc_ph2 = { xe_fa_params.workgroup_size[1], pipeline_state.pv_ratio, pipeline_state.HSV,
+ xe_fa_params.decode_params.out_row_per_wg_ph2, xe_fa_params.subgroup_size[1], xe_fa_params.decode_params.split_p_per_iter_ph2,
+ xe_fa_params.decode_params.p_chunk_ph1, xe_fa_params.decode_params.out_col_per_wg_ph2,
+ xe_fa_params.tm, xe_fa_params.tn, xe_fa_params.decode_params.pipeline_flag_ph2 };
+ ggml_vk_create_pipeline(device, pipelines.first, "xe_fa_decode_ph1", fa_decode_ph1_cm1_len, fa_decode_ph1_cm1_data, "main", 5, sizeof(vk_xe_fa_push_constants), { 1, xe_fa_params.decode_params.p_chunk_ph1, 1 }, sc_ph1, 1, false, true, xe_fa_params.subgroup_size[0]);
+ ggml_vk_create_pipeline(device, pipelines.second, "xe_fa_decode_ph2", fa_decode_ph2_cm1_len, fa_decode_ph2_cm1_data, "main", 5, sizeof(vk_xe_fa_push_constants), { 1, 1, 1 }, sc_ph2, 1, false, true, xe_fa_params.subgroup_size[0]);
+ }
+
+ for (auto& it : device->pipeline_xe_fa_prefill) {
+ auto pipeline_state = it.first;
+ auto& pipelines = it.second;
+ xe_fa_params = get_xe_fa_tuning_params(device.get(), pipeline_state, true);
+
+ std::vector<uint32_t> sc = { xe_fa_params.workgroup_size[0], xe_fa_params.prefill_params.wg_row, xe_fa_params.prefill_params.sg_col,
+ pipeline_state.HSK, pipeline_state.HSV, pipeline_state.qk_ratio, xe_fa_params.prefill_params.p_per_iter,
+ xe_fa_params.tm, xe_fa_params.tn, xe_fa_params.prefill_params.q_in_reg, shmem_sz, xe_fa_params.prefill_params.acc_count,
+ xe_fa_params.prefill_params.pipeline_flag, xe_fa_params.subgroup_size[0], xe_fa_params.prefill_params.gqa_fold };
+ ggml_vk_create_pipeline(device, pipelines, "xe_fa_prefill", fa_prefill_cm1_len, fa_prefill_cm1_data, "main", 6, sizeof(vk_xe_fa_push_constants), { xe_fa_params.prefill_params.gqa_fold, xe_fa_params.prefill_params.wg_row / xe_fa_params.prefill_params.gqa_fold, 1 }, sc, 1, false, true, xe_fa_params.subgroup_size[0]);
}
}
#endif
@@ -4125,6 +4298,7 @@ vk_device ggml_vk_get_device(size_t idx) {
bool dot2_f16_support = false;
bool ocp_microscaling_extension = false;
bool shader_float8_extension = false;
+ bool workgroup_memory_explicit_layout_support = false;
for (const auto& properties : ext_props) {
if (strcmp("VK_KHR_maintenance4", properties.extensionName) == 0) {
@@ -4180,6 +4354,10 @@ vk_device ggml_vk_get_device(size_t idx) {
dot2_f16_support = true;
} else if (strcmp("VK_KHR_pipeline_executable_properties", properties.extensionName) == 0) {
pipeline_executable_properties_support = true;
+#if defined(VK_KHR_workgroup_memory_explicit_layout)
+ } else if (strcmp("VK_KHR_workgroup_memory_explicit_layout", properties.extensionName) == 0) {
+ workgroup_memory_explicit_layout_support = true;
+#endif
} else if (strcmp("VK_EXT_memory_priority", properties.extensionName) == 0 &&
getenv("GGML_VK_ENABLE_MEMORY_PRIORITY")) {
device->memory_priority = true;
@@ -4532,6 +4710,16 @@ vk_device ggml_vk_get_device(size_t idx) {
device_extensions.push_back("VK_KHR_pipeline_executable_properties");
}
+#if defined(VK_KHR_workgroup_memory_explicit_layout)
+ VkPhysicalDeviceWorkgroupMemoryExplicitLayoutFeaturesKHR wg_explicit_layout_features {};
+ wg_explicit_layout_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_WORKGROUP_MEMORY_EXPLICIT_LAYOUT_FEATURES_KHR;
+ if (workgroup_memory_explicit_layout_support) {
+ last_struct->pNext = (VkBaseOutStructure *)&wg_explicit_layout_features;
+ last_struct = (VkBaseOutStructure *)&wg_explicit_layout_features;
+ device_extensions.push_back("VK_KHR_workgroup_memory_explicit_layout");
+ }
+#endif
+
if (device->external_memory_host) {
device_extensions.push_back("VK_EXT_external_memory_host");
}
@@ -4593,6 +4781,16 @@ vk_device ggml_vk_get_device(size_t idx) {
shader_float8_extension && shader_float8_features.shaderFloat8 &&
!getenv("GGML_VK_DISABLE_OCP_FP4");
+#if defined(VK_KHR_workgroup_memory_explicit_layout)
+ device->workgroup_memory_explicit_layout = workgroup_memory_explicit_layout_support &&
+ wg_explicit_layout_features.workgroupMemoryExplicitLayout;
+ device->workgroup_memory_explicit_layout_16bit = device->workgroup_memory_explicit_layout &&
+ wg_explicit_layout_features.workgroupMemoryExplicitLayout16BitAccess;
+#else
+ device->workgroup_memory_explicit_layout = false;
+ device->workgroup_memory_explicit_layout_16bit = false;
+#endif
+
device->pipeline_robustness = pl_robustness_features.pipelineRobustness;
device->multi_add = vk12_props.shaderRoundingModeRTEFloat16 &&
@@ -4958,6 +5156,13 @@ vk_device ggml_vk_get_device(size_t idx) {
descriptor_set_layout_create_info.setPNext(&dslbfci);
device->dsl = device->device.createDescriptorSetLayout(descriptor_set_layout_create_info);
+#if defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT)
+ device->xe_fa_support = device->coopmat_support && device->properties.limits.maxComputeSharedMemorySize >= 48 * 1024 &&
+ device->workgroup_memory_explicit_layout && device->workgroup_memory_explicit_layout_16bit &&
+ ((device->vendor_id == VK_VENDOR_ID_INTEL && (device->architecture == INTEL_XE2 || (device->architecture == INTEL_XE1 && device->uma))) ||
+ (device->subgroup_size == 32 && device->coopmat_support_16x16x16_f32acc && (getenv("GGML_VK_FORCE_XE_FA") != nullptr)));
+#endif
+
ggml_vk_load_shaders(device);
// Prefer a dedicated transfer queue on AMD dGPUs (non-GCN) when graphics queue use is disabled.
@@ -8403,32 +8608,98 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
vk_pipeline pipeline = nullptr;
+ auto tensor_aligned_with = [&](const ggml_tensor * t, uint32_t in) {
+ GGML_ASSERT(t != nullptr && in != 0);
+ return t->nb[1] % in == 0 && t->nb[2] % in == 0 && t->nb[3] % in == 0 && ggml_vk_tensor_buffer_offset(ctx, t) % in == 0;
+ };
+ auto vk_dev = ctx->device.get();
bool xe_fa_opt = false;
- bool fa_copy_qstate = false;
- bool xe_fa_supported_platform =
- (ctx->device.get()->architecture == INTEL_XE2 && ctx->device.get()->properties.deviceID != 0xFD80 && ctx->device.get()->properties.deviceID != 0xFD81) ||
- (ctx->device.get()->architecture == INTEL_XE1 && ctx->device.get()->coopmat_support && ctx->device.get()->uma);
- bool xe_fa_supported_usage = neq0 % 32 == 0 && nev0 % 16 == 0 && q->nb[1] > q->nb[2] && k->nb[1] > k->nb[2] && v->nb[1] > v->nb[2] && mask != nullptr;
- bool xe_fa_supported_dtype = q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16 && (mask != nullptr && mask->type == GGML_TYPE_F16);
+ bool xe_fa_use_2phases = false;
+ bool xe_fa_supported_platform = vk_dev->xe_fa_support && vk_dev->properties.deviceID != 0xFD80 && vk_dev->properties.deviceID != 0xFD81;
+ bool xe_fa_supported_usage = neq0 % 32 == 0 && nev0 % 16 == 0 && mask != nullptr && !use_sparse && max_bias == 0.0f && logit_softcap == 0.0f &&
+ tensor_aligned_with(k, 16u) && tensor_aligned_with(v, 16u);
+ bool xe_fa_supported_dtype = q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16 &&
+ mask != nullptr && mask->type == GGML_TYPE_F16;
+ xe_fa_opt = xe_fa_supported_platform && xe_fa_supported_usage && xe_fa_supported_dtype;
std::pair<vk_pipeline, vk_pipeline> xe_fa_pipeline_dual_phases = { nullptr , nullptr };
vk_pipeline xe_fa_pipeline = nullptr;
+ size_t size_q_staging = 0;
size_t size_p = 0;
size_t size_group_max = 0;
+ size_t size_decode_temp = 0;
+ uint32_t pv_ratio = uint32_t(neq2 / nev2);
+ vk_xe_fa_tuning_params xe_fa_params {};
{
std::lock_guard<std::mutex> guard(ctx->device->compile_mutex);
- auto &pipelines = ctx->device->pipeline_flash_attn_f32_f16;
- auto it = pipelines.find(fa_pipeline_state);
- if (it != pipelines.end()) {
- pipeline = it->second;
- } else {
- pipelines[fa_pipeline_state] = pipeline = std::make_shared<vk_pipeline_struct>();
+ if (xe_fa_opt) {
+ if (neq1 > 8 || pv_ratio > 64) {
+ uint32_t head_stride_q = uint32_t(nbq2 / ggml_type_size(q->type));
+ uint32_t batch_stride_q = uint32_t(nbq3 / ggml_type_size(q->type));
+ vk_xe_fa_pipeline_state prefill_state = { (uint32_t)neq0, (uint32_t)nev0, qk_ratio, pv_ratio, 0, false, false, false, false };
+ prefill_state.unaligned_seq_len = nek1 % 32 != 0;
+ prefill_state.unaligned_q = head_stride_q % 8 != 0 || q_stride % 8 != 0 || batch_stride_q % 8 != 0;
+ prefill_state.partial_row_mask = neq1 > nem1;
+ prefill_state.multi_head_mask = nem2 > 1;
+ xe_fa_params = get_xe_fa_tuning_params(vk_dev, prefill_state, true);
+ if (xe_fa_params.shmem_sz <= vk_dev->properties.limits.maxComputeSharedMemorySize) {
+ auto& xe_fa_pipelines = ctx->device->pipeline_xe_fa_prefill;
+ auto it = xe_fa_pipelines.find(prefill_state);
+ if (it != xe_fa_pipelines.end()) {
+ xe_fa_pipeline = it->second;
+ } else {
+ xe_fa_pipelines[prefill_state] = xe_fa_pipeline = std::make_shared<vk_pipeline_struct>();
+ }
+ } else {
+ xe_fa_opt = false;
+ }
+ } else {
+ vk_xe_fa_pipeline_state decode_state = { (uint32_t)neq0, (uint32_t)nev0, qk_ratio, pv_ratio, (uint32_t)neq1, false, false, false, false };
+ decode_state.unaligned_seq_len = nek1 % 256 != 0;
+ xe_fa_params = get_xe_fa_tuning_params(vk_dev, decode_state, false);
+ const uint32_t split_p_size = xe_fa_params.decode_params.p_chunk_ph1;
+ const size_t max_dim = std::max<size_t>(CEIL_DIV(nek1, split_p_size), 1);
+ const size_t p_dim = max_dim * split_p_size;
+ const size_t q_head_padding = 16;
+ size_q_staging = neq0 * neq1 * (neq2 * neq3 + q_head_padding) * sizeof(ggml_fp16_t);
+ size_p = neq1 * neq2 * p_dim * neq3 * sizeof(ggml_fp16_t);
+ size_group_max = neq1 * neq2 * max_dim * neq3 * sizeof(float);
+ size_decode_temp = size_q_staging + size_p + size_group_max;
+ uint32_t ph1_wg1 = CEIL_DIV(nek1, xe_fa_params.decode_params.p_chunk_ph1);
+ if (xe_fa_params.shmem_sz <= vk_dev->properties.limits.maxComputeSharedMemorySize &&
+ ph1_wg1 <= ctx->device->properties.limits.maxComputeWorkGroupCount[1] &&
+ size_decode_temp <= ctx->device->properties.limits.maxStorageBufferRange) {
+ auto& pipelines = ctx->device->pipeline_xe_fa_decode_dual_phases;
+ auto it = pipelines.find(decode_state);
+ if (it != pipelines.end()) {
+ xe_fa_pipeline_dual_phases = it->second;
+ } else {
+ pipelines[decode_state] = xe_fa_pipeline_dual_phases =
+ std::make_pair(std::make_shared<vk_pipeline_struct>(), std::make_shared<vk_pipeline_struct>());
+ }
+ xe_fa_use_2phases = true;
+ } else {
+ xe_fa_opt = false;
+ }
+ }
+ }
+
+ if (!xe_fa_opt) {
+ auto &pipelines = ctx->device->pipeline_flash_attn_f32_f16;
+ auto it = pipelines.find(fa_pipeline_state);
+ if (it != pipelines.end()) {
+ pipeline = it->second;
+ } else {
+ pipelines[fa_pipeline_state] = pipeline = std::make_shared<vk_pipeline_struct>();
+ }
}
}
- assert(pipeline);
- // Compile early to initialize wg_denoms.
- ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
+ if (!xe_fa_opt) {
+ assert(pipeline);
+ // Compile early to initialize wg_denoms.
+ ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1);
+ }
uint32_t split_kv = KV;
uint32_t split_k = 1;
@@ -8442,7 +8713,9 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
const uint32_t Br = fa_pipeline_state.Br;
const uint32_t Bc = fa_pipeline_state.Bc;
- GGML_ASSERT(Br == pipeline->wg_denoms[0]);
+ if (!xe_fa_opt) {
+ GGML_ASSERT(Br == pipeline->wg_denoms[0]);
+ }
const uint32_t Tr = CEIL_DIV(N, Br);
// Try to use split_k when KV is large enough to be worth the overhead.
@@ -8472,45 +8745,24 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
// of "align", so recompute split_k based on that.
split_kv = ROUNDUP_POW2(std::max(1u, KV / split_k), alignment);
split_k = CEIL_DIV(KV, split_kv);
- xe_fa_opt = xe_fa_supported_platform && xe_fa_supported_usage && xe_fa_supported_dtype;
- if (xe_fa_opt) {
- const uint32_t split_p_size = 32;
- const size_t max_dim = (nek1 + split_p_size - 1) / split_p_size;
- const size_t p_dim = max_dim * split_p_size;
-#if defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT)
- std::lock_guard<std::mutex> guard(ctx->device->compile_mutex);
- auto& pipelines = ctx->device->pipeline_xe_fa_decode_dual_phases;
- auto it = pipelines.find({ (uint32_t)neq0, (uint32_t)nev0, qk_ratio, (uint32_t)neq1 });
- if (it != pipelines.end()) {
- xe_fa_pipeline_dual_phases = it->second;
- } else {
- pipelines[{(uint32_t)neq0, (uint32_t)nev0, qk_ratio, (uint32_t)neq1}] = xe_fa_pipeline_dual_phases = std::make_pair(std::make_shared<vk_pipeline_struct>(), std::make_shared<vk_pipeline_struct>());
+ }
+
+ if (xe_fa_opt) {
+ // Bypass mask_opt and split_k buffer grow when they are not in use
+ use_mask_opt = false;
+ split_k = 1;
+ if (xe_fa_use_2phases) {
+ if (ctx->prealloc_size_x < size_decode_temp) {
+ ctx->prealloc_size_x = size_decode_temp;
+ ggml_vk_preallocate_buffers(ctx, subctx);
}
-#endif
- if (xe_fa_pipeline_dual_phases.first == nullptr || xe_fa_pipeline_dual_phases.second == nullptr) {
- xe_fa_opt = false;
- fa_copy_qstate = false;
- } else {
- size_p = neq1 * neq2 * p_dim * neq3 * sizeof(ggml_fp16_t);
- size_group_max = neq1 * neq2 * max_dim * neq3 * sizeof(float);
- size_t temp_size = ggml_nelements(q) * sizeof(ggml_fp16_t) + size_p + size_group_max;
- fa_copy_qstate = true;
- if (ctx->prealloc_size_x < temp_size) {
- ctx->prealloc_size_x = temp_size;
- ggml_vk_preallocate_buffers(ctx, subctx);
- }
- if (ctx->prealloc_x_need_sync) {
- ggml_vk_sync_buffers(ctx, subctx);
- }
+ if (ctx->prealloc_x_need_sync) {
+ ggml_vk_sync_buffers(ctx, subctx);
}
}
}
- if (xe_fa_opt == true) {
- use_mask_opt = false;
- }
-
// Reserve space for split_k temporaries. For each split x batch, we need to store the O matrix (D x ne1)
// and the per-row m and L values (ne1 rows). We store all the matrices first, followed by the rows.
// For matrices, the order is (inner to outer) [HSV, ne1, k, ne2, ne3].
@@ -8697,67 +8949,59 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
mask_n_head_log2, m0, m1,
gqa_ratio, split_kv, split_k };
- if (xe_fa_opt && split_k > 1) {
- auto upper_power_of_2 = [&](uint32_t in) {
- GGML_ASSERT(in != 0);
- if (in <= 1) return 1u;
- uint32_t ret = in - 1;
- ret |= ret >> 1;
- ret |= ret >> 2;
- ret |= ret >> 4;
- ret |= ret >> 8;
- ret |= ret >> 16;
- return ret + 1;
- };
- auto to_fp16_vk_0 = ggml_vk_get_to_fp16(ctx, q->type);
- const uint32_t out_dim_per_wg = qk_ratio > 16 ? 8 : 16;
- size_t x_ne = ggml_nelements(q);
- size_t temp_buf_offset = 0;
+ if (xe_fa_opt) {
+ uint32_t head_stride_q = uint32_t(nbq2 / ggml_type_size(q->type));
uint32_t head_stride_k = uint32_t(nbk2 / ggml_type_size(k->type));
uint32_t head_stride_v = uint32_t(nbv2 / ggml_type_size(v->type));
uint32_t batch_stride_q = uint32_t(nbq3 / ggml_type_size(q->type));
uint32_t batch_stride_k = uint32_t(nbk3 / ggml_type_size(k->type));
uint32_t batch_stride_v = uint32_t(nbv3 / ggml_type_size(v->type));
- uint32_t batch_stride_m = mask ? uint32_t(mask->nb[3] / ggml_type_size(mask->type)) : 0u;
uint32_t batch_stride_o = uint32_t(nb3 / ggml_type_size(dst->type));
- vk_fa_xe_opt_push_constants pc_ph1 = { (uint32_t)nek1, (uint32_t)neq1, (uint32_t)neq2, (uint32_t)nek2, qk_ratio, 1, (sinks != nullptr) ? 1u : 0u, (uint32_t)k_stride, head_stride_k,
- batch_stride_q, batch_stride_k, batch_stride_v, batch_stride_m, batch_stride_o, scale };
- vk_fa_xe_opt_push_constants pc_ph2 = pc_ph1;
- pc_ph2.nbkv_tok = v_stride;
- pc_ph2.nbkv_head = head_stride_v;
- vk_subbuffer q_temp_buf = fa_copy_qstate ? ggml_vk_subbuffer(ctx, ctx->prealloc_x, temp_buf_offset) : q_buf;
- temp_buf_offset += fa_copy_qstate ? x_ne * sizeof(ggml_fp16_t) : 0;
- vk_subbuffer p_temp_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_x, temp_buf_offset);
- temp_buf_offset += size_p;
- vk_subbuffer max_temp_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_x, temp_buf_offset);
- temp_buf_offset += size_group_max;
- uint32_t xe_native_sub_group_size = ctx->device.get()->architecture == INTEL_XE1 ? 8 : 16;
- uint32_t aligned_gqa_ratio = upper_power_of_2(qk_ratio);
- uint32_t out_per_wg_ph1 = std::min(256u / xe_native_sub_group_size, (uint32_t)neq1);
- uint32_t out_per_wg_ph2 = std::min(std::max(16u / aligned_gqa_ratio, 1u), (uint32_t)neq1);
- uint32_t ph1_wg = ((neq1 + out_per_wg_ph1 - 1) / out_per_wg_ph1) * nek2;
- uint32_t ph2_wg = ((neq1 + out_per_wg_ph2 - 1) / out_per_wg_ph2) * ne0 / out_dim_per_wg;
- if (fa_copy_qstate) {
- const std::vector<uint32_t> pc_cpy_fp16 =
- { (uint32_t)q->ne[0], (uint32_t)q->ne[1], (uint32_t)q->ne[2], (uint32_t)q->ne[3], (uint32_t)(x_ne) };
- ggml_vk_sync_buffers(ctx, subctx);
+
+ vk_xe_fa_push_constants pc_xe_fa = { (uint32_t)neq1, (uint32_t)neq2, (uint32_t)neq3,
+ (uint32_t)nek1, (uint32_t)nek2, (uint32_t)nek3, (uint32_t)nev2, (uint32_t)nev3,
+ (uint32_t)nem1, (uint32_t)nem2, (uint32_t)nem3,
+ q_stride, head_stride_q, k_stride, head_stride_k, v_stride, head_stride_v,
+ batch_stride_q, batch_stride_k, batch_stride_v, batch_stride_o,
+ (sinks != nullptr) ? (1u << 16) : 0u, scale };
+
+ if (xe_fa_use_2phases) {
+ auto to_fp16_vk_0 = ggml_vk_get_cpy_pipeline(ctx, q, nullptr, GGML_TYPE_F16);
+ size_t temp_buf_offset = 0;
+ vk_subbuffer q_temp_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_x, temp_buf_offset);
+ temp_buf_offset += size_q_staging;
+ vk_subbuffer p_temp_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_x, temp_buf_offset);
+ temp_buf_offset += size_p;
+ vk_subbuffer max_temp_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_x, temp_buf_offset);
+ temp_buf_offset += size_group_max;
+ uint32_t ph1_wg = CEIL_DIV(neq1, xe_fa_params.decode_params.out_per_wg_ph1) * nek2;
+ uint32_t ph2_wg = CEIL_DIV(neq1, xe_fa_params.decode_params.out_row_per_wg_ph2) * ne0 / xe_fa_params.decode_params.out_col_per_wg_ph2;
+
+ pc_xe_fa.nbq1 = uint32_t(neq0);
+ pc_xe_fa.nbq2 = uint32_t(neq0 * neq1);
+ pc_xe_fa.nbq3 = uint32_t(neq0 * neq1 * neq2);
+
ggml_pipeline_request_descriptor_sets(ctx, to_fp16_vk_0, 1);
- ggml_vk_dispatch_pipeline(ctx, subctx, to_fp16_vk_0, { q_buf, q_temp_buf }, pc_cpy_fp16, { (uint32_t)(x_ne), 1, 1 });
- }
+ ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_0, q, q_buf, q_temp_buf);
- ggml_vk_sync_buffers(ctx, subctx);
- ggml_pipeline_request_descriptor_sets(ctx, xe_fa_pipeline_dual_phases.first, 1);
- ggml_vk_dispatch_pipeline(ctx, subctx, xe_fa_pipeline_dual_phases.first,
- { q_temp_buf, k_buf, mask_buf, p_temp_buf, max_temp_buf },
- pc_ph1, { (uint32_t)ph1_wg, (uint32_t)nek1, (uint32_t)neq3 });
+ ggml_pipeline_request_descriptor_sets(ctx, xe_fa_pipeline_dual_phases.first, 1);
+ ggml_vk_dispatch_pipeline(ctx, subctx, xe_fa_pipeline_dual_phases.first,
+ { q_temp_buf, k_buf, mask_buf, p_temp_buf, max_temp_buf },
+ pc_xe_fa, { (uint32_t)ph1_wg, (uint32_t)nek1, (uint32_t)neq3 });
- ggml_vk_sync_buffers(ctx, subctx);
- ggml_pipeline_request_descriptor_sets(ctx, xe_fa_pipeline_dual_phases.second, 1);
- ggml_vk_dispatch_pipeline(ctx, subctx, xe_fa_pipeline_dual_phases.second,
- { p_temp_buf, v_buf, max_temp_buf, sinks_buf, dst_buf },
- pc_ph2, { (uint32_t)ph2_wg, (uint32_t)nev2, (uint32_t)neq3 });
+ ggml_vk_sync_buffers(ctx, subctx);
+ ggml_pipeline_request_descriptor_sets(ctx, xe_fa_pipeline_dual_phases.second, 1);
+ ggml_vk_dispatch_pipeline(ctx, subctx, xe_fa_pipeline_dual_phases.second,
+ { p_temp_buf, v_buf, max_temp_buf, sinks_buf, dst_buf },
+ pc_xe_fa, { (uint32_t)ph2_wg, (uint32_t)nev2, (uint32_t)neq3 });
- ctx->prealloc_x_need_sync = true;
+ ctx->prealloc_x_need_sync = true;
+ } else {
+ ggml_pipeline_request_descriptor_sets(ctx, xe_fa_pipeline, 1);
+ ggml_vk_dispatch_pipeline(ctx, subctx, xe_fa_pipeline,
+ { q_buf, k_buf, v_buf, mask_buf, sinks_buf, dst_buf },
+ pc_xe_fa, { (uint32_t)neq2, (uint32_t)neq1, (uint32_t)neq3 });
+ }
} else if (split_k > 1) {
ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_split_k_reduce, 1);
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_1.comp
index b5f95aaa0..53beab149 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_1.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_1.comp
@@ -3,261 +3,327 @@
#extension GL_EXT_control_flow_attributes : enable
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
-#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
#extension GL_KHR_memory_scope_semantics : enable
#extension GL_KHR_shader_subgroup_basic : enable
-#extension GL_KHR_shader_subgroup_ballot : enable
#extension GL_KHR_shader_subgroup_arithmetic : enable
#extension GL_KHR_cooperative_matrix : enable
#extension GL_EXT_shared_memory_block : enable
-layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+layout (local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
-layout (binding = 0) readonly buffer Q {float16_t qState[];};
-layout (binding = 1) readonly buffer K_VEC4 {f16vec4 kStateVec4[];};
-layout (binding = 2) buffer MASK_F16 {float16_t mState_f16[];};
-layout (binding = 3) buffer P_FP16 {float16_t matP_f16[];};
-layout (binding = 4) buffer OUT_MAX {float out_max_f32[];};
+layout (binding = 0) readonly buffer Q {float16_t data_q[];};
+layout (binding = 1) readonly buffer KV4 {f16vec4 data_kv4[];};
+layout (binding = 2) readonly buffer M {float16_t data_m[];};
+layout (binding = 3) buffer P {float16_t data_p[];};
+layout (binding = 4) buffer MAX {float data_max[];};
layout (push_constant) uniform parameter
{
- uint kvSeqLen;
- uint activationLength;
- uint qHead;
- uint kvHead;
- uint qkRatio;
- uint qkSubGroups;
- uint flag;
- uint kvStride1;
- uint kvStride2;
- uint batchStrideQ;
- uint batchStrideK;
- uint batchStrideV;
- uint batchStrideM;
- uint batchStrideO;
- float softMaxScale;
+ uint neq1;
+ uint neq2;
+ uint neq3;
+ uint nek1;
+ uint nek2;
+ uint nek3;
+ uint nev2;
+ uint nev3;
+ uint nem1;
+ uint nem2;
+ uint nem3;
+ uint nbq1;
+ uint nbq2;
+ uint nbk1;
+ uint nbk2;
+ uint nbv1;
+ uint nbv2;
+ uint nbq3;
+ uint nbk3;
+ uint nbv3;
+ uint nb3;
+ uint flag;
+ float scale;
} p;
-layout (constant_id = 0) const uint GROUPSIZE = 128;
-layout (constant_id = 1) const uint GQA_RATIO = 8;
-layout (constant_id = 2) const uint HEAD_DIM = 128;
-layout (constant_id = 3) const uint WARPSIZE = 16;
-layout (constant_id = 4) const uint MATP_REDUCE = 32;
-layout (constant_id = 5) const uint N_TOK = 1;
-layout (constant_id = 6) const uint COOP_MAT_P_PER_LOOP = 4;
-
-#define MAX_HEADS 8
-
-#define TN WARPSIZE
-#define TM 8
-#define TK 16
-#define SUBGROUP_COUNT (GROUPSIZE / WARPSIZE)
-#define MATP_PER_LOOP (COOP_MAT_P_PER_LOOP * TM)
-#define P_LOOP_COUNT (MATP_REDUCE / MATP_PER_LOOP)
-
-#define COOP_MAT_Q_PER_TOKEN ((GQA_RATIO + TN - 1) / TN)
-#define COOP_MAT_P_M COOP_MAT_Q_PER_TOKEN
-#define COOP_MAT_P_N (MATP_REDUCE / TM)
-#define SLM_PV_SIZE (MATP_REDUCE * COOP_MAT_P_M * TN)
-#define SLM_MASK_SIZE (N_TOK * MATP_REDUCE)
-#define SLM_POOL_SIZE_K (MATP_PER_LOOP * HEAD_DIM)
-#define K_LOAD_PER_LOOP (GROUPSIZE * 4)
-#define HEAD_DIM_VEC4 (HEAD_DIM / 4)
-#define SLM_CHUNK_SIZE (TK / 4)
-#define K_LOAD_LOOPS ((SLM_POOL_SIZE_K + K_LOAD_PER_LOOP - 1) / K_LOAD_PER_LOOP)
-#define O_COUNT ((GQA_RATIO + SUBGROUP_COUNT - 1) / SUBGROUP_COUNT)
-
-shared slm_pool_block {
- float slm_pool_pv[SLM_PV_SIZE + SLM_MASK_SIZE];
-} slm_pool_f32;
-
-shared slm_pool_alias_block {
- float16_t slm_pool_k[SLM_POOL_SIZE_K];
-} slm_pool_f16;
+layout (constant_id = 0) const uint WorkGroupSize = 128;
+layout (constant_id = 1) const uint gqa_ratio = 8;
+layout (constant_id = 2) const uint HSK = 128;
+layout (constant_id = 3) const uint SubGroupSize = 16;
+layout (constant_id = 4) const uint p_reduce = 32;
+layout (constant_id = 5) const uint Br = 1;
+layout (constant_id = 6) const uint pmat_count = 4;
+layout (constant_id = 7) const uint TM = 8;
+layout (constant_id = 8) const uint TN = 16;
+layout (constant_id = 9) const uint shmem_size = 49152;
+
+const uint TK = 16;
+const uint sg_count = WorkGroupSize / SubGroupSize;
+const uint p_per_loop = pmat_count * TM;
+const uint p_loops = p_reduce / p_per_loop;
+
+const uint pmat_row = (gqa_ratio + TN - 1) / TN;
+const uint pmat_col = p_reduce / TM;
+
+const uint k_per_loop = WorkGroupSize * 4;
+const uint hsk_v4 = HSK / 4;
+const uint slm_col_chunk = TK / 4;
+const uint k_loops = (p_per_loop * HSK + k_per_loop - 1) / k_per_loop;
+const uint warp_o_row = (gqa_ratio + sg_count - 1) / sg_count;
+
+#define SLM_PV_SIZE (p_reduce * pmat_row * TN)
+#define SLM_MASK_SIZE (Br * p_reduce)
+
+const bool k_in_shmem = k_loops * k_per_loop <= shmem_size / 2;
+const uint ksh_size = k_in_shmem ? (k_loops * k_per_loop) : (sg_count * TM * TK);
+
+shared sh_block {
+ float F32sh[SLM_PV_SIZE + SLM_MASK_SIZE];
+} F32Poolsh;
+
+shared shf16_alias {
+ float16_t F16sh[ksh_size];
+} F16Poolsh;
+
+shared shf16v4_alias {
+ f16vec4 F16V4sh[ksh_size / 4];
+} F16V4Poolsh;
void main() {
- const uint lane = gl_SubgroupInvocationID;
- const uint kHeadIdx = gl_WorkGroupID.x % p.kvHead;
- const uint outGroupIdx = gl_WorkGroupID.x / p.kvHead;
- const uint v = gl_WorkGroupID.y;
- const uint d = gl_WorkGroupID.z;
- const uint localLinearId = gl_SubgroupID;
- const uint wgLane = localLinearId * WARPSIZE + lane;
- const uint qDim = p.qHead * HEAD_DIM;
- const uint kvDim = p.kvStride1;
- const uint maskDim = p.kvSeqLen;
- const uint maxDim = (p.kvSeqLen + MATP_REDUCE - 1) / MATP_REDUCE;
- const uint pDim = maxDim * MATP_REDUCE;
- const uint tokFlatIdx = localLinearId + outGroupIdx * N_TOK;
- uint offsetBaseQ = min(tokFlatIdx, p.activationLength - 1) * qDim;
- offsetBaseQ = offsetBaseQ + d * p.batchStrideQ + kHeadIdx * HEAD_DIM * GQA_RATIO;
- const uint offsetBaseK = (d * p.batchStrideK + (v * MATP_REDUCE) * kvDim + kHeadIdx * p.kvStride2) / 4;
- uint offsetOut = d * p.qHead * p.activationLength * pDim + v * MATP_REDUCE + kHeadIdx * GQA_RATIO * pDim + (localLinearId * O_COUNT + outGroupIdx * N_TOK * p.qHead) * pDim + lane;
- uint offsetMax = d * p.qHead * p.activationLength * maxDim + v + kHeadIdx * GQA_RATIO * maxDim + (localLinearId * O_COUNT + outGroupIdx * N_TOK * p.qHead) * maxDim;
- const uint offsetSlmLoadPv = (localLinearId * O_COUNT * MATP_REDUCE + lane);
- const uint offsetBaseM = v * MATP_REDUCE + lane;
- const float fp32Min = uintBitsToFloat(0xFEFFFFFF);
-
- const uint loopCount = HEAD_DIM / TK;
- float maskFp32[MATP_REDUCE / WARPSIZE];
-
- if (tokFlatIdx < p.activationLength) {
- [[unroll]] for (uint mk = 0; mk < MATP_REDUCE / WARPSIZE; mk++) {
- const uint maskOffset = mk * WARPSIZE + offsetBaseM;
- if (maskOffset < maskDim) {
- maskFp32[mk] = float(mState_f16[d * p.batchStrideM + tokFlatIdx * maskDim + maskOffset]);
- } else {
- maskFp32[mk] = fp32Min;
- }
+ const uint lane = gl_SubgroupInvocationID;
+ const uint head_k_idx = gl_WorkGroupID.x % p.nek2;
+ const uint out_idx = gl_WorkGroupID.x / p.nek2;
+ const uint p_idx = gl_WorkGroupID.y;
+ const uint batch_k_ratio = p.neq3 / p.nek3;
+ const uint batch_idx = gl_WorkGroupID.z;
+ const uint batch_k_idx = batch_idx / batch_k_ratio;
+ const uint batch_mask_idx = batch_idx % p.nem3;
+ const uint batch_mask_stride = p.nek1 * p.nem1 * p.nem2;
+ const uint sg = gl_SubgroupID;
+ const uint wg_lane = sg * SubGroupSize + lane;
+ const uint k_dim = p.nbk1;
+ const uint max_dim = (p.nek1 + p_reduce - 1) / p_reduce;
+ const uint p_dim = max_dim * p_reduce;
+ const uint out_row_idx = sg + out_idx * Br;
+ const bool single_head_mask = p.nem2 == 1;
+ uint q_offset = min(out_row_idx, p.neq1 - 1) * p.nbq1;
+ q_offset = q_offset + batch_idx * p.nbq3 + head_k_idx * p.nbq2 * gqa_ratio;
+ const uint k_base = (batch_k_idx * p.nbk3 + (p_idx * p_reduce) * k_dim + head_k_idx * p.nbk2) / 4;
+ uint o_offset = batch_idx * p.neq2 * p.neq1 * p_dim + p_idx * p_reduce + head_k_idx * gqa_ratio * p_dim +
+ (sg * warp_o_row + out_idx * Br * p.neq2) * p_dim + lane;
+ uint max_offset = batch_idx * p.neq2 * p.neq1 * max_dim + p_idx + head_k_idx * gqa_ratio * max_dim +
+ (sg * warp_o_row + out_idx * Br * p.neq2) * max_dim;
+ const uint pvsh_load_offset = sg * warp_o_row * p_reduce + lane;
+ const uint mask_base = p_idx * p_reduce + lane;
+ const uint k_seq_boundary = ((p.nek1 - 1) * k_dim + batch_k_idx * p.nbk3 + head_k_idx * p.nbk2 + HSK) / 4 - 1;
+ const uint k_head_base = (batch_k_idx * p.nbk3 + head_k_idx * p.nbk2) / 4;
+ const float NEG_FLT_MAX_OVER_2 = uintBitsToFloat(0xFEFFFFFF);
+
+ float Maskf[p_reduce / SubGroupSize];
+
+ if (out_row_idx < p.neq1 && single_head_mask) {
+ if (out_row_idx < p.nem1) {
+ [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+ const uint mask_offset = idx * SubGroupSize + mask_base;
+ if (mask_offset < p.nek1) {
+ Maskf[idx] = float(data_m[batch_mask_idx * batch_mask_stride + out_row_idx * p.nek1 + mask_offset]);
+ } else {
+ Maskf[idx] = NEG_FLT_MAX_OVER_2;
+ }
+ }
+ } else {
+ [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+ const uint mask_offset = idx * SubGroupSize + mask_base;
+ Maskf[idx] = (mask_offset < p.nek1) ? 0.0f : NEG_FLT_MAX_OVER_2;
+ }
+ }
}
- }
- coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator> matP[COOP_MAT_P_M][COOP_MAT_P_N];
+ coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator> PMat[pmat_row][pmat_col];
- [[unroll]] for (uint mp = 0; mp < COOP_MAT_P_M; mp++) {
- [[unroll]] for (uint np = 0; np < COOP_MAT_P_N; np++) {
- matP[mp][np] = coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator>(0.0f);
+ [[unroll]] for (uint r = 0; r < pmat_row; ++r) {
+ [[unroll]] for (uint c = 0; c < pmat_col; ++c) {
+ PMat[r][c] = coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator>(0.0f);
+ }
}
- }
-
- [[unroll]] for (uint kLoad = 0; kLoad < K_LOAD_LOOPS; kLoad++) {
- const uint flatOffset = kLoad * GROUPSIZE + wgLane;
- const uint kRowIdx = flatOffset / HEAD_DIM_VEC4;
- const uint kColIdx = flatOffset % HEAD_DIM_VEC4;
- const uint slmChunkCol = kColIdx % SLM_CHUNK_SIZE;
- const uint slmChunkRow = kColIdx / SLM_CHUNK_SIZE;
- const uint offsetK = offsetBaseK + kRowIdx * kvDim / 4 + kColIdx;
- const uint offsetSlmK = kRowIdx * TK + slmChunkRow * TK * MATP_PER_LOOP + slmChunkCol * 4;
- slm_pool_f16.slm_pool_k[offsetSlmK + 0] = kStateVec4[offsetK].x;
- slm_pool_f16.slm_pool_k[offsetSlmK + 1] = kStateVec4[offsetK].y;
- slm_pool_f16.slm_pool_k[offsetSlmK + 2] = kStateVec4[offsetK].z;
- slm_pool_f16.slm_pool_k[offsetSlmK + 3] = kStateVec4[offsetK].w;
- }
-
- [[unroll]] for (uint pLoop = 0; pLoop < P_LOOP_COUNT; pLoop++) {
- f16vec4 kTemp[K_LOAD_LOOPS];
-
- if (pLoop + 1 < P_LOOP_COUNT) {
- [[unroll]] for (uint kLoad = 0; kLoad < K_LOAD_LOOPS; kLoad++) {
- const uint flatOffset = kLoad * GROUPSIZE + wgLane;
- const uint kRowIdx = flatOffset / HEAD_DIM_VEC4 + (pLoop + 1) * MATP_PER_LOOP;
- const uint kColIdx = flatOffset % HEAD_DIM_VEC4;
- const uint offsetK = offsetBaseK + kRowIdx * kvDim / 4 + kColIdx;
- kTemp[kLoad] = kStateVec4[offsetK];
- }
+
+ if (k_in_shmem) {
+ [[unroll]] for (uint idx = 0; idx < k_loops; ++idx) {
+ const uint flat_offset = idx * WorkGroupSize + wg_lane;
+ const uint row_idx = flat_offset / hsk_v4;
+ const uint col_idx = flat_offset % hsk_v4;
+ const uint sh_chunk_c = col_idx % slm_col_chunk;
+ const uint sh_chunk_r = col_idx / slm_col_chunk;
+ uint k_offset = k_base + row_idx * k_dim / 4 + col_idx;
+ k_offset = min(k_offset, k_seq_boundary);
+ const uint ksh_offset = (row_idx * TK + sh_chunk_r * TK * p_per_loop) / 4 + sh_chunk_c;
+ F16V4Poolsh.F16V4sh[ksh_offset] = data_kv4[k_offset];
+ }
}
- barrier();
- if (localLinearId < N_TOK) {
- [[unroll]] for (uint loop = 0; loop < loopCount; loop++) {
- coopmat<float16_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> matQ[COOP_MAT_P_M];
- coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> matK[COOP_MAT_P_PER_LOOP];
-
- [[unroll]] for (uint mq = 0; mq < COOP_MAT_P_M; mq++) {
- coopMatLoad(
- matQ[mq],
- qState,
- offsetBaseQ + mq * TN * HEAD_DIM + loop * TK,
- HEAD_DIM,
- gl_CooperativeMatrixLayoutColumnMajor);
+ [[unroll]] for (uint pl = 0; pl < p_loops; ++pl) {
+ f16vec4 kv4_temp[k_loops];
+
+ if (k_in_shmem) {
+ if (pl + 1 < p_loops) {
+ [[unroll]] for (uint idx = 0; idx < k_loops; ++idx) {
+ const uint flat_offset = idx * WorkGroupSize + wg_lane;
+ const uint row_idx = flat_offset / hsk_v4 + (pl + 1) * p_per_loop;
+ const uint col_idx = flat_offset % hsk_v4;
+ uint k_offset = k_base + row_idx * k_dim / 4 + col_idx;
+ k_offset = min(k_offset, k_seq_boundary);
+ kv4_temp[idx] = data_kv4[k_offset];
+ }
+ }
}
- [[unroll]] for (uint np = 0; np < COOP_MAT_P_PER_LOOP; np++) {
- coopMatLoad(
- matK[np],
- slm_pool_f16.slm_pool_k,
- loop * TK * MATP_PER_LOOP + np * TM * TK,
- TK,
- gl_CooperativeMatrixLayoutRowMajor);
+ barrier();
+ if (sg < Br) {
+ [[unroll]] for (uint loop = 0; loop < HSK / TK; ++loop) {
+ coopmat<float16_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> QMat[pmat_row];
+ coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> KMat[pmat_count];
+
+ [[unroll]] for (uint r = 0; r < pmat_row; ++r) {
+ const uint coord = q_offset + r * TN * p.nbq2 + loop * TK;
+ coopMatLoad(QMat[r], data_q, coord, p.nbq2, gl_CooperativeMatrixLayoutColumnMajor);
+ }
+
+ [[unroll]] for (uint c = 0; c < pmat_count; ++c) {
+ if (k_in_shmem) {
+ const uint coord = loop * TK * p_per_loop + c * TM * TK;
+ coopMatLoad(KMat[c], F16Poolsh.F16sh, coord, TK, gl_CooperativeMatrixLayoutRowMajor);
+ } else {
+ const uint k_base_row = (c * TM + pl * pmat_count * TM + p_idx * p_reduce);
+ const uint kv4_base_offset = k_head_base + k_base_row * p.nbk1 / 4 + loop * TK / 4;
+ [[unroll]] for (uint idx = 0; idx < TM * TK / 4 / SubGroupSize; ++idx) {
+ const uint flat_offset = idx * SubGroupSize + lane;
+ const uint row_idx = flat_offset / (TK / 4);
+ const uint col_idx = flat_offset % (TK / 4);
+ const uint k_temp_offset = kv4_base_offset + row_idx * p.nbk1 / 4 + col_idx;
+ const uint ksh_temp_offset = sg * TM * TK / 4 + idx * SubGroupSize + lane;
+ f16vec4 kv4_temp = (k_base_row + row_idx) < p.nek1 ? data_kv4[k_temp_offset] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+ F16V4Poolsh.F16V4sh[ksh_temp_offset] = kv4_temp;
+ }
+
+ subgroupBarrier();
+ coopMatLoad(KMat[c], F16Poolsh.F16sh, sg * TM * TK, TK, gl_CooperativeMatrixLayoutRowMajor);
+ subgroupBarrier();
+ }
+ }
+
+ [[unroll]] for (uint r = 0; r < pmat_row; ++r) {
+ [[unroll]] for (uint c = 0; c < pmat_count; ++c) {
+ PMat[r][pl * pmat_count + c] = coopMatMulAdd(KMat[c], QMat[r], PMat[r][pl * pmat_count + c]);
+ }
+ }
+ }
}
- [[unroll]] for (uint mp = 0; mp < COOP_MAT_P_M; mp++) {
- [[unroll]] for (uint np = 0; np < COOP_MAT_P_PER_LOOP; np++) {
- matP[mp][pLoop * COOP_MAT_P_PER_LOOP + np] = coopMatMulAdd(matK[np], matQ[mp], matP[mp][pLoop * COOP_MAT_P_PER_LOOP + np]);
- }
+ barrier();
+
+ if (k_in_shmem) {
+ if (pl + 1 < p_loops) {
+ [[unroll]] for (uint idx = 0; idx < k_loops; ++idx) {
+ const uint flat_offset = idx * WorkGroupSize + wg_lane;
+ const uint row_idx = flat_offset / hsk_v4;
+ const uint col_idx = flat_offset % hsk_v4;
+ const uint sh_chunk_c = col_idx % slm_col_chunk;
+ const uint sh_chunk_r = col_idx / slm_col_chunk;
+ const uint ksh_offset = (row_idx * TK + sh_chunk_r * TK * p_per_loop) / 4 + sh_chunk_c;
+ F16V4Poolsh.F16V4sh[ksh_offset] = kv4_temp[idx];
+ }
+ }
}
- }
}
barrier();
- if (pLoop + 1 < P_LOOP_COUNT) {
- [[unroll]] for (uint kLoad = 0; kLoad < K_LOAD_LOOPS; kLoad++) {
- const uint flatOffset = kLoad * GROUPSIZE + wgLane;
- const uint kRowIdx = flatOffset / HEAD_DIM_VEC4;
- const uint kColIdx = flatOffset % HEAD_DIM_VEC4;
- const uint slmChunkCol = kColIdx % SLM_CHUNK_SIZE;
- const uint slmChunkRow = kColIdx / SLM_CHUNK_SIZE;
- const uint offsetSlmK = kRowIdx * TK + slmChunkRow * TK * MATP_PER_LOOP + slmChunkCol * 4;
- slm_pool_f16.slm_pool_k[offsetSlmK + 0] = kTemp[kLoad].x;
- slm_pool_f16.slm_pool_k[offsetSlmK + 1] = kTemp[kLoad].y;
- slm_pool_f16.slm_pool_k[offsetSlmK + 2] = kTemp[kLoad].z;
- slm_pool_f16.slm_pool_k[offsetSlmK + 3] = kTemp[kLoad].w;
- }
- }
- }
-
- barrier();
-
- if (tokFlatIdx < p.activationLength) {
- [[unroll]] for (uint mk = 0; mk < MATP_REDUCE / WARPSIZE; mk++) {
- slm_pool_f32.slm_pool_pv[SLM_PV_SIZE + localLinearId * MATP_REDUCE + mk * WARPSIZE + lane] = maskFp32[mk];
- }
- }
-
- [[unroll]] for (uint oLoop = 0; oLoop < N_TOK; oLoop++) {
- if (oLoop + outGroupIdx * N_TOK < p.activationLength) {
- if (localLinearId == oLoop) {
- [[unroll]] for (uint mp = 0; mp < COOP_MAT_P_M; mp++) {
- [[unroll]] for (uint np = 0; np < COOP_MAT_P_N; np++) {
- coopMatStore(matP[mp][np], slm_pool_f32.slm_pool_pv, mp * MATP_REDUCE * TN + np * TM, MATP_REDUCE, gl_CooperativeMatrixLayoutColumnMajor);
- }
- }
- }
-
- barrier();
-
- [[unroll]] for (uint maskIdx = 0; maskIdx < MATP_REDUCE / WARPSIZE; maskIdx++) {
- maskFp32[maskIdx] = slm_pool_f32.slm_pool_pv[SLM_PV_SIZE + oLoop * MATP_REDUCE + maskIdx * WARPSIZE + lane];
- }
-
- float fp32O[O_COUNT][MATP_REDUCE / WARPSIZE];
- float maxOut[O_COUNT];
-
- [[unroll]] for (uint oc = 0; oc < O_COUNT; oc++) {
- [[unroll]] for (uint os = 0; os < MATP_REDUCE / WARPSIZE; os++) {
- fp32O[oc][os] = slm_pool_f32.slm_pool_pv[offsetSlmLoadPv + os * WARPSIZE + oc * MATP_REDUCE] * p.softMaxScale;
+ if (sg < Br && out_row_idx < p.neq1 && single_head_mask) {
+ [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+ F32Poolsh.F32sh[SLM_PV_SIZE + sg * p_reduce + idx * SubGroupSize + lane] = Maskf[idx];
}
+ }
- [[unroll]] for (uint os = 0; os < MATP_REDUCE / WARPSIZE; os++) {
- fp32O[oc][os] = fp32O[oc][os] + maskFp32[os];
- }
-
- float maxTemp = fp32Min;
- [[unroll]] for (uint os = 0; os < MATP_REDUCE / WARPSIZE; os++) {
- maxTemp = max(maxTemp, fp32O[oc][os]);
- }
- maxOut[oc] = subgroupMax(maxTemp);
- [[unroll]] for (uint os = 0; os < MATP_REDUCE / WARPSIZE; os++) {
- fp32O[oc][os] = exp(fp32O[oc][os] - maxOut[oc]);
- }
- }
-
- [[unroll]] for (uint oc = 0; oc < O_COUNT; oc++) {
- if (localLinearId * O_COUNT + oc < GQA_RATIO) {
- [[unroll]] for (uint os = 0; os < MATP_REDUCE / WARPSIZE; os++) {
- matP_f16[offsetOut + oc * pDim + os * WARPSIZE] = float16_t(fp32O[oc][os]);
- }
-
- if (lane == 0) {
- out_max_f32[offsetMax + oc * maxDim] = maxOut[oc];
- }
+ [[unroll]] for (uint out_r = 0; out_r < Br; ++out_r) {
+ if (out_r + out_idx * Br < p.neq1) {
+ const uint mask_global_base = batch_mask_idx * batch_mask_stride + (out_r + out_idx * Br) * p.nek1;
+ if (sg == out_r) {
+ [[unroll]] for (uint r = 0; r < pmat_row; ++r) {
+ [[unroll]] for (uint c = 0; c < pmat_col; ++c) {
+ const uint coord = r * p_reduce * TN + c * TM;
+ coopMatStore(PMat[r][c], F32Poolsh.F32sh, coord, p_reduce, gl_CooperativeMatrixLayoutColumnMajor);
+ }
+ }
+ }
+
+ barrier();
+
+ if (single_head_mask) {
+ [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+ Maskf[idx] = F32Poolsh.F32sh[SLM_PV_SIZE + out_r * p_reduce + idx * SubGroupSize + lane];
+ }
+ }
+
+ float Of[warp_o_row][p_reduce / SubGroupSize];
+ float Mf[warp_o_row];
+
+ [[unroll]] for (uint r = 0; r < warp_o_row; ++r) {
+ const uint head_q_idx = head_k_idx * gqa_ratio + sg * warp_o_row + r;
+ const uint mask_head_offset = (head_q_idx % p.nem2) * p.nek1 * p.nem1;
+
+ [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+ Of[r][idx] = F32Poolsh.F32sh[pvsh_load_offset + idx * SubGroupSize + r * p_reduce] * p.scale;
+ }
+
+ if (!single_head_mask) {
+ if (out_r + out_idx * Br < p.nem1) {
+ [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+ const uint mask_offset = idx * SubGroupSize + mask_base;
+ if (mask_offset < p.nek1) {
+ Maskf[idx] = float(data_m[mask_global_base + mask_head_offset + mask_offset]);
+ } else {
+ Maskf[idx] = NEG_FLT_MAX_OVER_2;
+ }
+ }
+ } else {
+ [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+ const uint mask_offset = idx * SubGroupSize + mask_base;
+ Maskf[idx] = (mask_offset < p.nek1) ? 0.0f : NEG_FLT_MAX_OVER_2;
+ }
+ }
+ }
+
+ [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+ Of[r][idx] = Of[r][idx] + Maskf[idx];
+ }
+
+ float Mtempf = NEG_FLT_MAX_OVER_2;
+ [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+ const uint p_col_offset = idx * SubGroupSize + mask_base;
+ if (p_col_offset < p.nek1) {
+ Mtempf = max(Mtempf, Of[r][idx]);
+ }
+ }
+ Mf[r] = subgroupMax(Mtempf);
+ [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+ const uint p_col_offset = idx * SubGroupSize + mask_base;
+ Of[r][idx] = (p_col_offset < p.nek1) ? exp(Of[r][idx] - Mf[r]) : 0.0f;
+ }
+ }
+
+ [[unroll]] for (uint r = 0; r < warp_o_row; ++r) {
+ if (sg * warp_o_row + r < gqa_ratio) {
+ [[unroll]] for (uint idx = 0; idx < p_reduce / SubGroupSize; ++idx) {
+ data_p[o_offset + r * p_dim + idx * SubGroupSize] = float16_t(Of[r][idx]);
+ }
+
+ if (lane == 0) {
+ data_max[max_offset + r * max_dim] = Mf[r];
+ }
+ }
+ }
+
+ o_offset = o_offset + p.neq2 * p_dim;
+ max_offset = max_offset + p.neq2 * max_dim;
+ barrier();
}
- }
-
- offsetOut = offsetOut + p.qHead * pDim;
- offsetMax = offsetMax + p.qHead * maxDim;
- barrier();
}
- }
}
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_2.comp
index 60a1c2ce7..42a5101d8 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_2.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_2.comp
@@ -3,406 +3,470 @@
#extension GL_EXT_control_flow_attributes : enable
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
-#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
#extension GL_KHR_memory_scope_semantics : enable
#extension GL_KHR_shader_subgroup_basic : enable
-#extension GL_KHR_shader_subgroup_ballot : enable
#extension GL_KHR_shader_subgroup_arithmetic : enable
#extension GL_KHR_cooperative_matrix : enable
#extension GL_EXT_shared_memory_block : enable
-layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+layout (local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
-layout (binding = 0) readonly buffer P {f16vec4 pStateVec4[];};
-layout (binding = 1) readonly buffer V {float16_t vState[];};
-layout (binding = 1) readonly buffer V_VEC4 {f16vec4 vStateVec4[];};
-layout (binding = 2) buffer MAX_FP32 {float max_f32[];};
-layout (binding = 3) buffer SINK_FP32 {float sink_f32[];};
-layout (binding = 4) buffer OUT_FP32 {float out_f32[];};
-layout (binding = 4) buffer OUT_VEC4 {vec4 out_f32_vec4[];};
-layout (binding = 4) buffer OUT_F16 {float16_t out_f16[];};
+layout (binding = 0) readonly buffer PV4 {f16vec4 data_pv4[];};
+layout (binding = 1) readonly buffer V {float16_t data_v[];};
+layout (binding = 1) readonly buffer VV4 {f16vec4 data_vv4[];};
+layout (binding = 2) readonly buffer MAX {float data_max[];};
+layout (binding = 3) readonly buffer S {float data_s[];};
+layout (binding = 4) buffer OV4 {vec4 data_ov4[];};
layout (push_constant) uniform parameter
{
- uint kvSeqLen;
- uint activationLength;
- uint qHead;
- uint kvHead;
- uint qkRatio;
- uint qkSubGroups;
- uint flag;
- uint kvStride1;
- uint kvStride2;
- uint batchStrideQ;
- uint batchStrideK;
- uint batchStrideV;
- uint batchStrideM;
- uint batchStrideO;
- float softMaxScale;
+ uint neq1;
+ uint neq2;
+ uint neq3;
+ uint nek1;
+ uint nek2;
+ uint nek3;
+ uint nev2;
+ uint nev3;
+ uint nem1;
+ uint nem2;
+ uint nem3;
+ uint nbq1;
+ uint nbq2;
+ uint nbk1;
+ uint nbk2;
+ uint nbv1;
+ uint nbv2;
+ uint nbq3;
+ uint nbk3;
+ uint nbv3;
+ uint nb3;
+ uint flag;
+ float scale;
} p;
-layout (constant_id = 0) const uint GROUPSIZE = 256;
-layout (constant_id = 1) const uint GQA_RATIO = 8;
-layout (constant_id = 2) const uint HEAD_DIM = 128;
-layout (constant_id = 3) const uint N_TOKS_PER_GROUP = 1;
-layout (constant_id = 4) const uint WARPSIZE = 16;
-layout (constant_id = 5) const uint MATP_PER_LOOP = 64;
-layout (constant_id = 6) const uint MATP_REDUCE = 32;
-layout (constant_id = 7) const uint WARP_V_DIM = 16;
-
-#define TN WARPSIZE
-#define TM 8
-#define TK 16
-#define MAT_O_N (WARP_V_DIM / TM)
-#define MAT_P_M (GQA_RATIO * N_TOKS_PER_GROUP)
-#define ALIGNED_P_M ((MAT_P_M + WARPSIZE - 1) / WARPSIZE)
-#define V_HEAD_GROUPS (HEAD_DIM / WARP_V_DIM)
-
-#define SUBGROUP_COUNT (GROUPSIZE / WARPSIZE)
-#define SPLIT_P_GROUPS (MATP_PER_LOOP / TK)
-
-#define SLM_POOL_SIZE_O (SUBGROUP_COUNT * ALIGNED_P_M * TN * MAT_O_N * TM)
-
-#define P_LOAD_PER_LOOP (GROUPSIZE * 4)
-#define P_LOAD_LOOPS ((MAT_P_M * MATP_PER_LOOP + P_LOAD_PER_LOOP - 1) / P_LOAD_PER_LOOP)
-#define SLM_POOL_SIZE_P (P_LOAD_LOOPS * P_LOAD_PER_LOOP)
-#define SIZE_LOCAL_MAX (MAT_P_M * MATP_PER_LOOP / MATP_REDUCE)
-#define MAX_LOAD_LOOPS ((SIZE_LOCAL_MAX + GROUPSIZE - 1) / GROUPSIZE)
-#define SLM_POOL_SIZE_LOCAL_MAX (MAX_LOAD_LOOPS * GROUPSIZE)
-#define MAX_REDUCE_COUNT ((MAT_P_M + SUBGROUP_COUNT - 1) / SUBGROUP_COUNT)
-#define GLOBAL_MAX_SIZE (MAX_REDUCE_COUNT * SUBGROUP_COUNT)
-
-#define SLM_POOL_SIZE_SOFTMAX_SUM (SUBGROUP_COUNT * P_LOAD_LOOPS)
-
-#define SLM_OFFSET_P (GLOBAL_MAX_SIZE * 2 + SLM_POOL_SIZE_SOFTMAX_SUM * 2 + SLM_POOL_SIZE_LOCAL_MAX * 2 * 2)
-
-#define SLM_OFFSET_GLOBAL_MAX 0
-#define SLM_OFFSET_SOFTMAX_SUM (SLM_OFFSET_GLOBAL_MAX + GLOBAL_MAX_SIZE)
-#define SLM_OFFSET_O (SLM_OFFSET_SOFTMAX_SUM + SLM_POOL_SIZE_SOFTMAX_SUM)
-#define SLM_OFFSET_LOCAL_MAX (GLOBAL_MAX_SIZE + SLM_POOL_SIZE_SOFTMAX_SUM)
-
-#define P_REDUCE_VEC4 (MATP_PER_LOOP / 4)
-#define MAX_PER_LOOP (MATP_PER_LOOP / MATP_REDUCE)
-#define SLM_MAX_STRIDE (MATP_REDUCE / 4)
-#define SUB_GROUPS_PER_LINE (MATP_PER_LOOP / WARPSIZE / 4)
-
-shared slm_pool_block {
- float slm_pool_o[GLOBAL_MAX_SIZE + SLM_POOL_SIZE_SOFTMAX_SUM + SLM_POOL_SIZE_O];
-} slm_pool_f32;
-
-shared slm_pool_alias_block {
- float16_t slm_pool_pv[GLOBAL_MAX_SIZE * 2 + SLM_POOL_SIZE_SOFTMAX_SUM * 2 + SLM_POOL_SIZE_LOCAL_MAX * 2 * 2 + SLM_POOL_SIZE_P * 2];
-} slm_pool_alias_f16;
+layout (constant_id = 0) const uint WorkGroupSize = 256;
+layout (constant_id = 1) const uint gqa_ratio = 8;
+layout (constant_id = 2) const uint HSV = 128;
+layout (constant_id = 3) const uint Wr = 1;
+layout (constant_id = 4) const uint SubGroupSize = 16;
+layout (constant_id = 5) const uint p_per_loop = 64;
+layout (constant_id = 6) const uint p_reduce = 32;
+layout (constant_id = 7) const uint Bc = 16;
+layout (constant_id = 8) const uint TM = 8;
+layout (constant_id = 9) const uint TN = 16;
+layout (constant_id = 10) const uint Flags = 0;
+
+#define UNALIGNED_SEQ_LEN_BIT (1 << 0)
+#define SINK_ENABLE_BIT (1 << 16)
+
+const uint TK = 16;
+const uint o_row = gqa_ratio * Wr;
+const uint omat_row = (o_row + TN - 1) / TN;
+const uint o_scalar_row = (o_row + SubGroupSize - 1) / SubGroupSize;
+
+const uint omat_col = Bc / TM;
+const uint v_head_wgs = HSV / Bc;
+
+const uint sg_count = WorkGroupSize / SubGroupSize;
+const uint split_p_sg = p_per_loop / TK;
+const uint p_loop_col = WorkGroupSize * 4;
+const uint p_loops = (o_row * p_per_loop + p_loop_col - 1) / p_loop_col;
+
+const uint size_max = o_row * p_per_loop / p_reduce;
+const uint max_loops = (size_max + WorkGroupSize - 1) / WorkGroupSize;
+const uint global_max_loops = (o_row + sg_count - 1) / sg_count;
+
+#define SLM_POOL_SIZE_GLOBAL_MAX ((global_max_loops * sg_count + SubGroupSize - 1) / SubGroupSize * SubGroupSize)
+#define SLM_POOL_SIZE_O (sg_count * omat_row * TN * omat_col * TM)
+#define SLM_POOL_SIZE_MAX (max_loops * WorkGroupSize)
+#define SLM_POOL_SIZE_SOFTMAX_SUM ((sg_count * p_loops + SubGroupSize - 1) / SubGroupSize * SubGroupSize)
+#define SLM_POOL_SIZE_OOB_STAGING (sg_count * TM * TK / 4)
+
+const uint psh_size = (p_loops * p_loop_col < omat_row * TN * p_per_loop) ? (omat_row * TN * p_per_loop) : (p_loops * p_loop_col);
+const uint psh_base = SLM_POOL_SIZE_GLOBAL_MAX * 2 + SLM_POOL_SIZE_SOFTMAX_SUM * 2 + SLM_POOL_SIZE_MAX * 2 * 2;
+
+const uint global_maxsh_offset = 0;
+const uint lfsh_offset = global_maxsh_offset + SLM_POOL_SIZE_GLOBAL_MAX;
+const uint osh_offset = lfsh_offset + SLM_POOL_SIZE_SOFTMAX_SUM;
+const uint maxsh_offset = SLM_POOL_SIZE_GLOBAL_MAX + SLM_POOL_SIZE_SOFTMAX_SUM;
+
+const uint p_v4 = p_per_loop / 4;
+const uint max_per_loop = p_per_loop / p_reduce;
+const uint maxsh_stride = p_reduce / 4;
+const uint sg_per_line = p_per_loop / SubGroupSize / 4;
+const bool seq_boundary_check = (Flags & UNALIGNED_SEQ_LEN_BIT) != 0;
+
+shared sh_block {
+ float F32sh[SLM_POOL_SIZE_GLOBAL_MAX + SLM_POOL_SIZE_SOFTMAX_SUM + SLM_POOL_SIZE_O];
+} F32Poolsh;
+
+shared shf16_alias {
+ float16_t F16sh[SLM_POOL_SIZE_GLOBAL_MAX * 2 + SLM_POOL_SIZE_SOFTMAX_SUM * 2 + SLM_POOL_SIZE_MAX * 2 * 2 + psh_size * 2];
+} F16Poolsh;
+
+shared shf16v4_alias {
+ f16vec4 F16V4sh[(SLM_POOL_SIZE_GLOBAL_MAX + SLM_POOL_SIZE_SOFTMAX_SUM + SLM_POOL_SIZE_MAX * 2) / 2 + SLM_POOL_SIZE_OOB_STAGING];
+} F16V4Poolsh;
void main() {
- const uint lane = gl_SubgroupInvocationID;
- const uint v = gl_WorkGroupID.y;
- const uint d = gl_WorkGroupID.z;
- const uint vWarpIdx = gl_WorkGroupID.x % V_HEAD_GROUPS;
- const uint outTokIdx = gl_WorkGroupID.x / V_HEAD_GROUPS;
- const uint localLinearId = gl_SubgroupID;
- const uint wgLane = localLinearId * WARPSIZE + lane;
- const uint splitIdx = localLinearId;
- const uint maxDim = (p.kvSeqLen + MATP_REDUCE - 1) / MATP_REDUCE;
- const uint pDim = maxDim * MATP_REDUCE;
- const uint kvDim = p.kvStride1;
- const uint oDim = p.qHead * HEAD_DIM;
- const uint offsetBaseP = (d * p.activationLength * p.qHead + v * GQA_RATIO + outTokIdx * N_TOKS_PER_GROUP * p.qHead) * pDim / 4;
- const uint offsetBaseMax = (d * p.activationLength * p.qHead + v * GQA_RATIO + outTokIdx * N_TOKS_PER_GROUP * p.qHead) * maxDim;
- const uint offsetBaseV = (d * p.batchStrideV + v * p.kvStride2 + vWarpIdx * WARP_V_DIM + splitIdx * TK * kvDim);
- const uint offsetSlmP = (SLM_OFFSET_P + wgLane * 4);
- const float fp32Min = uintBitsToFloat(0xFEFFFFFF);
- const float fp32Max = uintBitsToFloat(0x7EFFFFFF);
- uint offsetV = offsetBaseV;
-
- coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator> sums[ALIGNED_P_M][MAT_O_N];
- f16vec4 pStateTemp[P_LOAD_LOOPS];
-
- float fp32CompensationP[P_LOAD_LOOPS];
-
- uint loadRowBase[P_LOAD_LOOPS];
- uint loadColBase[P_LOAD_LOOPS];
- float fp32SoftMaxSum[P_LOAD_LOOPS];
- float fp32GlobalMaxP[P_LOAD_LOOPS];
- uint maxRowBase[MAX_LOAD_LOOPS];
- uint maxColBase[MAX_LOAD_LOOPS];
- uint outOffsets[ALIGNED_P_M];
- bool outputMask[ALIGNED_P_M];
- float fp32SinkCoeff[ALIGNED_P_M];
-
- [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
- const uint flatOffset = pm * WARPSIZE + lane;
- const uint inGroupTokIdx = flatOffset / GQA_RATIO;
- const uint inGroupHeadIdx = flatOffset % GQA_RATIO;
- outputMask[pm] = (N_TOKS_PER_GROUP * outTokIdx + inGroupTokIdx < p.activationLength) && (inGroupHeadIdx < GQA_RATIO) && (inGroupTokIdx < N_TOKS_PER_GROUP);
- outOffsets[pm] = (inGroupTokIdx * oDim + inGroupHeadIdx * HEAD_DIM) / 4;
- if ((0x1 & p.flag) != 0) {
- fp32SinkCoeff[pm] = sink_f32[inGroupHeadIdx + v * GQA_RATIO];
- }
- }
-
- [[unroll]] for (uint maxCount = 0; maxCount < MAX_REDUCE_COUNT; maxCount++) {
- const uint flatIdx = maxCount * SUBGROUP_COUNT + localLinearId;
- const uint rowIdx = flatIdx % GQA_RATIO;
- const uint tokIdx = flatIdx / GQA_RATIO;
-
- if (tokIdx < N_TOKS_PER_GROUP) {
- float fp32MaxReduce = fp32Min;
- const uint maxOffset = offsetBaseMax + (tokIdx * p.qHead + rowIdx) * maxDim;
- [[unroll]] for (uint maxReduce = 0; maxReduce < (maxDim + WARPSIZE - 1) / WARPSIZE; maxReduce++) {
- if (maxReduce * WARPSIZE + lane < maxDim) {
- fp32MaxReduce = max(fp32MaxReduce, max_f32[maxOffset + maxReduce * WARPSIZE + lane]);
+ const uint lane = gl_SubgroupInvocationID;
+ const uint head_v_idx = gl_WorkGroupID.y;
+ const uint batch_v_ratio = p.neq3 / p.nev3;
+ const uint batch_idx = gl_WorkGroupID.z;
+ const uint batch_v_idx = batch_idx / batch_v_ratio;
+ const uint v_idx = gl_WorkGroupID.x % v_head_wgs;
+ const uint out_row_idx = gl_WorkGroupID.x / v_head_wgs;
+ const uint sg = gl_SubgroupID;
+ const uint wg_lane = sg * SubGroupSize + lane;
+ const uint max_dim = (p.nek1 + p_reduce - 1) / p_reduce;
+ const uint p_dim = max_dim * p_reduce;
+ const uint v_dim = p.nbv1;
+ const uint o_dim = p.neq2 * HSV;
+ const uint p_base = (batch_idx * p.neq1 * p.neq2 + head_v_idx * gqa_ratio + out_row_idx * Wr * p.neq2) * p_dim / 4;
+ const uint max_base = (batch_idx * p.neq1 * p.neq2 + head_v_idx * gqa_ratio + out_row_idx * Wr * p.neq2) * max_dim;
+ const uint v_base = batch_v_idx * p.nbv3 + head_v_idx * p.nbv2 + v_idx * Bc + sg * TK * v_dim;
+ const uint o_base = (batch_idx * p.nb3 + v_idx * Bc + head_v_idx * gqa_ratio * HSV + out_row_idx * o_dim * Wr) / 4;
+ const uint psh_offset = psh_base + wg_lane * 4;
+ const float NEG_FLT_MAX_OVER_2 = uintBitsToFloat(0xFEFFFFFF);
+ const float POS_FLT_MAX_OVER_2 = uintBitsToFloat(0x7EFFFFFF);
+ uint v_offset = v_base;
+
+ coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator> OMat[omat_row][omat_col];
+ f16vec4 pv4_temp[p_loops];
+
+ float eMf[p_loops];
+
+ uint p_row_base[p_loops];
+ uint p_col_base[p_loops];
+ float Lf[p_loops];
+ float global_max[p_loops];
+ uint max_row_base[max_loops];
+ uint max_col_base[max_loops];
+ uint o_offset[o_scalar_row];
+ bool o_mask[o_scalar_row];
+ float sink[o_scalar_row];
+
+ [[unroll]] for (uint r = 0; r < o_scalar_row; ++r) {
+ const uint flat_offset = r * SubGroupSize + lane;
+ const uint row_idx = flat_offset / gqa_ratio;
+ const uint col_idx = flat_offset % gqa_ratio;
+ o_mask[r] = (Wr * out_row_idx + row_idx < p.neq1) && (col_idx < gqa_ratio) && (row_idx < Wr);
+ o_offset[r] = (row_idx * o_dim + col_idx * HSV) / 4;
+ if ((p.flag & SINK_ENABLE_BIT) != 0) {
+ sink[r] = data_s[col_idx + head_v_idx * gqa_ratio];
}
- }
- fp32MaxReduce = subgroupMax(fp32MaxReduce);
- if (lane == 0) {
- slm_pool_f32.slm_pool_o[SLM_OFFSET_GLOBAL_MAX + maxCount * SUBGROUP_COUNT + localLinearId] = fp32MaxReduce;
- }
- } else {
- if (lane == 0) {
- slm_pool_f32.slm_pool_o[SLM_OFFSET_GLOBAL_MAX + maxCount * SUBGROUP_COUNT + localLinearId] = fp32Max;
- }
- }
- }
-
- barrier();
-
- [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
- const uint flatOffset = (pLoad * GROUPSIZE + wgLane) / P_REDUCE_VEC4;
- const uint rowIdxFlat = flatOffset % GQA_RATIO;
- const uint tokenIdxFlat = min(flatOffset / GQA_RATIO, N_TOKS_PER_GROUP - 1);
- loadColBase[pLoad] = (pLoad * GROUPSIZE + wgLane) % P_REDUCE_VEC4;
- loadRowBase[pLoad] = (tokenIdxFlat * p.qHead + rowIdxFlat);
- fp32SoftMaxSum[pLoad] = 0.0f;
- fp32GlobalMaxP[pLoad] = slm_pool_f32.slm_pool_o[SLM_OFFSET_GLOBAL_MAX + flatOffset];
- }
-
- [[unroll]] for (uint maxLoad = 0; maxLoad < MAX_LOAD_LOOPS; maxLoad++) {
- const uint flatOffset = (maxLoad * GROUPSIZE + wgLane) / MAX_PER_LOOP;
- const uint rowIdxFlat = flatOffset % GQA_RATIO;
- const uint tokenIdxFlat = min(flatOffset / GQA_RATIO, N_TOKS_PER_GROUP - 1);
- maxColBase[maxLoad] = (maxLoad * GROUPSIZE + wgLane) % MAX_PER_LOOP;
- maxRowBase[maxLoad] = (tokenIdxFlat * p.qHead + rowIdxFlat);
- }
-
- [[unroll]] for (uint maxLoad = 0; maxLoad < MAX_LOAD_LOOPS; maxLoad++) {
- const uint flatMaxOffset = maxRowBase[maxLoad] * maxDim + maxColBase[maxLoad];
- slm_pool_f32.slm_pool_o[SLM_OFFSET_LOCAL_MAX + maxLoad * GROUPSIZE + wgLane] = max_f32[offsetBaseMax + flatMaxOffset];
- maxColBase[maxLoad] = maxColBase[maxLoad] + MATP_PER_LOOP / MATP_REDUCE;
- }
-
- [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
- const uint flatOffset = loadRowBase[pLoad] * pDim / 4 + loadColBase[pLoad];
- pStateTemp[pLoad] = pStateVec4[offsetBaseP + flatOffset];
- }
-
- [[unroll]] for (uint n = 0; n < ALIGNED_P_M; n++) {
- [[unroll]] for (uint i = 0; i < MAT_O_N; i++) {
- sums[n][i] = coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator>(0.0f);
- }
- }
-
- barrier();
-
- [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
- const uint maxOffset = (pLoad * GROUPSIZE + wgLane) / SLM_MAX_STRIDE;
- if (loadColBase[pLoad] < pDim / 4) {
- fp32CompensationP[pLoad] = slm_pool_f32.slm_pool_o[SLM_OFFSET_LOCAL_MAX + maxOffset];
- float pTemp[4] = float[4](pStateTemp[pLoad].x, pStateTemp[pLoad].y, pStateTemp[pLoad].z, pStateTemp[pLoad].w);
- float compTemp = exp(fp32CompensationP[pLoad] - fp32GlobalMaxP[pLoad]);
- [[unroll]] for (uint kk = 0; kk < 4; kk++) {
- pTemp[kk] = pTemp[kk] * compTemp;
- fp32SoftMaxSum[pLoad] = fp32SoftMaxSum[pLoad] + pTemp[kk];
- slm_pool_alias_f16.slm_pool_pv[offsetSlmP + pLoad * GROUPSIZE * 4 + kk] = float16_t(pTemp[kk]);
- }
- } else {
- [[unroll]] for (uint kk = 0; kk < 4; kk++) {
- slm_pool_alias_f16.slm_pool_pv[offsetSlmP + pLoad * GROUPSIZE * 4 + kk] = float16_t(0.0f);
- }
}
- loadColBase[pLoad] = loadColBase[pLoad] + P_REDUCE_VEC4;
- }
+ if (p.nek1 == 0) {
+ if (sg == 0) {
+ [[unroll]] for (uint c = 0; c < Bc / 4; ++c) {
+ [[unroll]] for (uint r = 0; r < o_scalar_row; ++r) {
+ if (o_mask[r]) {
+ data_ov4[o_base + o_offset[r] + c] = vec4(0.0f, 0.0f, 0.0f, 0.0f);
+ }
+ }
+ }
+ }
- const uint loopCount = (p.kvSeqLen + MATP_PER_LOOP - 1) / MATP_PER_LOOP;
+ return;
+ }
- for (uint loop = 0; loop < loopCount; loop++) {
- const uint slmPingPongLoad = (loop & 0x1);
- const uint slmPingPongStore = ((loop + 1) & 0x1);
+ [[unroll]] for (uint idx = 0; idx < global_max_loops; ++idx) {
+ const uint flat_offset = idx * sg_count + sg;
+ const uint row_idx = flat_offset % gqa_ratio;
+ const uint tok_idx = flat_offset / gqa_ratio;
+
+ if (tok_idx < Wr) {
+ float Maxf = NEG_FLT_MAX_OVER_2;
+ const uint max_offset = max_base + (tok_idx * p.neq2 + row_idx) * max_dim;
+ [[unroll]] for (uint max_reduce = 0; max_reduce < (max_dim + SubGroupSize - 1) / SubGroupSize; ++max_reduce) {
+ if (max_reduce * SubGroupSize + lane < max_dim) {
+ Maxf = max(Maxf, data_max[max_offset + max_reduce * SubGroupSize + lane]);
+ }
+ }
+ Maxf = subgroupMax(Maxf);
+ if (lane == 0) {
+ F32Poolsh.F32sh[global_maxsh_offset + idx * sg_count + sg] = Maxf;
+ }
+ } else {
+ if (lane == 0) {
+ F32Poolsh.F32sh[global_maxsh_offset + idx * sg_count + sg] = POS_FLT_MAX_OVER_2;
+ }
+ }
+ }
- if (loop + 1 < loopCount) {
- [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
- const uint flatOffset = loadRowBase[pLoad] * pDim / 4 + loadColBase[pLoad];
- pStateTemp[pLoad] = pStateVec4[offsetBaseP + flatOffset];
- }
+ barrier();
- [[unroll]] for (uint maxLoad = 0; maxLoad < MAX_LOAD_LOOPS; maxLoad++) {
- const uint flatMaxOffset = maxRowBase[maxLoad] * maxDim + maxColBase[maxLoad];
- slm_pool_f32.slm_pool_o[SLM_OFFSET_LOCAL_MAX + slmPingPongStore * SLM_POOL_SIZE_LOCAL_MAX + maxLoad * GROUPSIZE + wgLane] = max_f32[offsetBaseMax + flatMaxOffset];
- maxColBase[maxLoad] = maxColBase[maxLoad] + MATP_PER_LOOP / MATP_REDUCE;
- }
+ [[unroll]] for (uint idx = 0; idx < p_loops; ++idx) {
+ const uint flat_offset = (idx * WorkGroupSize + wg_lane) / p_v4;
+ const uint row_idx = flat_offset % gqa_ratio;
+ const uint tok_idx = min(flat_offset / gqa_ratio, Wr - 1);
+ p_col_base[idx] = (idx * WorkGroupSize + wg_lane) % p_v4;
+ p_row_base[idx] = (tok_idx * p.neq2 + row_idx);
+ Lf[idx] = 0.0f;
+ global_max[idx] = F32Poolsh.F32sh[global_maxsh_offset + flat_offset];
}
- barrier();
+ [[unroll]] for (uint idx = 0; idx < max_loops; ++idx) {
+ const uint flat_offset = (idx * WorkGroupSize + wg_lane) / max_per_loop;
+ const uint row_idx = flat_offset % gqa_ratio;
+ const uint tok_idx = min(flat_offset / gqa_ratio, Wr - 1);
+ max_col_base[idx] = (idx * WorkGroupSize + wg_lane) % max_per_loop;
+ max_row_base[idx] = (tok_idx * p.neq2 + row_idx);
+ }
- {
- const uint coopMatOffsetP = SLM_OFFSET_P + slmPingPongLoad * SLM_POOL_SIZE_P + splitIdx * TK;
- coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> matV[MAT_O_N];
- [[unroll]] for (uint cc = 0; cc < MAT_O_N; cc++) {
- coopMatLoad(
- matV[cc],
- vState,
- offsetV + TM * cc,
- kvDim,
- gl_CooperativeMatrixLayoutColumnMajor);
- }
- [[unroll]] for (uint mo = 0; mo < ALIGNED_P_M; mo++) {
- coopmat<float16_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> matP;
- coopMatLoad(
- matP,
- slm_pool_alias_f16.slm_pool_pv,
- coopMatOffsetP + mo * TN * MATP_PER_LOOP,
- MATP_PER_LOOP,
- gl_CooperativeMatrixLayoutColumnMajor);
-
- [[unroll]] for (uint no = 0; no < MAT_O_N; no++) {
- sums[mo][no] = coopMatMulAdd(matV[no], matP, sums[mo][no]);
+ [[unroll]] for (uint idx = 0; idx < max_loops; ++idx) {
+ const uint flat_offset = max_row_base[idx] * max_dim + max_col_base[idx];
+ if (seq_boundary_check) {
+ F32Poolsh.F32sh[maxsh_offset + idx * WorkGroupSize + wg_lane] =
+ (max_col_base[idx] < max_dim) ? data_max[max_base + flat_offset] : NEG_FLT_MAX_OVER_2;
+ } else {
+ F32Poolsh.F32sh[maxsh_offset + idx * WorkGroupSize + wg_lane] = data_max[max_base + flat_offset];
}
- }
+ max_col_base[idx] = max_col_base[idx] + p_per_loop / p_reduce;
}
- offsetV += MATP_PER_LOOP * kvDim;
- if (loop * MATP_PER_LOOP + splitIdx * TK >= p.kvSeqLen) {
- offsetV = 0;
- }
- if (loop + 1 < loopCount) {
- [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
- const uint maxOffset = (pLoad * GROUPSIZE + wgLane) / SLM_MAX_STRIDE;
- if (loadColBase[pLoad] < pDim / 4) {
- fp32CompensationP[pLoad] = slm_pool_f32.slm_pool_o[SLM_OFFSET_LOCAL_MAX + slmPingPongStore * SLM_POOL_SIZE_LOCAL_MAX + maxOffset];
- float pTemp[4] = float[4](pStateTemp[pLoad].x, pStateTemp[pLoad].y, pStateTemp[pLoad].z, pStateTemp[pLoad].w);
- float compTemp = exp(fp32CompensationP[pLoad] - fp32GlobalMaxP[pLoad]);
- [[unroll]] for (uint kk = 0; kk < 4; kk++) {
- pTemp[kk] = pTemp[kk] * compTemp;
- fp32SoftMaxSum[pLoad] = fp32SoftMaxSum[pLoad] + pTemp[kk];
- slm_pool_alias_f16.slm_pool_pv[offsetSlmP + slmPingPongStore * SLM_POOL_SIZE_P + pLoad * GROUPSIZE * 4 + kk] = float16_t(pTemp[kk]);
- }
+ [[unroll]] for (uint idx = 0; idx < p_loops; ++idx) {
+ const uint flat_offset = p_row_base[idx] * p_dim / 4 + p_col_base[idx];
+ if (seq_boundary_check) {
+ pv4_temp[idx] = (p_col_base[idx] < p_dim / 4) ? data_pv4[p_base + flat_offset] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
} else {
- [[unroll]] for (uint kk = 0; kk < 4; kk++) {
- slm_pool_alias_f16.slm_pool_pv[offsetSlmP + slmPingPongStore * SLM_POOL_SIZE_P + pLoad * GROUPSIZE * 4 + kk] = float16_t(0.0f);
- }
+ pv4_temp[idx] = data_pv4[p_base + flat_offset];
}
- loadColBase[pLoad] = loadColBase[pLoad] + P_REDUCE_VEC4;
- }
}
- }
-
- barrier();
-
- [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
- fp32SoftMaxSum[pLoad] = subgroupAdd(fp32SoftMaxSum[pLoad]);
- }
-
- [[unroll]] for (uint mo = 0; mo < ALIGNED_P_M; mo++) {
- [[unroll]] for (uint no = 0; no < MAT_O_N; no++) {
- coopMatStore(
- sums[mo][no],
- slm_pool_f32.slm_pool_o,
- SLM_OFFSET_O + mo * TN * WARP_V_DIM + TM * no + localLinearId * ALIGNED_P_M * TN * WARP_V_DIM,
- WARP_V_DIM,
- gl_CooperativeMatrixLayoutColumnMajor);
+
+ [[unroll]] for (uint r = 0; r < omat_row; ++r) {
+ [[unroll]] for (uint c = 0; c < omat_col; ++c) {
+ OMat[r][c] = coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator>(0.0f);
+ }
}
- }
- [[unroll]] for (uint pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) {
- slm_pool_f32.slm_pool_o[SLM_OFFSET_SOFTMAX_SUM + pLoad * SUBGROUP_COUNT + localLinearId] = fp32SoftMaxSum[pLoad];
- }
+ barrier();
- barrier();
+ [[unroll]] for (uint idx = 0; idx < p_loops; ++idx) {
+ const uint max_offset = (idx * WorkGroupSize + wg_lane) / maxsh_stride;
+ eMf[idx] = F32Poolsh.F32sh[maxsh_offset + max_offset];
+ float p_temp[4] = float[4](pv4_temp[idx].x, pv4_temp[idx].y, pv4_temp[idx].z, pv4_temp[idx].w);
+ float eMf_temp = exp(eMf[idx] - global_max[idx]);
+ [[unroll]] for (uint i = 0; i < 4; ++i) {
+ p_temp[i] = p_temp[i] * eMf_temp;
+ Lf[idx] = Lf[idx] + p_temp[i];
+ F16Poolsh.F16sh[psh_offset + idx * WorkGroupSize * 4 + i] = float16_t(p_temp[i]);
+ }
+ p_col_base[idx] = p_col_base[idx] + p_v4;
+ }
- if (localLinearId == 1) {
- const uint sumBase = SLM_OFFSET_SOFTMAX_SUM + lane * SUB_GROUPS_PER_LINE;
- float sumTemp[ALIGNED_P_M][SUB_GROUPS_PER_LINE];
- [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
- [[unroll]] for (uint reduce = 0; reduce < SUB_GROUPS_PER_LINE; reduce++) {
- sumTemp[pm][reduce] = slm_pool_f32.slm_pool_o[sumBase + pm * WARPSIZE * SUB_GROUPS_PER_LINE + reduce];
- }
+ const uint loop_count = seq_boundary_check ? ((p.nek1 + p_per_loop - 1) / p_per_loop - 1) : (p.nek1 / p_per_loop);
+ uint sh_pingpong_load = 0;
+ uint sh_pingpong_store = 0;
+
+ for (uint loop = 0; loop < loop_count; ++loop) {
+ sh_pingpong_load = loop & 0x1;
+ sh_pingpong_store = (loop + 1) & 0x1;
+
+ if (seq_boundary_check || loop + 1 < loop_count) {
+ [[unroll]] for (uint idx = 0; idx < p_loops; ++idx) {
+ const uint flat_offset = p_row_base[idx] * p_dim / 4 + p_col_base[idx];
+ if (seq_boundary_check) {
+ pv4_temp[idx] = (p_col_base[idx] < p_dim / 4) ? data_pv4[p_base + flat_offset] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+ } else {
+ pv4_temp[idx] = data_pv4[p_base + flat_offset];
+ }
+ }
+
+ [[unroll]] for (uint idx = 0; idx < max_loops; ++idx) {
+ const uint flat_offset = max_row_base[idx] * max_dim + max_col_base[idx];
+ if (seq_boundary_check) {
+ F32Poolsh.F32sh[maxsh_offset + sh_pingpong_store * SLM_POOL_SIZE_MAX + idx * WorkGroupSize + wg_lane] =
+ (max_col_base[idx] < max_dim) ? data_max[max_base + flat_offset] : NEG_FLT_MAX_OVER_2;
+ } else {
+ F32Poolsh.F32sh[maxsh_offset + sh_pingpong_store * SLM_POOL_SIZE_MAX + idx * WorkGroupSize + wg_lane] = data_max[max_base + flat_offset];
+ }
+ max_col_base[idx] = max_col_base[idx] + p_per_loop / p_reduce;
+ }
+ }
+
+ barrier();
+
+ {
+ const uint coord_p = psh_base + sh_pingpong_load * psh_size + sg * TK;
+ coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> VMat[omat_col];
+ [[unroll]] for (uint c = 0; c < omat_col; ++c) {
+ coopMatLoad(VMat[c], data_v, v_offset + TM * c, v_dim, gl_CooperativeMatrixLayoutColumnMajor);
+ }
+ [[unroll]] for (uint r = 0; r < omat_row; ++r) {
+ coopmat<float16_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> PMat;
+ coopMatLoad(PMat, F16Poolsh.F16sh, coord_p + r * TN * p_per_loop, p_per_loop, gl_CooperativeMatrixLayoutColumnMajor);
+
+ [[unroll]] for (uint c = 0; c < omat_col; ++c) {
+ OMat[r][c] = coopMatMulAdd(VMat[c], PMat, OMat[r][c]);
+ }
+ }
+ }
+
+ v_offset += p_per_loop * v_dim;
+
+ if (seq_boundary_check || loop + 1 < loop_count) {
+ [[unroll]] for (uint idx = 0; idx < p_loops; ++idx) {
+ const uint max_offset = (idx * WorkGroupSize + wg_lane) / maxsh_stride;
+ eMf[idx] = F32Poolsh.F32sh[maxsh_offset + sh_pingpong_store * SLM_POOL_SIZE_MAX + max_offset];
+ float p_temp[4] = float[4](pv4_temp[idx].x, pv4_temp[idx].y, pv4_temp[idx].z, pv4_temp[idx].w);
+ float eMf_temp = exp(eMf[idx] - global_max[idx]);
+ [[unroll]] for (uint i = 0; i < 4; ++i) {
+ p_temp[i] = p_temp[i] * eMf_temp;
+ Lf[idx] = Lf[idx] + p_temp[i];
+ F16Poolsh.F16sh[psh_offset + sh_pingpong_store * psh_size + idx * WorkGroupSize * 4 + i] = float16_t(p_temp[i]);
+ }
+ p_col_base[idx] = p_col_base[idx] + p_v4;
+ }
+ }
}
- [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
- [[unroll]] for (uint reduce = 1; reduce < SUB_GROUPS_PER_LINE; reduce++) {
- sumTemp[pm][0] = sumTemp[pm][0] + sumTemp[pm][reduce];
- }
+ barrier();
+
+ // last loop for oob handling
+ if (seq_boundary_check) {
+ const uint coord_p = psh_base + sh_pingpong_store * psh_size + sg * TK;
+ coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> VMat[omat_col];
+ coopmat<float16_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> PMat[omat_row];
+
+ [[unroll]] for (uint r = 0; r < omat_row; ++r) {
+ coopMatLoad(PMat[r], F16Poolsh.F16sh, coord_p + r * TN * p_per_loop, p_per_loop, gl_CooperativeMatrixLayoutColumnMajor);
+ }
+
+ barrier();
+ const uint v_global_coord_base = sg * TK + loop_count * p_per_loop;
+ const uint v_global_offset = (v_base + loop_count * p_per_loop * p.nbv1) / 4;
+ const uint vv4sh_base = psh_base / 4 + sg * TM * TK / 4;
+ [[unroll]] for (uint c = 0; c < omat_col; ++c) {
+ [[unroll]] for (uint idx = 0; idx < TM * TK / SubGroupSize / 4; ++idx) {
+ const uint vv4_idx = idx * SubGroupSize + lane;
+ const uint row_idx = vv4_idx / (TM / 4);
+ const uint col_idx = vv4_idx % (TM / 4);
+ const uint vv4_offset = v_global_offset + row_idx * p.nbv1 / 4 + c * TM / 4 + col_idx;
+
+ f16vec4 vv4_temp = (v_global_coord_base + row_idx < p.nek1) ? data_vv4[vv4_offset] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+ F16V4Poolsh.F16V4sh[vv4sh_base + vv4_idx] = vv4_temp;
+ }
+
+ subgroupBarrier();
+ coopMatLoad(VMat[c], F16V4Poolsh.F16V4sh, vv4sh_base, TM / 4, gl_CooperativeMatrixLayoutColumnMajor);
+ subgroupBarrier();
+ }
+
+ [[unroll]] for (uint r = 0; r < omat_row; ++r) {
+ [[unroll]] for (uint c = 0; c < omat_col; ++c) {
+ OMat[r][c] = coopMatMulAdd(VMat[c], PMat[r], OMat[r][c]);
+ }
+ }
}
- if ((0x1 & p.flag) != 0) {
- [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
- float fp32GlobalMax = slm_pool_f32.slm_pool_o[SLM_OFFSET_GLOBAL_MAX + pm * WARPSIZE + lane];
- float sinkCompensation = fp32GlobalMax - fp32SinkCoeff[pm];
- sinkCompensation = exp(sinkCompensation);
- float softmaxSumTemp = sumTemp[pm][0] * sinkCompensation;
- sumTemp[pm][0] = sumTemp[pm][0] + 1.0f / sinkCompensation;
- sumTemp[pm][0] = 1.0f / sumTemp[pm][0];
- sinkCompensation = sinkCompensation / (1.0f + softmaxSumTemp);
- sumTemp[pm][0] = fp32GlobalMax < fp32SinkCoeff[pm] ? sinkCompensation : sumTemp[pm][0];
- slm_pool_f32.slm_pool_o[SLM_OFFSET_SOFTMAX_SUM + pm * WARPSIZE + lane] = sumTemp[pm][0];
- }
- } else {
- [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
- slm_pool_f32.slm_pool_o[SLM_OFFSET_SOFTMAX_SUM + pm * WARPSIZE + lane] = 1.0f / sumTemp[pm][0];
- }
+ barrier();
+
+ [[unroll]] for (uint idx = 0; idx < p_loops; ++idx) {
+ Lf[idx] = subgroupAdd(Lf[idx]);
}
- }
-
- [[unroll]] for (uint reduce = 2; reduce < SPLIT_P_GROUPS; reduce = reduce << 1 ) {
- const uint stride = (reduce >> 1) * ALIGNED_P_M * TN * MAT_O_N * TM;
- if ((localLinearId % reduce) == 0) {
- const uint reduceBase = localLinearId * ALIGNED_P_M * TN * MAT_O_N * TM + SLM_OFFSET_O;
- float sumTemp0[4];
- float sumTemp1[4];
- const uint reduceVec4Count = ALIGNED_P_M * TN * MAT_O_N * TM / 4 / WARPSIZE;
- [[unroll]] for (uint totalLoads = 0; totalLoads < reduceVec4Count; totalLoads++) {
- [[unroll]] for (uint kk = 0; kk < 4; kk++) {
- sumTemp0[kk] = slm_pool_f32.slm_pool_o[reduceBase + totalLoads * 4 * WARPSIZE + 4 * lane + kk];
- sumTemp1[kk] = slm_pool_f32.slm_pool_o[reduceBase + stride + totalLoads * 4 * WARPSIZE + 4 * lane + kk];
- }
- [[unroll]] for (uint kk = 0; kk < 4; kk++) {
- sumTemp0[kk] = sumTemp0[kk] + sumTemp1[kk];
+ [[unroll]] for (uint r = 0; r < omat_row; ++r) {
+ [[unroll]] for (uint c = 0; c < omat_col; ++c) {
+ const uint coord = osh_offset + r * TN * Bc + TM * c + sg * omat_row * TN * Bc;
+ coopMatStore(OMat[r][c], F32Poolsh.F32sh, coord, Bc, gl_CooperativeMatrixLayoutColumnMajor);
}
+ }
- [[unroll]] for (uint kk = 0; kk < 4; kk++) {
- slm_pool_f32.slm_pool_o[reduceBase + totalLoads * 4 * WARPSIZE + 4 * lane + kk] = sumTemp0[kk];
+ [[unroll]] for (uint idx = 0; idx < p_loops; ++idx) {
+ if (lane == 0) {
+ F32Poolsh.F32sh[lfsh_offset + idx * sg_count + sg] = Lf[idx];
}
- }
}
+
barrier();
- }
- if (localLinearId == 0) {
- const uint slmBase0 = SLM_OFFSET_O + lane * WARP_V_DIM;
- const uint slmBase1 = slmBase0 + SPLIT_P_GROUPS / 2 * ALIGNED_P_M * TN * MAT_O_N * TM;
+ if (sg == 1) {
+ const uint lf_base = lfsh_offset + lane * sg_per_line;
+ float Lfrcp_temp[o_scalar_row][sg_per_line];
+ [[unroll]] for (uint r = 0; r < o_scalar_row; ++r) {
+ [[unroll]] for (uint reduce = 0; reduce < sg_per_line; ++reduce) {
+ Lfrcp_temp[r][reduce] = F32Poolsh.F32sh[lf_base + r * SubGroupSize * sg_per_line + reduce];
+ }
+ }
+
+ [[unroll]] for (uint pm = 0; pm < o_scalar_row; ++pm) {
+ [[unroll]] for (uint reduce = 1; reduce < sg_per_line; ++reduce) {
+ Lfrcp_temp[pm][0] = Lfrcp_temp[pm][0] + Lfrcp_temp[pm][reduce];
+ }
+ }
- const uint offsetOutBase = (d * p.batchStrideO + vWarpIdx * WARP_V_DIM + v * GQA_RATIO * HEAD_DIM + outTokIdx * oDim * N_TOKS_PER_GROUP) / 4;
- float fp32SoftMaxMul[ALIGNED_P_M];
- float fp32Output[ALIGNED_P_M][4];
- [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
- fp32SoftMaxMul[pm] = slm_pool_f32.slm_pool_o[SLM_OFFSET_SOFTMAX_SUM + pm * WARPSIZE + lane];
+ subgroupBarrier();
+
+ if ((p.flag & SINK_ENABLE_BIT) != 0) {
+ [[unroll]] for (uint r = 0; r < o_scalar_row; ++r) {
+ float Mf = F32Poolsh.F32sh[global_maxsh_offset + r * SubGroupSize + lane];
+ float sink_emf = Mf - sink[r];
+ sink_emf = exp(sink_emf);
+ float Lf_sink = Lfrcp_temp[r][0] * sink_emf;
+ Lfrcp_temp[r][0] = Lfrcp_temp[r][0] + 1.0f / sink_emf;
+ Lfrcp_temp[r][0] = 1.0f / Lfrcp_temp[r][0];
+ sink_emf = sink_emf / (1.0f + Lf_sink);
+ Lfrcp_temp[r][0] = (Mf < sink[r]) ? sink_emf : Lfrcp_temp[r][0];
+ F32Poolsh.F32sh[lfsh_offset + r * SubGroupSize + lane] = Lfrcp_temp[r][0];
+ }
+ } else {
+ [[unroll]] for (uint r = 0; r < o_scalar_row; ++r) {
+ F32Poolsh.F32sh[lfsh_offset + r * SubGroupSize + lane] = (Lfrcp_temp[r][0] == 0.0f) ? 0.0f : (1.0f / Lfrcp_temp[r][0]);
+ }
+ }
}
- [[unroll]] for (uint vg = 0; vg < WARP_V_DIM / 4; vg++) {
- [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
- [[unroll]] for (uint vc = 0; vc < 4; vc++) {
- fp32Output[pm][vc] = slm_pool_f32.slm_pool_o[slmBase0 + pm * WARPSIZE * WARP_V_DIM + vg * 4 + vc] * fp32SoftMaxMul[pm];
- fp32Output[pm][vc] = fp32Output[pm][vc] + slm_pool_f32.slm_pool_o[slmBase1 + pm * WARPSIZE * WARP_V_DIM + vg * 4 + vc] * fp32SoftMaxMul[pm];
+ [[unroll]] for (uint reduce = 2; reduce < split_p_sg; reduce = reduce << 1) {
+ const uint stride = (reduce >> 1) * omat_row * TN * omat_col * TM;
+ if ((sg % reduce) == 0) {
+ const uint reduce_base = sg * omat_row * TN * omat_col * TM + osh_offset;
+ float sum_temp0[4];
+ float sum_temp1[4];
+ const uint reduce_count = omat_row * TN * omat_col * TM / 4 / SubGroupSize;
+ [[unroll]] for (uint idx = 0; idx < reduce_count; ++idx) {
+ [[unroll]] for (uint i = 0; i < 4; ++i) {
+ sum_temp0[i] = F32Poolsh.F32sh[reduce_base + idx * 4 * SubGroupSize + 4 * lane + i];
+ sum_temp1[i] = F32Poolsh.F32sh[reduce_base + stride + idx * 4 * SubGroupSize + 4 * lane + i];
+ }
+
+ [[unroll]] for (uint i = 0; i < 4; ++i) {
+ sum_temp0[i] = sum_temp0[i] + sum_temp1[i];
+ }
+
+ [[unroll]] for (uint i = 0; i < 4; ++i) {
+ F32Poolsh.F32sh[reduce_base + idx * 4 * SubGroupSize + 4 * lane + i] = sum_temp0[i];
+ }
+ }
+ }
+ barrier();
+ }
+
+ if (sg == 0) {
+ const uint sh_base0 = osh_offset + lane * Bc;
+ const uint sh_base1 = sh_base0 + split_p_sg / 2 * omat_row * TN * omat_col * TM;
+
+ float Lfrcp[o_scalar_row];
+ float Of[o_scalar_row][4];
+ [[unroll]] for (uint r = 0; r < o_scalar_row; ++r) {
+ Lfrcp[r] = F32Poolsh.F32sh[lfsh_offset + r * SubGroupSize + lane];
}
- }
- [[unroll]] for (uint pm = 0; pm < ALIGNED_P_M; pm++) {
- if (outputMask[pm] == true) {
- out_f32_vec4[offsetOutBase + outOffsets[pm] + vg] = vec4(fp32Output[pm][0], fp32Output[pm][1], fp32Output[pm][2], fp32Output[pm][3]);
+ [[unroll]] for (uint c = 0; c < Bc / 4; ++c) {
+ [[unroll]] for (uint r = 0; r < o_scalar_row; ++r) {
+ [[unroll]] for (uint vc = 0; vc < 4; ++vc) {
+ Of[r][vc] = F32Poolsh.F32sh[sh_base0 + r * SubGroupSize * Bc + c * 4 + vc] * Lfrcp[r];
+ Of[r][vc] = Of[r][vc] + F32Poolsh.F32sh[sh_base1 + r * SubGroupSize * Bc + c * 4 + vc] * Lfrcp[r];
+ }
+ }
+
+ [[unroll]] for (uint r = 0; r < o_scalar_row; ++r) {
+ if (o_mask[r]) {
+ data_ov4[o_base + o_offset[r] + c] = vec4(Of[r][0], Of[r][1], Of[r][2], Of[r][3]);
+ }
+ }
}
- }
}
- }
}
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_prefill.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_prefill.comp
new file mode 100644
index 000000000..407654ca5
--- /dev/null
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_prefill.comp
@@ -0,0 +1,817 @@
+#version 450
+
+#extension GL_EXT_control_flow_attributes : enable
+#extension GL_EXT_shader_16bit_storage : require
+#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
+#extension GL_KHR_cooperative_matrix : enable
+#extension GL_KHR_memory_scope_semantics : enable
+#extension GL_KHR_shader_subgroup_basic : enable
+#extension GL_EXT_shared_memory_block : enable
+#extension GL_KHR_shader_subgroup_shuffle : enable
+#extension GL_KHR_shader_subgroup_shuffle_relative : enable
+
+layout (local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+
+layout (binding = 0) readonly buffer Q {float data_q[];};
+layout (binding = 0) readonly buffer QM2x4 {mat2x4 data_qm2x4[];};
+layout (binding = 1) readonly buffer K {float16_t data_k[];};
+layout (binding = 1) readonly buffer KV4 {f16vec4 data_kv4[];};
+layout (binding = 1) readonly buffer KM2x4 {f16mat2x4 data_km2x4[];};
+layout (binding = 2) readonly buffer V {float16_t data_v[];};
+layout (binding = 2) readonly buffer VV4 {f16vec4 data_vv4[];};
+layout (binding = 3) readonly buffer M {float16_t data_m[];};
+layout (binding = 3) readonly buffer MV2 {f16vec2 data_mv2[];};
+layout (binding = 3) readonly buffer MV4 {f16vec4 data_mv4[];};
+layout (binding = 4) readonly buffer S {float data_s[];};
+layout (binding = 5) buffer OV4 {vec4 data_ov4[];};
+
+layout (push_constant) uniform parameter
+{
+ uint neq1;
+ uint neq2;
+ uint neq3;
+ uint nek1;
+ uint nek2;
+ uint nek3;
+ uint nev2;
+ uint nev3;
+ uint nem1;
+ uint nem2;
+ uint nem3;
+ uint nbq1;
+ uint nbq2;
+ uint nbk1;
+ uint nbk2;
+ uint nbv1;
+ uint nbv2;
+ uint nbq3;
+ uint nbk3;
+ uint nbv3;
+ uint nb3;
+ uint flag;
+ float scale;
+} p;
+
+layout (constant_id = 0) const uint WorkGroupSize = 256;
+layout (constant_id = 1) const uint Wr = 128;
+layout (constant_id = 2) const uint Bc = 64;
+layout (constant_id = 3) const uint HSK = 128;
+layout (constant_id = 4) const uint HSV = 128;
+layout (constant_id = 5) const uint qk_ratio = 8;
+layout (constant_id = 6) const uint step_p = 32;
+layout (constant_id = 7) const uint TM = 8;
+layout (constant_id = 8) const uint TN = 16;
+layout (constant_id = 9) const uint q_reg = 96;
+layout (constant_id = 10) const uint shmem_size = 64;
+layout (constant_id = 11) const uint acc_count = 1;
+layout (constant_id = 12) const uint Flags = 0;
+layout (constant_id = 13) const uint SubGroupSize = 16;
+layout (constant_id = 14) const uint gqa_fold = 1;
+
+#define UNALIGNED_SEQ_LEN_BIT (1 << 0)
+#define UNALIGNED_Q_STRIDE_BIT (1 << 1)
+#define PARTIAL_ROW_MASK_BIT (1 << 2)
+#define INTEL_GPU_BIT (1 << 3)
+#define SINK_ENABLE_BIT (1 << 16)
+
+const bool aligned_q_stride = (Flags & UNALIGNED_Q_STRIDE_BIT) == 0;
+const bool seq_boundary_check = (Flags & UNALIGNED_SEQ_LEN_BIT) != 0;
+const bool partial_row_mask = (Flags & PARTIAL_ROW_MASK_BIT) != 0;
+const uint TK = 16;
+const uint sg_count = WorkGroupSize / SubGroupSize;
+const uint matp_count = step_p / TM;
+const uint mato_count = Bc / TM;
+const uint qk_sg = Wr / TN;
+const uint folded_row_wg = Wr / gqa_fold;
+const uint loadv_sg = (qk_sg == sg_count) ? sg_count : (sg_count - qk_sg);
+const uint k_col = (HSK % 64 == 0 && HSK <= 512 && WorkGroupSize <= 256) ? 8 : 4;
+const uint k_loop_col = WorkGroupSize * k_col;
+const uint k_loops = (step_p * HSK + k_loop_col - 1) / k_loop_col;
+const bool k_boundary_check = (step_p * HSK) % k_loop_col != 0;
+const uint v_vert_groups = step_p / TK;
+const uint v_hori_groups = loadv_sg / v_vert_groups;
+const uint v_loops = (HSV / TM + v_hori_groups - 1) / v_hori_groups;
+const bool v_boundary_check = HSV % (v_hori_groups * TM) != 0;
+const uint m_vert_groups = folded_row_wg / SubGroupSize;
+const uint m_hori_groups = sg_count / m_vert_groups;
+const uint m_loops = (step_p + m_hori_groups - 1) / m_hori_groups;
+
+const uint matq_count = q_reg / TK;
+const uint q_shmem = HSK - q_reg;
+const uint matq_shmem_count = q_shmem / TK;
+const uint max_q_reg_staging = shmem_size / 2 / Wr;
+
+const bool intel_gpu = (Flags & INTEL_GPU_BIT) != 0;
+const bool k_in_shmem = HSK <= 576;
+const bool v_in_shmem = HSK <= 128 && HSV <= 128;
+const bool mask_in_shmem = folded_row_wg <= 64;
+
+const uint oobsh_staging_size = sg_count * TM * TK;
+const uint qsh_size = Wr * q_shmem;
+const uint ksh_size = k_in_shmem ? (k_loops * k_loop_col) : oobsh_staging_size;
+const uint msh_size = mask_in_shmem ? (folded_row_wg * step_p) : 0;
+const uint psh_size = Wr * step_p;
+const uint vsh_size = v_in_shmem ? (step_p * HSV) : 0;
+const uint emfsh_size = (qk_sg == sg_count) ? 0 : Wr;
+const uint pvsh_size = intel_gpu ? 0 : (acc_count * TM * TN * sg_count);
+
+const uint qsh_base = 0;
+const uint emfsh_base = (qsh_base + qsh_size) / 2;
+const uint vsh_base = (emfsh_base + emfsh_size) * 2;
+const uint psh_base = (vsh_base + vsh_size) / 2;
+const uint ksh_base = (psh_base + psh_size) * 2;
+const uint pvsh_base = ksh_base / 2;
+const uint msh_base = ksh_base + ksh_size;
+const uint vsh_oob_staging_base = ksh_base / 4;
+
+const uint sg_mat_reduce = SubGroupSize / TN;
+const uint softmax_loops = step_p / sg_mat_reduce;
+const uint acc_loops = acc_count * TM / sg_mat_reduce;
+const uint CM1_ELEMS = (TM * TN) / SubGroupSize;
+const uint tk_sg_reduce = (TK + SubGroupSize - 1) / SubGroupSize;
+const uint o_sg_cols = (sg_count + qk_sg - 1) / qk_sg;
+const bool v_row_overshoot_check = o_sg_cols * Bc > HSV;
+
+shared sh_block {
+ float F32sh[shmem_size / 4];
+} F32Poolsh;
+
+shared shf16_alias {
+ float16_t F16sh[shmem_size / 2];
+} F16Poolsh;
+
+shared shf16v2_alias {
+ f16vec2 F16V2sh[shmem_size / 4];
+} F16V2Poolsh;
+
+shared shf16v4_alias {
+ f16vec4 F16V4sh[shmem_size / 8];
+} F16V4Poolsh;
+
+shared shf16m2x4_alias {
+ f16mat2x4 M2x4sh[shmem_size / 16];
+} F16M2x4Poolsh;
+
+void main() {
+ const uint lane = gl_SubgroupInvocationID;
+ const uint lane_c = lane % TN;
+ const uint lane_r = lane / TN;
+ const uint pv_ratio = p.neq2 / p.nev2;
+ const uint head_q_idx = gl_WorkGroupID.x;
+ const uint head_k_idx = head_q_idx * gqa_fold / qk_ratio;
+ const uint head_v_idx = head_q_idx * gqa_fold / pv_ratio;
+ const uint batch_k_ratio = p.neq3 / p.nek3;
+ const uint batch_v_ratio = p.neq3 / p.nev3;
+ const uint head_mask_idx = head_q_idx % p.nem2;
+ const uint out_r = gl_WorkGroupID.y;
+ const uint batch_idx = gl_WorkGroupID.z;
+ const uint batch_k_idx = batch_idx / batch_k_ratio;
+ const uint batch_v_idx = batch_idx / batch_v_ratio;
+ const uint batch_mask_idx = batch_idx % p.nem3;
+ const uint sg = gl_SubgroupID;
+ const uint batch_mask_stride = p.nek1 * p.nem1 * p.nem2;
+ uint sg_v_r = sg / v_vert_groups;
+ const uint sg_v_c = sg % v_vert_groups;
+ if (qk_sg < sg_count && sg >= qk_sg) {
+ sg_v_r -= qk_sg / v_vert_groups;
+ }
+ const uint sg_m_c = sg % m_hori_groups;
+ const uint sg_m_r = sg / m_hori_groups;
+ const uint sg_o_c = sg / qk_sg;
+ const uint sg_o_r = sg % qk_sg;
+ const uint sg_o_r_h = sg_o_r % gqa_fold;
+ const uint sg_o_r_row = sg_o_r / gqa_fold;
+ const uint o_dim = p.neq2 * HSV;
+ uint q_offset = out_r * folded_row_wg + sg_o_r_row * TN + lane_c;
+ q_offset = min(q_offset, p.neq1 - 1);
+ if (aligned_q_stride) {
+ q_offset = (q_offset * p.nbq1 + batch_idx * p.nbq3 + (head_q_idx * gqa_fold + sg_o_r_h) * p.nbq2) / 8 + lane_r;
+ } else {
+ q_offset = (q_offset * p.nbq1 + batch_idx * p.nbq3 + (head_q_idx * gqa_fold + sg_o_r_h) * p.nbq2) + lane_r;
+ }
+ uint k_base;
+ if (k_in_shmem) {
+ k_base = (batch_k_idx * p.nbk3 + head_k_idx * p.nbk2) / k_col;
+ } else {
+ k_base = batch_k_idx * p.nbk3 + head_k_idx * p.nbk2;
+ }
+ uint k_offset[k_loops];
+ const uint v_base = batch_v_idx * p.nbv3 + head_v_idx * p.nbv2;
+ uint v_offset;
+ if (v_in_shmem) {
+ v_offset = v_base + sg_v_r * TM + sg_v_c * p.nbv1 * TK;
+ } else {
+ v_offset = v_base + sg_o_c * Bc;
+ }
+ uint m_offset;
+ bool mask_in_boundary;
+ if (mask_in_shmem) {
+ m_offset = out_r * folded_row_wg + sg_m_r * SubGroupSize + lane;
+ mask_in_boundary = m_offset < p.nem1;
+ m_offset = min(m_offset, p.nem1 - 1);
+ m_offset = m_offset * p.nek1 + batch_mask_idx * batch_mask_stride + head_mask_idx * p.nek1 * p.nem1;
+ if (seq_boundary_check) {
+ m_offset += sg_m_c;
+ } else if (m_loops % 4 == 0) {
+ m_offset = m_offset / 4 + sg_m_c;
+ } else if (m_loops % 2 == 0) {
+ m_offset = m_offset / 2 + sg_m_c;
+ } else {
+ m_offset += sg_m_c;
+ }
+ } else {
+ m_offset = out_r * folded_row_wg + sg_o_r_row * TN + lane_c;
+ mask_in_boundary = m_offset < p.nem1;
+ m_offset = min(m_offset, p.nem1 - 1);
+ m_offset = m_offset * p.nek1 + batch_mask_idx * batch_mask_stride + lane_r * softmax_loops + head_mask_idx * p.nek1 * p.nem1;
+ if (!seq_boundary_check) {
+ m_offset /= 4;
+ }
+ }
+ const bool mask_row_valid = !partial_row_mask || mask_in_boundary;
+ const uint o_offset = (batch_idx * p.nb3 + out_r * folded_row_wg * o_dim + (head_q_idx * gqa_fold + sg_o_r_h) * HSV +
+ sg_o_r_row * TN * o_dim + sg_o_c * Bc + lane_c * o_dim) / 4 + lane_r;
+ const uint vsh_offset = sg_v_c * TK * HSV + sg_v_r * TM * TK + vsh_base;
+ const uint msh_store_offset = (sg_m_r * SubGroupSize + lane) * step_p + msh_base;
+ const uint msh_load_offset = sg_o_r_row * TN * step_p + lane_c * step_p + lane_r * softmax_loops + msh_base;
+ const uint psh_load_offset = psh_base + sg * TN * step_p + lane_c * step_p + lane_r * softmax_loops;
+ const uint psh_store_offset = (psh_base + sg * TN * step_p + lane_c * step_p) * 2 + lane_r * softmax_loops;
+ const float NEG_FLT_MAX_OVER_2 = uintBitsToFloat(0xFEFFFFFF);
+ const float16_t NEG_F16_MAX_OVER_2 = float16_t(NEG_FLT_MAX_OVER_2);
+ const float sink = ((p.flag & SINK_ENABLE_BIT) != 0) ? data_s[head_q_idx * gqa_fold + sg_o_r_h] : 0.0f;
+ coopmat<float16_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> QMat[matq_count];
+ uint ksh_offset[k_loops];
+ bool mask_load_k[k_loops];
+
+ if (q_reg <= max_q_reg_staging) {
+ if (sg < qk_sg) {
+ if (aligned_q_stride) {
+ [[dont_unroll]] for (uint i = 0; i < q_reg / sg_mat_reduce / 8; ++i) {
+ const uint qsh_offset = (qsh_base + sg * q_reg * TN + lane_c * q_reg) / 8 + i * sg_mat_reduce + lane_r;
+ mat2x4 q_temp = data_qm2x4[q_offset];
+ q_temp *= p.scale;
+ F16M2x4Poolsh.M2x4sh[qsh_offset] = f16mat2x4(q_temp);
+ q_offset += sg_mat_reduce;
+ }
+ } else {
+ [[dont_unroll]] for (uint i = 0; i < q_reg / sg_mat_reduce; ++i) {
+ const uint qsh_offset = (qsh_base + sg * q_reg * TN + lane_c * q_reg) + i * sg_mat_reduce + lane_r;
+ float q_temp = data_q[q_offset];
+ q_temp *= p.scale;
+ F16Poolsh.F16sh[qsh_offset] = float16_t(q_temp);
+ q_offset += sg_mat_reduce;
+ }
+ }
+ }
+
+ subgroupBarrier();
+
+ if (sg < qk_sg) {
+ const uint coord_q = sg * TN * q_reg + qsh_base;
+ [[unroll]] for (uint i = 0; i < matq_count; ++i) {
+ coopMatLoad(QMat[i], F16Poolsh.F16sh, coord_q + TK * i, q_reg, gl_CooperativeMatrixLayoutColumnMajor);
+ }
+ }
+ } else {
+ if (sg < qk_sg) {
+ const uint coord_q = sg * TN * TK + qsh_base;
+ [[unroll]] for (uint iq = 0; iq < q_reg / TK; ++iq) {
+ if (aligned_q_stride) {
+ [[dont_unroll]] for (uint idx = 0; idx < TK / sg_mat_reduce / 8; ++idx) {
+ const uint qsh_offset = (qsh_base + sg * TK * TN + lane_c * TK) / 8 + idx * sg_mat_reduce + lane_r;
+ mat2x4 q_temp = data_qm2x4[q_offset];
+ q_temp *= p.scale;
+ F16M2x4Poolsh.M2x4sh[qsh_offset] = f16mat2x4(q_temp);
+ q_offset += sg_mat_reduce;
+ }
+ } else {
+ [[dont_unroll]] for (uint idx = 0; idx < TK / sg_mat_reduce; ++idx) {
+ const uint qsh_offset = (qsh_base + sg * TK * TN + lane_c * TK) + idx * sg_mat_reduce + lane_r;
+ float q_temp = data_q[q_offset];
+ q_temp *= p.scale;
+ F16Poolsh.F16sh[qsh_offset] = float16_t(q_temp);
+ q_offset += sg_mat_reduce;
+ }
+ }
+
+ subgroupBarrier();
+ coopMatLoad(QMat[iq], F16Poolsh.F16sh, coord_q, TK, gl_CooperativeMatrixLayoutColumnMajor);
+ subgroupBarrier();
+ }
+ }
+ }
+
+ float Of[Bc / sg_mat_reduce];
+ [[unroll]] for (uint i = 0; i < Bc / sg_mat_reduce; ++i) {
+ Of[i] = 0;
+ }
+
+ barrier();
+
+ if (sg < qk_sg) {
+ if (aligned_q_stride) {
+ [[dont_unroll]] for (uint i = 0; i < q_shmem / sg_mat_reduce / 8; ++i) {
+ const uint qsh_offset = (qsh_base + sg * q_shmem * TN + lane_c * q_shmem) / 8 + i * sg_mat_reduce + lane_r;
+ mat2x4 q_temp = data_qm2x4[q_offset + i * sg_mat_reduce];
+ q_temp *= p.scale;
+ F16M2x4Poolsh.M2x4sh[qsh_offset] = f16mat2x4(q_temp);
+ }
+ } else {
+ [[dont_unroll]] for (uint i = 0; i < q_shmem / sg_mat_reduce; ++i) {
+ const uint qsh_offset = (qsh_base + sg * q_shmem * TN + lane_c * q_shmem) + i * sg_mat_reduce + lane_r;
+ float q_temp = data_q[q_offset + i * sg_mat_reduce];
+ q_temp *= p.scale;
+ F16Poolsh.F16sh[qsh_offset] = float16_t(q_temp);
+ }
+ }
+ }
+
+ if (k_in_shmem) {
+ [[unroll]] for (uint idx = 0; idx < k_loops; ++idx) {
+ const uint flat_offset = idx * WorkGroupSize + sg * SubGroupSize + lane;
+ const uint row_idx = flat_offset / (HSK / k_col);
+ const uint col_idx = flat_offset % (HSK / k_col);
+ const uint sh_chunk_c = col_idx % (TK / k_col);
+ const uint sh_chunk_r = col_idx / (TK / k_col);
+ k_offset[idx] = k_base + row_idx * p.nbk1 / k_col + col_idx;
+ ksh_offset[idx] = (ksh_base + row_idx * TK + sh_chunk_r * TK * step_p) / k_col + sh_chunk_c;
+ if (k_boundary_check) {
+ mask_load_k[idx] = row_idx < step_p;
+ }
+ }
+ }
+
+ const uint loop_count = (p.nek1 + step_p - 1) / step_p;
+ const uint k_seq_boundary = ((p.nek1 - 1) * p.nbk1 + batch_k_idx * p.nbk3 + head_k_idx * p.nbk2 + HSK) / k_col - 1;
+ float Mf = NEG_FLT_MAX_OVER_2;
+ float eMf = 0.0f;
+ float Lf = 0;
+ for (uint loop = 0; loop < loop_count; ++loop) {
+ if (k_in_shmem) {
+ if (k_boundary_check) {
+ [[unroll]] for (uint idx = 0; idx < k_loops; ++idx) {
+ if (seq_boundary_check && loop + 1 == loop_count) {
+ k_offset[idx] = min(k_offset[idx], k_seq_boundary);
+ }
+
+ if (mask_load_k[idx]) {
+ if (k_col == 4) {
+ F16V4Poolsh.F16V4sh[ksh_offset[idx]] = data_kv4[k_offset[idx]];
+ } else {
+ F16M2x4Poolsh.M2x4sh[ksh_offset[idx]] = data_km2x4[k_offset[idx]];
+ }
+ }
+ k_offset[idx] += step_p * p.nbk1 / k_col;
+ }
+ } else {
+ [[unroll]] for (uint idx = 0; idx < k_loops; ++idx) {
+ if (seq_boundary_check && loop + 1 == loop_count) {
+ k_offset[idx] = min(k_offset[idx], k_seq_boundary);
+ }
+
+ if (k_col == 4) {
+ F16V4Poolsh.F16V4sh[ksh_offset[idx]] = data_kv4[k_offset[idx]];
+ } else {
+ F16M2x4Poolsh.M2x4sh[ksh_offset[idx]] = data_km2x4[k_offset[idx]];
+ }
+ k_offset[idx] += step_p * p.nbk1 / k_col;
+ }
+ }
+ }
+
+ if (mask_in_shmem) {
+ if (seq_boundary_check) {
+ [[unroll]] for (uint idx = 0; idx < m_loops; ++idx) {
+ if (idx * m_hori_groups + loop * step_p + sg_m_c < p.nek1) {
+ F16Poolsh.F16sh[msh_store_offset + sg_m_c + idx * m_hori_groups] = mask_row_valid ? data_m[m_offset] : float16_t(0.0f);
+ m_offset += m_hori_groups;
+ } else {
+ F16Poolsh.F16sh[msh_store_offset + sg_m_c + idx * m_hori_groups] = NEG_F16_MAX_OVER_2;
+ }
+ }
+ } else if (m_loops % 4 == 0) {
+ [[unroll]] for (uint idx = 0; idx < m_loops / 4; ++idx) {
+ F16V4Poolsh.F16V4sh[msh_store_offset / 4 + sg_m_c + idx * m_hori_groups] =
+ mask_row_valid ? data_mv4[m_offset] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+ m_offset += m_hori_groups;
+ }
+ } else if (m_loops % 2 == 0) {
+ [[unroll]] for (uint idx = 0; idx < m_loops / 2; ++idx) {
+ F16V2Poolsh.F16V2sh[msh_store_offset / 2 + sg_m_c + idx * m_hori_groups] = mask_row_valid ? data_mv2[m_offset] : f16vec2(0.0f, 0.0f);
+ m_offset += m_hori_groups;
+ }
+ } else {
+ [[unroll]] for (uint idx = 0; idx < m_loops; ++idx) {
+ F16Poolsh.F16sh[msh_store_offset + sg_m_c + idx * m_hori_groups] = mask_row_valid ? data_m[m_offset] : float16_t(0.0f);
+ m_offset += m_hori_groups;
+ }
+ }
+ }
+
+ barrier();
+ if (qk_sg == sg_count || sg < qk_sg) {
+ {
+ coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator> PMat[matp_count];
+ coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> KMat[matp_count];
+ [[unroll]] for (uint pc = 0; pc < matp_count; ++pc) {
+ PMat[pc] = coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator>(0.0f);
+ }
+
+ [[unroll]] for (uint d0 = 0; d0 < matq_count; ++d0) {
+ if (k_in_shmem) {
+ [[unroll]] for (uint pc = 0; pc < matp_count; ++pc) {
+ const uint coord = ksh_base + step_p * TK * d0 + TM * TK * pc;
+ coopMatLoad(KMat[pc], F16Poolsh.F16sh, coord, TK, gl_CooperativeMatrixLayoutRowMajor);
+ }
+ } else {
+ [[unroll]] for (uint pc = 0; pc < matp_count; ++pc) {
+ const uint k_base_row = loop * step_p + pc * TM;
+ if (seq_boundary_check && loop + 1 == loop_count) {
+ const uint k_base_offset = (k_base + k_base_row * p.nbk1 + TK * d0) / 4;
+ [[unroll]] for (uint idx = 0; idx < TM * TK / 4 / SubGroupSize; ++idx) {
+ const uint flat_offset = idx * SubGroupSize + lane;
+ const uint row_idx = flat_offset / (TK / 4);
+ const uint col_idx = flat_offset % (TK / 4);
+ const uint k_temp_offset = k_base_offset + row_idx * p.nbk1 / 4 + col_idx;
+ const uint ksh_temp_offset = ksh_base / 4 + sg * TM * TK / 4 + idx * SubGroupSize + lane;
+ f16vec4 kv4_temp = (k_base_row + row_idx) < p.nek1 ? data_kv4[k_temp_offset] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+ F16V4Poolsh.F16V4sh[ksh_temp_offset] = kv4_temp;
+ }
+
+ subgroupBarrier();
+ coopMatLoad(KMat[pc], F16Poolsh.F16sh, ksh_base + sg * TM * TK, TK, gl_CooperativeMatrixLayoutRowMajor);
+ subgroupBarrier();
+ } else {
+ const uint coord = k_base + k_base_row * p.nbk1 + TK * d0;
+ coopMatLoad(KMat[pc], data_k, coord, p.nbk1, gl_CooperativeMatrixLayoutRowMajor);
+ }
+ }
+ }
+
+ [[unroll]] for (uint pc = 0; pc < matp_count; ++pc) {
+ PMat[pc] = coopMatMulAdd(KMat[pc], QMat[d0], PMat[pc]);
+ }
+ }
+
+ [[unroll]] for (uint d0 = 0; d0 < matq_shmem_count; ++d0) {
+ coopmat<float16_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> QMatsh;
+ if (k_in_shmem) {
+ [[unroll]] for (uint pc = 0; pc < matp_count; ++pc) {
+ const uint coord = ksh_base + matq_count * TK * step_p + step_p * TK * d0 + TM * TK * pc;
+ coopMatLoad(KMat[pc], F16Poolsh.F16sh, coord, TK, gl_CooperativeMatrixLayoutRowMajor);
+ }
+ } else {
+ [[unroll]] for (uint pc = 0; pc < matp_count; ++pc) {
+ const uint k_base_row = loop * step_p + pc * TM;
+ if (seq_boundary_check && loop + 1 == loop_count) {
+ const uint k_base_offset = (k_base + k_base_row * p.nbk1 + TK * d0 + TK * matq_count) / 4;
+ [[unroll]] for (uint idx = 0; idx < TM * TK / 4 / SubGroupSize; ++idx) {
+ const uint flat_offset = idx * SubGroupSize + lane;
+ const uint row_idx = flat_offset / (TK / 4);
+ const uint col_idx = flat_offset % (TK / 4);
+ const uint k_temp_offset = k_base_offset + row_idx * p.nbk1 / 4 + col_idx;
+ const uint ksh_temp_offset = ksh_base / 4 + sg * TM * TK / 4 + idx * SubGroupSize + lane;
+ f16vec4 kv4_temp = (k_base_row + row_idx) < p.nek1 ? data_kv4[k_temp_offset] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+ F16V4Poolsh.F16V4sh[ksh_temp_offset] = kv4_temp;
+ }
+
+ subgroupBarrier();
+ coopMatLoad(KMat[pc], F16Poolsh.F16sh, ksh_base + sg * TM * TK, TK, gl_CooperativeMatrixLayoutRowMajor);
+ subgroupBarrier();
+ } else {
+ const uint coord = k_base + k_base_row * p.nbk1 + TK * d0 + TK * matq_count;
+ coopMatLoad(KMat[pc], data_k, coord, p.nbk1, gl_CooperativeMatrixLayoutRowMajor);
+ }
+ }
+ }
+
+ const uint coord = qsh_base + sg * TN * q_shmem + TK * d0;
+ coopMatLoad(QMatsh, F16Poolsh.F16sh, coord, q_shmem, gl_CooperativeMatrixLayoutColumnMajor);
+ [[unroll]] for (uint pc = 0; pc < matp_count; ++pc) {
+ PMat[pc] = coopMatMulAdd(KMat[pc], QMatsh, PMat[pc]);
+ }
+ }
+
+ [[unroll]] for (uint pc = 0; pc < matp_count; ++pc) {
+ coopMatStore(PMat[pc], F32Poolsh.F32sh, psh_base + TM * pc + sg * TN * step_p, step_p, gl_CooperativeMatrixLayoutColumnMajor);
+ }
+ }
+
+ subgroupBarrier();
+
+ float Pf[softmax_loops];
+ float Mnewf = Mf;
+
+ [[unroll]] for (uint idx = 0; idx < softmax_loops; ++idx) {
+ Pf[idx] = F32Poolsh.F32sh[psh_load_offset + idx];
+ }
+
+ if (mask_in_shmem) {
+ [[unroll]] for (uint idx = 0; idx < softmax_loops; ++idx) {
+ Pf[idx] += float(F16Poolsh.F16sh[msh_load_offset + idx]);
+ }
+ } else {
+ if (seq_boundary_check) {
+ [[unroll]] for (uint idx = 0; idx < softmax_loops; ++idx) {
+ if (loop * step_p + idx + lane_r * softmax_loops < p.nek1) {
+ Pf[idx] += mask_row_valid ? float(data_m[m_offset + idx]) : 0.0f;
+ } else {
+ Pf[idx] += NEG_FLT_MAX_OVER_2;
+ }
+ }
+ m_offset += step_p;
+ } else {
+ [[unroll]] for (uint idx = 0; idx < softmax_loops / 4; ++idx) {
+ f16vec4 mv4_temp = mask_row_valid ? data_mv4[m_offset + idx] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+ float mask_f_temp[4] = float[4](float(mv4_temp.x), float(mv4_temp.y), float(mv4_temp.z), float(mv4_temp.w));
+ [[unroll]] for (uint vec_idx = 0; vec_idx < 4; ++vec_idx) {
+ Pf[4 * idx + vec_idx] += mask_f_temp[vec_idx];
+ }
+ }
+ m_offset += step_p / 4;
+ }
+ }
+
+ if (seq_boundary_check && loop + 1 == loop_count) {
+ [[unroll]] for (uint idx = 0; idx < softmax_loops; ++idx) {
+ if (loop * step_p + idx + lane_r * softmax_loops < p.nek1) {
+ Mnewf = max(Mnewf, Pf[idx]);
+ }
+ }
+ } else {
+ [[unroll]] for (uint idx = 0; idx < softmax_loops; ++idx) {
+ Mnewf = max(Mnewf, Pf[idx]);
+ }
+ }
+
+ [[unroll]] for (uint idx = 1; idx < sg_mat_reduce; ++idx) {
+ Mnewf = max(Mnewf, subgroupShuffleDown(Mnewf, idx * TN));
+ }
+
+ Mnewf = subgroupShuffle(Mnewf, lane & (TN - 1));
+ eMf = exp(Mf - Mnewf);
+ Lf *= eMf;
+
+ if (seq_boundary_check && loop + 1 == loop_count) {
+ [[unroll]] for (uint idx = 0; idx < softmax_loops; ++idx) {
+ if (loop * step_p + idx + lane_r * softmax_loops < p.nek1) {
+ Pf[idx] = exp(Pf[idx] - Mnewf);
+ } else {
+ Pf[idx] = 0.0f;
+ }
+ Lf += Pf[idx];
+ }
+ } else {
+ [[unroll]] for (uint idx = 0; idx < softmax_loops; ++idx) {
+ Pf[idx] = exp(Pf[idx] - Mnewf);
+ Lf += Pf[idx];
+ }
+ }
+
+ if (qk_sg != sg_count && lane_r == 0) {
+ F32Poolsh.F32sh[emfsh_base + sg * TN + lane_c] = eMf;
+ }
+
+ if (SubGroupSize > TN) {
+ subgroupBarrier();
+ }
+
+ [[unroll]] for (uint idx = 0; idx < softmax_loops; ++idx) {
+ F16Poolsh.F16sh[psh_store_offset + idx] = float16_t(Pf[idx]);
+ }
+ Mf = Mnewf;
+ }
+
+ if (v_in_shmem && (qk_sg == sg_count || sg >= sg_count - loadv_sg)) {
+ if (seq_boundary_check && loop + 1 == loop_count) {
+ const uint v_global_coord_base = loop * step_p + sg_v_c * TK + lane_c;
+ if (v_boundary_check) {
+ [[dont_unroll]] for (uint c = 0; c < tk_sg_reduce; ++c) {
+ const uint v_coord_temp = v_global_coord_base + c * SubGroupSize;
+ uint v_oob_offset = (v_base + v_coord_temp * p.nbv1 + sg_v_r * TM) / 4 + lane_r;
+ uint vsh_temp_offset = vsh_offset + c * SubGroupSize + lane_c + lane_r * 4 * TK;
+ [[dont_unroll]] for (uint r = 0; r < v_loops - 1; ++r) {
+ [[dont_unroll]] for (uint idx = 0; idx < TM / sg_mat_reduce / 4; ++idx) {
+ f16vec4 vv4_temp = (v_coord_temp < p.nek1) ? data_vv4[v_oob_offset + idx * sg_mat_reduce] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+ F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 0 * TK] = vv4_temp.x;
+ F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 1 * TK] = vv4_temp.y;
+ F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 2 * TK] = vv4_temp.z;
+ F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 3 * TK] = vv4_temp.w;
+ }
+ v_oob_offset += TM * v_hori_groups / 4;
+ vsh_temp_offset += TM * TK * v_hori_groups;
+ }
+ }
+
+ if ((sg_v_r + (v_loops - 1) * v_hori_groups) * TM < HSV) {
+ [[dont_unroll]] for (uint c = 0; c < tk_sg_reduce; ++c) {
+ const uint v_coord_temp = v_global_coord_base + c * SubGroupSize;
+ const uint v_oob_offset = (v_base + v_coord_temp * p.nbv1 + sg_v_r * TM + TM * (v_loops - 1) * v_hori_groups) / 4 + lane_r;
+ const uint vsh_temp_offset = vsh_offset + TM * (v_loops - 1) * TK * v_hori_groups + c * SubGroupSize + lane_c + lane_r * 4 * TK;
+ [[dont_unroll]] for (uint idx = 0; idx < TM / sg_mat_reduce / 4; ++idx) {
+ f16vec4 vv4_temp = (v_coord_temp < p.nek1) ? data_vv4[v_oob_offset + idx * sg_mat_reduce] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+ F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 0 * TK] = vv4_temp.x;
+ F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 1 * TK] = vv4_temp.y;
+ F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 2 * TK] = vv4_temp.z;
+ F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 3 * TK] = vv4_temp.w;
+ }
+ }
+ }
+ } else {
+ [[dont_unroll]] for (uint c = 0; c < tk_sg_reduce; ++c) {
+ const uint v_coord_temp = v_global_coord_base + c * SubGroupSize;
+ uint v_oob_offset = (v_base + v_coord_temp * p.nbv1 + sg_v_r * TM) / 4 + lane_r;
+ uint vsh_temp_offset = vsh_offset + c * SubGroupSize + lane_c + lane_r * 4 * TK;
+ [[dont_unroll]] for (uint r = 0; r < v_loops; ++r) {
+ [[dont_unroll]] for (uint idx = 0; idx < TM / sg_mat_reduce / 4; ++idx) {
+ f16vec4 vv4_temp = (v_coord_temp < p.nek1) ? data_vv4[v_oob_offset + idx * sg_mat_reduce] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+ F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 0 * TK] = vv4_temp.x;
+ F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 1 * TK] = vv4_temp.y;
+ F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 2 * TK] = vv4_temp.z;
+ F16Poolsh.F16sh[vsh_temp_offset + idx * sg_mat_reduce * TK * 4 + 3 * TK] = vv4_temp.w;
+ }
+ v_oob_offset += TM * v_hori_groups / 4;
+ vsh_temp_offset += TM * TK * v_hori_groups;
+ }
+ }
+ }
+ } else {
+ if (v_boundary_check) {
+ [[unroll]] for (uint r = 0; r < v_loops - 1; ++r) {
+ coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> VMat;
+ coopMatLoad(VMat, data_v, v_offset + TM * r * v_hori_groups, p.nbv1, gl_CooperativeMatrixLayoutColumnMajor);
+ coopMatStore(VMat, F16Poolsh.F16sh, vsh_offset + TM * r * TK * v_hori_groups, TK, gl_CooperativeMatrixLayoutRowMajor);
+ }
+
+ if ((sg_v_r + (v_loops - 1) * v_hori_groups) * TM < HSV) {
+ coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> VMat;
+ coopMatLoad(VMat, data_v, v_offset + TM * (v_loops - 1) * v_hori_groups, p.nbv1, gl_CooperativeMatrixLayoutColumnMajor);
+ coopMatStore(VMat, F16Poolsh.F16sh, vsh_offset + TM * (v_loops - 1) * TK * v_hori_groups, TK, gl_CooperativeMatrixLayoutRowMajor);
+ }
+ } else {
+ [[unroll]] for (uint r = 0; r < v_loops; ++r) {
+ coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> VMat;
+ coopMatLoad(VMat, data_v, v_offset + TM * r * v_hori_groups, p.nbv1, gl_CooperativeMatrixLayoutColumnMajor);
+ coopMatStore(VMat, F16Poolsh.F16sh, vsh_offset + TM * r * TK * v_hori_groups, TK, gl_CooperativeMatrixLayoutRowMajor);
+ }
+ }
+ }
+ }
+
+ barrier();
+ {
+ const uint coord_p = psh_base * 2 + sg_o_r * TN * step_p * 2;
+ const uint coord_v_base = vsh_base + sg_o_c * mato_count * TM * TK;
+ const uint coord_pv = pvsh_base + sg * acc_count * TN * TM;
+ const uint pvsh_offset = pvsh_base + sg * acc_count * TN * TM + lane_c * acc_count * TM + lane_r * 4;
+ if (qk_sg != sg_count) {
+ eMf = F32Poolsh.F32sh[emfsh_base + sg_o_r * TN + lane_c];
+ }
+ coopmat<float16_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> PMat[step_p / TK];
+ [[unroll]] for (uint d0 = 0; d0 < step_p / TK; ++d0) {
+ coopMatLoad(PMat[d0], F16Poolsh.F16sh, coord_p + TK * d0, 2 * step_p, gl_CooperativeMatrixLayoutColumnMajor);
+ }
+
+ [[unroll]] for (uint pv_col = 0; pv_col < mato_count / acc_count; ++pv_col) {
+ const uint coord_v = coord_v_base + pv_col * acc_count * TM * TK;
+ coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator> PVMat[acc_count];
+ coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> VMat[acc_count];
+ [[unroll]] for (uint idx = 0; idx < acc_count; ++idx) {
+ PVMat[idx] = coopmat<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator>(0.0f);
+ }
+
+ [[unroll]] for (uint d0 = 0; d0 < step_p / TK; ++d0) {
+ if (v_in_shmem) {
+ [[unroll]] for (uint idx = 0; idx < acc_count; ++idx) {
+ coopMatLoad(VMat[idx], F16Poolsh.F16sh, coord_v + HSV * TK * d0 + TM * TK * idx, TK, gl_CooperativeMatrixLayoutRowMajor);
+ }
+ } else {
+ if (seq_boundary_check && loop + 1 == loop_count) {
+ const uint v_global_coord_base = d0 * TK + loop * step_p + lane_c;
+ const uint vsh_staging_store = vsh_oob_staging_base + (sg * TM * TK + lane_c * TM) / 4 + lane_r;
+ const uint vsh_staging_load = vsh_oob_staging_base * 4 + sg * TM * TK;
+ [[unroll]] for (uint idx = 0; idx < acc_count; ++idx) {
+ subgroupBarrier();
+ [[dont_unroll]] for (uint c = 0; c < tk_sg_reduce; ++c) {
+ const uint v_coord_temp = v_global_coord_base + c * SubGroupSize;
+ const uint v_oob_offset = (v_base + v_coord_temp * p.nbv1 + sg_o_c * Bc + pv_col * acc_count * TM + idx * TM) / 4 + lane_r;
+
+ if (v_row_overshoot_check) {
+ const uint v_row_temp = sg_o_c * Bc + pv_col * acc_count * TM + idx * TM;
+ if (v_row_temp < HSV) {
+ [[dont_unroll]] for (uint r = 0; r < TM / sg_mat_reduce / 4; ++r) {
+ f16vec4 vv4_temp = (v_coord_temp < p.nek1) ? data_vv4[v_oob_offset + sg_mat_reduce * r] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+ F16V4Poolsh.F16V4sh[vsh_staging_store + c * SubGroupSize * TM / 4 + sg_mat_reduce * r] = vv4_temp;
+ }
+ } else {
+ [[dont_unroll]] for (uint r = 0; r < TM / sg_mat_reduce / 4; ++r) {
+ F16V4Poolsh.F16V4sh[vsh_staging_store + c * SubGroupSize * TM / 4 + sg_mat_reduce * r] = f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+ }
+ }
+ } else {
+ [[dont_unroll]] for (uint r = 0; r < TM / sg_mat_reduce / 4; ++r) {
+ f16vec4 vv4_temp = (v_coord_temp < p.nek1) ? data_vv4[v_oob_offset + sg_mat_reduce * r] : f16vec4(0.0f, 0.0f, 0.0f, 0.0f);
+ F16V4Poolsh.F16V4sh[vsh_staging_store + c * SubGroupSize * TM / 4 + sg_mat_reduce * r] = vv4_temp;
+ }
+ }
+ }
+ subgroupBarrier();
+ coopMatLoad(VMat[idx], F16Poolsh.F16sh, vsh_staging_load, TM, gl_CooperativeMatrixLayoutColumnMajor);
+ }
+ } else {
+ const uint v_coord_global = v_offset + pv_col * acc_count * TM + d0 * TK * p.nbv1;
+ if (v_row_overshoot_check) {
+ [[unroll]] for (uint idx = 0; idx < acc_count; ++idx) {
+ const uint v_row_temp = sg_o_c * Bc + pv_col * acc_count * TM + idx * TM;
+ if (v_row_temp < HSV) {
+ coopMatLoad(VMat[idx], data_v, v_coord_global + TM * idx, p.nbv1, gl_CooperativeMatrixLayoutColumnMajor);
+ } else {
+ VMat[idx] = coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA>(0.0f);
+ }
+ }
+ } else {
+ [[unroll]] for (uint idx = 0; idx < acc_count; ++idx) {
+ coopMatLoad(VMat[idx], data_v, v_coord_global + TM * idx, p.nbv1, gl_CooperativeMatrixLayoutColumnMajor);
+ }
+ }
+ }
+ }
+
+ [[unroll]] for (uint idx = 0; idx < acc_count; ++idx) {
+ PVMat[idx] = coopMatMulAdd(VMat[idx], PMat[d0], PVMat[idx]);
+ }
+ }
+
+ if (intel_gpu) {
+ [[unroll]] for (uint idx = 0; idx < acc_count; ++idx) {
+ [[unroll]] for (uint e = 0; e < CM1_ELEMS; ++e) {
+ Of[pv_col * acc_loops + idx * CM1_ELEMS + e] = fma(Of[pv_col * acc_loops + idx * CM1_ELEMS + e],
+ eMf,
+ PVMat[idx][e]);
+ }
+ }
+ } else {
+ barrier();
+
+ [[unroll]] for (uint idx = 0; idx < acc_count; ++idx) {
+ coopMatStore(PVMat[idx], F32Poolsh.F32sh, coord_pv + TM * idx, acc_count * TM, gl_CooperativeMatrixLayoutColumnMajor);
+ }
+
+ subgroupBarrier();
+
+ [[unroll]] for (uint idx = 0; idx < acc_loops / 4; ++idx) {
+ [[unroll]] for (uint vec_idx = 0; vec_idx < 4; ++vec_idx) {
+ Of[pv_col * acc_loops + idx * 4 + vec_idx] = fma(Of[pv_col * acc_loops + idx * 4 + vec_idx],
+ eMf,
+ F32Poolsh.F32sh[pvsh_offset + idx * sg_mat_reduce * 4 + vec_idx]);
+ }
+ }
+ barrier();
+ }
+ }
+ }
+
+ v_offset += p.nbv1 * step_p;
+ barrier();
+ }
+
+ if (sg < qk_sg) {
+ [[unroll]] for (uint idx = 1; idx < sg_mat_reduce; ++idx) {
+ Lf += subgroupShuffleDown(Lf, idx * TN);
+ }
+ if (lane_r == 0) {
+ F32Poolsh.F32sh[sg * TN + lane_c] = Mf;
+ F32Poolsh.F32sh[qk_sg * TN + sg * TN + lane_c] = Lf;
+ }
+ }
+
+ barrier();
+
+ Mf = F32Poolsh.F32sh[sg_o_r * TN + lane_c];
+ Lf = F32Poolsh.F32sh[qk_sg * TN + sg_o_r * TN + lane_c];
+
+ float Lfrcp;
+ if ((p.flag & SINK_ENABLE_BIT) != 0) {
+ float sink_emf = Mf - sink;
+ sink_emf = exp(sink_emf);
+ float Lf_sink = Lf * sink_emf;
+ Lf += 1.0f / sink_emf;
+ Lf = 1.0f / Lf;
+ sink_emf /= (1.0f + Lf_sink);
+ Lfrcp = (Mf < sink) ? sink_emf : Lf;
+ } else {
+ Lfrcp = (Lf == 0.0f) ? 0.0f : (1.0f / Lf);
+ }
+
+ [[unroll]] for (uint out_c = 0; out_c < Bc / sg_mat_reduce; ++out_c) {
+ Of[out_c] *= Lfrcp;
+ }
+
+ const uint sg_out_r = out_r * folded_row_wg + sg_o_r_row * TN + lane_c;
+ [[unroll]] for (uint o_chunk = 0; o_chunk < Bc / sg_mat_reduce / 4; ++o_chunk) {
+ const uint out_c = sg_mat_reduce * 4 * o_chunk + sg_o_c * Bc + lane_r;
+ if (sg_out_r < p.neq1 && out_c < HSV) {
+ data_ov4[o_offset + o_chunk * sg_mat_reduce] = vec4(Of[4 * o_chunk], Of[4 * o_chunk + 1], Of[4 * o_chunk + 2], Of[4 * o_chunk + 3]);
+ }
+ }
+}
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
index c1696e4cd..f9ae46ef8 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
@@ -939,6 +939,7 @@ void process_shaders() {
#if defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT)
string_to_spv("fa_decode_ph1", "flash_attn_decode_phase_1.comp", {}, true, true, false, false);
string_to_spv("fa_decode_ph2", "flash_attn_decode_phase_2.comp", {}, true, true, false, false);
+ string_to_spv("fa_prefill", "flash_attn_prefill.comp", {}, true, true, false, false);
#endif
string_to_spv("fa_sparse_compact", "flash_attn_sparse_compact.comp", {});