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