diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 9d336ec373d3..5042d2e69ae2 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -7003,6 +7003,24 @@ static void ggml_vk_mul_mat_vec_q_f16_cols(ggml_backend_vk_context * ctx, vk_con const uint64_t y_sz = quantize_y ? (ggml_vk_align_size(y_ne, 128) * ggml_type_size(GGML_TYPE_Q8_1) / ggml_blck_size(GGML_TYPE_Q8_1)) : (f16_f32_kernel ? sizeof(float) * y_ne : sizeof(ggml_fp16_t) * y_ne); + // Swapped inputs with more than one output row: the mat-vec writes dst transposed + // ([ne1, ne0] row-major, the vector index is the slow dimension), so it goes to the + // split_k scratch buffer and a strided copy puts it in place. dst is ne0 x ne1 with + // ne0 <= mul_mat_vec_max_cols, so the copy is a few KB. + const bool transpose_d = swap_inputs && dst->ne[0] > 1; + ggml_tensor dst_t; + vk_pipeline cpy_d = nullptr; + if (transpose_d) { + GGML_ASSERT(dst->type == GGML_TYPE_F32 && ggml_is_contiguous(dst) && ctx->num_additional_fused_ops == 0); + dst_t = *dst; + dst_t.nb[0] = dst->ne[1] * sizeof(float); + dst_t.nb[1] = sizeof(float); + dst_t.nb[2] = ggml_nbytes(dst); + dst_t.nb[3] = ggml_nbytes(dst); + cpy_d = ggml_vk_get_cpy_pipeline(ctx, &dst_t, nullptr, GGML_TYPE_F32); + GGML_ASSERT(cpy_d != nullptr); + } + { if ( (qx_needs_dequant && x_sz > ctx->device->properties.limits.maxStorageBufferRange) || @@ -7017,6 +7035,10 @@ static void ggml_vk_mul_mat_vec_q_f16_cols(ggml_backend_vk_context * ctx, vk_con ctx->prealloc_size_y = y_sz; ggml_vk_preallocate_buffers(ctx, subctx); } + if (transpose_d && ctx->prealloc_size_split_k < ggml_nbytes(dst)) { + ctx->prealloc_size_split_k = ggml_nbytes(dst); + ggml_vk_preallocate_buffers(ctx, subctx); + } // Request descriptor sets if (qx_needs_dequant) { @@ -7028,6 +7050,9 @@ static void ggml_vk_mul_mat_vec_q_f16_cols(ggml_backend_vk_context * ctx, vk_con if (quantize_y) { ggml_pipeline_request_descriptor_sets(ctx, to_q8_1, 1); } + if (transpose_d) { + ggml_pipeline_request_descriptor_sets(ctx, cpy_d, 1); + } } vk_subbuffer d_D = ggml_vk_tensor_subbuffer(ctx, cgraph->nodes[node_idx + ctx->num_additional_fused_ops]); @@ -7102,6 +7127,14 @@ static void ggml_vk_mul_mat_vec_q_f16_cols(ggml_backend_vk_context * ctx, vk_con stride_batch_y = src1->nb[0] / ggml_type_size(src1->type); } + vk_subbuffer d_Out = d_D; + if (transpose_d) { + if (ctx->prealloc_split_k_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } + d_Out = { ctx->prealloc_split_k, 0, ggml_nbytes(dst) }; + } + const uint32_t max_groups_x = ctx->device->properties.limits.maxComputeWorkGroupCount[0]; uint32_t groups_x = ne01; @@ -7148,7 +7181,7 @@ static void ggml_vk_mul_mat_vec_q_f16_cols(ggml_backend_vk_context * ctx, vk_con { d_X, d_Y, - d_D, + d_Out, d_F0, d_F1, }, @@ -7156,6 +7189,12 @@ static void ggml_vk_mul_mat_vec_q_f16_cols(ggml_backend_vk_context * ctx, vk_con base_work_group_y += groups_y; } + if (transpose_d) { + ggml_vk_sync_buffers(ctx, subctx); + ggml_vk_cpy_to_contiguous(ctx, subctx, cpy_d, &dst_t, d_Out, d_D); + ctx->prealloc_split_k_need_sync = true; + } + if (x_non_contig) { ctx->prealloc_x_need_sync = true; } @@ -7643,11 +7682,17 @@ void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, const st ggml_vk_mul_mat_vec_nc_f16_f32(ctx, subctx, cgraph, node_idx); // mul_mat_vec supports batching ne12*ne13 when ne11==1, or treating ne11 as the batch size (up to four) // when ne12 and ne13 are one. - // With one output row, B^T*A has the same flat output as A^T*B. + // With one output row, B^T*A has the same flat output as A^T*B. With a few output rows + // (up to mul_mat_vec_max_cols, e.g. a [k, 4] f32 projection over a whole ubatch), the + // src0 rows become the mat-vec columns and the result is transposed into dst, instead + // of padding a tiny M to a full mul_mm tile. } else if (ctx->num_additional_fused_ops == 0 && (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16) && (src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16 || src1->type == GGML_TYPE_BF16 || ggml_is_quantized(src1->type)) && - dst->ne[0] == 1 && dst->ne[1] > mul_mat_vec_max_cols && + dst->ne[0] <= mul_mat_vec_max_cols && dst->ne[1] > mul_mat_vec_max_cols && + // more than one output row: transposed write, F32 dst only, and src1 (the swapped mat-vec matrix) + // limited to the f32/f16 that test-backend-ops can check against the CPU reference + (dst->ne[0] == 1 || (dst->type == GGML_TYPE_F32 && (src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16))) && src0->ne[2] == 1 && src0->ne[3] == 1 && src1->ne[2] == 1 && src1->ne[3] == 1 && ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst) && diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 5d8befe159b7..9e5bee849149 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -12481,6 +12481,21 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 32, 509, 2112, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 32, 509, 2112, {1, 1}, {1, 1})); + // small m (2..8) with n above MMVF_MAX_BATCH_SIZE: operand swap with a transposed result. + // m = 4, k = 10240 is the Flash-Next hyper-connection inject projection ([n_embd*hc, hc] f32); + // n = 1 and 4 are its decode / draft-verify sizes, which stay on the plain mat-vec path. + for (int64_t m : {2, 3, 4, 5, 8, 9}) { + for (int64_t n : {9, 16, 509, 2048}) { + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, m, n, 2048, {1, 1}, {1, 1})); + } + } + for (int64_t n : {1, 4, 8, 9, 512, 2048}) { + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 4, n, 10240, {1, 1}, {1, 1})); + } + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 4, 509, 2051, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 4, 509, 2051, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 4, 509, 2051, {2, 1}, {1, 1})); + #if 0 { // Test paths in OpenCL @@ -14168,6 +14183,8 @@ static std::vector> make_test_cases_perf() { } test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 248320, 1, 2048, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 256, 1, 2048, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 4, 2048, 10240, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 4, 512, 10240, {1, 1}, {1, 1})); for (int K : {3, 5}) { for (int IC : {256, 2560}) {