diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h b/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h index 68b3200b3..8446e313c 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h +++ b/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h @@ -124,6 +124,24 @@ 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; + 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; +}; + struct vk_op_push_constants { uint32_t KX; uint32_t KY; diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-types.h b/ggml/src/ggml-vulkan/ggml-vulkan-types.h index 67e3361ed..5df1c3900 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-types.h +++ b/ggml/src/ggml-vulkan/ggml-vulkan-types.h @@ -996,6 +996,7 @@ struct vk_device_struct { bool fa_sparse_compact_use_subgroups; vk_pipeline pipeline_flash_attn_split_k_reduce; + std::map, std::pair> pipeline_xe_fa_decode_dual_phases; vk_pipeline pipeline_count_experts; // [2] is for whether to take n_experts from spec constant (0) or push constant (1) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index f7e27703f..1a83ac320 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -2973,6 +2973,46 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_matmul_split_k_reduce, "split_k_reduce", split_k_reduce_len, split_k_reduce_data, "main", 2, 2 * sizeof(uint32_t), {256 * 4, 1, 1}, {}, 1); 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 (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; + } + + for (auto& it : device->pipeline_xe_fa_decode_dual_phases) { + const uint32_t split_p_chunk = 32; + auto HdQk = 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); + } + } + for (auto &it : device->pipeline_fa_mask_opt) { auto BrBc = it.first; ggml_vk_create_pipeline(device, it.second, "fa_mask_opt", fa_mask_opt_len, fa_mask_opt_data, "main", 2, sizeof(vk_op_flash_attn_mask_opt_push_constants), {1, 1, 1}, {128, 128 / device->subgroup_size, BrBc.first, BrBc.second}, 1, true, true, device->subgroup_size); @@ -7899,6 +7939,18 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const vk_pipeline pipeline = nullptr; + 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); + std::pair xe_fa_pipeline_dual_phases = { nullptr , nullptr }; + vk_pipeline xe_fa_pipeline = nullptr; + size_t size_p = 0; + size_t size_group_max = 0; + { std::lock_guard guard(ctx->device->compile_mutex); auto &pipelines = ctx->device->pipeline_flash_attn_f32_f16; @@ -7956,6 +8008,37 @@ 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) { + std::lock_guard guard(ctx->device->compile_mutex); + 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; + 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(), std::make_shared()); + } + + 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 (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) @@ -8111,7 +8194,68 @@ 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 (split_k > 1) { + 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; + 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 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); + 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_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_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 }); + + ctx->prealloc_x_need_sync = true; + } else if (split_k > 1) { ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_split_k_reduce, 1); if (ctx->prealloc_split_k_need_sync) { 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 new file mode 100644 index 000000000..b5f95aaa0 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_1.comp @@ -0,0 +1,263 @@ +#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_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 (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 (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; +} 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; + +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; + } + } + } + + coopmat matP[COOP_MAT_P_M][COOP_MAT_P_N]; + + [[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(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]; + } + } + + barrier(); + if (localLinearId < N_TOK) { + [[unroll]] for (uint loop = 0; loop < loopCount; loop++) { + coopmat matQ[COOP_MAT_P_M]; + coopmat 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 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); + } + + [[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 (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; + } + + [[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]; + } + } + } + + 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 new file mode 100644 index 000000000..60a1c2ce7 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_2.comp @@ -0,0 +1,408 @@ +#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_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 (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 (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; +} 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; + +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 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]); + } + } + 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(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; + } + + const uint loopCount = (p.kvSeqLen + MATP_PER_LOOP - 1) / MATP_PER_LOOP; + + for (uint loop = 0; loop < loopCount; loop++) { + const uint slmPingPongLoad = (loop & 0x1); + const uint slmPingPongStore = ((loop + 1) & 0x1); + + 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]; + } + + [[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; + } + } + + barrier(); + + { + const uint coopMatOffsetP = SLM_OFFSET_P + slmPingPongLoad * SLM_POOL_SIZE_P + splitIdx * TK; + coopmat 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 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]); + } + } + } + + 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]); + } + } 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); + } + } + 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 pLoad = 0; pLoad < P_LOAD_LOOPS; pLoad++) { + slm_pool_f32.slm_pool_o[SLM_OFFSET_SOFTMAX_SUM + pLoad * SUBGROUP_COUNT + localLinearId] = fp32SoftMaxSum[pLoad]; + } + + barrier(); + + 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]; + } + } + + [[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]; + } + } + + 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]; + } + } + } + + [[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 kk = 0; kk < 4; kk++) { + slm_pool_f32.slm_pool_o[reduceBase + totalLoads * 4 * WARPSIZE + 4 * lane + kk] = sumTemp0[kk]; + } + } + } + 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; + + 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]; + } + + [[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 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]); + } + } + } + } +} 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 73cef00b0..5b2479da2 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -922,6 +922,10 @@ void process_shaders() { string_to_spv("fa_split_k_reduce", "flash_attn_split_k_reduce.comp", {}); string_to_spv("fa_mask_opt", "flash_attn_mask_opt.comp", {}); + + 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_sparse_compact", "flash_attn_sparse_compact.comp", {}); string_to_spv("fa_sparse_compact_subgroup", "flash_attn_sparse_compact.comp", {{"USE_SUBGROUPS", "1"}});