From 1a07c39e7cf43e4006b014ffcdb25cda9a416690 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 20 Oct 2026 10:40:30 +0300 Subject: [PATCH 1/3] 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) --- ggml/src/ggml-backend.cpp | 1 + ggml/src/ggml-metal/ggml-metal-device.cpp | 22 ++++ ggml/src/ggml-metal/ggml-metal-device.h | 2 + ggml/src/ggml-metal/ggml-metal-impl.h | 7 +- ggml/src/ggml-metal/ggml-metal-ops.cpp | 79 +++++++++++-- ggml/src/ggml-metal/ggml-metal-ops.h | 1 + ggml/src/ggml-metal/ggml-metal.cpp | 4 + ggml/src/ggml-metal/kernels/misc.metal | 134 ++++++++++++++++++++++ ggml/src/ggml-metal/kernels/pool.metal | 122 ++++++++++++++------ 9 files changed, 328 insertions(+), 44 deletions(-) diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index ffbcba3661cf..22e5fd9c7d99 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 2c7972b0a5d6..65611314adb1 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 4012b9d5f69b..5ccaec6a34f2 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 4ffb5a77ddfb..46b56c9d9cd4 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 53c981618511..fc856b5b24af 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(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(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 583d1156bf27..d5b792304445 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops.h @@ -95,6 +95,7 @@ int ggml_metal_op_pad_reflect_1d (ggml_metal_op_t ctx, int idx); int ggml_metal_op_roll (ggml_metal_op_t ctx, int idx); int ggml_metal_op_arange (ggml_metal_op_t ctx, int idx); int ggml_metal_op_timestep_embedding(ggml_metal_op_t ctx, int idx); +size_t ggml_metal_op_argmax_extra (const struct ggml_tensor * op); int ggml_metal_op_argmax (ggml_metal_op_t ctx, int idx); int ggml_metal_op_argsort (ggml_metal_op_t ctx, int idx); int ggml_metal_op_top_k (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 344882e044f3..074a9d8e69dd 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 fb27983db3e6..90d3fbe211af 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 13d355b9deb3..3dc0261c3376 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; + } } From bbd743f86a6f10880bbdfc0cf9799713ab4f19f7 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 20 Oct 2026 11:06:48 +0300 Subject: [PATCH 2/3] 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) --- tests/test-backend-ops.cpp | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 8fc90e037ba5..c8e5c4406565 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -12139,6 +12139,17 @@ static std::vector> 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)); From 157796b6c62b66ea1822d11360022b0998d7d8ca Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 20 Oct 2026 12:06:46 +0300 Subject: [PATCH 3/3] cont : minor fix --- ggml/src/ggml-metal/ggml-metal-ops.h | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h index d5b792304445..b8296bf2b481 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); @@ -95,7 +97,6 @@ int ggml_metal_op_pad_reflect_1d (ggml_metal_op_t ctx, int idx); int ggml_metal_op_roll (ggml_metal_op_t ctx, int idx); int ggml_metal_op_arange (ggml_metal_op_t ctx, int idx); int ggml_metal_op_timestep_embedding(ggml_metal_op_t ctx, int idx); -size_t ggml_metal_op_argmax_extra (const struct ggml_tensor * op); int ggml_metal_op_argmax (ggml_metal_op_t ctx, int idx); int ggml_metal_op_argsort (ggml_metal_op_t ctx, int idx); int ggml_metal_op_top_k (ggml_metal_op_t ctx, int idx);