From d68216aa3ece8bf56a3ddc80a0a238e09028ab8c Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Thu, 20 Aug 2026 21:31:29 +0300 Subject: [PATCH] metal : clamp K extent in tensor API mat-mat kernel for K not a multiple of 32 (llama/27450) The Tensor API mat-mat path of kernel_mul_mm (GGML_METAL_HAS_TENSOR) fed a static K=32 tile to the matmul2d op on every iteration. On the last, partial K tile (ne00 % 32 != 0) the src1 slice extends past the K extent of the tensor, and the op reads those out-of-bounds elements (undefined behavior per the MSL specification, section 2.22.2). Depending on stale memory contents, this corrupted the result or produced NaN. Make the matmul2d op use dynamic_extent for K, and clamp the K extent of both operand tensor views to the remaining valid K range (min(32, K - loop_k)) per iteration, so the op reads exactly the valid K range on every iteration (mirroring the tail handling of the MPP matmul2d examples). On K-aligned inputs the clamp degenerates to the full 32-wide tile: the only difference from the static-K op is that the dynamic-K op derives K from the operand extents and edge-checks the tile against the tensor extents (a handful of integer ops per iteration). Add test-backend-ops MUL_MAT cases with K not a multiple of 32 to exercise the unaligned K path. Assisted-by: pi:llama.cpp/Qwen3.8-27B --- ggml/src/ggml-metal/ggml-metal.metal | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index 949931c8d..27f97b5e0 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -10365,9 +10365,12 @@ kernel void kernel_mul_mm( auto tB = tensor(ptrB, dextents(K, N), array({1, strideB})); // Configure matmul operation + // note: K is dynamic_extent (clamped to the valid range in PHASE 2), since a static + // N_MM_NK_TOTAL K tile would read src1 out of bounds when K % N_MM_NK_TOTAL != 0 + // ref: https://github.com/ggml-org/llama.cpp/pull/27064 mpp::tensor_ops::matmul2d< mpp::tensor_ops::matmul2d_descriptor( - NRB, NRA, N_MM_NK_TOTAL, false, true, true, + NRB, NRA, static_cast(dynamic_extent), false, true, true, mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate), execution_simdgroups> mm; @@ -10419,10 +10422,14 @@ kernel void kernel_mul_mm( threadgroup_barrier(mem_flags::mem_threadgroup); // === PHASE 2: Tensor matmul === - auto mA = tA.slice(0, 0); - auto mB = tB.slice(loop_k, rb); + // Clamp the K extent of both operand tensors to the remaining valid K range so + // the dynamic-K op never reads past the K extent of src1 (or the staged A tile). + const int kExt = min(N_MM_NK_TOTAL, K - loop_k); - mm.run(mB, mA, cT); + auto tAv = tensor(sa, dextents(kExt, NRA), array({1, N_MM_NK_TOTAL})); + auto tBv = tensor(ptrB + loop_k + rb * strideB, dextents(kExt, N - rb), array({1, strideB})); + + mm.run(tBv, tAv, cT); threadgroup_barrier(mem_flags::mem_threadgroup); }