QVAC-21921 ggml-speech: ACE-Step Oobleck VAE custom ops (snake, col2im_1d) — CPU + Metal

include/ggml.h

@@ -226,7 +226,19 @@
 #define GGML_MAX_OP_PARAMS      64
 
 #ifndef GGML_MAX_NAME
-#   define GGML_MAX_NAME        64
+// ACE-Step DiT tensor names exceed the historical 64-char limit (e.g. 67 chars),
+// which would truncate/mismatch on load, so this fork bumps the default to 128
+// (matching upstream acestep.cpp).
+//
+// ABI WARNING: name[] is an inline array in struct ggml_tensor, so changing this
+// alters the tensor's size and field offsets. Linking objects built with 64
+// against a lib built with 128 is an ODR/ABI mismatch (silent memory
+// corruption). This is safe here ONLY because the whole QVAC speech stack
+// consumes ggml exclusively through the ggml-speech vcpkg port and is rebuilt
+// from THIS header in the same build (audiogen-cpp, tts-cpp, parakeet-cpp, …) —
+// there are no prebuilt 64-byte-layout artifacts in the link graph. Any change
+// to this value is a HARD rebuild-everything requirement for all downstreams.
+#   define GGML_MAX_NAME        128
 #endif
 
 #define GGML_DEFAULT_N_THREADS  4
@@ -604,6 +616,14 @@ extern "C" {
         // out = x*aw + ab + max(x,0) + slope*min(x,0).
         GGML_OP_AFFINE_PRELU,
 
+        // col2im for 1D transpose-conv (ACE-Step Oobleck VAE):
+        // scatter-add GEMM columns back into a 1D signal.
+        GGML_OP_COL2IM_1D,
+
+        // Snake activation y = x + sin^2(a*x) * inv_b, per-channel a / inv_b
+        // (ACE-Step Oobleck VAE).
+        GGML_OP_SNAKE,
+
         GGML_OP_COUNT,
     };
 
@@ -2497,6 +2517,25 @@ extern "C" {
             struct ggml_tensor  * ab,
             struct ggml_tensor  * slope);
 
+    // col2im_1d: scatter-add GEMM columns back to a 1D signal (GEMM-based
+    // conv_transpose_1d, ACE-Step Oobleck VAE).
+    // a: [K*OC, T_in]  (columns from the matmul, K = a->ne[0]/oc)
+    // result: [T_out, OC]  where T_out = (T_in - 1)*s0 + K - 2*p0
+    GGML_API struct ggml_tensor * ggml_col2im_1d(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,   // columns [K*OC, T_in]
+            int                   s0,  // stride
+            int                   oc,  // output channels
+            int                   p0); // padding to crop from both sides
+
+    // snake activation: y = x + sin^2(a*x) * inv_b, per-channel a / inv_b
+    // (ACE-Step Oobleck VAE).  x: [T, C]; a, inv_b: one per channel.
+    GGML_API struct ggml_tensor * ggml_snake(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * x,      // [T, C]
+            struct ggml_tensor  * a,      // per-channel scale inside sin(), F32
+            struct ggml_tensor  * inv_b); // per-channel output scale, F32
+
     // Move tensor elements by an offset given for each dimension. Elements that
     // are shifted beyond the last position are wrapped around to the beginning.
     GGML_API struct ggml_tensor * ggml_roll(

src/ggml-cpu/ggml-cpu.c

@@ -1962,6 +1962,14 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_supertonic_edge_pad_1d(params, tensor);
             } break;
+        case GGML_OP_COL2IM_1D:
+            {
+                ggml_compute_forward_col2im_1d(params, tensor);
+            } break;
+        case GGML_OP_SNAKE:
+            {
+                ggml_compute_forward_snake(params, tensor);
+            } break;
         case GGML_OP_ZERO_UPSAMPLE:
             {
                 ggml_compute_forward_zero_upsample(params, tensor);
@@ -2392,6 +2400,8 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_ZERO_UPSAMPLE:
         case GGML_OP_CHANNEL_SHUFFLE:
         case GGML_OP_AFFINE_PRELU:
+        case GGML_OP_COL2IM_1D:
+        case GGML_OP_SNAKE:
         case GGML_OP_ROLL:
         case GGML_OP_ARANGE:
         case GGML_OP_TIMESTEP_EMBEDDING:

src/ggml-cpu/ops.cpp

@@ -8136,6 +8136,111 @@ void ggml_compute_forward_supertonic_edge_pad_1d(
     }
 }
 
+// ggml_compute_forward_col2im_1d
+//
+// Scatter-add columns [K*OC, T_in] -> signal [T_out, OC], where
+// T_out = (T_in - 1)*s0 + K - 2*p0.  Implemented as a gather: each output time
+// reads the (at most ceil(K/s0)) columns that land on it, so threads write
+// disjoint outputs and no atomics/locks are needed.  Parallelized over the time
+// axis so the split stays balanced whatever OC is (down to OC = 1 mono audio).
+// CPU bring-up: F32 only.
+void ggml_compute_forward_col2im_1d(
+        const ggml_compute_params * params,
+              ggml_tensor * dst) {
+
+    const ggml_tensor * src = dst->src[0]; // [K*OC, T_in]
+
+    GGML_ASSERT(src->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type == GGML_TYPE_F32);
+    GGML_ASSERT(ggml_is_contiguous(src));
+    GGML_ASSERT(ggml_is_contiguous(dst));
+
+    const int32_t s0 = ggml_get_op_params_i32(dst, 0);
+    const int32_t OC = ggml_get_op_params_i32(dst, 1);
+    const int32_t p0 = ggml_get_op_params_i32(dst, 2);
+
+    const int64_t K_OC  = src->ne[0];
+    const int64_t T_in  = src->ne[1];
+    const int64_t K     = K_OC / OC;
+    const int64_t T_out = dst->ne[0];
+
+    const float * col_data = (const float *) src->data;
+          float * dst_data = (float *)       dst->data;
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    // Threads own disjoint output-time bands.
+    const int64_t dr  = (T_out + nth - 1) / nth;
+    const int64_t it0 = dr * ith;
+    const int64_t it1 = it0 + dr < T_out ? it0 + dr : T_out;
+
+    for (int64_t oc = 0; oc < OC; oc++) {
+        for (int64_t t_out = it0; t_out < it1; t_out++) {
+            const int64_t t_abs = t_out + p0; // position in the uncropped signal
+            // Gather every (t_in, k) with t_in*s0 + k == t_abs, 0 <= k < K.
+            int64_t t_in_min = (t_abs - K + 1 + s0 - 1) / s0; // ceil((t_abs-K+1)/s0)
+            if (t_in_min < 0) t_in_min = 0;
+            int64_t t_in_max = t_abs / s0;
+            if (t_in_max >= T_in) t_in_max = T_in - 1;
+
+            float sum = 0.0f;
+            for (int64_t t_in = t_in_min; t_in <= t_in_max; t_in++) {
+                const int64_t k = t_abs - t_in * s0;
+                if (k >= 0 && k < K) {
+                    sum += col_data[(oc * K + k) + t_in * K_OC];
+                }
+            }
+            dst_data[t_out + oc * T_out] = sum;
+        }
+    }
+}
+
+// ggml_compute_forward_snake
+//
+// Snake activation y = x + sin^2(a*x) * inv_b, with per-channel a and inv_b.
+// x / dst are [T, C] (T = ne0, contiguous); a and inv_b hold one F32 value per
+// channel.  Parallelized over channels (threads own disjoint [c*T, (c+1)*T)
+// bands).  ACE-Step Oobleck VAE.  CPU bring-up: F32 only.
+void ggml_compute_forward_snake(
+        const ggml_compute_params * params,
+              ggml_tensor * dst) {
+
+    const ggml_tensor * x     = dst->src[0];
+    const ggml_tensor * a     = dst->src[1];
+    const ggml_tensor * inv_b = dst->src[2];
+
+    GGML_ASSERT(x->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type == GGML_TYPE_F32);
+    GGML_ASSERT(a->type == GGML_TYPE_F32 && inv_b->type == GGML_TYPE_F32);
+    GGML_ASSERT(ggml_is_contiguous(x));
+    GGML_ASSERT(ggml_is_contiguous(dst));
+
+    const int64_t T = x->ne[0];
+    const int64_t C = x->ne[1];
+
+    const float * xd = (const float *) x->data;
+    const float * ad = (const float *) a->data;
+    const float * bd = (const float *) inv_b->data;
+          float * yd = (float *)       dst->data;
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    // Split over channels: each thread owns disjoint time bands.
+    for (int64_t c = ith; c < C; c += nth) {
+        const float ac = ad[c];
+        const float bc = bd[c];
+        const float * xc = xd + c * T;
+              float * yc = yd + c * T;
+        for (int64_t t = 0; t < T; t++) {
+            const float xi = xc[t];
+            const float si = sinf(ac * xi);
+            yc[t] = xi + si * si * bc;
+        }
+    }
+}
+
 // ggml_compute_forward_roll
 
 static int64_t ggml_wrap_index(int64_t i, int64_t ne) {

src/ggml-cpu/ops.h

@@ -86,6 +86,8 @@ void ggml_compute_forward_gru(const struct ggml_compute_params * params, struct
 void ggml_compute_forward_zero_upsample(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_channel_shuffle(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_affine_prelu(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_col2im_1d(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_roll(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_arange(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_timestep_embedding(const struct ggml_compute_params * params, struct ggml_tensor * dst);

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

@@ -1746,6 +1746,48 @@ 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_col2im_1d(ggml_metal_library_t lib, const ggml_tensor * op) {
+    assert(op->op == GGML_OP_COL2IM_1D);
+
+    GGML_ASSERT(ggml_is_contiguous(op->src[0]));
+    GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32);
+    GGML_ASSERT(op->type         == GGML_TYPE_F32);
+
+    char base[256];
+    char name[256];
+
+    snprintf(base, 256, "kernel_col2im_1d_%s", ggml_type_name(op->type));
+    snprintf(name, 256, "%s", base);
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
+    }
+
+    return res;
+}
+
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_snake(ggml_metal_library_t lib, const ggml_tensor * op) {
+    assert(op->op == GGML_OP_SNAKE);
+
+    GGML_ASSERT(ggml_is_contiguous(op->src[0]));
+    GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32);
+    GGML_ASSERT(op->type         == GGML_TYPE_F32);
+
+    char base[256];
+    char name[256];
+
+    snprintf(base, 256, "kernel_snake_%s", ggml_type_name(op->type));
+    snprintf(name, 256, "%s", base);
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, base, 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);
 

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

@@ -148,6 +148,8 @@ 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_col2im_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);

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

@@ -1111,6 +1111,16 @@ 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_COL2IM_1D:
+            return ggml_is_contiguous(op->src[0]) &&
+                   op->src[0]->type == GGML_TYPE_F32 &&
+                   op->type         == GGML_TYPE_F32;
+        case GGML_OP_SNAKE:
+            return ggml_is_contiguous(op->src[0]) &&
+                   op->src[0]->type == GGML_TYPE_F32 &&
+                   op->src[1]->type == GGML_TYPE_F32 &&
+                   op->src[2]->type == GGML_TYPE_F32 &&
+                   op->type         == GGML_TYPE_F32;
         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) &&

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

@@ -603,6 +603,24 @@ typedef struct {
     uint64_t nb1;
 } ggml_metal_kargs_conv_transpose_1d;
 
+// col2im for 1D transpose-conv (ACE-Step Oobleck VAE): scatter-add
+// GEMM columns [K*OC, T_in] back into a 1D signal [T_out, OC].
+typedef struct {
+    int32_t T_out;
+    int32_t T_in;
+    int32_t K;
+    int32_t OC;
+    int32_t s0;
+    int32_t p0;
+} ggml_metal_kargs_col2im_1d;
+
+// snake activation y = x + sin^2(a*x) * inv_b, per-channel a / inv_b
+// (ACE-Step Oobleck VAE).  x / y are [T, C] contiguous.
+typedef struct {
+    int32_t L; // T (time steps, contiguous inner dim)
+    int32_t C; // channels (one a / inv_b per channel)
+} ggml_metal_kargs_snake;
+
 typedef struct {
     int32_t  IC;
     int32_t  IH;