Skip to content

vulkan : copy strided f16 KV once on AMD so reads spread across memory channels #27703

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Nathanw1014 wants to merge 4 commits into
base: master
Choose a base branch
from
+62 −14
Open

vulkan : copy strided f16 KV once on AMD so reads spread across memory channels #27703

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
07c113e
vulkan : copy strided f16 KV once so reads spread across memory channels
Nathanw1014 Jul 28, 2026
9f3e5fa
tests : cover the f16 KV copy to scratch
Nathanw1014 Jul 28, 2026
0e57398
vulkan : gate the f16 KV scratch copy to AMD
Nathanw1014 Aug 22, 2026
4d9e62f
vulkan : reuse the generic copy shader for the f16 KV scratch fill
Nathanw1014 Aug 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 56 additions & 14 deletions ggml/src/ggml-vulkan/ggml-vulkan.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10834,12 +10834,22 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
};
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 &&
// f16 K/V in the strided KV-cache layout goes through the same scratch as a plain copy.
// A plain copy amortizes later than the fused dequant, so it needs a larger batch.
// Gated to AMD, where the strided-load penalty was measured.
const bool kv_f16_strided = ctx->device->vendor_id == VK_VENDOR_ID_AMD &&
neq1 >= 256 &&
k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16 &&
(k->nb[1] != (uint64_t)HSK * sizeof(ggml_fp16_t) ||
v->nb[1] != (uint64_t)HSV * sizeof(ggml_fp16_t)) &&
(HSK % 8) == 0 && (HSV % 8) == 0;
const bool use_dequant_kv = ((k_quant && v_quant) || kv_f16_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 &&
(kv_f16_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
Expand Down Expand Up @@ -11026,21 +11036,53 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
ctx->prealloc_size_x = k_f16_sz + v_f16_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 };
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 (k->type == GGML_TYPE_F16) {
// f16 K/V need only a strided copy, so reuse the generic copy shader, iterating
// in source memory order so the strided side is the write.
vk_pipeline cp_k = ggml_vk_get_cpy_pipeline(ctx, k, nullptr, GGML_TYPE_F16);
vk_pipeline cp_v = ggml_vk_get_cpy_pipeline(ctx, v, nullptr, GGML_TYPE_F16);
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;
// read [HS, NH, KV, NS] linearly, write the scratch layout [HS, KV, NH, NS]
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;
Expand Down
6 changes: 6 additions & 0 deletions tests/test-backend-ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9979,6 +9979,12 @@ 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));

// dense-allocated f16 K/V in the cache layout at batch >= 256 (engages the f16 copy into the FA scratch)
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_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3}));
test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1}));
test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, { 10, 5, 4, 3}));
Expand Down