From 5f310fa112bd700e5b6f3f53b2da5b3939923724 Mon Sep 17 00:00:00 2001 From: Foad Abo Dahood <32059146+masterFoad@users.noreply.github.com> Date: Fri, 11 Sep 2026 14:12:55 +0300 Subject: [PATCH] metal : skip the empty half of the mul_mm_id token tile (llama/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. --- ggml/src/ggml-metal/kernels/mul_mm.metal | 99 +++++++++++++++--------- 1 file changed, 64 insertions(+), 35 deletions(-) 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(0.f); } + + // simdgroups 2,3 own rows NR1H..NR1-1 + const bool sg_active = has_hi || sgitg < 2; #else - auto tA = tensor, tensor_inline>(sa, dextents(NK, NR0)); - auto tB = tensor, tensor_inline>(sb, dextents(NR1, NK )); + auto tA = tensor, tensor_inline>(sa, dextents(NK, NR0)); + + // sb is [NR1][NK] row-major + auto tB0 = tensor, tensor_inline>(sb, dextents(NK, NR1H)); + auto tB1 = tensor, tensor_inline>(sb + NR1H*NK, dextents(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(); + auto cT0 = mm.get_destination_cooperative_tensor(); + auto cT1 = mm.get_destination_cooperative_tensor(); #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); + + FOR_UNROLL (short i = 0; i < 2; i++) { + simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false); + } + + 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]); + } + + lsma += 8*64; + lsmb += 4*64; } - - simdgroup_barrier(mem_flags::mem_none); - - FOR_UNROLL (short i = 0; i < 2; i++) { - simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false); - } - - 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]); - } - - 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, tensor_inline>(sc, dextents(NR0, NR1)); - cT.store(tC); -#else - threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0; + auto tC0 = tensor, tensor_inline>(sc, dextents(NR0, NR1H)); + cT0.store(tC0); - for (short i = 0; i < 8; i++) { - simdgroup_store(mc[i], temp_str + 8*(i%4) + 8*NR0*(i/4), NR0, 0, false); + if (has_hi) { + auto tC1 = tensor, tensor_inline>(sc + NR1H*NR0, dextents(NR0, NR1H)); + cT1.store(tC1); + } +#else + 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); + } } #endif