metal: mul_mat: fix src1 row stride edge cases with has_tensor

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

tests/test-backend-ops.cpp

@@ -9082,6 +9082,14 @@ static std::vector<std::unique_ptr<test_case>> 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<std::unique_ptr<test_case>> 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;