Commit 9b8446b91 for llama.cpp
commit 9b8446b91434e0c0cdc293baf0b29c099eaaa0c7
Author: feni6 <fenissova@proton.me>
Date: Mon Oct 26 02:13:39 2026 -0600
metal : stop overflow from happening in large matrix output offsets (#30072)
* Metal : stop overflow from happening in large matrix output offsets
Assisted-by: Codex
* metal : widen batch offsets and MUL_MAT_ID token indices
Assisted-by: Codex
---------
Co-authored-by: feni6 <270037094+feni6@users.noreply.github.com>
diff --git a/ggml/src/ggml-metal/kernels/mul_mm.metal b/ggml/src/ggml-metal/kernels/mul_mm.metal
index 2c8492037..edcc061d5 100644
--- a/ggml/src/ggml-metal/kernels/mul_mm.metal
+++ b/ggml/src/ggml-metal/kernels/mul_mm.metal
@@ -323,10 +323,10 @@ kernel void kernel_mul_mm(
// if no bounds checks on the output are needed, we can directly write to device memory
device float * C = (device float *) dst +
(r0 + 32*(sgitg & 1)) + \
- (r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0;
+ (uint64_t) (r1 + 16*(sgitg >> 1))*args.ne0 + (uint64_t) im*args.ne1*args.ne0;
for (short i = 0; i < 8; i++) {
- simdgroup_store(mc[i], C + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false);
+ simdgroup_store(mc[i], C + 8*(i%4) + (uint64_t) 8*args.ne0*(i/4), args.ne0, 0, false);
}
} else {
// block is smaller than 64x32, we should avoid writing data outside of the matrix
@@ -342,7 +342,7 @@ kernel void kernel_mul_mm(
if (sgitg == 0) {
for (int j = tiitg; j < nr1; j += NR1) {
- device float * D = (device float *) dst + r0 + (r1 + j)*args.ne0 + im*args.ne1*args.ne0;
+ device float * D = (device float *) dst + r0 + (uint64_t) (r1 + j)*args.ne0 + (uint64_t) im*args.ne1*args.ne0;
device float4 * D4 = (device float4 *) D;
threadgroup float * C = temp_str + (j*NR0);
@@ -633,7 +633,7 @@ kernel void kernel_mul_mm_id(
const int id = ids_i32[idb + lr1];
const short i11 = (id % args.ne20) % args.ne11;
- const short i12 = (id / args.ne20);
+ const int i12 = (id / args.ne20);
device const block_q * x = (device const block_q *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1;
@@ -866,9 +866,9 @@ kernel void kernel_mul_mm_id(
const int id = ids_i32[idb + j];
const short ide = id % args.ne20;
- const short idt = id / args.ne20;
+ const int idt = id / args.ne20;
- device float * D = (device float *) dst + r0 + ide*args.ne0 + idt*args.ne1*args.ne0;
+ device float * D = (device float *) dst + r0 + (uint64_t) ide*args.ne0 + (uint64_t) idt*args.ne1*args.ne0;
device float4 * D4 = (device float4 *) D;
threadgroup float * C = (threadgroup float *) shmem + j*NR0;
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 3e9439252..9a4ddaa29 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -11080,6 +11080,17 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_mul_mat_id_w4a8(GGML_TYPE_MXFP4, GGML_TYPE_F32, 8, 2, false, 32, 32, 256));
test_cases.emplace_back(new test_mul_mat_id_w4a4(GGML_TYPE_MXFP4, GGML_TYPE_F32, 8, 2, false, 32, 32, 256));
+#if 0
+ // > 2^31-element dst (8.9 GB), about 45 GB peak memory. Too large to be enabled by default.
+ // kernel_mul_mm dst offset of the last batch is past INT32_MAX: rows 0-31 use full tiles, rows 32-39 use edge tiles.
+ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 1048576, 40, 64, {1, 1}, {53, 1}));
+
+ // kernel_mul_mm_id: the index of tokens 32768-33791 does not fit in 16 bits (about 0.3 GB peak memory).
+ test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F32, GGML_TYPE_F32, 4, 2, false, 64, 33792, 64));
+ // kernel_mul_mm_id: > 2^31-element dst (8.7 GB), about 45 GB peak memory. Tokens 32514+ have a dst offset past INT32_MAX, tokens 32768+ also need a 32-bit token index.
+ test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F32, GGML_TYPE_F32, 16, 8, false, 8256, 32832, 64));
+#endif
+
#if 0
// > 4GB A matrix. Too slow to be enabled by default.
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 900000, 3, 2592, {1, 1}, {1, 1}));