Public/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-20 01:31:31 +02:00
Code Issues Packages Projects Releases Wiki Activity

metal: fuse snake activation (mul, sin, sqr, mul, add) (#25459)

Browse Source
This commit is contained in:
Georgi Gerganov
2026-07-28 19:02:22 +03:00
committed by YiChen Lv
parent 0b6cbb572c
commit b2d79a7e7f
1 changed files with 28 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-metal/kernels/misc.metal
+28
View File
@@ -345,3 +345,31 @@ kernel void kernel_count_equal(
typedef decltype(kernel_count_equal<int32_t>) kernel_count_equal_t;
template [[host_name("kernel_count_equal_i32")]] kernel kernel_count_equal_t kernel_count_equal<int32_t>;
template <typename T>
kernel void kernel_snake(
constant ggml_metal_kargs_snake & args,
device const T * x,
device const float * a,
device const float * inv_b,
device T * dst,
uint tgpig [[threadgroup_position_in_grid]],
uint tpitg [[thread_position_in_threadgroup]],
uint ntg [[threads_per_threadgroup]]) {
const int idx = tgpig * ntg + tpitg;
if (idx >= args.T * args.C) {
return;
}
const int c = idx / args.T; // x is [T, C], a / inv_b collapse to [1, C]
const float xi = float(x[idx]);
const float si = sin(a[c] * xi);
dst[idx] = T(xi + si * si * inv_b[c]);
}
template [[host_name("kernel_snake_f32")]] kernel void kernel_snake<float>(constant ggml_metal_kargs_snake &, device const float *, device const float *, device const float *, device float *, uint, uint, uint);
template [[host_name("kernel_snake_f16")]] kernel void kernel_snake<half>(constant ggml_metal_kargs_snake &, device const half *, device const float *, device const float *, device half *, uint, uint, uint);
#if defined(GGML_METAL_HAS_BF16)
template [[host_name("kernel_snake_bf16")]] kernel void kernel_snake<bfloat>(constant ggml_metal_kargs_snake &, device const bfloat *, device const float *, device const float *, device bfloat *, uint, uint, uint);
#endif