From bcc7c630f8db39add0922df597eeaee8f0fd4f9e Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Fri, 14 Aug 2026 14:28:08 +0200 Subject: [PATCH] metal: mul_mat: fix src1 row stride edge cases with has_tensor --- ggml/src/ggml-metal/ggml-metal-ops.cpp | 10 +++++++++- tests/test-backend-ops.cpp | 11 ++++++++++- 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 6d324056dd..f2c069267c 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -2437,7 +2437,15 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) { !ggml_is_transposed(op->src[1]) && // for now the matrix-matrix multiplication kernel only works on A14+/M1+ SoCs // AMD GPU and older A-chips will reuse matrix-vector multiplication kernel - props_dev->has_simdgroup_mm && ne00 >= 64 && ne11 > ne11_mm_min) { + props_dev->has_simdgroup_mm && ne00 >= 64 && ne11 > ne11_mm_min && + // the tensor API variant of kernel_mul_mm reads src1 directly from device memory through + // a cooperative tensor, which has two limitations - use the mat-vec kernels in these cases: + // - the src1 row stride is limited to 16 bits + // example: the permuted src1 produced by the MLA + FA attention epilogue in llama.cpp + // - the last K tile is read out of bounds when ne00 is not a multiple of the tile size, + // unlike src0 which is staged through threadgroup memory with zero padding + // example: the im2col src1 of a conv_2d with K = 14*14*3 (see #25652) + (!props_dev->has_tensor || (nb11/ggml_type_size(op->src[1]->type) < 65536 && ne00 % 32 == 0))) { //GGML_LOG_INFO("matrix: ne00 = %6d, ne01 = %6d, ne02 = %6d, ne11 = %6d, ne12 = %6d\n", ne00, ne01, ne02, ne11, ne12); // some Metal matrix data types require aligned pointers diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 08c29eec63..1011265310 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9082,6 +9082,14 @@ static std::vector> make_test_cases_eval() { } } + // permuted src1 with a row stride >= 2^16 elements, as produced by the MLA + FA attention + // epilogue with 128 heads (e.g. deepseek32, dots3note) + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 128, 32, 512, {128, 1}, {1, 1}, {0, 2, 1, 3})); + + // K not a multiple of the mat-mat tile size, as produced by conv_2d im2col with K = 14*14*3 + // (vision patch embedding, see #25652) + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 64, 32, 588, {1, 1}, {1, 1})); + // BF16 is absent from base_types: add the 3 standard non-contig permutations explicitly test_cases.emplace_back(new test_mul_mat(GGML_TYPE_BF16, GGML_TYPE_F32, 16, 1, 256, {2, 3}, {1, 1}, {0, 2, 1, 3})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_BF16, GGML_TYPE_F32, 16, 1, 256, {2, 3}, {1, 1}, {0, 1, 3, 2})); @@ -9711,12 +9719,13 @@ static std::vector> make_test_cases_eval() { if (nh == 1 && hsk != 320 && hsk != 576) continue; for (int nr3 : { 1, 3, }) { if (hsk > 64 && nr3 > 1) continue; // skip broadcast for large head sizes - for (int nr2 : { 1, 4, 8, 12, 16, 20, 32 }) { + for (int nr2 : { 1, 4, 8, 12, 16, 20, 32, 128 }) { if (nr2 == 8 && hsk != 192) continue; if (nr2 == 12 && hsk != 128) continue; if (nr2 == 16 && hsk != 192) continue; if (nr2 == 20 && (nh != 1 || hsk != 576)) continue; if (nr2 == 32 && (nh != 1 || hsk != 320)) continue; + if (nr2 == 128 && (nh != 1 || hsk != 576)) continue; // deepseek32/dots3note MLA-as-MQA (128 q heads, 1 kv head) //for (int kv : { 1, 17, 31, 33, 61, 113, 65, 127, 129, 130, 255, 260, 371, 380, 407, 512, 1024, }) { for (int kv : { 113, 512, 1024, }) { if (nr2 != 1 && kv != 512) continue;