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
This commit is contained in:
Georgi Gerganov authored and 
GitHub committed
1 parent
9df0b73c70
commit
df11702136
10 files changed
+341
-45
No files matched your search
@@ -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;
|
||||
|
||||
@@ -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);
|
||||
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 base = ow * args.s0 - args.p0;
|
||||
const int src_off = row * args.IW;
|
||||
const int dst_off = row * args.OW;
|
||||
|
||||
float acc = -INFINITY;
|
||||
float lmax = -INFINITY;
|
||||
|
||||
const int src_off = row * args.IW;
|
||||
const int dst_off = row * args.OW;
|
||||
|
||||
for (int ki = 0; ki < args.k0; ++ki) {
|
||||
int j = base + ki;
|
||||
if (j < 0 || j >= args.IW){
|
||||
continue;
|
||||
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]);
|
||||
}
|
||||
float v = src[src_off + j];
|
||||
acc = max(acc, v);
|
||||
}
|
||||
|
||||
dst[dst_off + ow] = acc;
|
||||
float max_val = simd_max(lmax);
|
||||
|
||||
threadgroup float * shared_max = (threadgroup float *) shmem;
|
||||
|
||||
if (ntg > N_SIMDWIDTH) {
|
||||
if (sgitg == 0) {
|
||||
shared_max[tiisg] = -INFINITY;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
if (tiisg == 0) {
|
||||
shared_max[sgitg] = max_val;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
max_val = shared_max[tiisg];
|
||||
max_val = simd_max(max_val);
|
||||
}
|
||||
|
||||
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 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 base = ow * args.s0 - args.p0;
|
||||
const int src_off = row * args.IW;
|
||||
const int dst_off = row * args.OW;
|
||||
|
||||
float acc = 0.0f;
|
||||
int cnt = 0;
|
||||
float lsum = 0.0f;
|
||||
int lcnt = 0;
|
||||
|
||||
const int src_off = row * args.IW;
|
||||
const int dst_off = row * args.OW;
|
||||
|
||||
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;
|
||||
}
|
||||
acc += src[src_off + j];
|
||||
cnt += 1;
|
||||
}
|
||||
|
||||
dst[dst_off + ow] = (cnt > 0) ? (acc / (float)cnt) : 0.0f;
|
||||
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;
|
||||
}
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
sum = shared_sum[tiisg];
|
||||
cnt = shared_cnt[tiisg];
|
||||
|
||||
sum = simd_sum(sum);
|
||||
cnt = simd_sum(cnt);
|
||||
}
|
||||
|
||||
if (tpitg == 0) {
|
||||
dst[dst_off + ow] = (cnt > 0.0f) ? (sum / cnt) : 0.0f;
|
||||
}
|
||||
}
|
||||
@@ -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));
|
||||
|
||||
Reference in new issue
Block a user