diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index f6cbaecb7f54..8bdd61f00768 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -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 @@ -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 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 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 { + 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 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 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/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 8e3b273a1e46..e7ccd44b9c2e 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9979,6 +9979,12 @@ static std::vector> 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}));