diff --git a/ggml/src/ggml-metal/kernels/mul_mm.metal b/ggml/src/ggml-metal/kernels/mul_mm.metal index ee848eed6d..0a45bb1bbe 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 diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 15c42a1081..b63b3773ee 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 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 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> 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) {