public_git/whisper.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/whisper.cpp.git synced 2026-09-30 03:21:18 +02:00
Code Issues Packages Projects Releases Wiki Activity

metal : skip the empty half of the mul_mm_id token tile (llama/28301)

Browse Source 
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.
This commit is contained in:
Foad Abo Dahood
2026-09-14 20:45:06 +03:00
committed by Georgi Gerganov
parent 76da3529d9
commit 5f310fa112
1 changed files with 64 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-metal/kernels/mul_mm.metal
+64 -35
View File
@@ -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);
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<threadgroup float, dextents<int32_t, 2>, tensor_inline>(sc, dextents<int32_t, 2>(NR0, NR1));
cT.store(tC);
#else
threadgroup float * temp_str = ((threadgroup float *) shmem) + 32*(sgitg&1) + (16*(sgitg >> 1))*NR0;
auto tC0 = tensor<threadgroup float, dextents<int32_t, 2>, tensor_inline>(sc, dextents<int32_t, 2>(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<threadgroup float, dextents<int32_t, 2>, tensor_inline>(sc + NR1H*NR0, dextents<int32_t, 2>(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