metal: fuse snake activation (mul, sin, sqr, mul, add) by ServeurpersoCom · Pull Request #25459 · ggml-org/llama.cpp

ServeurpersoCom · 2026-07-08T18:12:16Z

Overview

Metal counterpart of #22667 (CUDA) and #22855 (Vulkan): same matcher on the naive 5-op chain (mul, sin, sqr, mul, add), same broadcast contract (a / inv_b in F32, shaped [1, C] over x [T, C]), single kernel templated on x's type with F32 compute.

Follows the Metal backend conventions: the bf16 instantiation is gated behind GGML_METAL_HAS_BF16, and the fused dispatch checks concurrency ranges on the remaining chain nodes before encoding, as the bin fusion does, since the pre-check in the encode loop only covers the leading mul.

Covered by the existing backend-agnostic SNAKE_FUSE tests from #22667.

Additional information

Requirements

I have read and agree with the contributing guidelines
AI usage disclosure: YES Fable 5 / Rootless dev containers

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.

forforever73 · 2026-07-14T09:26:15Z

F32 looks good, but the kernel_snake_f16/bf16 are never dispatched. Binary Metal ops currently require src0 == F32 in ggml_metal_device_supports_op. the head of the Snake fusion chain mul(x, a) will fall back to CPU, and the 5-op fusion chain never materializes on Metal.

ServeurpersoCom · 2026-07-14T10:16:34Z

F32 looks good, but the kernel_snake_f16/bf16 are never dispatched. Binary Metal ops currently require src0 == F32 in ggml_metal_device_supports_op. the head of the Snake fusion chain mul(x, a) will fall back to CPU, and the 5-op fusion chain never materializes on Metal.

Right, the leading mul is gated to F32 by supports_op, so the f16/bf16 variants never dispatch today. The kernel body is a single template shared with the F32 path, so this is just future proofing, and every backend carries such variants: f16/bf16 on Vulkan behind the SIN/SQR gate (#22855), bf16 on CUDA behind the bin_bcast filter (#22912). They light up when the missing ops are added.

ServeurpersoCom · 2026-07-14T10:57:13Z

Snake fusion, fused chain and missing ops inventory :

Fused chain           CUDA        Vulkan      Metal
---------------------------------------------------
F32                   live        live        live
F16                   live        live        dead
BF16                  dead        dead        dead

Chain ops             CUDA        Vulkan      Metal
---------------------------------------------------
bin F32               ok          ok          ok
bin F16               ok          ok          MISSING (kernel + gate)
bin BF16              MISSING     MISSING     MISSING
sin/sqr F32           ok          ok          ok
sin/sqr F16           ok          ok          ok
sin/sqr BF16          ok          MISSING     MISSING

PRs to do (this would double performance for vocoders) :

cuda: bin_bcast BF16 -> to enables BF16
metal: bin F16 -> to enables F16
metal: bin BF16 + sin/sqr BF16 -> to enables BF16
vulkan: BF16 skipped, device reports bf16: 0 -> unsupported by driver tested on my RTX PRO 6000 (I am not sure if that is possible.)

forforever73 · 2026-07-14T11:17:32Z

Got it, I'm ok keeping the dormant f16/bf16 variants here for now.

ServeurpersoCom · 2026-07-14T11:21:22Z

I’ll do some follow-up PRs, starting tests with my GGML fork, to get real benchmarks and user feedback on my small *.cpp projects.

ggerganov · 2026-07-14T11:59:59Z

+                // Snake activation autofuse: mul -> sin -> sqr -> mul -> add
+                bool fused = false;
+                if (node->op == GGML_OP_MUL && ctx->use_fusion) {
+                    static constexpr ggml_op snake_ops[5] = { GGML_OP_MUL, GGML_OP_SIN, GGML_OP_SQR, GGML_OP_MUL, GGML_OP_ADD };
+                    if (ctx->can_fuse(idx, snake_ops, 5)) {
+                        const ggml_tensor * mul0     = ctx->node(idx + 0);
+                        const ggml_tensor * sin_node = ctx->node(idx + 1);
+                        const ggml_tensor * sqr      = ctx->node(idx + 2);
+                        const ggml_tensor * mul1     = ctx->node(idx + 3);
+                        const ggml_tensor * add      = ctx->node(idx + 4);
+
+                        // x carries the full activation shape, a is the broadcast operand
+                        const ggml_tensor * x = ggml_are_same_shape(mul0, mul0->src[0]) ? mul0->src[0] : mul0->src[1];
+                        const ggml_tensor * a = (x == mul0->src[0]) ? mul0->src[1] : mul0->src[0];
+
+                        // mul1 reads sqr and inv_b in either operand order
+                        const ggml_tensor * inv_b    = (mul1->src[0] == sqr) ? mul1->src[1] : mul1->src[0];
+
+                        // closure check: the trailing add reads the same x as the leading mul
+                        const ggml_tensor * x_in_add = (add->src[0] == mul1) ? add->src[1] : add->src[0];
+
+                        // x is in the supported whitelist and every chain intermediate shares x's type.
+                        // a and inv_b bind as device const float * in the kernel, so they stay F32.
+                        const bool types_ok =
+                            (x->type == GGML_TYPE_F32 || x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) &&
+                            (a->type    == GGML_TYPE_F32) && (inv_b->type    == GGML_TYPE_F32) &&
+                            (mul0->type == x->type)       && (sin_node->type == x->type) &&
+                            (sqr->type  == x->type)       && (mul1->type     == x->type) &&
+                            (add->type  == x->type);
+                        // a / inv_b collapse to [1, C, 1, 1], x and add stay 2D
+                        const bool shape_ok = ggml_are_same_shape(a, inv_b) && a->ne[0] == 1 && a->ne[1] == x->ne[1];
+                        const bool dim_ok =
+                            (x->ne[2]     == 1) && (x->ne[3]     == 1) &&
+                            (add->ne[2]   == 1) && (add->ne[3]   == 1) &&
+                            (a->ne[2]     == 1) && (a->ne[3]     == 1) &&
+                            (inv_b->ne[2] == 1) && (inv_b->ne[3] == 1);
+                        // kernel reads x[idx] and a[c] / inv_b[c] linearly, so every operand is contiguous
+                        const bool contig_ok =
+                            ggml_is_contiguous(x) && ggml_is_contiguous(add) &&
+                            ggml_is_contiguous(a) && ggml_is_contiguous(inv_b);
+                        if (types_ok && shape_ok && dim_ok && contig_ok && x_in_add == x) {
+                            n_fuse = ggml_metal_op_snake_fused(ctx, idx);
+                            fused = true;
+                        }
+                    }
+                }
+                if (!fused) {
+                    n_fuse = ggml_metal_op_bin(ctx, idx);
+                }


This logic should be absorbed inside the ggml_metal_op_bin(). We want to keep this switch statement clean.

Done, matcher extracted to ggml_metal_op_can_fuse_snake and dispatch absorbed into ggml_metal_op_bin, the switch is back to a single call.

@ggerganov

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

ggerganov · 2026-07-14T12:53:06Z

+    // closure check: the trailing add reads the same x as the leading mul
+    const ggml_tensor * x_in_add = (add->src[0] == mul1) ? add->src[1] : add->src[0];
+
+        // x is in the supported whitelist and every chain intermediate shares x's type.
+        // a and inv_b bind as device const float * in the kernel, so they stay F32.
+        const bool types_ok =
+            (x->type == GGML_TYPE_F32 || x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) &&
+            (a->type    == GGML_TYPE_F32) && (inv_b->type    == GGML_TYPE_F32) &&
+            (mul0->type == x->type)       && (sin_node->type == x->type) &&
+            (sqr->type  == x->type)       && (mul1->type     == x->type) &&
+            (add->type  == x->type);
+        // a / inv_b collapse to [1, C, 1, 1], x and add stay 2D
+        const bool shape_ok = ggml_are_same_shape(a, inv_b) && a->ne[0] == 1 && a->ne[1] == x->ne[1];
+        const bool dim_ok =
+            (x->ne[2]     == 1) && (x->ne[3]     == 1) &&
+            (add->ne[2]   == 1) && (add->ne[3]   == 1) &&
+            (a->ne[2]     == 1) && (a->ne[3]     == 1) &&
+            (inv_b->ne[2] == 1) && (inv_b->ne[3] == 1);
+        // kernel reads x[idx] and a[c] / inv_b[c] linearly, so every operand is contiguous
+        const bool contig_ok =
+            ggml_is_contiguous(x) && ggml_is_contiguous(add) &&
+            ggml_is_contiguous(a) && ggml_is_contiguous(inv_b);
+
+    return types_ok && shape_ok && dim_ok && contig_ok && x_in_add == x;


Indentation is wrong in this section.

ServeurpersoCom requested a review from a team as a code owner July 8, 2026 18:12

github-actions Bot added ggml changes relating to the ggml tensor library for machine learning Apple Metal https://en.wikipedia.org/wiki/Metal_(API) labels Jul 8, 2026

ggerganov self-assigned this Jul 13, 2026

forforever73 approved these changes Jul 14, 2026

View reviewed changes

ggerganov reviewed Jul 14, 2026

View reviewed changes

metal: absorb snake fusion into ggml_metal_op_bin

cf4c3d2

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

ggerganov reviewed Jul 14, 2026

View reviewed changes

metal: fix indentation in ggml_metal_op_can_fuse_snake

ec507ba

ggerganov approved these changes Jul 14, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

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

metal: fuse snake activation (mul, sin, sqr, mul, add)#25459
ServeurpersoCom wants to merge 3 commits into
ggml-org:masterfrom
ServeurpersoCom:ggml/metal-snake-fusion

ServeurpersoCom commented Jul 8, 2026

Uh oh!

forforever73 commented Jul 14, 2026

Uh oh!

ServeurpersoCom commented Jul 14, 2026

Uh oh!

ServeurpersoCom commented Jul 14, 2026

Uh oh!

forforever73 commented Jul 14, 2026

Uh oh!

ServeurpersoCom commented Jul 14, 2026

Uh oh!

ggerganov Jul 14, 2026

Uh oh!

ServeurpersoCom Jul 14, 2026

Uh oh!

ggerganov Jul 14, 2026

Uh oh!

ServeurpersoCom Jul 14, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

ServeurpersoCom commented Jul 8, 2026

Overview

Additional information

Requirements

Uh oh!

forforever73 commented Jul 14, 2026

Uh oh!

ServeurpersoCom commented Jul 14, 2026

Uh oh!

ServeurpersoCom commented Jul 14, 2026

Uh oh!

forforever73 commented Jul 14, 2026

Uh oh!

ServeurpersoCom commented Jul 14, 2026

Uh oh!

ggerganov Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

ServeurpersoCom Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

ggerganov Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

ServeurpersoCom Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants