Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source 
* metal: fuse snake activation (mul, sin, sqr, mul, add)

Mirror the CUDA, Vulkan and CPU snake fusion: same matcher on the naive
5-op chain, same F32 contract on a and inv_b, same F32/F16/BF16 kernel
with F32 compute. Follows the Metal backend idioms: bf16 instantiation
gated behind GGML_METAL_HAS_BF16 and concurrency ranges checked on the
remaining chain nodes before encoding, as done by the bin fusion.

Covered by the existing backend-agnostic SNAKE_FUSE tests.

* metal: absorb snake fusion into ggml_metal_op_bin

Extract the matcher to ggml_metal_op_can_fuse_snake, mirroring the
Vulkan naming, and dispatch the fused path from ggml_metal_op_bin.
The encode loop switch is back to a single call per case.

Address review from ggerganov

* metal: fix indentation in ggml_metal_op_can_fuse_snake
This commit is contained in:
Pascal
2026-07-15 13:53:31 +03:00
committed by GitHub
parent f955e394bf
commit b3c9d1b846
6 changed files with 153 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-metal/ggml-metal.metal
+29
View File
@@ -5406,6 +5406,35 @@ template [[host_name("kernel_col2im_1d_bf16")]] kernel void kernel_col2im_1d<bfl
#endif
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
typedef void (conv_transpose_2d_t)(
constant ggml_metal_kargs_conv_transpose_2d & args,
device const float * src0,