Commit 5bda51bfb for llama.cpp
commit 5bda51bfbc62e64193221e639f6ad4e08767d760
Author: Foad Abo Dahood <32059146+masterFoad@users.noreply.github.com>
Date: Fri Sep 11 14:12:55 2026 +0300
metal : skip the empty half of the mul_mm_id token tile (#28301)
kernel_mul_mm_id splits its NR1 = 32 token tile into two 16-row halves and skips
the upper half when the expert did not fill it, on both the tensor and simdgroup
paths. The tB extents are corrected to (NK, NR1H) for the [NR1][NK] row-major tile.
The B tile is staged unconditionally, as on master: rows past nr1 restage a clamped
duplicate of a valid row, lie in the output-row dimension so they never contribute
to a valid row, and are dropped by the final store loop.
test-backend-ops: re-draw the expert ids between perf iterations of test_mul_mat_id
so MoE perf numbers are not warm-cache, and add token-tile boundary coverage using
n_used == n_mats, which routes every token to every expert so each expert receives
exactly n rows; n = 32, 33, 47, 48, 49 reach mul_mm_id and leave a last tile of 32,
1, 15, 16 and 17 rows.
diff --git a/ggml/src/ggml-metal/kernels/mul_mm.metal b/ggml/src/ggml-metal/kernels/mul_mm.metal
index ee848eed6..0a45bb1bb 100644
--- a/ggml/src/ggml-metal/kernels/mul_mm.metal
+++ b/ggml/src/ggml-metal/kernels/mul_mm.metal
@@ -496,6 +496,13 @@ kernel void kernel_mul_mm_id(
+ args.nb11*i11
+ args.nb10*iy);
+ // skip the upper half of the token tile when the expert did not fill it
+ constexpr short NR1H = NR1/2;
+
+ const bool has_hi = nr1 > NR1H;
+
+ const short lb1 = (short) tiitg/NL1; // 0 .. NR1-1, this thread's row of the B tile
+
#ifndef GGML_METAL_HAS_TENSOR
S0_8x8 ma[4];
S1_8x8 mb[2];
@@ -505,15 +512,22 @@ kernel void kernel_mul_mm_id(
for (short i = 0; i < 8; i++){
mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
}
+
+ // simdgroups 2,3 own rows NR1H..NR1-1
+ const bool sg_active = has_hi || sgitg < 2;
#else
- auto tA = tensor<threadgroup S0, dextents<int32_t, 2>, tensor_inline>(sa, dextents<int32_t, 2>(NK, NR0));
- auto tB = tensor<threadgroup S1, dextents<int32_t, 2>, tensor_inline>(sb, dextents<int32_t, 2>(NR1, NK ));
+ auto tA = tensor<threadgroup S0, dextents<int32_t, 2>, tensor_inline>(sa, dextents<int32_t, 2>(NK, NR0));
+
+ // sb is [NR1][NK] row-major
+ auto tB0 = tensor<threadgroup S1, dextents<int32_t, 2>, tensor_inline>(sb, dextents<int32_t, 2>(NK, NR1H));
+ auto tB1 = tensor<threadgroup S1, dextents<int32_t, 2>, tensor_inline>(sb + NR1H*NK, dextents<int32_t, 2>(NK, NR1H));
mpp::tensor_ops::matmul2d<
- mpp::tensor_ops::matmul2d_descriptor(NR1, NR0, NK, false, true, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate),
+ mpp::tensor_ops::matmul2d_descriptor(NR1H, NR0, NK, false, true, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate),
execution_simdgroups<4>> mm;
- auto cT = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB), float>();
+ auto cT0 = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB0), float>();
+ auto cT1 = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB1), float>();
#endif
for (int loop_k = 0; loop_k < args.ne00; loop_k += NK) {
@@ -656,37 +670,45 @@ kernel void kernel_mul_mm_id(
threadgroup_barrier(mem_flags::mem_threadgroup);
#ifndef GGML_METAL_HAS_TENSOR
- // load matrices from threadgroup memory and conduct outer products
- threadgroup const S0 * lsma = (sa + 4*64*(sgitg%2));
- threadgroup const S1 * lsmb = (sb + 2*64*(sgitg/2));
+ if (sg_active) {
+ // load matrices from threadgroup memory and conduct outer products
+ threadgroup const S0 * lsma = (sa + 4*64*(sgitg%2));
+ threadgroup const S1 * lsmb = (sb + 2*64*(sgitg/2));
- FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
- simdgroup_barrier(mem_flags::mem_none);
+ FOR_UNROLL (short ik = 0; ik < NK/8; ik++) {
+ simdgroup_barrier(mem_flags::mem_none);
- FOR_UNROLL (short i = 0; i < 4; i++) {
- simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);
- }
+ FOR_UNROLL (short i = 0; i < 4; i++) {
+ simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);
+ }
- simdgroup_barrier(mem_flags::mem_none);
+ simdgroup_barrier(mem_flags::mem_none);
- FOR_UNROLL (short i = 0; i < 2; i++) {
- simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
- }
+ FOR_UNROLL (short i = 0; i < 2; i++) {
+ simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);
+ }
- simdgroup_barrier(mem_flags::mem_none);
+ simdgroup_barrier(mem_flags::mem_none);
- FOR_UNROLL (short i = 0; i < 8; i++){
- simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
- }
+ FOR_UNROLL (short i = 0; i < 8; i++){
+ simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
+ }
- lsma += 8*64;
- lsmb += 4*64;
+ lsma += 8*64;
+ lsmb += 4*64;
+ }
}
#else
- auto sA = tA.slice(0, 0);
- auto sB = tB.slice(0, 0);
+ auto sA = tA.slice(0, 0);
+ auto sB0 = tB0.slice(0, 0);
- mm.run(sB, sA, cT);
+ mm.run(sB0, sA, cT0);
+
+ if (has_hi) {
+ auto sB1 = tB1.slice(0, 0);
+
+ mm.run(sB1, sA, cT1);
+ }
#endif
}
@@ -694,13 +716,20 @@ kernel void kernel_mul_mm_id(
threadgroup_barrier(mem_flags::mem_threadgroup);
#ifdef GGML_METAL_HAS_TENSOR
- auto tC = tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline>(sc, dextents<int32_t, 2>(NR0, NR1));
- cT.store(tC);
+ auto tC0 = tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline>(sc, dextents<int32_t, 2>(NR0, NR1H));
+ cT0.store(tC0);
+
+ if (has_hi) {
+ auto tC1 = tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline>(sc + NR1H*NR0, dextents<int32_t, 2>(NR0, NR1H));
+ cT1.store(tC1);
+ }
#else
- threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
+ if (sg_active) {
+ threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
- for (short i = 0; i < 8; i++) {
- simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
+ for (short i = 0; i < 8; i++) {
+ simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false);
+ }
}
#endif
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 15c42a108..b63b3773e 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -1219,6 +1219,11 @@ struct test_case {
}
}
+ // re-draw data-dependent inputs between timed perf iterations
+ virtual void reinit_perf_iter(ggml_context * ctx) {
+ GGML_UNUSED(ctx);
+ }
+
virtual size_t op_size(ggml_tensor * t) {
size_t size = ggml_nbytes(t);
// add source tensors
@@ -1653,6 +1658,9 @@ struct test_case {
total_time_us += end_time - start_time;
total_mem += mem;
total_runs += n_runs;
+
+ // re-draw any data-dependent inputs (expert ids) outside the timed region
+ reinit_perf_iter(ctx.get());
} while (total_time_us < 1000*1000); // run for at least 1 second
// Create test result
@@ -5000,25 +5008,31 @@ struct test_mul_mat_hadamard : public test_mul_mat {
}
};
-static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats) {
+static void init_mul_mat_id_ids(ggml_context * ctx, int n_mats) {
std::random_device rd;
std::default_random_engine rng(rd());
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
- if (t->type == GGML_TYPE_I32) {
- if (ggml_is_view_op(t->op)) { continue; }
- // ids
- for (int64_t r = 0; r < ggml_nrows(t); r++) {
- std::vector<int32_t> data(t->ne[0]);
- for (int i = 0; i < t->ne[0]; i++) {
- data[i] = i % n_mats;
- }
- std::shuffle(data.begin(), data.end(), rng);
- ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t));
+ if (t->type != GGML_TYPE_I32 || ggml_is_view_op(t->op)) {
+ continue;
+ }
+ for (int64_t r = 0; r < ggml_nrows(t); r++) {
+ std::vector<int32_t> data(t->ne[0]);
+ for (int i = 0; i < t->ne[0]; i++) {
+ data[i] = i % n_mats;
}
- } else {
+ std::shuffle(data.begin(), data.end(), rng);
+ ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t));
+ }
+ }
+}
+
+static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats) {
+ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
+ if (t->type != GGML_TYPE_I32) {
init_tensor_uniform(t);
}
}
+ init_mul_mat_id_ids(ctx, n_mats);
}
// GGML_OP_MUL_MAT_ID
@@ -5085,6 +5099,10 @@ struct test_mul_mat_id : public test_case {
void initialize_tensors(ggml_context * ctx) override {
init_mul_mat_id_tensors(ctx, n_mats);
}
+
+ void reinit_perf_iter(ggml_context * ctx) override {
+ init_mul_mat_id_ids(ctx, n_mats);
+ }
};
// GGML_OP_MUL_MAT_ID + GGML_OP_ADD or GGML_OP_MUL
@@ -9890,6 +9908,19 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_BF16, GGML_TYPE_F32, 16, 16, 256, {2, 3}, {1, 1}, {0, 1, 3, 2}));
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_BF16, GGML_TYPE_F32, 16, 16, 256, {2, 3}, {1, 1}, {0, 3, 2, 1}));
+ // token-tile boundary coverage. With n_used == n_mats every token routes to every expert, so
+ // each expert receives exactly n rows, with no dependence on the random draw. mul_mm_id is used
+ // from 32 tokens up: n = 32, 33, 47, 48, 49 reach it, leaving a last tile of 32, 1, 15, 16 and
+ // 17 rows - 16 and 17 straddle the point where the upper half stops being skipped. The smaller
+ // n cover the same row counts on the mat-vec path.
+ for (ggml_type type_a : {GGML_TYPE_Q4_K, GGML_TYPE_IQ2_XS, GGML_TYPE_F16}) {
+ for (int n : {1, 15, 16, 17, 31, 32, 33, 47, 48, 49}) {
+ test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 4, 4, false, 512, n, 256));
+ }
+ // experts that receive no rows at all
+ test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 8, 1, false, 512, 1, 256));
+ }
+
for (ggml_type type_a : other_types) {
for (ggml_type type_b : {GGML_TYPE_F32}) {
if (ggml_blck_size(type_a) != 256) {