Skip to content

vulkan : run small-M mul_mat as a swapped mat-vec #129

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

Merged
gaetan-puleo merged 4 commits into from
Oct 3, 2026
Merged

vulkan : run small-M mul_mat as a swapped mat-vec #129

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
4fa5cc3
vulkan: swap operands for small-M mul_mat with a transposed result
Nathanw1014 Oct 1, 2026
5fe3be2
Merge branch 'master' into hb/vk-hc-inject
gaetan-puleo Oct 3, 2026
109119a
vulkan : keep the small-M swapped mat-vec to F32 dst and f32/f16 src1
Nathanw1014 Oct 3, 2026
eeeeb88
Merge branch 'master' into hb/vk-hc-inject
gaetan-puleo Oct 3, 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
51 changes: 48 additions & 3 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 @@ -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) ||
Expand All @@ -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) {
Expand All @@ -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]);
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -7148,14 +7181,20 @@ 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,
},
pc, { groups_x, groups_y, groups_z });
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;
}
Expand Down Expand Up @@ -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) &&
Expand Down
17 changes: 17 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 @@ -12481,6 +12481,21 @@ static std::vector<std::unique_ptr<test_case>> 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
Expand Down Expand Up @@ -14168,6 +14183,8 @@ static std::vector<std::unique_ptr<test_case>> 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}) {
Expand Down
Loading