Commit df1170213 for llama.cpp
commit df11702136d02b9e8fa37971de0bf0e95a91b2de
Author: Georgi Gerganov <ggerganov@gmail.com>
Date: Tue Oct 20 16:46:00 2026 +0300
metal : optimize argmax and pool_1d (#30862)
* metal : optimize argmax and pool_1d
Add a two-pass partial reduction for large-row argmax and use cooperative threadgroup reductions for pool_1d. Allocate scratch space for the argmax partials through the Metal buffer allocator.
Assisted-by: pi:llama.cpp/ggml-org/DeepSeek-V4-Flash-Vision-Exp-GGUF (Mac Studio, M2 Ultra, 192GB)
* tests : add pool_1d perf cases
Add representative POOL_1D max cases for the MoE down-projection amax shapes, reducing the whole intermediate dim.
Assisted-by: pi:llama.cpp/ggml-org/DeepSeek-V4-Flash-Vision-Exp-GGUF (Mac Studio, M2 Ultra, 192GB)
* cont : minor fix
diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp
index ffbcba366..22e5fd9c7 100644
--- a/ggml/src/ggml-backend.cpp
+++ b/ggml/src/ggml-backend.cpp
@@ -2226,6 +2226,7 @@ bool ggml_op_alloc_size_may_expand(enum ggml_op op) {
case GGML_OP_CUMSUM:
case GGML_OP_ARGSORT:
case GGML_OP_TOP_K:
+ case GGML_OP_ARGMAX:
return true;
default:
return false;
diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp
index 2c7972b0a..65611314a 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-device.cpp
@@ -1567,6 +1567,28 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax(ggml_meta
return res;
}
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax_part(ggml_metal_library_t lib) {
+ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, "kernel_argmax_part_f32");
+ if (!res.pipeline) {
+ res = ggml_metal_library_compile_pipeline(lib, "kernel_argmax_part_f32", "kernel_argmax_part_f32", nullptr);
+ }
+
+ res.smem = 32*(sizeof(float) + sizeof(int32_t));
+
+ return res;
+}
+
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax_final(ggml_metal_library_t lib) {
+ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, "kernel_argmax_final_f32");
+ if (!res.pipeline) {
+ res = ggml_metal_library_compile_pipeline(lib, "kernel_argmax_final_f32", "kernel_argmax_final_f32", nullptr);
+ }
+
+ res.smem = 32*(sizeof(float) + sizeof(int32_t));
+
+ return res;
+}
+
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort(ggml_metal_library_t lib, const ggml_tensor * op) {
assert(op->op == GGML_OP_ARGSORT);
diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h
index 4012b9d5f..5ccaec6a3 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.h
+++ b/ggml/src/ggml-metal/ggml-metal-device.h
@@ -145,6 +145,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax_part (ggml_metal_library_t lib);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax_final (ggml_metal_library_t lib);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort_merge (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht (ggml_metal_library_t lib, int n, enum ggml_type tsrc);
diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h
index 4ffb5a77d..46b56c9d9 100644
--- a/ggml/src/ggml-metal/ggml-metal-impl.h
+++ b/ggml/src/ggml-metal/ggml-metal-impl.h
@@ -16,6 +16,9 @@
#define N_MM_NPART_AMAX 256
+#define N_ARGMAX_NPART_MAX 256
+#define N_ARGMAX_CHUNK 4096
+
// kernel parameters for mat-vec threadgroups
//
// N_R0: number of src0 rows to process per simdgroup
@@ -1379,8 +1382,10 @@ typedef struct {
} ggml_metal_kargs_pool_1d;
typedef struct {
- int64_t ne00;
+ int64_t ne00;
uint64_t nb01;
+ int64_t npart;
+ int64_t ne01;
} ggml_metal_kargs_argmax;
typedef struct {
diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp
index 53c981618..fc856b5b2 100644
--- a/ggml/src/ggml-metal/ggml-metal-ops.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp
@@ -2319,13 +2319,18 @@ int ggml_metal_op_pool_1d(ggml_metal_op_t ctx, int idx) {
auto pipeline = ggml_metal_library_get_pipeline_pool_1d(lib, op, op_pool);
- const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), (int) np);
- const int ntg = (np + nth - 1) / nth;
+ const int nth = std::min(256, std::max(1, k0));
+ const int ntg = (int) np;
+
+ const size_t smem = op_pool == GGML_OP_POOL_AVG ?
+ 2*32*sizeof(float) :
+ 32*sizeof(float);
ggml_metal_encoder_set_pipeline(enc, pipeline);
ggml_metal_encoder_set_bytes (enc, &args_pool_1d, sizeof(args_pool_1d), 0);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2);
+ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
ggml_metal_encoder_dispatch_threadgroups(enc, ntg, 1, 1, nth, 1, 1);
@@ -5421,6 +5426,25 @@ int ggml_metal_op_timestep_embedding(ggml_metal_op_t ctx, int idx) {
return 1;
}
+size_t ggml_metal_op_argmax_extra(const ggml_tensor * op) {
+ assert(op->op == GGML_OP_ARGMAX);
+
+ const int64_t ne00 = op->src[0]->ne[0];
+ const int64_t ne01 = op->src[0]->ne[1];
+
+ int npart = 1;
+ if (ne00 > N_ARGMAX_CHUNK) {
+ npart = (int) std::min<int64_t>(N_ARGMAX_NPART_MAX, (ne00 + N_ARGMAX_CHUNK - 1)/N_ARGMAX_CHUNK);
+ }
+
+ if (npart == 1) {
+ return 0;
+ }
+
+ // partial max + partial argmax for every row
+ return (sizeof(float) + sizeof(int32_t)) * npart * ne01;
+}
+
int ggml_metal_op_argmax(ggml_metal_op_t ctx, int idx) {
ggml_tensor * op = ctx->node(idx);
@@ -5432,22 +5456,61 @@ int ggml_metal_op_argmax(ggml_metal_op_t ctx, int idx) {
GGML_TENSOR_LOCALS( int32_t, ne, op, ne);
GGML_TENSOR_LOCALS(uint64_t, nb, op, nb);
+ const int64_t nrows = ggml_nrows(op->src[0]);
+
+ int npart = 1;
+ if (ne00 > N_ARGMAX_CHUNK) {
+ npart = (int) std::min<int64_t>(N_ARGMAX_NPART_MAX, (ne00 + N_ARGMAX_CHUNK - 1)/N_ARGMAX_CHUNK);
+ }
+
ggml_metal_kargs_argmax args = {
- /*.ne00 = */ ne00,
- /*.nb01 = */ nb01,
+ /*.ne00 = */ ne00,
+ /*.nb01 = */ nb01,
+ /*.npart = */ npart,
+ /*.ne01 = */ ne01,
};
- auto pipeline = ggml_metal_library_get_pipeline_argmax(lib, op);
+ const size_t smem = 32*(sizeof(float) + sizeof(int32_t));
- const int64_t nrows = ggml_nrows(op->src[0]);
+ if (npart > 1) {
+ auto pipeline_part = ggml_metal_library_get_pipeline_argmax_part(lib);
+ auto pipeline_final = ggml_metal_library_get_pipeline_argmax_final(lib);
+
+ ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op);
+ ggml_metal_buffer_id bid_tmp = bid_dst;
+ bid_tmp.offs += ggml_nbytes(op);
+
+ ggml_metal_encoder_set_pipeline(enc, pipeline_part);
+ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
+ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
+ ggml_metal_encoder_set_buffer (enc, bid_tmp, 2);
+ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
+ ggml_metal_encoder_dispatch_threadgroups(enc, npart, nrows, 1, 256, 1, 1);
+
+ ggml_metal_op_concurrency_reset(ctx);
+
+ int nth = 32; // SIMD width
+ while (nth < npart && nth < 256) {
+ nth *= 2;
+ }
+
+ ggml_metal_encoder_set_pipeline(enc, pipeline_final);
+ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
+ ggml_metal_encoder_set_buffer (enc, bid_tmp, 1);
+ ggml_metal_encoder_set_buffer (enc, bid_dst, 2);
+ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
+ ggml_metal_encoder_dispatch_threadgroups(enc, nrows, 1, 1, nth, 1, 1);
+
+ return 1;
+ }
+
+ auto pipeline = ggml_metal_library_get_pipeline_argmax(lib, op);
int nth = 32; // SIMD width
while (nth < ne00 && nth*ne01*ne02*ne03 < 256) {
nth *= 2;
}
- const size_t smem = pipeline.smem;
-
ggml_metal_encoder_set_pipeline(enc, pipeline);
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1);
diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h
index 583d1156b..b8296bf2b 100644
--- a/ggml/src/ggml-metal/ggml-metal-ops.h
+++ b/ggml/src/ggml-metal/ggml-metal-ops.h
@@ -47,6 +47,8 @@ size_t ggml_metal_op_flash_attn_ext_extra_tmp(const struct ggml_tensor * op);
size_t ggml_metal_op_flash_attn_ext_extra_kv_f16(const struct ggml_tensor * op);
size_t ggml_metal_op_flash_attn_ext_extra_idx(const struct ggml_tensor * op);
+size_t ggml_metal_op_argmax_extra(const struct ggml_tensor * op);
+
int ggml_metal_op_concat (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_repeat (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_acc (ggml_metal_op_t ctx, int idx);
diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp
index 344882e04..074a9d8e6 100644
--- a/ggml/src/ggml-metal/ggml-metal.cpp
+++ b/ggml/src/ggml-metal/ggml-metal.cpp
@@ -245,6 +245,10 @@ static size_t ggml_backend_metal_buffer_type_get_alloc_size(ggml_backend_buffer_
{
res = 2*sizeof(int32_t)*ggml_nelements(tensor->src[0]);
} break;
+ case GGML_OP_ARGMAX:
+ {
+ res += ggml_metal_op_argmax_extra(tensor);
+ } break;
default:
break;
}
diff --git a/ggml/src/ggml-metal/kernels/misc.metal b/ggml/src/ggml-metal/kernels/misc.metal
index fb27983db..90d3fbe21 100644
--- a/ggml/src/ggml-metal/kernels/misc.metal
+++ b/ggml/src/ggml-metal/kernels/misc.metal
@@ -60,6 +60,140 @@ kernel void kernel_argmax_f32(
dst_i32[tgpig] = arg_val;
}
+kernel void kernel_argmax_part_f32(
+ constant ggml_metal_kargs_argmax & args,
+ device const char * src0,
+ device char * tmp,
+ threadgroup char * shmem [[threadgroup(0)]],
+ uint3 tgpig[[threadgroup_position_in_grid]],
+ uint3 tpitg[[thread_position_in_threadgroup]],
+ uint sgitg[[simdgroup_index_in_threadgroup]],
+ uint tiisg[[thread_index_in_simdgroup]],
+ uint3 ntg[[threads_per_threadgroup]]) {
+ const int row = tgpig.y;
+ const int part = tgpig.x;
+ const int ne00 = (int) args.ne00;
+ const int npart = (int) args.npart;
+ const int chunk = (ne00 + npart - 1) / npart;
+ const int start = part * chunk;
+ const int end = min(start + chunk, ne00);
+
+ device const float * x_row = (device const float *) ((device const char *) src0 + row * args.nb01);
+
+ float lmax = -INFINITY;
+ int32_t larg = -1;
+
+ for (int i00 = start + tpitg.x; i00 < end; i00 += ntg.x) {
+ if (x_row[i00] > lmax) {
+ lmax = x_row[i00];
+ larg = i00;
+ }
+ }
+
+ // find the argmax value in the block
+ float max_val = simd_max(lmax);
+ int32_t arg_val = simd_max(select(-1, larg, lmax == max_val));
+
+ threadgroup float * shared_maxval = (threadgroup float *) shmem;
+ threadgroup int32_t * shared_argmax = (threadgroup int32_t *) shmem + N_SIMDWIDTH;
+
+ if (ntg.x > N_SIMDWIDTH) {
+ if (sgitg == 0) {
+ shared_maxval[tiisg] = -INFINITY;
+ shared_argmax[tiisg] = -1;
+ }
+
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+
+ if (tiisg == 0) {
+ shared_maxval[sgitg] = max_val;
+ shared_argmax[sgitg] = arg_val;
+ }
+
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+
+ max_val = shared_maxval[tiisg];
+ arg_val = shared_argmax[tiisg];
+
+ float max_val_reduced = simd_max(max_val);
+ int32_t arg_val_reduced = simd_max(select(-1, arg_val, max_val == max_val_reduced));
+
+ max_val = max_val_reduced;
+ arg_val = arg_val_reduced;
+ }
+
+ if (tpitg.x == 0) {
+ device float * tmp_max = (device float *) tmp;
+ device int32_t * tmp_arg = (device int32_t *) tmp + npart * (int) args.ne01;
+
+ tmp_max[row*npart + part] = max_val;
+ tmp_arg[row*npart + part] = arg_val;
+ }
+}
+
+kernel void kernel_argmax_final_f32(
+ constant ggml_metal_kargs_argmax & args,
+ device const char * tmp,
+ device char * dst,
+ threadgroup char * shmem [[threadgroup(0)]],
+ uint3 tgpig[[threadgroup_position_in_grid]],
+ uint3 tpitg[[thread_position_in_threadgroup]],
+ uint sgitg[[simdgroup_index_in_threadgroup]],
+ uint tiisg[[thread_index_in_simdgroup]],
+ uint3 ntg[[threads_per_threadgroup]]) {
+ const int row = tgpig.x;
+ const int npart = (int) args.npart;
+
+ device const float * tmp_max = (device const float *) tmp;
+ device const int32_t * tmp_arg = (device const int32_t *) tmp + npart * (int) args.ne01;
+
+ float lmax = -INFINITY;
+ int32_t larg = -1;
+
+ for (int i = tpitg.x; i < npart; i += ntg.x) {
+ if (tmp_max[row*npart + i] > lmax) {
+ lmax = tmp_max[row*npart + i];
+ larg = tmp_arg[row*npart + i];
+ }
+ }
+
+ // find the argmax value in the block
+ float max_val = simd_max(lmax);
+ int32_t arg_val = simd_max(select(-1, larg, lmax == max_val));
+
+ threadgroup float * shared_maxval = (threadgroup float *) shmem;
+ threadgroup int32_t * shared_argmax = (threadgroup int32_t *) shmem + N_SIMDWIDTH;
+
+ if (ntg.x > N_SIMDWIDTH) {
+ if (sgitg == 0) {
+ shared_maxval[tiisg] = -INFINITY;
+ shared_argmax[tiisg] = -1;
+ }
+
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+
+ if (tiisg == 0) {
+ shared_maxval[sgitg] = max_val;
+ shared_argmax[sgitg] = arg_val;
+ }
+
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+
+ max_val = shared_maxval[tiisg];
+ arg_val = shared_argmax[tiisg];
+
+ float max_val_reduced = simd_max(max_val);
+ int32_t arg_val_reduced = simd_max(select(-1, arg_val, max_val == max_val_reduced));
+
+ max_val = max_val_reduced;
+ arg_val = arg_val_reduced;
+ }
+
+ if (tpitg.x == 0) {
+ ((device int32_t *) dst)[row] = arg_val;
+ }
+}
+
kernel void kernel_diag_f32(
constant ggml_metal_kargs_diag & args,
device const char * src0,
diff --git a/ggml/src/ggml-metal/kernels/pool.metal b/ggml/src/ggml-metal/kernels/pool.metal
index 13d355b9d..3dc0261c3 100644
--- a/ggml/src/ggml-metal/kernels/pool.metal
+++ b/ggml/src/ggml-metal/kernels/pool.metal
@@ -84,65 +84,117 @@ kernel void kernel_pool_1d_max_f32(
constant ggml_metal_kargs_pool_1d & args,
device const float * src,
device float * dst,
- uint gid [[thread_position_in_grid]]
-) {
-
- if (gid >= args.np) {
+ threadgroup char * shmem [[threadgroup(0)]],
+ uint tgpig[[threadgroup_position_in_grid]],
+ uint tpitg[[thread_position_in_threadgroup]],
+ uint sgitg[[simdgroup_index_in_threadgroup]],
+ uint tiisg[[thread_index_in_simdgroup]],
+ uint ntg[[threads_per_threadgroup]]) {
+ if (tgpig >= args.np) {
return;
}
- const int ow = (int)gid % args.OW;
- const int row = (int)gid / args.OW;
+ const int out = (int) tgpig;
+ const int ow = out % (int) args.OW;
+ const int row = out / (int) args.OW;
+
+ const int base = ow * args.s0 - args.p0;
+ const int src_off = row * args.IW;
+ const int dst_off = row * args.OW;
+
+ float lmax = -INFINITY;
+
+ for (int ki = tpitg; ki < args.k0; ki += ntg) {
+ const int j = base + ki;
+ if (j >= 0 && j < args.IW) {
+ lmax = max(lmax, src[src_off + j]);
+ }
+ }
- const int base = ow * args.s0 - args.p0;
+ float max_val = simd_max(lmax);
- float acc = -INFINITY;
+ threadgroup float * shared_max = (threadgroup float *) shmem;
- const int src_off = row * args.IW;
- const int dst_off = row * args.OW;
+ if (ntg > N_SIMDWIDTH) {
+ if (sgitg == 0) {
+ shared_max[tiisg] = -INFINITY;
+ }
+ threadgroup_barrier(mem_flags::mem_threadgroup);
- for (int ki = 0; ki < args.k0; ++ki) {
- int j = base + ki;
- if (j < 0 || j >= args.IW){
- continue;
+ if (tiisg == 0) {
+ shared_max[sgitg] = max_val;
}
- float v = src[src_off + j];
- acc = max(acc, v);
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+
+ max_val = shared_max[tiisg];
+ max_val = simd_max(max_val);
}
- dst[dst_off + ow] = acc;
+ if (tpitg == 0) {
+ dst[dst_off + ow] = max_val;
+ }
}
kernel void kernel_pool_1d_avg_f32(
constant ggml_metal_kargs_pool_1d & args,
device const float * src,
device float * dst,
- uint gid [[thread_position_in_grid]]
-) {
-
- if (gid >= args.np) {
+ threadgroup char * shmem [[threadgroup(0)]],
+ uint tgpig[[threadgroup_position_in_grid]],
+ uint tpitg[[thread_position_in_threadgroup]],
+ uint sgitg[[simdgroup_index_in_threadgroup]],
+ uint tiisg[[thread_index_in_simdgroup]],
+ uint ntg[[threads_per_threadgroup]]) {
+ if (tgpig >= args.np) {
return;
}
- const int ow = (int)gid % args.OW;
- const int row = (int)gid / args.OW;
-
- const int base = ow * args.s0 - args.p0;
+ const int out = (int) tgpig;
+ const int ow = out % (int) args.OW;
+ const int row = out / (int) args.OW;
- float acc = 0.0f;
- int cnt = 0;
+ const int base = ow * args.s0 - args.p0;
+ const int src_off = row * args.IW;
+ const int dst_off = row * args.OW;
- const int src_off = row * args.IW;
- const int dst_off = row * args.OW;
+ float lsum = 0.0f;
+ int lcnt = 0;
- for (int ki = 0; ki < args.k0; ++ki) {
+ for (int ki = tpitg; ki < args.k0; ki += ntg) {
const int j = base + ki;
- if (j < 0 || j >= args.IW) {
- continue;
+ if (j >= 0 && j < args.IW) {
+ lsum += src[src_off + j];
+ lcnt += 1;
+ }
+ }
+
+ float sum = simd_sum(lsum);
+ float cnt = simd_sum((float) lcnt);
+
+ threadgroup float * shared_sum = (threadgroup float *) shmem;
+ threadgroup float * shared_cnt = (threadgroup float *) shmem + N_SIMDWIDTH;
+
+ if (ntg > N_SIMDWIDTH) {
+ if (sgitg == 0) {
+ shared_sum[tiisg] = 0.0f;
+ shared_cnt[tiisg] = 0.0f;
+ }
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+
+ if (tiisg == 0) {
+ shared_sum[sgitg] = sum;
+ shared_cnt[sgitg] = cnt;
}
- acc += src[src_off + j];
- cnt += 1;
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+
+ sum = shared_sum[tiisg];
+ cnt = shared_cnt[tiisg];
+
+ sum = simd_sum(sum);
+ cnt = simd_sum(cnt);
}
- dst[dst_off + ow] = (cnt > 0) ? (acc / (float)cnt) : 0.0f;
+ if (tpitg == 0) {
+ dst[dst_off + ow] = (cnt > 0.0f) ? (sum / cnt) : 0.0f;
+ }
}
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 1df1129cd..485629cf6 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -12199,6 +12199,17 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
}
}
+ // POOL_1D max over MoE down-projection amax shapes:
+ // amax is [1, n_expert, n_tokens] after reducing the whole intermediate dim.
+ for (int64_t n_ff : { 2048, 4096, 8192 }) {
+ for (int64_t n_expert : { 8, 16 }) {
+ for (int64_t n_tokens : { 1, 512 }) {
+ test_cases.emplace_back(new test_pool1d(
+ GGML_OP_POOL_MAX, GGML_TYPE_F32, {n_ff, n_expert, n_tokens, 1}, n_ff, n_ff, 0));
+ }
+ }
+ }
+
// LEAKY_RELU at FFN activation width, for direct comparison with RELU
for (int64_t n_tokens : {512, 2048}) {
test_cases.emplace_back(new test_leaky_relu(GGML_TYPE_F32, { 17408, n_tokens, 1, 1 }, 0.1f));