ggml: add GGML_OP_SNAKE for fused Snake activation

[ServeurpersoCom]

Overview

Adds "GGML_OP_SNAKE", a fused snake activation "y = x + sin(a * x)^2 * inv_b" for audio decoders. The naive decomposition needs 12 ggml ops including two "ggml_repeat" broadcasts that materialize the full tensor twice; the fused op cuts that to one read + one write per backend.

Ships with backend coverage on CPU, CUDA, Metal and Vulkan (F32, F16, BF16) and a naive vs fused equivalence test that proves the fused op produces the same result as a decomposition built from already validated ggml ops ("mul", "sin", "sqr", "mul", "add").

Additional information

Used by acestep.cpp on the SEANet decoder of the VAE, and ready to drop into koboldcpp's qwen3-tts (currently the naive 12 op form with "ggml_repeat" broadcasts) and ace-step (currently a 5 op form with "exp(alpha)" and "1/exp(beta)" precomputed at load) audio decoders.

Beyond music generation, this is the same "y = x + sin(a*x)^2 * inv_b" activation introduced by Ziyin et al. (NeurIPS 2020) and adopted as the standard non-linearity in the BigVGAN family of neural vocoders (Lee et al., ICLR 2023), which is the audio decoder backbone shared by Qwen3-TTS, Qwen3-Omni, OmniVoice, DAC (Descript Audio Codec) and most current speech and music generation models.

Validation

The op is exercised across CPU, CUDA, ROCm, Metal and Vulkan via "acestep.cpp" real-world music generation, with active backend-specific tuning in the upstream repo. Coverage of exotic architectures (Apple Silicon, ARM Linux, AMD discrete) has been gathered through rented Metal instances and user SSH access. On this PR's machine (RTX PRO 6000 Blackwell + Ryzen 9 9950X3D) the test suite passes 14/14:

Backend	SNAKE coverage	SNAKE naive vs fused
CPU SIMD	13/13	1/1
CUDA	13/13	1/1
Vulkan NVIDIA	13/13	1/1
Vulkan llvmpipe	13/13	1/1

Remaining backends will be covered by the standard ggml CI matrix.

The equivalence test uses mean absolute error between backends (override of the virtual "err()", already used by other tests in the file) instead of the default NMSE. NMSE is "mse(a, b) / mse(a, 0)" and collapses when the reference signal is ~0 by construction (delta of two equivalent paths); empirically NMSE blows up to 1.49 on a clean run versus a 1e-7 threshold. MAE is well defined near zero with a 1e-5 tolerance.

A companion PR for "GGML_OP_COL2IM_1D" (memory efficient transposed 1d convolution, used by the same audio decoders to avoid the dilated buffer that "ggml_conv_transpose_1d" allocates) will follow once this lands.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES Opus 4.7 + rootless pod

[am17an]

This can just a be fusion inside the CUDA backend, as it's a pretty straight forward one.

[ServeurpersoCom]

This can just a be fusion inside the CUDA backend, as it's a pretty straight forward one.

Thanks for the look. Quick questions before deciding, just want to make sure I understand the fusion path:

Does it cover all backends or only CUDA? Snake is exercised on CUDA, Vulkan, Metal, WebGPU and CPU SIMD via acestep.cpp, so I'd want non-CUDA users to get the fused path too.

Does it handle multiple input shapes of the same logical op? Users write Snake differently depending on precomputation: the full form y = x + sin(exp(a)x)^2 * exp(-b) (12 ops, used by koboldcpp/qwen3tts) versus the reduced form y = x + sin(ax)^2 * inv_b with a and inv_b precomputed at load (5 ops, used by koboldcpp/acestep).

And can downstream code still call something like ggml_snake(x, a, inv_b) directly, or does every user have to emit the decomposition and trust the matcher?

[am17an]

No each backend will need to detect these nodes in the graph and do the fusion. It's what we do for frankly way common operations like topk-moe and gated ffn. So I don't see why this should be any different. Supporting a new OP is a maintenance burden, and as such adding new ops must be drastically effective. Just as an example, acestep would fall back to the CPU if any backend doesn't implement this op, making it way slower than the unfused version.

[0cc4m]

Maybe it would be good to have some common fusion infrastructure. It is a lot of duplicated work to detect and fuse the same pattern in each backend.

[am17an]

We do have something like this for FA and now GDN, where ggml looks for the fused op. However if you look at the machinery involved in doing that, it has to be really worth it to do this.

[ServeurpersoCom]

I add the missing abstraction layer...

[ServeurpersoCom]

This is great! The fusion machinery works perfectly. I'm seeing less than 1% overhead compared to a direct kernel call, within run-to-run variance, while the naive 5-op path loses 30% on CUDA (20% on Vulkan) on my Oobleck VAE, so the difference between fused and naive is clearly measurable (I observed the full loss when my matching was incorrectly written, before I fixed it). I'm checking the correctness for Vulkan and how to perform clean tests like existing fused ops. then remove the public GGML_OP_SNAKE op as you requested.

[ServeurpersoCom]

ggml_flash_attn_ext under Vulkan causes non-deterministic BF16 divergences, and work some time!

For documentation purpose :

    GGML graph for snake activation

           y = x + sin(a*x)^2 * inv_b

       +------+    +-------+    +-------+
       |  x   |    |   a   |    | inv_b |
       +--+---+    +---+---+    +---+---+
          |            |            |
          +-----+------+            |
                |                   |
                v                   |
          +-----------+             |
 +------> | n[0]  MUL |             |
 |        |  ax=x*a   |             |
 |        +-----+-----+             |
 |              |                   |
 |              v                   |
 |        +-----------+             |
 |        | n[1]  SIN |             |
 |        | s=sin(ax) |             |
 |        +-----+-----+             |
 |              |                   |
 |              v                   |
 |        +-----------+             |
 |        | n[2]  SQR |             |
 |        |  s2=s*s   |             |
 |        +-----+-----+             |
 |              |                   |
 |              v                   |
 |        +-----------+             |
 |        | n[3]  MUL | <-----------+
 |        | d=s2*invb |
 |        +-----+-----+
 |              |
 |              v
 |        +-----------+
 |        | n[4]  ADD |
 |        |  y=x+d    |
 |        +-----+-----+
 |              |
 |              v
 |
 |              y
 |
 |
 |       fusion machinery scans the graph and detects
 |       the chain MUL > SIN > SQR > MUL > ADD with
 |       single use intermediates, contiguous shapes,
 |       and closure x_in_add == x_in_mul0
 |
 |       on match, the driver advances i by 4 and
 |       routes the whole chain to one specialized
 |       kernel that reads x, a, inv_b once and
 |       writes y once
 |
 |              +---------------------------+
 +------------> | op_snake_fused(x,a,invb)  |
                |                           |
     bypass     |  one register pass:       |
     5 launches |    xi = load x[idx]       |
     collapse   |    s  = sin(a[c] * xi)    |
     into 1     |    y  = xi + s*s*invb[c]  |
                |    store y[idx]           |
                +-------------+-------------+
                              |
                              v

                              y


    Public op (old direct kernel) vs. GGML fusion (diff)

    The project's math can remain decomposed;
    GGML will do the optimizations automatically

 // Graph building
-// Snake activation (fused): y = x + sin^2(a * x) * inv_b
+// Snake activation (5-op naive decomposition for backend pattern fusion)
+// y = x + sin(a * x)^2 * inv_b
 // x: [T, C], exp_a: [1, C], inv_b: [1, C] (pre-computed at load)
+// Backends pattern-match (mul -> sin -> sqr -> mul -> add) and dispatch
+// to their fused snake kernel under the hood.
 static struct ggml_tensor * vae_snake(struct ggml_context * ctx,
                                       struct ggml_tensor *  x,
                                       struct ggml_tensor *  exp_a,
                                       struct ggml_tensor *  inv_b) {
-    return ggml_snake(ctx, x, exp_a, inv_b);
+    struct ggml_tensor * ax = ggml_mul(ctx, x, exp_a);
+    struct ggml_tensor * s  = ggml_sin(ctx, ax);
+    struct ggml_tensor * s2 = ggml_sqr(ctx, s);
+    struct ggml_tensor * d  = ggml_mul(ctx, s2, inv_b);
+    return ggml_add(ctx, x, d);
 }

[0cc4m]

ggml_flash_attn_ext under Vulkan causes non-deterministic BF16 divergences, and work some time!

What do you mean?

[ServeurpersoCom]

ggml_flash_attn_ext under Vulkan causes non-deterministic BF16 divergences, and work some time!

What do you mean?

On acestep.cpp I spent the day narrowing this down: cosine similarities were blowing up in my DiT, only in BF16, mostly on the deeper / longer models (SFT, Xl-sft, 24-32 layers x 50 steps), and non-deterministically (roughly 50% per run). I eventually tried --no-fa and everything became bit-exact good across all 4 models, all runs.

Root cause: under Vulkan, ggml_flash_attn_ext only has f32_f16 pipelines, no BF16 variant. K/V in BF16 get silently down-cast to F16 before the kernel, so any activation above 65504 becomes ±Inf and explodes through soft_max. CUDA has native BF16 FA pipelines, so it didn't trigger there.

That detour delayed me on finishing the Vulkan pattern matching work.

[jeffbolznv]

Please file a bug for that, with instructions for how to reproduce it. If you can catch it with GGML_VULKAN_CHECK_RESULTS, even better.

[ServeurpersoCom]

Please file a bug for that, with instructions for how to reproduce it. If you can catch it with GGML_VULKAN_CHECK_RESULTS, even better.

Of course, I'll file a separate bug for it with a proper repro. I'm planning to isolate it, possibly via a small standalone C++
Either a simple BF16 test, or something more complex if necessary. because on my case I need a test that pushes the activations through enough layers / iterations to reproduce,

[ServeurpersoCom]

Measured on acestep, which I migrated to fully use the GGML autofuse pass for all snake kernels. The test script runs sequentially across all DiT quants (BF16, Q8_0, Q6_K, Q5_K_M) while the VAE itself stays BF16. Metal is still pending measurement, where the speedup should be even more pronounced thanks to the threadgroup-dispatched fused kernel avoiding the 5 separate command-encoder roundtrips of the naive chain.

Env var used for the NAIVE pass
GGML_CUDA_DISABLE_FUSION=1
GGML_VK_DISABLE_FUSION=1
GGML_CPU_DISABLE_FUSION=1 (added for testing!)

backend pass   FUSION (ms) NAIVE (ms) speedup
---------------------------------------------
CUDA0   BF16         442.7      615.4   1.39x
CUDA0   Q8_0         428.9      615.1   1.43x
CUDA0   Q6_K         429.4      615.1   1.43x
CUDA0   Q5_K_M       436.5      615.0   1.41x

Vulkan0 BF16         766.4      953.2   1.24x
Vulkan0 Q8_0         778.0      963.5   1.24x
Vulkan0 Q6_K         775.0      969.0   1.25x
Vulkan0 Q5_K_M       771.2      959.4   1.24x

CPU     BF16       30437.6    34572.9   1.14x
CPU     Q8_0       31106.8    36405.0   1.17x
CPU     Q6_K       29990.0    35059.1   1.17x
CPU     Q5_K_M     29662.2    34304.3   1.16x

The metric is the total VAE execution time; only the snake activation path differs between the two columns (autofused kernel vs naive mul/sin/sqr/mul/add chain), everything else including the optimized col2im_1d kernel is identical.

[ServeurpersoCom]

Therefore, this approach would allow us to separate the kernels we want upstream into separate PR since everything is decoupled

[ServeurpersoCom]

Creation of decoupled mini-PRs

#22667

[ServeurpersoCom]

Noting two BF16 coverage gaps on the Vulkan backend: (1) FlashAttention has no BF16 pipeline variant, so K/V get silently downcast to F16 and overflow to ±Inf above 65504, and (2) the CPY matcher is missing the bf16 -> f32 pipeline, which aborts hard (hit downstream in acestep.cpp on LoRA merges). Both stem from the same root cause: BF16 was added to Vulkan piecewise without a full op-by-op audit.

[0cc4m]

Why would bf16 FA get swapped with f16? It should just fall back to CPU if the op is not available.

[ServeurpersoCom]

Why would bf16 FA get swapped with f16? It should just fall back to CPU if the op is not available.

You're right, the fallback path does work. supports_op correctly returns false for FLASH_ATTN_EXT with BF16 K/V, so it dispatches to CPU.
In my DiT K/V actually reach FA in F32 because mul_mat hardcodes F32 output, and I tend to avoid casts unless I'm sure they're needed (minimalist by habit).
I just checked llama.cpp (BERT, T5 encoder, LLaDA diffusion models): it does cast K/V to F16 explicitly before FA for graphs without a KV cache. I'll fix it on my side!

Thanks, this kind of pushback genuinely helps.