Commit d8f30a5b6 for llama.cpp
commit d8f30a5b6db907110b6f8d67c2f84c1443b871ef
Author: Ruben Ortlam <rortlam@redhat.com>
Date: Sat Nov 14 23:36:58 2026 +0100
vulkan: optimize prompt processing Flash Attention for RDNA3 and RDNA4 (#31323)
* vulkan: optimize Flash Attention cm1 for RDNA3 and RDNA4
* reduce shmem use through buffer reuse
* clean up and fixes
* limit q_in_regs to prefill
* vulkan: add flipped FA path, each subgroup owns 16 Q rows
* cleanup
* move flip shader into own file
* replace kvsep with direct load mode
* fix nvidia hang, tune AMD
* only store pv in sfsh if necessary
* fix skip condition
* nvidia tuning
* reduce occupancy on RDNA4 for large attention head sizes
* fix sparse FA fp16 underflow bug
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-common.h b/ggml/src/ggml-vulkan/ggml-vulkan-common.h
index b74062931..dfd324134 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan-common.h
+++ b/ggml/src/ggml-vulkan/ggml-vulkan-common.h
@@ -34,7 +34,7 @@ bool ggml_vk_intel_windows_driver_in_range(uint32_t driver_version, uint32_t low
// shaders
void ggml_vk_destroy_pipeline(vk::Device& device, vk_pipeline& pipeline);
-vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_t hsk, uint32_t hsv, uint32_t n_rows, uint32_t n_kv, ggml_type k_type, ggml_type v_type, bool f32acc);
+vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_t hsk, uint32_t hsv, uint32_t n_rows, uint32_t n_kv, ggml_type k_type, ggml_type v_type, bool f32acc, uint32_t qk_ratio, bool allow_flip = true);
vk_fa_pipeline_state get_fa_pipeline_state(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool aligned, bool f32acc, bool use_mask, bool use_mask_opt, bool use_logit_softcap, bool use_sparse, ggml_type k_type, ggml_type v_type);
uint32_t get_subgroup_size(const std::string &pipeline_name, const vk_device_architecture &arch);
void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested = nullptr);
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-types.h b/ggml/src/ggml-vulkan/ggml-vulkan-types.h
index ac9d05e6e..64a39bf25 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan-types.h
+++ b/ggml/src/ggml-vulkan/ggml-vulkan-types.h
@@ -1427,13 +1427,16 @@ struct vk_fa_tuning_params {
uint32_t row_split;
bool shmem_staging;
bool disable_subgroups;
+ bool q_in_regs;
+ bool flip;
+ bool flip_direct;
uint32_t limit_occupancy_shmem;
void print() const {
std::cerr << "path=" << path << " workgroup_size=" << workgroup_size << " subgroup_size=" << subgroup_size <<
" block_rows=" << block_rows << " block_cols=" << block_cols << " d_split=" << d_split <<
" row_split=" << row_split << " shmem_staging=" << shmem_staging << " disable_subgroups=" << disable_subgroups <<
- " limit_occupancy_shmem=" << limit_occupancy_shmem << std::endl;
+ " q_in_regs=" << q_in_regs << " flip=" << flip << " flip_direct=" << flip_direct << " limit_occupancy_shmem=" << limit_occupancy_shmem << std::endl;
}
};
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
index b692f2f8b..1af36cdca 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp
+++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
@@ -1287,12 +1287,8 @@ static vk_fa_tuning_params get_fa_tuning_params_scalar(const vk_device& device,
return result;
}
-static vk_fa_tuning_params get_fa_tuning_params_coopmat1(const vk_device& device, uint32_t hsk, uint32_t hsv, uint32_t n_rows, uint32_t n_kv, ggml_type k_type, ggml_type v_type, bool f32acc) {
- GGML_UNUSED(n_rows);
+static vk_fa_tuning_params get_fa_tuning_params_coopmat1(const vk_device& device, uint32_t hsk, uint32_t hsv, uint32_t n_rows, uint32_t n_kv, ggml_type k_type, ggml_type v_type, bool f32acc, uint32_t qk_ratio, bool allow_flip) {
GGML_UNUSED(n_kv);
- GGML_UNUSED(k_type);
- GGML_UNUSED(v_type);
- GGML_UNUSED(f32acc);
vk_fa_tuning_params result{};
result.path = FA_COOPMAT1;
@@ -1304,7 +1300,10 @@ static vk_fa_tuning_params get_fa_tuning_params_coopmat1(const vk_device& device
const uint32_t num_subgroups = 4;
- result.block_rows = coopmat_block_rows;
+ // Halves K/V re-reads; only wins on UMA, regresses on discrete RDNA4
+ const bool double_block_rows = device->uma &&
+ (device->architecture == AMD_RDNA3 || device->architecture == AMD_RDNA4) && n_rows >= 64;
+ result.block_rows = double_block_rows ? 2 * coopmat_block_rows : coopmat_block_rows;
result.block_cols = coopmat_block_cols * num_subgroups;
result.row_split = num_subgroups;
result.subgroup_size = device->subgroup_size;
@@ -1315,6 +1314,41 @@ static vk_fa_tuning_params get_fa_tuning_params_coopmat1(const vk_device& device
result.shmem_staging = (device->vendor_id == VK_VENDOR_ID_NVIDIA && hsk < 256 && hsv < 256) ? 1 : 0;
+ if (double_block_rows && !ggml_vk_flash_attn_coopmat_shmem_support(device, result, hsk, hsv, f32acc, k_type, v_type)) {
+ result.block_rows = coopmat_block_rows;
+ }
+
+ // Br=32 spills registers, only prefill benefits
+ result.q_in_regs = result.block_rows == coopmat_block_rows && device->vendor_id == VK_VENDOR_ID_AMD && n_rows >= 64;
+
+ // Flipped ownership: each subgroup owns 16 Q rows.
+ // uma, nvidia: K/V tiles staged in separate LDS buffers; discrete AMD GQA: K/V loaded directly from global memory.
+ static const bool disable_flip = getenv("GGML_VK_DISABLE_FA_FLIP") != nullptr;
+ static const char * force_flip = getenv("GGML_VK_FORCE_FA_FLIP");
+ const bool flip_arch = force_flip != nullptr || device->vendor_id == VK_VENDOR_ID_NVIDIA ||
+ (device->vendor_id == VK_VENDOR_ID_AMD && (device->architecture == AMD_RDNA3 || device->architecture == AMD_RDNA4));
+ if (allow_flip && !disable_flip && flip_arch && n_rows >= 64 &&
+ (result.subgroup_size == 32 || result.subgroup_size == 64) &&
+ hsk % 16 == 0 && hsv % 16 == 0 && hsk <= 256 && hsv <= 256 &&
+ k_type == v_type && (k_type == GGML_TYPE_F16 || k_type == GGML_TYPE_BF16)) {
+ result.flip = true;
+ result.flip_direct = !device->uma && qk_ratio > 1 && device->vendor_id != VK_VENDOR_ID_NVIDIA;
+ if (force_flip != nullptr && strcmp(force_flip, "kvsep") == 0) {
+ result.flip_direct = false;
+ } else if (force_flip != nullptr && strcmp(force_flip, "direct") == 0) {
+ result.flip_direct = true;
+ }
+ result.block_rows = num_subgroups * coopmat_block_rows;
+ result.block_cols = (device->uma && (hsk >= 256 || hsv >= 256)) ? 16 : 32;
+ result.q_in_regs = false;
+ result.shmem_staging = 0;
+ }
+
+ // q_in_regs LDS lands in the 12-wave band, ACO caps VGPRs at 128 and serializes K loads; pad into the 8-wave band
+ if (device->architecture == AMD_RDNA4 && result.q_in_regs && hsk > 256) {
+ result.limit_occupancy_shmem = 14 * 1024 / 16;
+ }
+
return result;
}
@@ -1346,7 +1380,7 @@ static vk_fa_tuning_params get_fa_tuning_params_coopmat2(const vk_device& device
return result;
}
-vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_t hsk, uint32_t hsv, uint32_t n_rows, uint32_t n_kv, ggml_type k_type, ggml_type v_type, bool f32acc) {
+vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_t hsk, uint32_t hsv, uint32_t n_rows, uint32_t n_kv, ggml_type k_type, ggml_type v_type, bool f32acc, uint32_t qk_ratio, bool allow_flip) {
FaCodePath path = device->coopmat2 ? FA_COOPMAT2 :
device->coopmat1_fa_support ? FA_COOPMAT1 : FA_SCALAR;
@@ -1365,8 +1399,13 @@ vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_t hsk,
if (path == FA_COOPMAT1) {
bool shape_ok = (f32acc && device->coopmat_support_16x16x16_f32acc) ||
(!f32acc && device->coopmat_support_16x16x16_f16acc);
- const vk_fa_tuning_params params = get_fa_tuning_params_coopmat1(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc);
+ vk_fa_tuning_params params = get_fa_tuning_params_coopmat1(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc, qk_ratio, allow_flip);
bool shmem_ok = ggml_vk_flash_attn_coopmat_shmem_support(device, params, hsk, hsv, f32acc, k_type, v_type);
+ if (!shmem_ok && params.flip) {
+ allow_flip = false;
+ params = get_fa_tuning_params_coopmat1(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc, qk_ratio, allow_flip);
+ shmem_ok = ggml_vk_flash_attn_coopmat_shmem_support(device, params, hsk, hsv, f32acc, k_type, v_type);
+ }
if (!shape_ok || !shmem_ok) {
path = FA_SCALAR;
@@ -1382,7 +1421,7 @@ vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_t hsk,
case FA_SCALAR:
return get_fa_tuning_params_scalar(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc);
case FA_COOPMAT1:
- return get_fa_tuning_params_coopmat1(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc);
+ return get_fa_tuning_params_coopmat1(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc, qk_ratio, allow_flip);
case FA_COOPMAT2:
return get_fa_tuning_params_coopmat2(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc);
default:
@@ -1399,7 +1438,10 @@ vk_fa_pipeline_state get_fa_pipeline_state(const vk_device& device, const vk_fa_
(use_mask ? 2 : 0) |
(use_logit_softcap ? 4 : 0) |
(old_amd_windows ? 8 : 0) |
- (use_sparse ? 16 : 0);
+ (use_sparse ? 16 : 0) |
+ (params.q_in_regs ? 32 : 0) |
+ (params.flip ? 64 : 0) |
+ (params.flip_direct ? 128 : 0);
const uint32_t subgroup_size = params.disable_subgroups ? 0 : params.subgroup_size;
@@ -2237,7 +2279,7 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
align, disable_robustness, require_full_subgroups, required_subgroup_size);
};
- // FA scalar has two SPIR-V modules (MMQ vs non-MMQ); FA cm1 has one. K/V
+ // FA scalar has two SPIR-V modules (MMQ vs non-MMQ); FA cm1 has two (flip vs non-flip). K/V
// quant type is selected at runtime via the FaTypeK / FaTypeV spec constants.
for (auto &fa : device->pipeline_flash_attn_f32_f16) {
@@ -2302,6 +2344,7 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
const bool fa_ds = fa.first.subgroup_size == 0;
const bool bf16_kv = fa.first.k_type == GGML_TYPE_BF16;
+ const bool flip = (fa.first.flags & 64) != 0;
const void * spv_data;
size_t spv_size;
@@ -2309,12 +2352,22 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
if (bf16_kv) {
#if defined(VK_KHR_shader_bfloat16) && defined(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT)
if (!device->coopmat_bf16_support) continue;
- spv_data = flash_attn_f32_f16_bf16_cm1_data;
- spv_size = flash_attn_f32_f16_bf16_cm1_len;
- name = aligned ? "flash_attn_f32_bf16_aligned_cm1" : "flash_attn_f32_bf16_cm1";
+ if (flip) {
+ spv_data = flash_attn_f32_f16_bf16_flip_cm1_data;
+ spv_size = flash_attn_f32_f16_bf16_flip_cm1_len;
+ name = aligned ? "flash_attn_f32_bf16_aligned_flip_cm1" : "flash_attn_f32_bf16_flip_cm1";
+ } else {
+ spv_data = flash_attn_f32_f16_bf16_cm1_data;
+ spv_size = flash_attn_f32_f16_bf16_cm1_len;
+ name = aligned ? "flash_attn_f32_bf16_aligned_cm1" : "flash_attn_f32_bf16_cm1";
+ }
#else
continue;
#endif
+ } else if (flip) {
+ if (f32acc) { spv_data = flash_attn_f32_f16_flip_cm1_data; spv_size = flash_attn_f32_f16_flip_cm1_len; }
+ else { spv_data = flash_attn_f32_f16_flip_f16acc_cm1_data; spv_size = flash_attn_f32_f16_flip_f16acc_cm1_len; }
+ name = aligned ? "flash_attn_f32_f16_aligned_flip_cm1" : "flash_attn_f32_f16_flip_cm1";
} else {
if (f32acc) { spv_data = flash_attn_f32_f16_cm1_data; spv_size = flash_attn_f32_f16_cm1_len; }
else { spv_data = flash_attn_f32_f16_f16acc_cm1_data; spv_size = flash_attn_f32_f16_f16acc_cm1_len; }
@@ -8496,10 +8549,11 @@ bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_
const uint32_t iq_shmem = 16 * sizeof(ggml_fp16_t);
const uint32_t qstride = hsk_pad / 4 + 2;
- const uint32_t Qf = Br * qstride * f16vec4;
+ const uint32_t Qf = (params.q_in_regs ? 1 : Br * qstride) * f16vec4;
const uint32_t psh_stride = Br / 4 + 2;
- const uint32_t Psh = Bc * psh_stride * f16vec4;
+ const uint32_t qchunk_stride = 4 * 16 / 4 + 2;
+ const uint32_t Psh = std::max(Bc * psh_stride, params.q_in_regs ? Br * qchunk_stride : 0) * f16vec4;
const uint32_t sfshstride = (hsk <= 128) ? (Br + 8) : Br;
const uint32_t sfsh = Bc * sfshstride * acctype;
@@ -8508,14 +8562,30 @@ bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_
const uint32_t vsh_stride = MatBc / 4 * row_split;
const uint32_t ksh = ((kvshstride >= vsh_stride) ? (Bc * kvshstride) : (Bc * vsh_stride)) * f16vec4;
- // BF16 PVMat accumulator is f32 (no bf16 accumulator support), so pvsh is vec4 (16 bytes)
- const uint32_t pvsh_elem_size = (k_type == GGML_TYPE_BF16) ? 16u : f16vec4;
- const uint32_t osh_stride = params.row_split * MatBr / 4;
- const uint32_t pvsh = MatBc * osh_stride * pvsh_elem_size;
-
const uint32_t slope = Br * acctype;
- const uint32_t total_size = tmpsh + iq_shmem + Qf + Psh + sfsh + ksh + pvsh + slope;
+ uint32_t pvsh = 0;
+ if (Br <= MatBr) {
+ const uint32_t pvsh_elem_size = (k_type == GGML_TYPE_BF16) ? 16u : f16vec4;
+ const uint32_t osh_stride = params.row_split * MatBr / 4;
+ pvsh = Br * osh_stride * pvsh_elem_size;
+ }
+
+ uint32_t total_size = tmpsh + iq_shmem + Qf + Psh + sfsh + ksh + pvsh + slope + params.limit_occupancy_shmem * 16;
+
+ if (params.flip) {
+ const uint32_t d_pad = std::max(hsk_pad, hsv_pad);
+ const uint32_t wg = params.workgroup_size;
+ const uint32_t vld = CEIL_DIV(Bc * hsv_pad / 8, wg);
+ const bool vt = vld % 4 == 0 && wg % (hsv_pad / 8) == 0 && vld * (wg / (hsv_pad / 8)) == Bc;
+ const uint32_t kvbuf = std::max(Bc * (d_pad / 4 + 2), vt ? hsv_pad * (Bc / 4 + 2) : 0u) * 8;
+ const uint32_t kvshf = params.flip_direct ? 16 : 2 * kvbuf;
+ const uint32_t sshf = params.row_split * 16 * (Bc / 4 + (f32acc ? 1 : 2)) * 4 * acctype;
+ // f16/quant K/V stage P and Q through the S slab; only bf16 keeps a separate pshf.
+ const bool p_in_sshf = k_type != GGML_TYPE_BF16;
+ const uint32_t pshf = p_in_sshf ? 0 : params.row_split * 16 * (Bc / 4 + 2) * f16vec4;
+ total_size = kvshf + sshf + pshf;
+ }
const bool supported = total_size <= device->properties.limits.maxComputeSharedMemorySize;
VK_LOG_DEBUG("ggml_vk_flash_attn_coopmat_shmem_support(HSK=" << hsk << ", HSV=" << hsv << ", f32acc=" << f32acc << ", total_size=" << total_size << ", supported=" << supported);
@@ -8623,9 +8693,8 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
const ggml_type k_type_eff = (use_dequant_kv && !kv_raw_strided) ? GGML_TYPE_F16 : k->type;
const ggml_type v_type_eff = (use_dequant_kv && !kv_raw_strided) ? GGML_TYPE_F16 : v->type;
- // For scalar/coopmat1 FA, we can use the "large" size to accommodate qga.
- // For coopmat2 FA, we always use the small size (which is still pretty large for gqa).
- vk_fa_tuning_params tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, 512, KV, k_type_eff, v_type_eff, f32acc);
+ // N=32 avoids the cm1 prefill block_rows doubling
+ vk_fa_tuning_params tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, 32, KV, k_type_eff, v_type_eff, f32acc, qk_ratio);
const uint32_t max_gqa = std::min(tuning_params.block_rows, 32u);
if (N <= 8 && qk_ratio > 1 && qk_ratio <= max_gqa &&
@@ -8638,7 +8707,7 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
workgroups_y /= gqa_ratio;
}
- tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, N, KV, k_type_eff, v_type_eff, f32acc);
+ tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, N, KV, k_type_eff, v_type_eff, f32acc, qk_ratio);
float scale = 1.0f;
float max_bias = 0.0f;
@@ -8694,14 +8763,28 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
nbv2_eff = (uint32_t)((uint64_t)HSV * KV * sizeof(ggml_fp16_t));
nbv3_eff = (uint32_t)((uint64_t)HSV * KV * nev2 * sizeof(ggml_fp16_t));
}
- const uint32_t alignment = tuning_params.block_cols;
- // CM2 decoded K/V strides count whole blocks and do not need to be multiples of eight blocks.
- const bool cm2_decode_k = tuning_params.path == FA_COOPMAT2 &&
- k_type_eff != GGML_TYPE_F16 && k_type_eff != GGML_TYPE_BF16;
- const bool cm2_decode_v = tuning_params.path == FA_COOPMAT2 &&
- v_type_eff != GGML_TYPE_F16 && v_type_eff != GGML_TYPE_BF16;
- bool aligned = (KV % alignment) == 0 && (q_stride_eff & 7) == 0 &&
- (cm2_decode_k || (k_stride & 7) == 0) && (cm2_decode_v || (v_stride & 7) == 0);
+ auto fa_aligned = [&](const vk_fa_tuning_params & tp) {
+ // CM2 decoded K/V strides count whole blocks and do not need to be multiples of eight blocks.
+ const bool cm2_decode_k = tp.path == FA_COOPMAT2 &&
+ k_type_eff != GGML_TYPE_F16 && k_type_eff != GGML_TYPE_BF16;
+ const bool cm2_decode_v = tp.path == FA_COOPMAT2 &&
+ v_type_eff != GGML_TYPE_F16 && v_type_eff != GGML_TYPE_BF16;
+ return (KV % tp.block_cols) == 0 && (q_stride_eff & 7) == 0 &&
+ (cm2_decode_k || (k_stride & 7) == 0) && (cm2_decode_v || (v_stride & 7) == 0);
+ };
+ uint32_t alignment = tuning_params.block_cols;
+ bool aligned = fa_aligned(tuning_params);
+
+ if (tuning_params.flip) {
+ // flipped coopmat1 path reads K/V rows with 16-byte loads
+ const bool kv_16b = use_dequant_kv ||
+ (((uintptr_t)k->data | (uintptr_t)v->data | nbk2_eff | nbk3_eff | nbv2_eff | nbv3_eff) % 16) == 0;
+ if (!aligned || use_sparse || gqa_ratio != 1 || !kv_16b) {
+ tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, N, KV, k_type_eff, v_type_eff, f32acc, qk_ratio, false);
+ alignment = tuning_params.block_cols;
+ aligned = fa_aligned(tuning_params);
+ }
+ }
// Need to use the coopmat2 variant that clamps loads when HSK/HSV aren't sufficiently aligned.
if (((HSK | HSV) % 16) != 0 && tuning_params.path == FA_COOPMAT2) {
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp
index 2c5e85908..c1e4606cd 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp
@@ -95,13 +95,13 @@ void main() {
const uint32_t col_tid = (gl_LocalInvocationIndex % threads_per_rowgroup) / D_split;
if (LIMIT_OCCUPANCY_SHMEM > 0) {
- // This just exists to avoid the occupancy_limiter array getting optimized out
- occupancy_limiter[tid] = vec4(tid);
-
+ // Runtime-dependent index, otherwise RADV forwards the store and drops the array
+ const uint32_t li = (gl_WorkGroupID.x * gl_WorkGroupSize.x + tid) % LIMIT_OCCUPANCY_SHMEM;
+ occupancy_limiter[li] = vec4(tid);
barrier();
-
- if (occupancy_limiter[tid] == vec4(99999.0)) {
- data_ov4[0] = D_TYPEV4(occupancy_limiter[tid]);
+ const uint32_t lr = (li + 1) % LIMIT_OCCUPANCY_SHMEM;
+ if (occupancy_limiter[lr] == vec4(99999.0)) {
+ data_ov4[0] = D_TYPEV4(occupancy_limiter[lr]);
}
}
@@ -154,8 +154,6 @@ void main() {
float Lf[rows_per_thread], Mf[rows_per_thread];
- // Use -FLT_MAX/2 rather than -inf to reduce the possibility of NaNs, e.g. when computing Mold-M.
- const float NEG_FLT_MAX_OVER_2 = uintBitsToFloat(0xFEFFFFFF);
[[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
Lf[r] = 0;
@@ -220,7 +218,11 @@ void main() {
masksh[c * masksh_stride + r] = m;
max_mask = max(max_mask, float(m));
} else {
+#if defined(FLOAT16)
+ masksh[c * masksh_stride + r] = USE_SPARSE ? FLOAT_TYPE(NEG_FLT16_MAX) : FLOAT_TYPE(0);
+#else
masksh[c * masksh_stride + r] = USE_SPARSE ? FLOAT_TYPE(NEG_FLT_MAX_OVER_2) : FLOAT_TYPE(0);
+#endif
}
}
}
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl
index 2e0e23bc1..072edfa68 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl
@@ -26,6 +26,10 @@ const bool LOGIT_SOFTCAP = (Flags & 4) != 0;
const bool OLD_AMD_WINDOWS = (Flags & 8) != 0;
// Sparse: gather binding-7 indices instead of scanning [0,KV); p.split_kv = n_kv_max.
const bool USE_SPARSE = (Flags & 16) != 0;
+const bool Q_IN_REGS = (Flags & 32) != 0;
+// coopmat1: subgroups own Q rows, O in coopmat accumulators. K/V staged in LDS, or loaded from global memory with FLIP_DIRECT.
+const bool FLIP = (Flags & 64) != 0;
+const bool FLIP_DIRECT = (Flags & 128) != 0;
// Round up head sizes to a multiple of 16, for coopmat1/coopmat2 paths
const uint32_t HSK_pad = (HSK + 15) & ~15;
@@ -86,6 +90,11 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];};
layout (binding = 7) readonly buffer SP {int32_t data_sparse[];};
+// Use -FLT_MAX/2 rather than -inf to reduce the possibility of NaNs, e.g. when computing Mold-M.
+#define NEG_FLT_MAX_OVER_2 uintBitsToFloat(0xFEFFFFFF)
+// Lowest finite f16, mask value for inactive sparse KV slots
+#define NEG_FLT16_MAX -65504.0
+
#define MASK_OPT_ALL_NEG_INF 1
#define MASK_OPT_ALL_ZERO 2
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp
index aa9dd624b..bb9755ab8 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp
@@ -42,10 +42,14 @@ layout (binding = 3) readonly buffer M {float16_t data_m[];};
shared float tmpsh[row_split];
const uint32_t qstride = HSK_pad / 4 + 2;
-shared FLOAT_TYPEV4 Qf[Br * qstride];
+shared FLOAT_TYPEV4 Qf[Q_IN_REGS ? 1 : (Br * qstride)];
const uint psh_stride = Br / 4 + 2;
-shared FLOAT_TYPEV4 Psh[Bc * psh_stride];
+const uint32_t qchunk_dtiles = 4;
+const uint32_t qchunk_stride = qchunk_dtiles * 16 / 4 + 2;
+const uint32_t psh_size = Bc * psh_stride;
+const uint32_t qchunk_size = Q_IN_REGS ? Br * qchunk_stride : 0;
+shared FLOAT_TYPEV4 Psh[psh_size >= qchunk_size ? psh_size : qchunk_size];
// Avoid padding for hsk==256 to make it fit in 48KB shmem.
const uint32_t sfshstride = (HSK <= 128) ? (Br / 4 + 2) : Br / 4;
@@ -57,11 +61,33 @@ const uint v_cols = MatBc / 4 * row_split; // total cols, 4 vec4s per MatBc * nu
const uint vsh_stride = v_cols;
shared FLOAT_TYPEV4 kvsh[(kvsh_stride >= vsh_stride) ? (Bc * kvsh_stride) : (Bc * vsh_stride)];
-const uint32_t osh_stride = row_split * MatBr / 4;
-shared O_TYPEV4 pvsh[MatBc * osh_stride];
-
shared ACC_TYPE slope[Br];
+// row-tiled builds stage the PV result through sfsh
+const bool PV_IN_SFSH = Br > MatBr;
+const uint32_t osh_stride = row_split * MatBr / 4;
+shared O_TYPEV4 pvsh[PV_IN_SFSH ? 2 : Br * osh_stride];
+
+shared vec4 occupancy_limiter[LIMIT_OCCUPANCY_SHMEM > 0 ? LIMIT_OCCUPANCY_SHMEM : 1];
+#if defined(FLOAT16) && !defined(ACC_TYPE_MAX)
+const uint32_t pv_per_acc = 2;
+O_TYPEV4 pvsh_load(const uint32_t i) {
+ if (!PV_IN_SFSH) {
+ return pvsh[i];
+ }
+ const uvec2 u = floatBitsToUint(vec2(sfsh[i / 2][(i % 2) * 2], sfsh[i / 2][(i % 2) * 2 + 1]));
+ return O_TYPEV4(unpackFloat2x16(u.x), unpackFloat2x16(u.y));
+}
+#else
+const uint32_t pv_per_acc = 1;
+O_TYPEV4 pvsh_load(const uint32_t i) {
+ if (!PV_IN_SFSH) {
+ return pvsh[i];
+ }
+ return O_TYPEV4(sfsh[i]);
+}
+#endif
+
void main() {
#ifdef NEEDS_INIT_IQ_SHMEM
if (fa_type_needs_shmem(FaTypeK) || fa_type_needs_shmem(FaTypeV)) {
@@ -78,29 +104,69 @@ void main() {
const uint32_t row_tid = gl_LocalInvocationIndex / threads_per_rowgroup;
const uint32_t col_tid = gl_LocalInvocationIndex % threads_per_rowgroup;
-#define tile_row(r) (row_tid * rows_per_thread + (r))
-
- // Zero-initialize shared memory for Q/K when HSK is not a multiple of 16 (HSK_pad > HSK).
- if ((HSK % 16) != 0) {
- [[unroll]] for (uint i = 0; i < Br * qstride; i += gl_WorkGroupSize.x) {
- if (i + tid < Br * qstride) {
- Qf[i + tid] = FLOAT_TYPEV4(0);
- }
- }
+ if (LIMIT_OCCUPANCY_SHMEM > 0) {
+ // Runtime-dependent index, otherwise RADV forwards the store and drops the array
+ const uint32_t li = (gl_WorkGroupID.x * gl_WorkGroupSize.x + tid) % LIMIT_OCCUPANCY_SHMEM;
+ occupancy_limiter[li] = vec4(tid);
barrier();
+ const uint32_t lr = (li + 1) % LIMIT_OCCUPANCY_SHMEM;
+ if (occupancy_limiter[lr] == vec4(99999.0)) {
+ data_ov4[0] = D_TYPEV4(occupancy_limiter[lr]);
+ }
}
+#define tile_row(r) (row_tid * rows_per_thread + (r))
+
+ const uint32_t q_tiles = Br / MatBr;
+ const uint32_t d_tiles = HSK_pad / 16;
+
uint32_t q_offset = gqa_iq1*p.nb01 + (iq2*p.nb02+iq3*p.nb03) / 4;
- [[unroll]] for (uint32_t idx = 0; idx < Br * HSK / 4; idx += gl_WorkGroupSize.x) {
- uint32_t d = (idx + tid) % (HSK / 4);
- uint32_t r = (idx + tid) / (HSK / 4);
- if (r < Br && d < HSK / 4 &&
- i * Br + r < N) {
- Qf[r * qstride + d] = FLOAT_TYPEV4(data_qv4[q_offset / 4 + (i * Br + r) * q_stride / 4 + d] * p.scale);
+ coopmat<FLOAT_TYPE, gl_ScopeSubgroup, 16, MatBr, gl_MatrixUseB> Qreg[Q_IN_REGS ? (d_tiles * q_tiles) : 1];
+
+ if (Q_IN_REGS) {
+ [[unroll]] for (uint32_t dbase = 0; dbase < d_tiles; dbase += qchunk_dtiles) {
+ [[unroll]] for (uint32_t idx = 0; idx < Br * qchunk_dtiles * 16 / 4; idx += gl_WorkGroupSize.x) {
+ uint32_t dvec = (idx + tid) % (qchunk_dtiles * 16 / 4);
+ uint32_t r = (idx + tid) / (qchunk_dtiles * 16 / 4);
+ if (r < Br) {
+ uint32_t d = dbase * 4 + dvec;
+ FLOAT_TYPEV4 val = FLOAT_TYPEV4(0);
+ if (d < HSK / 4 && i * Br + r < N) {
+ val = FLOAT_TYPEV4(data_qv4[q_offset / 4 + (i * Br + r) * q_stride / 4 + d] * p.scale);
+ }
+ Psh[r * qchunk_stride + dvec] = val;
+ }
+ }
+ barrier();
+ [[unroll]] for (uint32_t dd = 0; dd < qchunk_dtiles && dbase + dd < d_tiles; ++dd) {
+ [[unroll]] for (uint32_t q = 0; q < q_tiles; ++q) {
+ coopMatLoad(Qreg[(dbase + dd) * q_tiles + q], Psh, dd * 16 / 4 + q * MatBr * qchunk_stride, qchunk_stride, gl_CooperativeMatrixLayoutColumnMajor);
+ }
+ }
+ barrier();
+ }
+ } else {
+ // Zero-initialize Qf when HSK is not a multiple of 16 (HSK_pad > HSK).
+ if ((HSK % 16) != 0) {
+ [[unroll]] for (uint ii = 0; ii < Br * qstride; ii += gl_WorkGroupSize.x) {
+ if (ii + tid < Br * qstride) {
+ Qf[ii + tid] = FLOAT_TYPEV4(0);
+ }
+ }
+ barrier();
}
+
+ [[unroll]] for (uint32_t idx = 0; idx < Br * HSK / 4; idx += gl_WorkGroupSize.x) {
+ uint32_t d = (idx + tid) % (HSK / 4);
+ uint32_t r = (idx + tid) / (HSK / 4);
+ if (r < Br && d < HSK / 4 &&
+ i * Br + r < N) {
+ Qf[r * qstride + d] = FLOAT_TYPEV4(data_qv4[q_offset / 4 + (i * Br + r) * q_stride / 4 + d] * p.scale);
+ }
+ }
+ barrier();
}
- barrier();
O_TYPEV4 Of[rows_per_thread][d_per_thread];
[[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
@@ -111,8 +177,6 @@ void main() {
float Lf[rows_per_thread], Mf[rows_per_thread];
- // Use -FLT_MAX/2 rather than -inf to reduce the possibility of NaNs, e.g. when computing Mold-M.
- const float NEG_FLT_MAX_OVER_2 = uintBitsToFloat(0xFEFFFFFF);
[[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
Lf[r] = 0;
@@ -214,7 +278,7 @@ void main() {
}
mask_cache[idx / WorkGroupSize] = m;
} else if (USE_SPARSE) {
- mask_cache[idx / WorkGroupSize] = f16vec4(NEG_FLT_MAX_OVER_2);
+ mask_cache[idx / WorkGroupSize] = f16vec4(NEG_FLT16_MAX);
}
}
}
@@ -265,11 +329,14 @@ void main() {
// K * Q^T -> S^T: Bc x HSK_pad * HSK_pad x Br -> Bc x Br
// Bc split across workgroup (four subgroups), loop over HSK in chunks of 16: 16 x 16 * 16 x 16 -> 16 x 16
// This is written transposed in order to allow for N being 8 if implementations need it
- coopmat<ACC_TYPE, gl_ScopeSubgroup, MatBc, MatBr, gl_MatrixUseAccumulator> SfMat = coopmat<ACC_TYPE, gl_ScopeSubgroup, MatBc, MatBr, gl_MatrixUseAccumulator>(0);
+ coopmat<ACC_TYPE, gl_ScopeSubgroup, MatBc, MatBr, gl_MatrixUseAccumulator> SfMat[q_tiles];
+ [[unroll]] for (uint32_t q = 0; q < q_tiles; ++q) {
+ SfMat[q] = coopmat<ACC_TYPE, gl_ScopeSubgroup, MatBc, MatBr, gl_MatrixUseAccumulator>(0);
+ }
coopmat<FLOAT_TYPE, gl_ScopeSubgroup, MatBc, 16, gl_MatrixUseA> KMat;
coopmat<FLOAT_TYPE, gl_ScopeSubgroup, 16, MatBr, gl_MatrixUseB> QMat;
- [[unroll]] for (uint32_t d = 0; d < HSK_pad / 16; ++d) {
+ [[unroll]] for (uint32_t d = 0; d < d_tiles; ++d) {
// If SHMEM_STAGING is set, a Bc * HSK_pad size tile of K is loaded to shmem
// If not, K is loaded directly from global memory if aligned, otherwise
// staged through a Bc * MatBr size staging buffer.
@@ -319,13 +386,20 @@ void main() {
coopMatLoad(KMat, kvsh, coord, kvsh_stride, gl_CooperativeMatrixLayoutRowMajor);
}
- coopMatLoad(QMat, Qf, d * 16 / 4, qstride, gl_CooperativeMatrixLayoutColumnMajor);
-
- SfMat = coopMatMulAdd(KMat, QMat, SfMat);
+ [[unroll]] for (uint32_t q = 0; q < q_tiles; ++q) {
+ if (Q_IN_REGS) {
+ SfMat[q] = coopMatMulAdd(KMat, Qreg[d * q_tiles + q], SfMat[q]);
+ } else {
+ coopMatLoad(QMat, Qf, d * 16 / 4 + q * MatBr * qstride, qstride, gl_CooperativeMatrixLayoutColumnMajor);
+ SfMat[q] = coopMatMulAdd(KMat, QMat, SfMat[q]);
+ }
+ }
}
- uint coord = gl_SubgroupID * MatBc * sfshstride;
- coopMatStore(SfMat, sfsh, coord, sfshstride, gl_CooperativeMatrixLayoutRowMajor);
+ [[unroll]] for (uint32_t q = 0; q < q_tiles; ++q) {
+ uint coord = gl_SubgroupID * MatBc * sfshstride + q * (MatBr / 4);
+ coopMatStore(SfMat[q], sfsh, coord, sfshstride, gl_CooperativeMatrixLayoutRowMajor);
+ }
barrier();
if (LOGIT_SOFTCAP) {
@@ -442,7 +516,10 @@ void main() {
[[unroll]] for (uint32_t hsv_tile = 0; hsv_tile < num_hsv_tiles; ++hsv_tile) {
const uint hsv_offset = (hsv_tile * row_split + gl_SubgroupID) * 16;
- coopmat<O_TYPE, gl_ScopeSubgroup, MatBc, MatBr, gl_MatrixUseAccumulator> PVMat = coopmat<O_TYPE, gl_ScopeSubgroup, MatBc, MatBr, gl_MatrixUseAccumulator>(0);
+ coopmat<O_TYPE, gl_ScopeSubgroup, MatBc, MatBr, gl_MatrixUseAccumulator> PVMat[q_tiles];
+ [[unroll]] for (uint32_t q = 0; q < q_tiles; ++q) {
+ PVMat[q] = coopmat<O_TYPE, gl_ScopeSubgroup, MatBc, MatBr, gl_MatrixUseAccumulator>(0);
+ }
// Preload V tiles for [Bc, 16 * num subgroups]
const uint v_rows = Bc;
@@ -492,8 +569,6 @@ void main() {
if (hsv_offset < HSV_pad) {
[[unroll]] for (uint32_t bc_chunk = 0; bc_chunk < Bc / MatBc; ++bc_chunk) {
- coopMatLoad(KMat, Psh, bc_chunk * MatBc * psh_stride, psh_stride, gl_CooperativeMatrixLayoutColumnMajor);
-
if (SHMEM_STAGING == 0) {
if (!USE_DECODE_V && !KV_bounds_check && !USE_SPARSE) {
// F16/BF16 values can be loaded directly from global memory
@@ -509,11 +584,19 @@ void main() {
coopMatLoad(QMat, kvsh, v_tile_offset, kvsh_stride, gl_CooperativeMatrixLayoutRowMajor);
}
- PVMat = coopMatMulAdd(KMat, QMat, PVMat);
+ [[unroll]] for (uint32_t q = 0; q < q_tiles; ++q) {
+ coopMatLoad(KMat, Psh, bc_chunk * MatBc * psh_stride + q * (MatBc / 4), psh_stride, gl_CooperativeMatrixLayoutColumnMajor);
+ PVMat[q] = coopMatMulAdd(KMat, QMat, PVMat[q]);
+ }
}
- // Store PVMat to pvsh and load into Of
- coopMatStore(PVMat, pvsh, o_offset, osh_stride, gl_CooperativeMatrixLayoutRowMajor);
+ [[unroll]] for (uint32_t q = 0; q < q_tiles; ++q) {
+ if (PV_IN_SFSH) {
+ coopMatStore(PVMat[q], sfsh, (o_offset + q * MatBc * osh_stride) / pv_per_acc, osh_stride / pv_per_acc, gl_CooperativeMatrixLayoutRowMajor);
+ } else {
+ coopMatStore(PVMat[q], pvsh, o_offset + q * MatBc * osh_stride, osh_stride, gl_CooperativeMatrixLayoutRowMajor);
+ }
+ }
}
barrier();
@@ -534,7 +617,7 @@ void main() {
if (hsv_col >= hsv_base && hsv_col < hsv_base + hsv_per_tile && hsv_col < HSV) {
const uint local_hsv = (hsv_col - hsv_base) / 4;
- Of[r][d_local] += pvsh[row * osh_stride + local_hsv];
+ Of[r][d_local] += pvsh_load(row * osh_stride + local_hsv);
}
}
}
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1_flip.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1_flip.comp
new file mode 100644
index 000000000..6d8641a91
--- /dev/null
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1_flip.comp
@@ -0,0 +1,441 @@
+#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_int32 : require
+
+#if defined(BFLOAT16)
+#extension GL_EXT_bfloat16 : enable
+#endif
+
+#extension GL_KHR_shader_subgroup_basic : enable
+#extension GL_KHR_shader_subgroup_vote : enable
+#extension GL_KHR_shader_subgroup_shuffle : enable
+#extension GL_KHR_memory_scope_semantics : enable
+#extension GL_KHR_cooperative_matrix : enable
+
+#include "types.glsl"
+#include "flash_attn_base.glsl"
+
+layout (binding = 0) readonly buffer QV4 {vec4 data_qv4[];};
+layout (binding = 1) readonly buffer KU {uvec4 data_ku[];};
+layout (binding = 1) readonly buffer KF4 {FLOAT_TYPEV4 data_kf4[];};
+layout (binding = 2) readonly buffer VU {uvec4 data_vu[];};
+layout (binding = 2) readonly buffer VF4 {FLOAT_TYPEV4 data_vf4[];};
+layout (binding = 3) readonly buffer M {float16_t data_m[];};
+
+const uint32_t D_pad = HSK_pad > HSV_pad ? HSK_pad : HSV_pad;
+
+// Flipped coopmat1 FA: each subgroup owns 16 Q rows and the full head dim. K/V tiles are either
+// staged once per workgroup in LDS or, with FLIP_DIRECT, loaded by each subgroup from global memory.
+// Requires gqa_ratio == 1, aligned KV, f16/bf16 K/V, HSK/HSV multiples of 16, Br == row_split * 16.
+
+const uint32_t FL_LPR = SubGroupSize / 16;
+const uint32_t FL_CPL = Bc / FL_LPR;
+const uint32_t FL_OCPL = 16 / FL_LPR;
+
+const uint32_t FL_KLD = CEIL_DIV(Bc * HSK_pad / 8, WorkGroupSize);
+const uint32_t FL_VLD = CEIL_DIV(Bc * HSV_pad / 8, WorkGroupSize);
+
+// V^T staging: each thread packs FL_VLD consecutive KV rows (multiple of 4) of 8 d values
+const bool FL_VT = (uint32_t((FL_VLD % 4) == 0) * uint32_t((WorkGroupSize % (HSV_pad / 8)) == 0) *
+ uint32_t(FL_VLD * (WorkGroupSize / (HSV_pad / 8)) == Bc)) != 0;
+
+// LDS in uvec2 (4 halfs) units
+const uint32_t fl_kstride = D_pad / 4 + 2;
+const uint32_t fl_vtstride = Bc / 4 + 2;
+const uint32_t fl_vtsize = FL_VT ? HSV_pad * fl_vtstride : 0;
+const uint32_t fl_kvbuf = fl_vtsize > Bc * fl_kstride ? fl_vtsize : Bc * fl_kstride;
+// 16 B keeps the following shared arrays 16-B aligned for coopmat load/store
+shared uvec2 kvshf[FLIP_DIRECT ? 2 : 2 * fl_kvbuf];
+
+#if defined(ACC_TYPE_MAX)
+const uint32_t fl_sstride = Bc / 4 + 2;
+#else
+const uint32_t fl_sstride = Bc / 4 + 1;
+#endif
+shared ACC_TYPEV4 sshf[row_split * 16 * fl_sstride];
+
+// P and Q (the coopmat A operands) are staged through the S slab, which is idle
+// between the S read and the PV matmul. f16 K/V aliases into sshf (f32acc packs two
+// f16 rows per element, FL_PACK==2; f16acc stores natively, FL_PACK==1); bf16 keeps pshf.
+#if !defined(BFLOAT16)
+#define FL_P_IN_SSHF 1
+#endif
+#if !defined(BFLOAT16) && !defined(ACC_TYPE_MAX)
+#define FL_PACK 2
+#else
+#define FL_PACK 1
+#endif
+
+#if defined(FL_P_IN_SSHF)
+const uint32_t fl_astride = fl_sstride;
+#else
+const uint32_t fl_astride = Bc / 4 + 2;
+shared FLOAT_TYPEV4 pshf[row_split * 16 * fl_astride];
+#endif
+
+// Store one vec4 (column group cfv) of the A slab.
+void fl_store_a(const uint32_t e_base, const uint32_t cfv, const FLOAT_TYPEV4 val) {
+#if FL_PACK == 2
+ const uint32_t e = e_base + cfv / 2;
+ const uint32_t h = (cfv & 1u) * 2u;
+ sshf[e][h] = uintBitsToFloat(packFloat2x16(f16vec2(val.xy)));
+ sshf[e][h + 1] = uintBitsToFloat(packFloat2x16(f16vec2(val.zw)));
+#elif defined(FL_P_IN_SSHF)
+ sshf[e_base + cfv] = ACC_TYPEV4(val);
+#else
+ pshf[e_base + cfv] = val;
+#endif
+}
+
+void fl_load_a(out coopmat<FLOAT_TYPE, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA> a, const uint32_t e) {
+#if defined(FL_P_IN_SSHF)
+ coopMatLoad(a, sshf, e, fl_sstride, gl_CooperativeMatrixLayoutRowMajor);
+#else
+ coopMatLoad(a, pshf, e, fl_astride, gl_CooperativeMatrixLayoutRowMajor);
+#endif
+}
+
+#if defined(ACC_TYPE_MAX)
+const float FL_TAU = 0.0;
+#else
+const float FL_TAU = 8.0;
+#endif
+
+uvec4 fl_kr[FL_KLD];
+uvec4 fl_vr[FL_VLD];
+
+uint32_t fl_mo_offset, fl_mo_idx, fl_mo_word;
+
+uint32_t fl_mask_bits(const uint32_t j) {
+ if (!MASK_ENABLE || !USE_MASK_OPT) {
+ return 0;
+ }
+ if (fl_mo_idx != j / 16) {
+ fl_mo_idx = j / 16;
+ fl_mo_word = data_mask_opt[fl_mo_offset + fl_mo_idx];
+ }
+ return (fl_mo_word >> ((j % 16) * 2)) & 0x3;
+}
+
+uint32_t fl_next_tile(uint32_t j) {
+ while (j < end_j && fl_mask_bits(j) == MASK_OPT_ALL_NEG_INF) {
+ ++j;
+ }
+ return j;
+}
+
+void fl_load_k(const uint32_t j, const uint32_t k_offset) {
+ [[unroll]] for (uint32_t l = 0; l < FL_KLD; ++l) {
+ const uint32_t idx = l * WorkGroupSize + gl_LocalInvocationIndex;
+ if (idx < Bc * HSK_pad / 8) {
+ const uint32_t c = idx / (HSK_pad / 8);
+ const uint32_t d = idx % (HSK_pad / 8);
+ fl_kr[l] = data_ku[(k_offset + (j * Bc + c) * k_stride) / 8 + d];
+ }
+ }
+}
+
+void fl_load_v(const uint32_t j, const uint32_t v_offset) {
+ [[unroll]] for (uint32_t l = 0; l < FL_VLD; ++l) {
+ const uint32_t idx = l * WorkGroupSize + gl_LocalInvocationIndex;
+ if (idx < Bc * HSV_pad / 8) {
+ uint32_t c = idx / (HSV_pad / 8);
+ const uint32_t d = idx % (HSV_pad / 8);
+ if (FL_VT) {
+ c = (gl_LocalInvocationIndex / (HSV_pad / 8)) * FL_VLD + l;
+ }
+ fl_vr[l] = data_vu[(v_offset + (j * Bc + c) * v_stride) / 8 + d];
+ }
+ }
+}
+
+void fl_store_k() {
+ [[unroll]] for (uint32_t l = 0; l < FL_KLD; ++l) {
+ const uint32_t idx = l * WorkGroupSize + gl_LocalInvocationIndex;
+ if (idx < Bc * HSK_pad / 8) {
+ const uint32_t a = (idx / (HSK_pad / 8)) * fl_kstride + 2 * (idx % (HSK_pad / 8));
+ kvshf[a] = fl_kr[l].xy;
+ kvshf[a + 1] = fl_kr[l].zw;
+ }
+ }
+}
+
+void fl_store_v(const uint32_t base) {
+ if (FL_VT) {
+ const uint32_t d0 = (gl_LocalInvocationIndex % (HSV_pad / 8)) * 8;
+ const uint32_t c0 = (gl_LocalInvocationIndex / (HSV_pad / 8)) * FL_VLD;
+ [[unroll]] for (uint32_t dd = 0; dd < 8; ++dd) {
+ const uint32_t ddv = (dd + gl_LocalInvocationIndex) & 7u; // rotate store order to spread banks
+ const uint32_t sh = (ddv % 2) * 16;
+ [[unroll]] for (uint32_t l = 0; l < FL_VLD; l += 4) {
+ const uvec2 w = uvec2(((fl_vr[l ][ddv / 2] >> sh) & 0xFFFF) | (((fl_vr[l + 1][ddv / 2] >> sh) & 0xFFFF) << 16),
+ ((fl_vr[l + 2][ddv / 2] >> sh) & 0xFFFF) | (((fl_vr[l + 3][ddv / 2] >> sh) & 0xFFFF) << 16));
+ kvshf[base + (d0 + ddv) * fl_vtstride + (c0 + l) / 4] = w;
+ }
+ }
+ return;
+ }
+ [[unroll]] for (uint32_t l = 0; l < FL_VLD; ++l) {
+ const uint32_t idx = l * WorkGroupSize + gl_LocalInvocationIndex;
+ if (idx < Bc * HSV_pad / 8) {
+ const uint32_t a = base + (idx / (HSV_pad / 8)) * fl_kstride + 2 * (idx % (HSV_pad / 8));
+ kvshf[a] = fl_vr[l].xy;
+ kvshf[a + 1] = fl_vr[l].zw;
+ }
+ }
+}
+
+void fl_sg_barrier() {
+ controlBarrier(gl_ScopeSubgroup, gl_ScopeSubgroup, gl_StorageSemanticsShared, gl_SemanticsAcquireRelease);
+}
+
+float fl_row_max(float v) {
+ [[unroll]] for (uint32_t o = 1; o < FL_LPR; o *= 2) {
+ v = max(v, subgroupShuffleXor(v, o));
+ }
+ return v;
+}
+
+float fl_row_sum(float v) {
+ [[unroll]] for (uint32_t o = 1; o < FL_LPR; o *= 2) {
+ v += subgroupShuffleXor(v, o);
+ }
+ return v;
+}
+
+void fl_rescale(inout coopmat<ACC_TYPE, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator> o, const uint32_t ss_base, const uint32_t lr, const uint32_t lc, const float alpha) {
+ fl_sg_barrier();
+ coopMatStore(o, sshf, ss_base, fl_sstride, gl_CooperativeMatrixLayoutRowMajor);
+ fl_sg_barrier();
+ [[unroll]] for (uint32_t v = 0; v < FL_OCPL / 4; ++v) {
+ const uint32_t a = ss_base + lr * fl_sstride + lc * (FL_OCPL / 4) + v;
+ sshf[a] = sshf[a] * ACC_TYPE(alpha);
+ }
+ fl_sg_barrier();
+ coopMatLoad(o, sshf, ss_base, fl_sstride, gl_CooperativeMatrixLayoutRowMajor);
+}
+
+void main() {
+ init_indices();
+
+ const uint32_t sg = gl_SubgroupID;
+ const uint32_t lane = gl_SubgroupInvocationID;
+ const uint32_t lr = lane / FL_LPR;
+ const uint32_t lc = lane % FL_LPR;
+ const uint32_t row0 = i * Br + sg * 16;
+ const uint32_t myrow = row0 + lr;
+
+ const uint32_t ss_base = sg * 16 * fl_sstride;
+ const uint32_t a_base = sg * 16 * fl_astride;
+
+ const uint32_t q_offset = (iq2 * p.nb02 + iq3 * p.nb03) / 4;
+
+ // Q (scaled) -> coopmat A registers, staged through this subgroup's A slab
+ coopmat<FLOAT_TYPE, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA> Qa[HSK_pad / 16];
+ const uint32_t qdt = fl_astride * FL_PACK / 4;
+ [[unroll]] for (uint32_t dbase = 0; dbase < HSK_pad / 16; dbase += qdt) {
+ [[unroll]] for (uint32_t idx0 = 0; idx0 < 16 * qdt * 4; idx0 += SubGroupSize) {
+ const uint32_t idx = idx0 + lane;
+ const uint32_t r = idx / (qdt * 4);
+ const uint32_t dv = idx % (qdt * 4);
+ const uint32_t d = dbase * 4 + dv;
+ if (r < 16) {
+ FLOAT_TYPEV4 val = FLOAT_TYPEV4(0);
+ if (d < HSK / 4 && row0 + r < N) {
+ val = FLOAT_TYPEV4(data_qv4[q_offset / 4 + (row0 + r) * p.nb01 / 4 + d] * p.scale);
+ }
+ fl_store_a(a_base + r * fl_astride, dv, val);
+ }
+ }
+ fl_sg_barrier();
+ [[unroll]] for (uint32_t dd = 0; dd < qdt; ++dd) {
+ if (dbase + dd < HSK_pad / 16) {
+ fl_load_a(Qa[dbase + dd], a_base + dd * (4 / FL_PACK));
+ }
+ }
+ fl_sg_barrier();
+ }
+
+ coopmat<ACC_TYPE, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator> Oacc[HSV_pad / 16];
+ [[unroll]] for (uint32_t t = 0; t < HSV_pad / 16; ++t) {
+ Oacc[t] = coopmat<ACC_TYPE, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>(0);
+ }
+
+ float M = NEG_FLT_MAX_OVER_2;
+ float L = 0;
+
+ float slope_v = 1.0;
+ if (p.max_bias > 0.0f) {
+ slope_v = float(perElemOpComputeSlope(0, 0, ACC_TYPE(0), iq2));
+ }
+
+ const uint32_t k_offset = (ik2 * p.nb12 + ik3 * p.nb13) / FaBlockBytesK;
+ const uint32_t v_offset = (iv2 * p.nb22 + iv3 * p.nb23) / FaBlockBytesV;
+ uint32_t m_offset = 0;
+ const uint32_t mo_stride = CEIL_DIV(KV, 16 * Bc);
+ fl_mo_offset = mo_stride * i;
+ fl_mo_idx = ~0u;
+ fl_mo_word = 0;
+ if (p.nem2 != 1 || p.nem3 != 1) {
+ m_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 * KV;
+ fl_mo_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * CEIL_DIV(p.nem1, Br) * mo_stride;
+ }
+
+ uint32_t j = (FLIP_DIRECT && row0 >= N) ? end_j : fl_next_tile(start_j);
+ if (!FLIP_DIRECT && j < end_j) {
+ fl_load_k(j, k_offset);
+ fl_load_v(j, v_offset);
+ }
+
+ [[dont_unroll]]
+ while (j < end_j) {
+ const uint32_t jn = fl_next_tile(j + 1);
+ const uint32_t mask_bits = fl_mask_bits(j);
+
+ if (!FLIP_DIRECT) {
+ barrier();
+ fl_store_k();
+ fl_store_v(fl_kvbuf);
+ barrier();
+ if (jn < end_j) {
+ fl_load_k(jn, k_offset);
+ fl_load_v(jn, v_offset);
+ }
+ }
+
+ coopmat<ACC_TYPE, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator> Sacc[Bc / 16];
+ [[unroll]] for (uint32_t c = 0; c < Bc / 16; ++c) {
+ Sacc[c] = coopmat<ACC_TYPE, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>(0);
+ }
+ [[unroll]] for (uint32_t d = 0; d < HSK_pad / 16; ++d) {
+ [[unroll]] for (uint32_t c = 0; c < Bc / 16; ++c) {
+ coopmat<FLOAT_TYPE, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB> Kb;
+ if (FLIP_DIRECT) {
+ coopMatLoad(Kb, data_kf4, (k_offset + (j * Bc + c * 16) * k_stride) / 4 + d * 4, k_stride / 4, gl_CooperativeMatrixLayoutColumnMajor);
+ } else {
+ coopMatLoad(Kb, kvshf, c * 16 * fl_kstride + d * 4, fl_kstride, gl_CooperativeMatrixLayoutColumnMajor);
+ }
+ Sacc[c] = coopMatMulAdd(Qa[d], Kb, Sacc[c]);
+ }
+ }
+
+ fl_sg_barrier();
+ [[unroll]] for (uint32_t c = 0; c < Bc / 16; ++c) {
+ coopMatStore(Sacc[c], sshf, ss_base + c * 4, fl_sstride, gl_CooperativeMatrixLayoutRowMajor);
+ }
+ fl_sg_barrier();
+ float s[FL_CPL];
+ [[unroll]] for (uint32_t v = 0; v < FL_CPL / 4; ++v) {
+ const vec4 sv = vec4(sshf[ss_base + lr * fl_sstride + lc * (FL_CPL / 4) + v]);
+ [[unroll]] for (uint32_t k = 0; k < 4; ++k) {
+ s[v * 4 + k] = sv[k];
+ }
+ }
+
+ if (LOGIT_SOFTCAP) {
+ [[unroll]] for (uint32_t k = 0; k < FL_CPL; ++k) {
+ s[k] = p.logit_softcap * tanh(s[k]);
+ }
+ }
+
+ if (MASK_ENABLE && mask_bits != MASK_OPT_ALL_ZERO && myrow < p.nem1) {
+ const uint32_t moff = m_offset + myrow * KV + j * Bc + lc * FL_CPL;
+ [[unroll]] for (uint32_t k = 0; k < FL_CPL; ++k) {
+ s[k] += slope_v * float(data_m[moff + k]);
+ }
+ }
+
+ float mx = s[0];
+ [[unroll]] for (uint32_t k = 1; k < FL_CPL; ++k) {
+ mx = max(mx, s[k]);
+ }
+ mx = fl_row_max(mx) + FATTN_KQ_MAX_OFFSET;
+
+ if (subgroupAny(mx > M + FL_TAU)) {
+ const float Mn = max(M, mx);
+ const float alpha = exp(M - Mn);
+ const bool o_empty = subgroupAll(M == NEG_FLT_MAX_OVER_2);
+ M = Mn;
+ L *= alpha;
+ if (!o_empty) {
+ [[unroll]] for (uint32_t t = 0; t < HSV_pad / 16; ++t) {
+ fl_rescale(Oacc[t], ss_base, lr, lc, alpha);
+ }
+ }
+ }
+
+ fl_sg_barrier();
+ [[unroll]] for (uint32_t v = 0; v < FL_CPL / 4; ++v) {
+ const vec4 pv = exp(vec4(s[v * 4], s[v * 4 + 1], s[v * 4 + 2], s[v * 4 + 3]) - M);
+ L += pv.x + pv.y + pv.z + pv.w;
+ fl_store_a(a_base + lr * fl_astride, lc * (FL_CPL / 4) + v, FLOAT_TYPEV4(pv));
+ }
+ fl_sg_barrier();
+
+ coopmat<FLOAT_TYPE, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA> Pa[Bc / 16];
+ [[unroll]] for (uint32_t c = 0; c < Bc / 16; ++c) {
+ fl_load_a(Pa[c], a_base + c * (4 / FL_PACK));
+ }
+
+ [[unroll]] for (uint32_t t = 0; t < HSV_pad / 16; ++t) {
+ [[unroll]] for (uint32_t c = 0; c < Bc / 16; ++c) {
+ coopmat<FLOAT_TYPE, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB> Vb;
+ if (FLIP_DIRECT) {
+ coopMatLoad(Vb, data_vf4, (v_offset + (j * Bc + c * 16) * v_stride) / 4 + t * 4, v_stride / 4, gl_CooperativeMatrixLayoutRowMajor);
+ } else if (FL_VT) {
+ coopMatLoad(Vb, kvshf, fl_kvbuf + t * 16 * fl_vtstride + c * 4, fl_vtstride, gl_CooperativeMatrixLayoutColumnMajor);
+ } else {
+ coopMatLoad(Vb, kvshf, fl_kvbuf + c * 16 * fl_kstride + t * 4, fl_kstride, gl_CooperativeMatrixLayoutRowMajor);
+ }
+ Oacc[t] = coopMatMulAdd(Pa[c], Vb, Oacc[t]);
+ }
+ }
+
+ j = jn;
+ }
+
+ L = fl_row_sum(L);
+
+ float oscale = 1.0;
+ uint32_t o_row_offset;
+ if (p.k_num > 1) {
+ o_row_offset = HSV * p.ne1 * (split_k_index + p.k_num * (myrow + p.ne2 * iq3)) / 4 + iq2 * HSV / 4;
+ if (myrow < N && lc == 0) {
+ const uint32_t lm_offset = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + p.k_num * (myrow + p.ne2 * iq3));
+ data_o[lm_offset + iq2] = D_TYPE(L);
+ data_o[lm_offset + p.ne1 + iq2] = D_TYPE(M);
+ }
+ } else {
+ if ((p.mask_n_head_log2 & SINK_ENABLE_BIT) != 0) {
+ const float sink = float(perElemOpGetSink(0, 0u, ACC_TYPE(0), iq2));
+ float ms = 1.0f;
+ float vs = 1.0f;
+ if (sink > M) {
+ ms = exp(M - sink);
+ } else {
+ vs = exp(sink - M);
+ }
+ L = L * ms + vs;
+ oscale = ms;
+ }
+ oscale *= (L == 0.0) ? 0.0 : (1.0 / L);
+ o_row_offset = (iq3 * p.ne2 * p.ne1 * HSV + iq2 * HSV + myrow * p.ne1 * HSV) / 4;
+ }
+
+ [[unroll]] for (uint32_t t = 0; t < HSV_pad / 16; ++t) {
+ fl_sg_barrier();
+ coopMatStore(Oacc[t], sshf, ss_base, fl_sstride, gl_CooperativeMatrixLayoutRowMajor);
+ fl_sg_barrier();
+ [[unroll]] for (uint32_t v = 0; v < FL_OCPL / 4; ++v) {
+ const uint32_t d4 = t * 4 + lc * (FL_OCPL / 4) + v;
+ const vec4 o = vec4(sshf[ss_base + lr * fl_sstride + lc * (FL_OCPL / 4) + v]) * oscale;
+ if (myrow < N && d4 < HSV / 4) {
+ data_ov4[o_row_offset + d4] = D_TYPEV4(o);
+ }
+ }
+ }
+}
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp
index c7253784c..9ddb4eb82 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp
@@ -193,12 +193,17 @@ FLOAT_TYPEV4 faGatherVVector(const decodeBufFA_V bl_in, const uint32_t blockCoor
#endif
// Add gathered mask to S (slope==1 since sparse requires max_bias==0). col = slot in block jblk.
+#if defined(ACC_TYPE_MAX)
+#define ACC_TYPE_NEG (-ACC_TYPE_MAX)
+#else
+#define ACC_TYPE_NEG ACC_TYPE(NEG_FLT_MAX_OVER_2)
+#endif
+
ACC_TYPE faAddSparseMask(const uint32_t row, const uint32_t col, const ACC_TYPE elem, const uint32_t jblk) {
- const float NEG = uintBitsToFloat(0xFEFFFFFF);
const uint32_t kvslot = jblk * Bc + col;
- if (kvslot >= p.split_kv) { return ACC_TYPE(NEG); }
+ if (kvslot >= p.split_kv) { return ACC_TYPE_NEG; }
const int r = data_sparse[sparse_base + kvslot];
- return r < 0 ? ACC_TYPE(NEG) : elem + ACC_TYPE(data_mf16[g_m_off_elem + row * m_stride + uint(r)]);
+ return r < 0 ? ACC_TYPE_NEG : elem + ACC_TYPE(data_mf16[g_m_off_elem + row * m_stride + uint(r)]);
}
ACC_TYPE maxReduce(const in ACC_TYPE x, const in ACC_TYPE y) {
@@ -319,8 +324,6 @@ void main() {
coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator> L, M;
- // Use -FLT_MAX/2 rather than -inf to reduce the possibility of NaNs, e.g. when computing Mold-M.
- const float NEG_FLT_MAX_OVER_2 = uintBitsToFloat(0xFEFFFFFF);
L = coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator>(0);
#if defined(ACC_TYPE_MAX)
@@ -443,7 +446,7 @@ void main() {
uint R = ((i + 1) * Br > N) ? (N % Br) : Br;
uint C = ((j + 1) * Bc > KV) ? (KV % Bc) : Bc;
- coopMatPerElementNV(S, S, replacePadding, ACC_TYPE(NEG_FLT_MAX_OVER_2), R, C);
+ coopMatPerElementNV(S, S, replacePadding, ACC_TYPE_NEG, R, C);
}
coopmat<ACC_TYPE, gl_ScopeWorkgroup, Br, Bc, gl_MatrixUseAccumulator> rowmax, P, rowsum, eM;
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 d26357498..52e5f0440 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
@@ -738,6 +738,11 @@ void process_shaders() {
#if defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT)
string_to_spv("flash_attn_f32_f16", "flash_attn_cm1.comp",
merge_maps(fa_base_dict, {{"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"D_TYPEV4", "vec4"}, {"COOPMAT", "1"}}), fp16, true, false, f16acc);
+
+ std::map<std::string, std::string> fa_flip_dict = fa_base_dict;
+ fa_flip_dict.erase("DATA_A_IQ4_NL");
+ string_to_spv("flash_attn_f32_f16_flip", "flash_attn_cm1_flip.comp",
+ merge_maps(fa_flip_dict, {{"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"D_TYPEV4", "vec4"}, {"COOPMAT", "1"}}), fp16, true, false, f16acc);
#endif
}
@@ -770,6 +775,9 @@ void process_shaders() {
string_to_spv("flash_attn_f32_f16_bf16", "flash_attn_cm1.comp",
merge_maps(fa_bf16_dict, {{"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"D_TYPEV4", "vec4"}, {"COOPMAT", "1"}}),
true, true, false, false);
+ string_to_spv("flash_attn_f32_f16_bf16_flip", "flash_attn_cm1_flip.comp",
+ merge_maps(fa_bf16_dict, {{"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"D_TYPEV4", "vec4"}, {"COOPMAT", "1"}}),
+ true, true, false, false);
#endif
#if defined(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT)