Commit 733f57a99 for llama.cpp
commit 733f57a9994f7a237176aa9f4cdc31bc466d056d
Author: Nathan Wilson <67372905+Nathanw1014@users.noreply.github.com>
Date: Thu Oct 15 17:39:02 2026 +1030
vulkan : copy strided f16 KV once on AMD so reads spread across memory channels (#27703)
* vulkan : copy strided f16 KV once so reads spread across memory channels
Assisted-by: Claude (Opus 4.8 / Opus 5)
* tests : cover the f16 KV copy to scratch
Assisted-by: Claude (Opus 4.8 / Opus 5)
* vulkan : gate the f16 KV scratch copy to AMD
Assisted-by: Claude (Opus 4.8 / Opus 5)
* vulkan : reuse the generic copy shader for the f16 KV scratch fill
The f16 clause needs only a strided copy, not a dequant, so drop the
dedicated shader and dispatch the existing copy pipeline instead. The
copy is described in source memory order so the strided access lands on
the writes, which stream, instead of the reads, which would concentrate
on a few memory channels.
Assisted-by: Claude (Opus 4.8 / Opus 5 / Fable 5)
* vulkan : extend the strided KV scratch copy to bf16
Assisted-by: Claude Opus 5
* vulkan : clean up comments
Assisted-by: Claude Opus 5
* consolidate copy shaders
* disable on RDNA4
---------
Co-authored-by: Ruben Ortlam <rortlam@redhat.com>
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-types.h b/ggml/src/ggml-vulkan/ggml-vulkan-types.h
index 7794c6310..a41ce5966 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan-types.h
+++ b/ggml/src/ggml-vulkan/ggml-vulkan-types.h
@@ -835,8 +835,8 @@ struct vk_device_struct {
vk_pipeline pipeline_roll_f32;
vk_pipeline pipeline_repeat_i32, pipeline_repeat_back_f32;
vk_pipeline pipeline_repeat_i16;
- vk_pipeline pipeline_cpy_f32_f32, pipeline_cpy_f32_f16, pipeline_cpy_f16_f16, pipeline_cpy_f16_f32, pipeline_cpy_f32_bf16, pipeline_cpy_bf16_f32, pipeline_cpy_bf16_f16, pipeline_cpy_f32_i32, pipeline_cpy_i32_f32;
- vk_pipeline pipeline_contig_cpy_f32_f32, pipeline_contig_cpy_f32_f16, pipeline_contig_cpy_f16_f16, pipeline_contig_cpy_f16_f32, pipeline_contig_cpy_f32_bf16, pipeline_contig_cpy_bf16_f32, pipeline_contig_cpy_bf16_f16, pipeline_contig_cpy_f32_i32, pipeline_contig_cpy_i32_f32;
+ vk_pipeline pipeline_cpy_32, pipeline_cpy_f32_f16, pipeline_cpy_16, pipeline_cpy_f16_f32, pipeline_cpy_f32_bf16, pipeline_cpy_bf16_f32, pipeline_cpy_bf16_f16, pipeline_cpy_f32_i32, pipeline_cpy_i32_f32;
+ vk_pipeline pipeline_contig_cpy_32, pipeline_contig_cpy_f32_f16, pipeline_contig_cpy_16, pipeline_contig_cpy_f16_f32, pipeline_contig_cpy_f32_bf16, pipeline_contig_cpy_bf16_f32, pipeline_contig_cpy_bf16_f16, pipeline_contig_cpy_f32_i32, pipeline_contig_cpy_i32_f32;
vk_pipeline pipeline_cpy_f32_quant[GGML_TYPE_COUNT];
vk_pipeline pipeline_cpy_quant_f32[GGML_TYPE_COUNT];
vk_pipeline pipeline_cpy_transpose_16, pipeline_cpy_transpose_32;
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
index 8357f01fe..420571140 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp
+++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
@@ -3273,9 +3273,9 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_rms_norm_back_f32, "rms_norm_back_f32", rms_norm_back_f32_len, rms_norm_back_f32_data, "main", 3, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_l2_norm_f32, "l2_norm_f32", l2_norm_f32_len, l2_norm_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1);
- ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_f32, "cpy_f32_f32", cpy_f32_f32_len, cpy_f32_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
+ ggml_vk_create_pipeline(device, device->pipeline_cpy_32, "cpy_32", cpy_32_len, cpy_32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_f16, "cpy_f32_f16", cpy_f32_f16_len, cpy_f32_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
- ggml_vk_create_pipeline(device, device->pipeline_cpy_f16_f16, "cpy_f16_f16", cpy_f16_f16_len, cpy_f16_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
+ ggml_vk_create_pipeline(device, device->pipeline_cpy_16, "cpy_16", cpy_16_len, cpy_16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_cpy_f16_f32, "cpy_f16_f32", cpy_f16_f32_len, cpy_f16_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_bf16,"cpy_f32_bf16",cpy_f32_bf16_len,cpy_f32_bf16_data,"main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_cpy_bf16_f32,"cpy_bf16_f32",cpy_bf16_f32_len,cpy_bf16_f32_data,"main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
@@ -3283,9 +3283,9 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
ggml_vk_create_pipeline(device, device->pipeline_cpy_i32_f32, "cpy_i32_f32", cpy_i32_f32_len, cpy_i32_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_i32, "cpy_f32_i32", cpy_f32_i32_len, cpy_f32_i32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
- ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f32_f32, "contig_cpy_f32_f32", contig_cpy_f32_f32_len, contig_cpy_f32_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
+ ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_32, "contig_cpy_32", contig_cpy_32_len, contig_cpy_32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f32_f16, "contig_cpy_f32_f16", contig_cpy_f32_f16_len, contig_cpy_f32_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
- ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f16_f16, "contig_cpy_f16_f16", contig_cpy_f16_f16_len, contig_cpy_f16_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
+ ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_16, "contig_cpy_16", contig_cpy_16_len, contig_cpy_16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f16_f32, "contig_cpy_f16_f32", contig_cpy_f16_f32_len, contig_cpy_f16_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f32_bf16,"contig_cpy_f32_bf16",contig_cpy_f32_bf16_len,contig_cpy_f32_bf16_data,"main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_bf16_f32,"contig_cpy_bf16_f32",contig_cpy_bf16_f32_len,contig_cpy_bf16_f32_data,"main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
@@ -6100,13 +6100,6 @@ vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const ggml_t
}
}
- if (src->type == GGML_TYPE_F32 && to == GGML_TYPE_F32) {
- if (contig) {
- return ctx->device->pipeline_contig_cpy_f32_f32;
- } else {
- return ctx->device->pipeline_cpy_f32_f32;
- }
- }
if (src->type == GGML_TYPE_F32 && to == GGML_TYPE_F16) {
if (contig) {
return ctx->device->pipeline_contig_cpy_f32_f16;
@@ -6114,13 +6107,6 @@ vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const ggml_t
return ctx->device->pipeline_cpy_f32_f16;
}
}
- if (src->type == GGML_TYPE_F16 && to == GGML_TYPE_F16) {
- if (contig) {
- return ctx->device->pipeline_contig_cpy_f16_f16;
- } else {
- return ctx->device->pipeline_cpy_f16_f16;
- }
- }
if (src->type == GGML_TYPE_F16 && to == GGML_TYPE_F32) {
if (contig) {
return ctx->device->pipeline_contig_cpy_f16_f32;
@@ -6203,15 +6189,15 @@ vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const ggml_t
GGML_ASSERT(ggml_is_quantized(to) || ggml_type_size(src->type) == 2 || ggml_type_size(src->type) == 4);
if ((ggml_type_size(src->type) % 4) == 0) {
if (contig) {
- return ctx->device->pipeline_contig_cpy_f32_f32;
+ return ctx->device->pipeline_contig_cpy_32;
} else {
- return ctx->device->pipeline_cpy_f32_f32;
+ return ctx->device->pipeline_cpy_32;
}
} else {
if (contig) {
- return ctx->device->pipeline_contig_cpy_f16_f16;
+ return ctx->device->pipeline_contig_cpy_16;
} else {
- return ctx->device->pipeline_cpy_f16_f16;
+ return ctx->device->pipeline_cpy_16;
}
}
}
@@ -8134,19 +8120,30 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
};
const bool k_quant = k->type != GGML_TYPE_F16 && k->type != GGML_TYPE_BF16 && k->type != GGML_TYPE_F32;
const bool v_quant = v->type != GGML_TYPE_F16 && v->type != GGML_TYPE_BF16 && v->type != GGML_TYPE_F32;
- const bool use_dequant_kv = k_quant && v_quant && neq1 >= 64 &&
+ const bool kv_raw = (k->type == GGML_TYPE_F16 || k->type == GGML_TYPE_BF16) && k->type == v->type;
+ // gated to AMD, except RDNA4, which regressed in benchmarks
+ const bool kv_raw_strided = ctx->device->vendor_id == VK_VENDOR_ID_AMD &&
+ ctx->device->architecture != vk_device_architecture::AMD_RDNA4 &&
+ neq1 >= 256 &&
+ kv_raw &&
+ (k->nb[1] != (uint64_t)HSK * ggml_type_size(k->type) ||
+ v->nb[1] != (uint64_t)HSV * ggml_type_size(v->type)) &&
+ (HSK % 8) == 0 && (HSV % 8) == 0;
+ const size_t kv_scratch_ts = kv_raw_strided ? ggml_type_size(k->type) : sizeof(ggml_fp16_t);
+ const bool use_dequant_kv = ((k_quant && v_quant) || kv_raw_strided) && neq1 >= 64 &&
is_dense_kv_cache(k) && is_dense_kv_cache(v) &&
- (uint64_t)ggml_nelements(k) * sizeof(ggml_fp16_t) <= ctx->device->properties.limits.maxStorageBufferRange &&
- (uint64_t)ggml_nelements(v) * sizeof(ggml_fp16_t) <= ctx->device->properties.limits.maxStorageBufferRange &&
- ctx->device->pipeline_dequant_transpose[k->type] != nullptr &&
- ctx->device->pipeline_dequant_transpose[v->type] != nullptr &&
+ (uint64_t)ggml_nelements(k) * kv_scratch_ts <= ctx->device->properties.limits.maxStorageBufferRange &&
+ (uint64_t)ggml_nelements(v) * kv_scratch_ts <= ctx->device->properties.limits.maxStorageBufferRange &&
+ (kv_raw_strided ||
+ (ctx->device->pipeline_dequant_transpose[k->type] != nullptr &&
+ ctx->device->pipeline_dequant_transpose[v->type] != nullptr)) &&
// coopmat2 path does not benefit from the f16 scratch
!ctx->device->coopmat2 &&
// Intel Xe1 regresses, see PR 25494
(ctx->device->vendor_id != VK_VENDOR_ID_INTEL ||
(ctx->device->coopmat_support && ctx->device->architecture != vk_device_architecture::INTEL_XE1));
- const ggml_type k_type_eff = use_dequant_kv ? GGML_TYPE_F16 : k->type;
- const ggml_type v_type_eff = use_dequant_kv ? GGML_TYPE_F16 : v->type;
+ 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).
@@ -8422,28 +8419,56 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
}
if (use_dequant_kv) {
- const uint64_t fp = sizeof(ggml_fp16_t);
- const uint64_t k_f16_sz = (uint64_t)ggml_nelements(k) * fp;
- const uint64_t v_f16_sz = (uint64_t)ggml_nelements(v) * fp;
- if (ctx->prealloc_size_x < k_f16_sz + v_f16_sz) {
- ctx->prealloc_size_x = k_f16_sz + v_f16_sz;
+ const uint64_t k_sz = (uint64_t)ggml_nelements(k) * kv_scratch_ts;
+ const uint64_t v_sz = (uint64_t)ggml_nelements(v) * kv_scratch_ts;
+ if (ctx->prealloc_size_x < k_sz + v_sz) {
+ ctx->prealloc_size_x = k_sz + v_sz;
ggml_vk_preallocate_buffers(ctx, subctx);
}
- vk_pipeline tr_k = ctx->device->pipeline_dequant_transpose[k->type];
- vk_pipeline tr_v = ctx->device->pipeline_dequant_transpose[v->type];
- ggml_pipeline_request_descriptor_sets(ctx, tr_k, 1);
- ggml_pipeline_request_descriptor_sets(ctx, tr_v, 1);
- if (ctx->prealloc_x_need_sync) {
- ggml_vk_sync_buffers(ctx, subctx);
- }
- vk_subbuffer k_dst = vk_subbuffer{ ctx->prealloc_x, 0, k_f16_sz };
- vk_subbuffer v_dst = vk_subbuffer{ ctx->prealloc_x, k_f16_sz, v_f16_sz };
+ vk_subbuffer k_dst = vk_subbuffer{ ctx->prealloc_x, 0, k_sz };
+ vk_subbuffer v_dst = vk_subbuffer{ ctx->prealloc_x, k_sz, v_sz };
const uint32_t k_nel = (uint32_t)ggml_nelements(k);
const uint32_t v_nel = (uint32_t)ggml_nelements(v);
- { const std::vector<uint32_t> pc = { (uint32_t)HSK, (uint32_t)nek2, (uint32_t)KV, 0, k_nel };
- ggml_vk_dispatch_pipeline(ctx, subctx, tr_k, { k_buf, k_dst }, pc, { k_nel, 1, 1 }); }
- { const std::vector<uint32_t> pc = { (uint32_t)HSV, (uint32_t)nev2, (uint32_t)KV, 0, v_nel };
- ggml_vk_dispatch_pipeline(ctx, subctx, tr_v, { v_buf, v_dst }, pc, { v_nel, 1, 1 }); }
+ if (kv_raw_strided) {
+ vk_pipeline cp_k = ggml_vk_get_cpy_pipeline(ctx, k, nullptr, k->type);
+ vk_pipeline cp_v = ggml_vk_get_cpy_pipeline(ctx, v, nullptr, v->type);
+ ggml_pipeline_request_descriptor_sets(ctx, cp_k, 1);
+ ggml_pipeline_request_descriptor_sets(ctx, cp_v, 1);
+ if (ctx->prealloc_x_need_sync) {
+ ggml_vk_sync_buffers(ctx, subctx);
+ }
+ auto make_pc = [](uint32_t hs, uint32_t nh, uint32_t kv, uint32_t nel) {
+ vk_op_unary_push_constants pc{};
+ pc.ne = nel;
+ pc.ne00 = hs; pc.ne01 = nh; pc.ne02 = kv; pc.ne03 = nel / (hs * nh * kv);
+ pc.nb00 = 1; pc.nb01 = hs; pc.nb02 = hs * nh; pc.nb03 = hs * nh * kv;
+ pc.ne10 = hs; pc.ne11 = nh; pc.ne12 = kv; pc.ne13 = pc.ne03;
+ pc.nb10 = 1; pc.nb11 = hs * kv; pc.nb12 = hs; pc.nb13 = hs * kv * nh;
+ init_pushconst_fastdiv(pc);
+ return pc;
+ };
+ auto cpy_elems = [](uint32_t ne) -> std::array<uint32_t, 3> {
+ if (ne > 262144) { return { 512, 512, CEIL_DIV(ne, 262144) }; }
+ if (ne > 512) { return { 512, CEIL_DIV(ne, 512), 1 }; }
+ return { ne, 1, 1 };
+ };
+ ggml_vk_dispatch_pipeline(ctx, subctx, cp_k, { k_buf, k_dst },
+ make_pc((uint32_t)HSK, (uint32_t)nek2, (uint32_t)KV, k_nel), cpy_elems(k_nel));
+ ggml_vk_dispatch_pipeline(ctx, subctx, cp_v, { v_buf, v_dst },
+ make_pc((uint32_t)HSV, (uint32_t)nev2, (uint32_t)KV, v_nel), cpy_elems(v_nel));
+ } else {
+ vk_pipeline tr_k = ctx->device->pipeline_dequant_transpose[k->type];
+ vk_pipeline tr_v = ctx->device->pipeline_dequant_transpose[v->type];
+ ggml_pipeline_request_descriptor_sets(ctx, tr_k, 1);
+ ggml_pipeline_request_descriptor_sets(ctx, tr_v, 1);
+ if (ctx->prealloc_x_need_sync) {
+ ggml_vk_sync_buffers(ctx, subctx);
+ }
+ { const std::vector<uint32_t> pc = { (uint32_t)HSK, (uint32_t)nek2, (uint32_t)KV, 0, k_nel };
+ ggml_vk_dispatch_pipeline(ctx, subctx, tr_k, { k_buf, k_dst }, pc, { k_nel, 1, 1 }); }
+ { const std::vector<uint32_t> pc = { (uint32_t)HSV, (uint32_t)nev2, (uint32_t)KV, 0, v_nel };
+ ggml_vk_dispatch_pipeline(ctx, subctx, tr_v, { v_buf, v_dst }, pc, { v_nel, 1, 1 }); }
+ }
ggml_vk_sync_buffers(ctx, subctx);
k_buf = k_dst;
v_buf = v_dst;
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 42ddf4858..f54501546 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp
@@ -864,18 +864,18 @@ void process_shaders() {
string_to_spv("rms_norm_back_f32", "rms_norm_back.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}));
string_to_spv("l2_norm_f32", "l2_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
- string_to_spv("cpy_f32_f32", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});
+ string_to_spv("cpy_32", "copy.comp", {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}});
string_to_spv("cpy_f32_f16", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}});
- string_to_spv("cpy_f16_f16", "copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}});
+ string_to_spv("cpy_16", "copy.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}});
string_to_spv("cpy_f16_f32", "copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}});
string_to_spv("cpy_f32_bf16","copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "uint16_t"}, {"DATA_D_BF16", "1"}});
string_to_spv("cpy_bf16_f32","copy.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "float"}, {"DATA_A_BF16", "1"}});
string_to_spv("cpy_bf16_f16","copy.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "float16_t"}, {"DATA_A_BF16", "1"}});
- string_to_spv("contig_cpy_f32_f32", "contig_copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}});
+ string_to_spv("contig_cpy_32", "contig_copy.comp", {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}});
string_to_spv("contig_cpy_f32_i32", "contig_copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "int"}});
string_to_spv("contig_cpy_i32_f32", "contig_copy.comp", {{"A_TYPE", "int"}, {"D_TYPE", "float"}});
string_to_spv("contig_cpy_f32_f16", "contig_copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}});
- string_to_spv("contig_cpy_f16_f16", "contig_copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}});
+ string_to_spv("contig_cpy_16", "contig_copy.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}});
string_to_spv("contig_cpy_f16_f32", "contig_copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}});
string_to_spv("contig_cpy_f32_bf16","contig_copy.comp",{{"A_TYPE", "float"}, {"D_TYPE", "uint16_t"}, {"DATA_D_BF16", "1"}});
string_to_spv("contig_cpy_bf16_f32","contig_copy.comp",{{"A_TYPE", "uint16_t"}, {"D_TYPE", "float"}, {"DATA_A_BF16", "1"}});
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 93b30d8d2..ba7f7eb4e 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -11817,6 +11817,13 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 1024, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false));
test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, false));
+ test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {8, 1}, 1024, 320, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3}, false));
+ test_cases.emplace_back(new test_flash_attn_ext(96, 96, 8, {4, 1}, 512, 256, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3}, false));
+ test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {8, 1}, 512, 259, true, false, 8.0f, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3}, false));
+ test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {8, 1}, 512, 288, true, false, 0, 30.0f, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3}, false));
+ test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {8, 1}, 1024, 320, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_BF16, GGML_TYPE_BF16, {0, 2, 1, 3}, false));
+ test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {8, 1}, 512, 259, true, false, 8.0f, 0, GGML_PREC_F32, GGML_TYPE_BF16, GGML_TYPE_BF16, {0, 2, 1, 3}, false));
+
// FLASH_ATTN_EXT MMA: non-pow2 head size and MLA K/V view.
test_cases.emplace_back(new test_flash_attn_ext(192, 128, 8, {8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true));