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); }