Commit 073b74034 for llama.cpp
commit 073b74034b8bfe652e410ec8bd1e0d95a7d46fea
Author: Georgi Gerganov <ggerganov@gmail.com>
Date: Mon Oct 12 13:31:41 2026 +0300
metal : fix mul_mat_id with duplicate ids (#30373)
* metal : fix mul_mat_id limits and kernel name
* metal : fix mul_mat_id with duplicate ids
kernel_mul_mm_id_map0 assumed at most one match per token and a fixed
capacity of ne21 entries per expert, which produced wrong or dropped
output rows and could overflow into the slice of the next expert when
the ids contain duplicates.
build a compact per-expert list of the (token, slot) entries instead:
- pass 1 counts the entries of every expert, thread 0 computes the
offsets of the expert slices (tpe now holds [base, count])
- pass 2 writes the entries with one unconditional store per token
(conditional stores of divergent lanes are very slow), the dummy
value of the tokens without a match lands on a padding slot of the
slice that is never read - duplicates are handled by a rare branch
- kernel_mul_mm_id reads the base/count of its expert and loops over
the entry tiles that do not fit into the grid
The ids buffer is sized as ne20*ne21 + ne02 (entries + padding).
883/883 test-backend-ops MUL_MAT_ID tests pass, llama-bench is within
run-to-run variance of the previous behavior.
Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD
* metal : make mul_mat_id map0 ne02 a function constant
use FOR_UNROLL for the id loops of kernel_mul_mm_id_map0 and turn ne02
into the function constant FC_mul_mm_id_map0_ne02 (FC_MUL_MM + 7), so
that the threadgroup offsets of cnt/bas/sids and the prefix-sum bound
fold at compile time.
the pipeline name now encodes ne02 (same convention as the ssm_conv
nc/silu constants) and the map0 pipeline is compiled with the constant
value set. ne02 is no longer part of ggml_metal_kargs_mul_mm_id_map0 -
the dispatch grid still comes from the tensor ne[2] on the host.
no measurable perf change: 883/883 test-backend-ops MUL_MAT_ID tests
pass, op-level and llama-bench results are within run variance.
Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD
* cont : clean-up + TODO
* metal : dispatch mul_mat_id entry tiles from the prepass
follow the approach of the vulkan implementation (#29998): instead of
launching a grid over experts x ne21 entry tiles and letting most of
the threadgroups early exit, kernel_mul_mm_id_map0 now also emits a
row map - the total tile count plus one (expert, first entry, valid
rows) descriptor per 32-entry tile that has work.
kernel_mul_mm_id is launched with a static upper bound on the number
of tiles (full tiles over all entries plus at most one partial tile
per expert), reads its descriptor and processes a single tile - the
tail early exits. This removes the round loop that was needed to cover
duplicate ids past the grid (register pressure: th_max 768 -> back to
the plain mul_mm level) and duplicates are now handled implicitly,
since they only produce more descriptors within the bound.
The grid shrinks by ~8-16x on typical shapes (e.g. 128 experts,
8 used, 512 tokens: 128 tiles instead of 2048 entry tiles, of which
94% only launched to exit) and ne21 is no longer needed by the mm
kernel kargs.
883/883 test-backend-ops MUL_MAT_ID tests pass. llama-bench
(Qwen3.6-35B-A3B Q8_0) is now faster than master on all pp sizes
(pp2048/4096/8192 +2%, pp512 +0.5%), tg32 unchanged.
Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD
* cont : add smem assert
diff --git a/ggml/src/ggml-metal/ggml-metal-common.cpp b/ggml/src/ggml-metal/ggml-metal-common.cpp
index a00f87c99..70fd551e3 100644
--- a/ggml/src/ggml-metal/ggml-metal-common.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-common.cpp
@@ -43,9 +43,10 @@ bool ggml_metal_op_mul_mat_use_mm(const struct ggml_tensor * op, bool has_simdgr
bool ggml_metal_op_mul_mat_id_use_mm(const struct ggml_tensor * op, bool has_simdgroup_mm) {
const int64_t ne00 = op->src[0]->ne[0];
+ const int64_t ne02 = op->src[0]->ne[2];
const int64_t ne21 = op->src[2]->ne[1];
- return has_simdgroup_mm && ne00 >= 64 && ne21 >= 32;
+ return has_simdgroup_mm && ne00 >= 64 && ne02 <= 1024 && ne21 >= 32;
}
// the most src1 rows of the few-row MMA kernels
diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp
index 8ac6e0830..b70129de8 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-device.cpp
@@ -1268,10 +1268,16 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0(g
ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
if (!res.pipeline) {
- res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
+ ggml_metal_cv_t cv = ggml_metal_cv_init();
+ ggml_metal_cv_set_int32(cv, ne02, FC_MUL_MM + 7);
+
+ res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
+
+ ggml_metal_cv_free(cv);
}
- res.smem = (size_t) ne02*ne20*sizeof(uint16_t);
+ // per expert: entry count + base offset, followed by the staging of the ids of ntg tokens
+ res.smem = 2*(size_t) ne02*sizeof(uint32_t) + (size_t) ne02*ne20*sizeof(uint16_t);
res.smem = GGML_PAD(res.smem, 16);
return res;
diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h
index e7c895631..4ffb5a77d 100644
--- a/ggml/src/ggml-metal/ggml-metal-impl.h
+++ b/ggml/src/ggml-metal/ggml-metal-impl.h
@@ -560,7 +560,6 @@ typedef struct {
} ggml_metal_kargs_mul_mv_ext;
typedef struct {
- int32_t ne02;
int32_t ne10;
int32_t ne11; // n_expert_used (bcast)
uint64_t nb11;
@@ -590,7 +589,6 @@ typedef struct {
uint64_t nb12;
uint64_t nb13;
int32_t ne20;
- int32_t ne21;
int32_t ne0;
int32_t ne1;
int16_t r2;
diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp
index 86b9fe4ef..115d73695 100644
--- a/ggml/src/ggml-metal/ggml-metal-ops.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp
@@ -2772,21 +2772,39 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) {
return 1;
}
+// upper bound on the number of entry tiles emitted by kernel_mul_mm_id_map0: the full tiles
+// over all entries plus at most one partial tile per expert
+static int64_t ggml_metal_op_mul_mat_id_max_tiles(const ggml_tensor * op) {
+ const int64_t ne02 = op->src[0]->ne[2]; // n_expert
+ const int64_t ne20 = op->src[2]->ne[0]; // n_expert_used
+ const int64_t ne21 = op->src[2]->ne[1]; // n_token
+
+ const int64_t n_rows = ne20*ne21;
+ const int64_t tiles = (n_rows + 31)/32 + (ne02 < n_rows ? ne02 : n_rows) - 1;
+
+ return n_rows < tiles ? n_rows : tiles;
+}
+
size_t ggml_metal_op_mul_mat_id_extra_tpe(const ggml_tensor * op) {
assert(op->op == GGML_OP_MUL_MAT_ID);
- const int64_t ne02 = op->src[0]->ne[2]; // n_expert
+ // row map of kernel_mul_mm_id: total tile count + one
+ // (expert, first entry, valid rows) descriptor per tile
+ const int64_t max_tiles = ggml_metal_op_mul_mat_id_max_tiles(op);
- return ggml_type_size(GGML_TYPE_I32)*ne02;
+ return ggml_type_size(GGML_TYPE_I32)*(1 + 3*max_tiles);
}
size_t ggml_metal_op_mul_mat_id_extra_ids(const ggml_tensor * op) {
assert(op->op == GGML_OP_MUL_MAT_ID);
const int64_t ne02 = op->src[0]->ne[2]; // n_expert
+ const int64_t ne20 = op->src[2]->ne[0]; // n_expert_used
const int64_t ne21 = op->src[2]->ne[1]; // n_token
- return ggml_type_size(GGML_TYPE_I32)*ne02*ne21;
+ // one entry per (expert used, token) pair of ids plus one padding slot per expert
+ // (see kernel_mul_mm_id_map0)
+ return ggml_type_size(GGML_TYPE_I32)*(ne20*ne21 + ne02);
}
size_t ggml_metal_op_mul_mat_id_extra_amax(const ggml_tensor * op) {
@@ -2884,7 +2902,6 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
{
ggml_metal_kargs_mul_mm_id_map0 args = {
- ne02,
ne10,
ne11, // n_expert_used (bcast)
nb11,
@@ -2942,7 +2959,6 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
/*.nb12 =*/ nb12,
/*.nb13 =*/ nb13,
/*.ne20 =*/ ne20, // n_expert_used
- /*.ne21 =*/ ne21, // n_tokens
/*.ne0 =*/ ne0,
/*.ne1 =*/ ne1,
/*.r2 =*/ r2,
@@ -2960,9 +2976,13 @@ int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) {
const size_t smem = pipeline.smem;
+ GGML_ASSERT(smem <= props_dev->max_theadgroup_memory_size);
+
ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
- ggml_metal_encoder_dispatch_threadgroups(enc, (ne21 + 31)/32, (ne01 + 63)/64, ne02, 128, 1, 1);
+ const int max_tiles = (int) ggml_metal_op_mul_mat_id_max_tiles(op);
+
+ ggml_metal_encoder_dispatch_threadgroups(enc, max_tiles, (ne01 + 63)/64, 1, 128, 1, 1);
}
} else {
auto pipeline = ggml_metal_library_get_pipeline_mul_mv_id(lib, op);
diff --git a/ggml/src/ggml-metal/kernels/mul_mm.metal b/ggml/src/ggml-metal/kernels/mul_mm.metal
index 9d262da1d..dee13d878 100644
--- a/ggml/src/ggml-metal/kernels/mul_mm.metal
+++ b/ggml/src/ggml-metal/kernels/mul_mm.metal
@@ -8,6 +8,7 @@ constant short FC_mul_mm_ne13 [[function_constant(FC_MUL_MM + 3)]];
constant short FC_mul_mm_r2 [[function_constant(FC_MUL_MM + 4)]];
constant short FC_mul_mm_r3 [[function_constant(FC_MUL_MM + 5)]];
constant bool FC_mul_mm_id_amax [[function_constant(FC_MUL_MM + 6)]];
+constant int32_t FC_mul_mm_id_map0_ne02 [[function_constant(FC_MUL_MM + 7)]];
// each block_q contains 16*nl weights
#ifdef GGML_METAL_HAS_TENSOR
@@ -371,19 +372,90 @@ kernel void kernel_mul_mm_id_map0(
ushort ntg[[threads_per_threadgroup]]) {
const short ide = tpitg; // expert id
- uint32_t n_all = 0;
+ // per expert: entry count (pass 1), replaced by the base offset of the slice (pass 2)
+ threadgroup uint32_t * cnt = (threadgroup uint32_t *) (shmem);
+ threadgroup uint32_t * bas = cnt + FC_mul_mm_id_map0_ne02;
+
+ // staging of the ids of ntg tokens
+ threadgroup uint16_t * sids = (threadgroup uint16_t *) (bas + FC_mul_mm_id_map0_ne02);
- device int32_t * ids_i32 = (device int32_t *) hids + ide*args.ne21;
+ // pass 1: count the entries of this expert - every occurrence of this expert within the
+ // ids is an entry (duplicate ids produce multiple entries for the same token)
+ uint32_t n_all = 0;
for (int i21 = 0; i21 < args.ne21; i21 += ntg) { // n_tokens
if (i21 + tpitg < args.ne21) {
device const int32_t * src2_i32 = (device const int32_t *) (src2 + (i21 + tpitg)*args.nb21);
- threadgroup uint16_t * sids = (threadgroup uint16_t *) shmem + tpitg*ne20;
+ FOR_UNROLL (short i20 = 0; i20 < ne20; i20++) {
+ sids[tpitg*ne20 + i20] = src2_i32[i20];
+ }
+ }
+
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+
+ for (short t = 0; t < ntg; t++) {
+ if (i21 + t >= args.ne21) {
+ break;
+ }
+
+ threadgroup const uint16_t * sids_t = sids + t*ne20;
- #pragma unroll(ne20)
- for (short i20 = 0; i20 < ne20; i20++) {
- sids[i20] = src2_i32[i20];
+ FOR_UNROLL (short i20 = 0; i20 < ne20; i20++) {
+ n_all += sids_t[i20] == ide;
+ }
+ }
+
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+ }
+
+ cnt[ide] = n_all;
+
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+
+ // exclusive prefix sum over the counts: the entries of this expert start at bas[ide] -
+ // every slice is padded by one entry for the uniform dummy store of pass 2 - and emit the
+ // entry tile descriptors of kernel_mul_mm_id: (expert, first entry, valid rows) per tile
+ if (tpitg == 0) {
+ constexpr uint32_t NTILE = 32; // must match NR1 of kernel_mul_mm_id
+
+ uint32_t base = 0;
+ uint32_t tile = 0;
+
+ device uint32_t * tpe_u32 = (device uint32_t *) (htpe);
+
+ for (int e = 0; e < FC_mul_mm_id_map0_ne02; ++e) {
+ const uint32_t n = cnt[e];
+
+ bas[e] = base;
+
+ for (uint32_t k = 0; k < n; k += NTILE) {
+ tpe_u32[1 + 3*tile ] = e;
+ tpe_u32[1 + 3*tile + 1] = base + k;
+ tpe_u32[1 + 3*tile + 2] = n - k < NTILE ? n - k : NTILE;
+
+ tile++;
+ }
+
+ base += n + 1;
+ }
+
+ tpe_u32[0] = tile;
+ }
+
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+
+ // pass 2: write the entries of this expert into its slice of the compact list
+ device int32_t * ids_i32 = (device int32_t *) (hids);
+
+ uint32_t ient = bas[ide];
+
+ for (int i21 = 0; i21 < args.ne21; i21 += ntg) { // n_tokens
+ if (i21 + tpitg < args.ne21) {
+ device const int32_t * src2_i32 = (device const int32_t *) (src2 + (i21 + tpitg)*args.nb21);
+
+ FOR_UNROLL (short i20 = 0; i20 < ne20; i20++) {
+ sids[tpitg*ne20 + i20] = src2_i32[i20];
}
}
@@ -394,24 +466,40 @@ kernel void kernel_mul_mm_id_map0(
break;
}
- threadgroup const uint16_t * sids = (threadgroup const uint16_t *) shmem + t*ne20;
+ threadgroup const uint16_t * sids_t = sids + t*ne20;
+
+ short first = -1;
+ short n_ents = 0;
+
+ FOR_UNROLL (short i20 = 0; i20 < ne20; i20++) {
+ const short hit = sids_t[i20] == ide;
- short sel = 0;
- #pragma unroll(ne20)
- for (short i20 = 0; i20 < ne20; i20++) {
- sel += (sids[i20] == ide)*(i20 + 1);
+ first = hit && first < 0 ? i20 : first;
+ n_ents += hit;
}
- ids_i32[n_all] = (i21 + t)*ne20 + sel - 1;
+ // one store per token for every thread, regardless of the match - conditional
+ // stores of the divergent lanes are very slow - the dummy value of the tokens
+ // without a match lands on the padding slot of this slice and is never read
+ ids_i32[ient] = first >= 0 ? (i21 + t)*ne20 + first : 0;
- n_all += sel > 0;
+ if (n_ents > 1) {
+ // duplicate ids: the remaining matches of this token
+ short j = 1;
+
+ FOR_UNROLL (short i20 = 0; i20 < ne20; i20++) {
+ if (sids_t[i20] == ide && i20 != first) {
+ ids_i32[ient + j] = (i21 + t)*ne20 + i20;
+ j += 1;
+ }
+ }
+ }
+
+ ient += n_ents;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
-
- device uint32_t * tpe_u32 = (device uint32_t *) (htpe);
- tpe_u32[ide] = n_all;
}
kernel void kernel_mul_mm_id_amax_part_f32(
@@ -505,6 +593,7 @@ template [[host_name("kernel_mul_mm_id_map0_ne20_10")]] kernel kernel_mul_mm_id_
template [[host_name("kernel_mul_mm_id_map0_ne20_16")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<16>;
template [[host_name("kernel_mul_mm_id_map0_ne20_22")]] kernel kernel_mul_mm_id_map0_t kernel_mul_mm_id_map0<22>;
+// TODO: refactor into two seperate functions to avoid ifdef branching: with and without tensor support
template<typename S0, typename S0_4x4, typename S0_8x8, typename S1, typename S1_2x4, typename S1_8x8, typename block_q, short nl, void (*dequantize_func)(device const block_q *, short, thread S0_4x4 &), typename T0, typename T0_4x4, typename T1, typename T1_2x4>
kernel void kernel_mul_mm_id(
constant ggml_metal_kargs_mul_mm_id & args,
@@ -533,55 +622,42 @@ kernel void kernel_mul_mm_id(
constexpr int NL0 = NK/16;
constexpr int NL1 = NK/8;
- const int im = tgpig.z; // expert
- const int r0 = tgpig.y*NR0;
- const int r1 = tgpig.x*NR1;
+ const int tile = tgpig.x; // entry tile, described by the prepass
+ const int r0 = tgpig.y*NR0;
device const uint32_t * tpe_u32 = (device const uint32_t *) (htpe);
device const int32_t * ids_i32 = (device const int32_t *) (hids);
- const int32_t neh1 = tpe_u32[im];
-
- if (r1 >= neh1) {
+ // the prepass emits one (expert, first entry, valid rows) descriptor per tile that has
+ // work - the grid is sized by a static upper bound, so the tail early exits
+ if (tile >= (int) tpe_u32[0]) {
return;
}
+ const int im = tpe_u32[1 + 3*tile]; // expert
+ const int idb = tpe_u32[1 + 3*tile + 1]; // first entry of this tile in hids
+
+ // number of entries in this tile, always <= NR1
+ const short nr1 = (short) ((tpe_u32[1 + 3*tile + 2] < NR1) ? tpe_u32[1 + 3*tile + 2] : NR1);
+
// if this block is of 64x32 shape or smaller
const short nr0 = (args.ne0 - r0 < NR0) ? (args.ne0 - r0) : NR0;
- const short nr1 = ( neh1 - r1 < NR1) ? ( neh1 - r1) : NR1;
// a thread shouldn't load data outside of the matrix
const short lr0 = ((short)tiitg/NL0) < nr0 ? ((short)tiitg/NL0) : nr0 - 1; // 0 .. 63
- const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; // 0 .. 31
const short il0 = (tiitg % NL0);
- short il = il0;
-
- const int id = ids_i32[im*args.ne21 + r1 + lr1];
-
- const short i11 = (id % args.ne20) % args.ne11;
- const short i12 = (id / args.ne20);
const short i13 = 0;
const uint64_t offset0 = im*args.nb02 + i13*args.nb03;
const short offset1 = il0/nl;
- device const block_q * x = (device const block_q *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1;
-
const short iy = 8*(tiitg % NL1);
- device const T1 * y = (device const T1 *)(src1
- + args.nb13*i13
- + args.nb12*i12
- + 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;
-
// power-of-two rescaling
const float s1_inv = FC_mul_mm_id_amax ? ((device const float *) amax)[0] : 1.0f;
const float s1_scale = FC_mul_mm_id_amax ? ((device const float *) amax)[1] : 1.0f;
@@ -591,13 +667,6 @@ kernel void kernel_mul_mm_id(
S1_8x8 mb[2];
simdgroup_float8x8 mc[8];
-
- 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));
@@ -609,6 +678,35 @@ kernel void kernel_mul_mm_id(
mpp::tensor_ops::matmul2d_descriptor(NR1H, NR0, NK, false, true, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate),
execution_simdgroups<4>> mm;
+#endif
+
+ const short lr1 = ((short)tiitg/NL1) < nr1 ? ((short)tiitg/NL1) : nr1 - 1; // 0 .. 31
+
+ short il = il0;
+
+ const int id = ids_i32[idb + lr1];
+
+ const short i11 = (id % args.ne20) % args.ne11;
+ const short i12 = (id / args.ne20);
+
+ device const block_q * x = (device const block_q *)(src0 + args.nb01*(r0 + lr0) + offset0) + offset1;
+
+ device const T1 * y = (device const T1 *)(src1
+ + args.nb13*i13
+ + args.nb12*i12
+ + args.nb11*i11
+ + args.nb10*iy);
+
+ const bool has_hi = nr1 > NR1H;
+
+#ifndef GGML_METAL_HAS_TENSOR
+ 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 cT0 = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB0), float>();
auto cT1 = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB1), float>();
#endif
@@ -819,7 +917,7 @@ kernel void kernel_mul_mm_id(
threadgroup_barrier(mem_flags::mem_threadgroup);
for (short j = sgitg; j < nr1; j += 4) {
- const int id = ids_i32[im*args.ne21 + r1 + j];
+ const int id = ids_i32[idb + j];
const short ide = id % args.ne20;
const short idt = id / args.ne20;