diff --git a/ggml/src/ggml-metal/kernels/mul_mm.metal b/ggml/src/ggml-metal/kernels/mul_mm.metal index ee848eed6d6a..509bb53b1b01 100644 --- a/ggml/src/ggml-metal/kernels/mul_mm.metal +++ b/ggml/src/ggml-metal/kernels/mul_mm.metal @@ -133,11 +133,15 @@ kernel void kernel_mul_mm( } // Store result tile to output matrix (with batch offset) - // cT.store handles bounds checking via tD's extents (M, N) - device float * dstBatch = (device float *)dst + im * N * M; + // a slice is not used because its byte offset is int32 and wraps at 2 GiB + // ref: https://github.com/ggml-org/llama.cpp/pull/28748 + const int tileM = min(NRA, M - ra); + const int tileN = min(NRB, N - rb); - auto tD = tensor(dstBatch, dextents(M, N), array({1, M})); - cT.store(tD.slice(ra, rb)); + device float * dstTile = (device float *)dst + (uint64_t) im * N * M + (uint64_t) rb * M + ra; + + auto tD = tensor(dstTile, dextents(tileM, tileN), array({1, M})); + cT.store(tD); } #else