public_git/whisper.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/whisper.cpp.git synced 2026-09-30 11:36:38 +02:00
Code Issues Packages Projects Releases Wiki Activity

vulkan: add Intel Xe flash attention optimization kernels (2/3, Xe-LPG Plus/Xe2/Xe3) (llama/24406)

Browse Source 
* vulkan : Intel FA kernel optimization for split k path

* vulkan : Host code update for Intel split k FA kernel path selection, fix A770 Linux op test failures

* vulkan : use symmetric coopMatMulAdd() in flash_attn_decode_phase_1 shader to resolve test op failre on A770 Linux with 26.2.3 mesa driver

* vulkan : fix editorconfig issue in flash_attn_decode_phase_2.comp

---------

Co-authored-by: Liu, Russell <russell.liu@intel.com>
This commit is contained in:
Jiang, Fish
2026-09-23 20:46:47 +03:00
committed by Georgi Gerganov
co-authored by Liu, Russell
parent ad0058be6e
commit e900c88838
6 changed files with 839 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h
+18
View File
@@ -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;
ggml/src/ggml-vulkan/ggml-vulkan-types.h
+1
View File
@@ -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::tuple<uint32_t, uint32_t, uint32_t, uint32_t>, std::pair<vk_pipeline, vk_pipeline>> 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)
ggml/src/ggml-vulkan/ggml-vulkan.cpp
+145 -1
View File
@@ -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<vk_pipeline, vk_pipeline> 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<std::mutex> 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<std::mutex> 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<vk_pipeline_struct>(), std::make_shared<vk_pipeline_struct>());
}
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<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);
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) {
ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_1.comp
+263
View File
@@ -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<float, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator> 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<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];
}
}
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 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();
}
}
}
ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_decode_phase_2.comp
+408
View File
@@ -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<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]);
}
}
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;
}
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<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]);
}
}
}
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]);
}
}
}
}
}
ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
+4
View File
@@ -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"}});