ggml: add GGML_OP_SNAKE for fused Snake activation

ggml/include/ggml-rpc.h

@@ -8,10 +8,10 @@ extern "C" {
 
 #define RPC_PROTO_MAJOR_VERSION    4
 #define RPC_PROTO_MINOR_VERSION    0
-#define RPC_PROTO_PATCH_VERSION    0
+#define RPC_PROTO_PATCH_VERSION    1
 
 #ifdef  __cplusplus
-static_assert(GGML_OP_COUNT == 96, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
+static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
 #endif
 
 #define GGML_RPC_MAX_SERVERS       16

ggml/include/ggml.h

@@ -529,6 +529,7 @@ extern "C" {
         GGML_OP_IM2COL,
         GGML_OP_IM2COL_BACK,
         GGML_OP_IM2COL_3D,
+        GGML_OP_SNAKE,
         GGML_OP_CONV_2D,
         GGML_OP_CONV_3D,
         GGML_OP_CONV_2D_DW,
@@ -1996,6 +1997,14 @@ extern "C" {
         int                   d1, // dilation dimension 1
         bool                  is_2D);
 
+    // Fused Snake activation: y = x + sin^2(a * x) * inv_b
+    // x: [T, C], a: [1, C] or [C], inv_b: [1, C] or [C]
+    GGML_API struct ggml_tensor * ggml_snake(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * x,
+        struct ggml_tensor  * a,
+        struct ggml_tensor  * inv_b);
+
     GGML_API struct ggml_tensor * ggml_conv_1d(
             struct ggml_context * ctx,
             struct ggml_tensor  * a,   // convolution kernel

ggml/src/ggml-cpu/ggml-cpu.c

@@ -1902,6 +1902,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_im2col_3d(params, tensor);
             } break;
+        case GGML_OP_SNAKE:
+            {
+                ggml_compute_forward_snake(params, tensor);
+            } break;
         case GGML_OP_CONV_2D:
             {
                 ggml_compute_forward_conv_2d(params, tensor);
@@ -2330,6 +2334,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_IM2COL:
         case GGML_OP_IM2COL_BACK:
         case GGML_OP_IM2COL_3D:
+        case GGML_OP_SNAKE:
         case GGML_OP_CONV_2D:
         case GGML_OP_CONV_3D:
         case GGML_OP_CONV_2D_DW:

ggml/src/ggml-cpu/ops.cpp

@@ -5924,6 +5924,19 @@ void ggml_compute_forward_rope_back(
     }
 }
 
+// Type-generic load/store helpers for F32, F16, BF16 snake activation
+template <typename T>
+static inline float snake_load(T x);
+template <> inline float snake_load(float x)       { return x; }
+template <> inline float snake_load(ggml_fp16_t x) { return GGML_FP16_TO_FP32(x); }
+template <> inline float snake_load(ggml_bf16_t x) { return GGML_BF16_TO_FP32(x); }
+
+template <typename T>
+static inline T snake_store(float x);
+template <> inline float       snake_store<float>(float x)       { return x; }
+template <> inline ggml_fp16_t snake_store<ggml_fp16_t>(float x) { return GGML_FP32_TO_FP16(x); }
+template <> inline ggml_bf16_t snake_store<ggml_bf16_t>(float x) { return GGML_FP32_TO_BF16(x); }
+
 // ggml_compute_forward_conv_transpose_1d
 
 static void ggml_compute_forward_conv_transpose_1d_f16_f32(
@@ -6650,8 +6663,53 @@ static inline int64_t ggml_wrap_around(int64_t coord, int64_t size) {
     return (coord  + size) % size; // adding size avoids negative number weirdness
 }
 
-// ggml_compute_forward_conv_2d
+// ggml_compute_forward_snake
+// Fused: y = x + sin^2(a * x) * inv_b
+// Supports F32, F16, BF16 input/output (same type), params always F32.
+template <typename elem_t>
+static void ggml_compute_forward_snake_impl(
+        const struct ggml_compute_params * params,
+        struct ggml_tensor * dst) {
+    const struct ggml_tensor * src0 = dst->src[0]; // x: [T, C]
+    const struct ggml_tensor * src1 = dst->src[1]; // a: [1, C] or [C]
+    const struct ggml_tensor * src2 = dst->src[2]; // inv_b: [1, C] or [C]
+
+    const int64_t T = src0->ne[0];
+    const int64_t C = src0->ne[1];
 
+    const elem_t * xd = (const elem_t *)src0->data;
+    const float  * ad = (const float  *)src1->data;
+    const float  * bd = (const float  *)src2->data;
+    elem_t       * yd = (elem_t       *)dst->data;
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    for (int64_t c = ith; c < C; c += nth) {
+        const float ac = ad[c];
+        const float bc = bd[c];
+        const elem_t * xc = xd + c * T;
+        elem_t       * yc = yd + c * T;
+        for (int64_t t = 0; t < T; t++) {
+            const float xi = snake_load(xc[t]);
+            const float s  = sinf(ac * xi);
+            yc[t] = snake_store<elem_t>(xi + s * s * bc);
+        }
+    }
+}
+
+void ggml_compute_forward_snake(
+        const struct ggml_compute_params * params,
+        struct ggml_tensor * dst) {
+    switch (dst->src[0]->type) {
+        case GGML_TYPE_F32:  ggml_compute_forward_snake_impl<float>      (params, dst); break;
+        case GGML_TYPE_F16:  ggml_compute_forward_snake_impl<ggml_fp16_t>(params, dst); break;
+        case GGML_TYPE_BF16: ggml_compute_forward_snake_impl<ggml_bf16_t>(params, dst); break;
+        default: GGML_ABORT("snake: unsupported type %d", dst->src[0]->type);
+    }
+}
+
+// ggml_compute_forward_conv_2d
 
 static void ggml_compute_forward_conv_2d_impl(const ggml_compute_params * params,
                                               const ggml_tensor *         kernel,  // [KW, KH, IC, OC]

ggml/src/ggml-cpu/ops.h

@@ -67,6 +67,7 @@ void ggml_compute_forward_conv_transpose_1d(const struct ggml_compute_params * p
 void ggml_compute_forward_im2col(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_im2col_back_f32(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_im2col_3d(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_snake(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_conv_2d(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_conv_3d(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_conv_transpose_2d(const struct ggml_compute_params * params, struct ggml_tensor * dst);

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -12,6 +12,7 @@
 #include "ggml-cuda/clamp.cuh"
 #include "ggml-cuda/concat.cuh"
 #include "ggml-cuda/conv-transpose-1d.cuh"
+#include "ggml-cuda/snake.cuh"
 #include "ggml-cuda/conv2d.cuh"
 #include "ggml-cuda/conv2d-dw.cuh"
 #include "ggml-cuda/conv2d-transpose.cuh"
@@ -2877,6 +2878,10 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
         case GGML_OP_IM2COL_3D:
             ggml_cuda_op_im2col_3d(ctx, dst);
             break;
+            break;
+        case GGML_OP_SNAKE:
+            ggml_cuda_op_snake(ctx, dst);
+            break;
         case GGML_OP_CONV_2D:
             ggml_cuda_op_conv2d(ctx, dst);
             break;
@@ -5159,6 +5164,10 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
         case GGML_OP_ROPE_BACK: {
             return op->src[0]->nb[0] == ggml_type_size(op->src[0]->type) && ggml_is_contiguous_2(op->src[0]);
         }
+        case GGML_OP_SNAKE:
+            return op->src[0]->type == GGML_TYPE_F32 ||
+                   op->src[0]->type == GGML_TYPE_F16 ||
+                   op->src[0]->type == GGML_TYPE_BF16;
         case GGML_OP_IM2COL:
         case GGML_OP_IM2COL_3D:
         case GGML_OP_CONV_2D:

ggml/src/ggml-cuda/snake.cu

@@ -0,0 +1,70 @@
+#include "snake.cuh"
+
+// Fused Snake activation: y = x + sin^2(a * x) * inv_b
+// x: [T, C] (T contiguous), a: [C], inv_b: [C]
+// Supports F32, F16, BF16 data with F32 compute.
+
+template <typename T>
+static __device__ __forceinline__ float snake_load(T x);
+template <> __device__ __forceinline__ float snake_load(float x)          { return x; }
+template <> __device__ __forceinline__ float snake_load(half x)           { return __half2float(x); }
+template <> __device__ __forceinline__ float snake_load(nv_bfloat16 x)    { return __bfloat162float(x); }
+
+template <typename T>
+static __device__ __forceinline__ T snake_store(float x);
+template <> __device__ __forceinline__ float       snake_store<float>(float x)       { return x; }
+template <> __device__ __forceinline__ half        snake_store<half>(float x)        { return __float2half(x); }
+template <> __device__ __forceinline__ nv_bfloat16 snake_store<nv_bfloat16>(float x) { return __float2bfloat16(x); }
+
+template <typename T>
+static __global__ void kernel_snake(
+        const T     * __restrict__ x,
+        const float * __restrict__ a,
+        const float * __restrict__ inv_b,
+        T           * __restrict__ dst,
+        const int T_len,
+        const int C) {
+    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= T_len * C) return;
+
+    const int c = idx / T_len;
+
+    const float xi = snake_load(x[idx]);
+    const float s  = sinf(a[c] * xi);
+    dst[idx] = snake_store<T>(xi + s * s * inv_b[c]);
+}
+
+void ggml_cuda_op_snake(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0];
+    const ggml_tensor * src1 = dst->src[1];
+    const ggml_tensor * src2 = dst->src[2];
+
+    const float * a_d     = (const float *)src1->data;
+    const float * inv_b_d = (const float *)src2->data;
+
+    const int T = (int)src0->ne[0];
+    const int C = (int)src0->ne[1];
+    const int total = T * C;
+
+    const int block_size = 256;
+    const int grid_size  = (total + block_size - 1) / block_size;
+
+    cudaStream_t stream = ctx.stream();
+
+    switch (src0->type) {
+        case GGML_TYPE_F32: {
+            kernel_snake<<<grid_size, block_size, 0, stream>>>(
+                (const float *)src0->data, a_d, inv_b_d, (float *)dst->data, T, C);
+        } break;
+        case GGML_TYPE_F16: {
+            kernel_snake<<<grid_size, block_size, 0, stream>>>(
+                (const half *)src0->data, a_d, inv_b_d, (half *)dst->data, T, C);
+        } break;
+        case GGML_TYPE_BF16: {
+            kernel_snake<<<grid_size, block_size, 0, stream>>>(
+                (const nv_bfloat16 *)src0->data, a_d, inv_b_d, (nv_bfloat16 *)dst->data, T, C);
+        } break;
+        default:
+            GGML_ABORT("snake: unsupported type");
+    }
+}

ggml/src/ggml-cuda/snake.cuh

@@ -0,0 +1,3 @@
+#include "common.cuh"
+
+void ggml_cuda_op_snake(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

ggml/src/ggml-metal/ggml-metal-device.cpp

@@ -1738,6 +1738,25 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_1
     return res;
 }
 
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_snake(ggml_metal_library_t lib, const ggml_tensor * op) {
+    const char * suffix;
+    switch (op->src[0]->type) {
+        case GGML_TYPE_F32:  suffix = "f32";  break;
+        case GGML_TYPE_F16:  suffix = "f16";  break;
+        case GGML_TYPE_BF16: suffix = "bf16"; break;
+        default: GGML_ABORT("snake: unsupported type");
+    }
+
+    char name[64];
+    snprintf(name, sizeof(name), "kernel_snake_%s", suffix);
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, name, name, nullptr);
+    }
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_2d(ggml_metal_library_t lib, const ggml_tensor * op) {
     assert(op->op == GGML_OP_CONV_TRANSPOSE_2D);
 

ggml/src/ggml-metal/ggml-metal-device.h

@@ -148,6 +148,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_norm
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rope              (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_im2col            (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_1d (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_snake             (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_2d (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_2d           (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_3d           (ggml_metal_library_t lib, const struct ggml_tensor * op);

ggml/src/ggml-metal/ggml-metal-device.m

@@ -1081,6 +1081,10 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
         case GGML_OP_REPEAT:
         case GGML_OP_CONV_TRANSPOSE_1D:
             return true;
+        case GGML_OP_SNAKE:
+            return op->src[0]->type == GGML_TYPE_F32 ||
+                   op->src[0]->type == GGML_TYPE_F16 ||
+                   op->src[0]->type == GGML_TYPE_BF16;
         case GGML_OP_CONV_TRANSPOSE_2D:
             return ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]) &&
                 (op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32) &&

ggml/src/ggml-metal/ggml-metal-impl.h

@@ -603,6 +603,11 @@ typedef struct {
     uint64_t nb1;
 } ggml_metal_kargs_conv_transpose_1d;
 
+typedef struct {
+    int32_t T;
+    int32_t C;
+} ggml_metal_kargs_snake;
+
 typedef struct {
     int32_t  IC;
     int32_t  IH;

ggml/src/ggml-metal/ggml-metal-ops.cpp

@@ -390,6 +390,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) {
             {
                 n_fuse = ggml_metal_op_conv_transpose_1d(ctx, idx);
             } break;
+        case GGML_OP_SNAKE:
+            {
+                n_fuse = ggml_metal_op_snake(ctx, idx);
+            } break;
         case GGML_OP_CONV_TRANSPOSE_2D:
             {
                 n_fuse = ggml_metal_op_conv_transpose_2d(ctx, idx);
@@ -3834,6 +3838,38 @@ int ggml_metal_op_conv_transpose_1d(ggml_metal_op_t ctx, int idx) {
     return 1;
 }
 
+int ggml_metal_op_snake(ggml_metal_op_t ctx, int idx) {
+    ggml_tensor * dst = ctx->node(idx);
+
+    ggml_metal_library_t lib = ctx->lib;
+    ggml_metal_encoder_t enc = ctx->enc;
+
+    const ggml_tensor * src0 = dst->src[0];
+    const ggml_tensor * src1 = dst->src[1];
+    const ggml_tensor * src2 = dst->src[2];
+
+    const int T = (int)src0->ne[0];
+    const int C = (int)src0->ne[1];
+    const int total = T * C;
+
+    auto pipeline = ggml_metal_library_get_pipeline_snake(lib, dst);
+
+    ggml_metal_kargs_snake args = { T, C };
+
+    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(src0), 1);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(src1), 2);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(src2), 3);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(dst),  4);
+
+    const int nth = 256;
+    const int ntg = (total + nth - 1) / nth;
+    ggml_metal_encoder_dispatch_threadgroups(enc, ntg, 1, 1, nth, 1, 1);
+
+    return 1;
+}
+
 int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) {
     ggml_tensor * op = ctx->node(idx);
 

ggml/src/ggml-metal/ggml-metal-ops.h

@@ -77,6 +77,7 @@ int ggml_metal_op_im2col            (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_conv_2d           (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_conv_3d           (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_conv_transpose_1d (ggml_metal_op_t ctx, int idx);
+int ggml_metal_op_snake             (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_conv_transpose_2d (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_upscale           (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_pad               (ggml_metal_op_t ctx, int idx);

ggml/src/ggml-metal/ggml-metal.metal

@@ -4887,6 +4887,29 @@ kernel void kernel_conv_transpose_1d<half>(
     uint3    tgpg[[threadgroups_per_grid]]);
 
 
+template <typename T>
+kernel void kernel_snake(
+        device const ggml_metal_kargs_snake & args [[buffer(0)]],
+        device const T     * x      [[buffer(1)]],
+        device const float * a      [[buffer(2)]],
+        device const float * inv_b  [[buffer(3)]],
+        device       T     * dst    [[buffer(4)]],
+        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;
+    const float xi = float(x[idx]);
+    const float s  = sin(a[c] * xi);
+    dst[idx] = T(xi + s * s * inv_b[c]);
+}
+
+template [[host_name("kernel_snake_f32")]]  kernel void kernel_snake<float>(device const 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>(device const ggml_metal_kargs_snake &, device const half *, device const float *, device const float *, device half *, uint, uint, uint);
+template [[host_name("kernel_snake_bf16")]] kernel void kernel_snake<bfloat>(device const ggml_metal_kargs_snake &, device const bfloat *, device const float *, device const float *, device bfloat *, uint, uint, uint);
+
 typedef void (conv_transpose_2d_t)(
         constant ggml_metal_kargs_conv_transpose_2d & args,
         device const float * src0,