ggml : add ggml_conv_1d_grouped

ggml/include/ggml.h

@@ -2059,6 +2059,21 @@ extern "C" {
             int                   s0,  // stride
             int                   d0); // dilation
 
+    // grouped 1D convolution
+    // a: [K, IC/G, OC]   convolution kernel
+    // b: [L, IC,   N]    data
+    // groups must divide both IC and OC evenly
+    // when groups == 1, equivalent to ggml_conv_1d
+    // when groups == IC, equivalent to ggml_conv_1d_dw
+    GGML_API struct ggml_tensor * ggml_conv_1d_grouped(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a,      // convolution kernel
+            struct ggml_tensor  * b,      // data
+            int                   s0,     // stride
+            int                   p0,     // padding
+            int                   d0,     // dilation
+            int                   groups); // number of groups
+
     GGML_API struct ggml_tensor * ggml_conv_transpose_1d(
             struct ggml_context * ctx,
             struct ggml_tensor  * a,   // convolution kernel

ggml/src/ggml.c

@@ -4521,10 +4521,16 @@ struct ggml_tensor * ggml_conv_1d(
         int                   d0) {
     struct ggml_tensor * im2col = ggml_im2col(ctx, a, b, s0, 0, p0, 0, d0, 0, false, a->type == GGML_TYPE_BF16 ? GGML_TYPE_F32 : GGML_TYPE_F16); // [N, OL, IC * K]
 
+    // convert BF16 kernel to F32 for mul_mat compatibility
+    struct ggml_tensor * a_op = ggml_reshape_2d(ctx, a, (a->ne[0] * a->ne[1]), a->ne[2]); // [OC, IC * K]
+    if (a->type == GGML_TYPE_BF16) {
+        a_op = ggml_cpy(ctx, a_op, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, a_op->ne[0], a_op->ne[1]));
+    }
+
     struct ggml_tensor * result =
         ggml_mul_mat(ctx,
                 ggml_reshape_2d(ctx, im2col, im2col->ne[0], (im2col->ne[2] * im2col->ne[1])), // [N, OL, IC * K] => [N*OL, IC * K]
-                ggml_reshape_2d(ctx, a, (a->ne[0] * a->ne[1]), a->ne[2]));                    // [OC，IC, K] => [OC, IC * K]
+                a_op);                                                                         // [OC, IC * K]
 
     result = ggml_reshape_3d(ctx, result, im2col->ne[1], a->ne[2], im2col->ne[2]); // [N, OC, OL]
 
@@ -4555,7 +4561,13 @@ struct ggml_tensor * ggml_conv_1d_dw(
 
     struct ggml_tensor * im2col = ggml_im2col(ctx, a, new_b, s0, 0, p0, 0, d0, 0, false, a->type == GGML_TYPE_BF16 ? GGML_TYPE_F32 : GGML_TYPE_F16);
 
-    struct ggml_tensor * result = ggml_mul_mat(ctx, im2col, a);
+    // convert BF16 kernel to F32 for mul_mat compatibility
+    struct ggml_tensor * a_op = a;
+    if (a->type == GGML_TYPE_BF16) {
+        a_op = ggml_cpy(ctx, a_op, ggml_new_tensor_3d(ctx, GGML_TYPE_F32, a->ne[0], a->ne[1], a->ne[2]));
+    }
+
+    struct ggml_tensor * result = ggml_mul_mat(ctx, im2col, a_op);
 
     result = ggml_reshape_3d(ctx, result, result->ne[0], result->ne[2], 1);
 
@@ -4573,6 +4585,63 @@ struct ggml_tensor * ggml_conv_1d_dw_ph(
     return ggml_conv_1d_dw(ctx, a, b, s0, a->ne[0] / 2, d0);
 }
 
+// ggml_conv_1d_grouped
+
+struct ggml_tensor * ggml_conv_1d_grouped(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * a,
+        struct ggml_tensor  * b,
+        int                   s0,
+        int                   p0,
+        int                   d0,
+        int                   groups) {
+    GGML_ASSERT(groups > 0);
+
+    const int64_t OC   = a->ne[2];   // total output channels
+    const int64_t IC_G = a->ne[1];   // input channels per group (kernel dim)
+    const int64_t IC   = b->ne[1];   // total input channels
+
+    GGML_ASSERT(IC % groups == 0);
+    GGML_ASSERT(OC % groups == 0);
+    GGML_ASSERT(IC_G == IC / groups);
+
+    // degenerate cases: fall back to existing implementations
+    if (groups == 1) {
+        return ggml_conv_1d(ctx, a, b, s0, p0, d0);
+    }
+    if (groups == IC && groups == OC) {
+        return ggml_conv_1d_dw(ctx, a, b, s0, p0, d0);
+    }
+
+    const int64_t OC_G = OC / groups;
+
+    struct ggml_tensor * result = NULL;
+
+    for (int g = 0; g < groups; g++) {
+        // slice kernel for group g: [K, IC_G, OC_G]
+        struct ggml_tensor * a_g = ggml_view_3d(ctx, a,
+            a->ne[0], IC_G, OC_G,
+            a->nb[1], a->nb[2],
+            g * OC_G * a->nb[2]);
+
+        // slice input for group g: [L, IC_G, N]
+        struct ggml_tensor * b_g = ggml_view_3d(ctx, b,
+            b->ne[0], IC_G, b->ne[2],
+            b->nb[1], b->nb[2],
+            g * IC_G * b->nb[1]);
+
+        struct ggml_tensor * out_g = ggml_conv_1d(ctx, a_g, b_g, s0, p0, d0);
+
+        if (result == NULL) {
+            result = out_g;
+        } else {
+            result = ggml_concat(ctx, result, out_g, 1);
+        }
+    }
+
+    return result;
+}
+
 // ggml_col2im_1d
 
 struct ggml_tensor * ggml_col2im_1d(

tests/CMakeLists.txt

@@ -264,6 +264,7 @@ if (NOT GGML_BACKEND_DL)
     llama_build_and_test(test-quantize-perf.cpp)
     llama_build_and_test(test-rope.cpp)
     llama_build_and_test(test-col2im-1d.cpp)
+    llama_build_and_test(test-conv-1d-grouped.cpp)
 endif()
 
 # libmtmd

tests/test-conv-1d-grouped.cpp

@@ -0,0 +1,178 @@
+// Test for ggml_conv_1d_grouped
+//
+// Verifies grouped 1D convolution by comparing against manual per-group computation.
+
+#include "ggml.h"
+#include "ggml-backend.h"
+#include "ggml-cpu.h"
+
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <cmath>
+#include <vector>
+
+static void fill_random_f32(float * data, int n) {
+    for (int i = 0; i < n; i++) {
+        data[i] = ((float)rand() / RAND_MAX) * 2.0f - 1.0f;
+    }
+}
+
+static bool all_close(const float * a, const float * b, int n, float eps = 2e-2f) {
+    for (int i = 0; i < n; i++) {
+        if (fabsf(a[i] - b[i]) > eps) {
+            fprintf(stderr, "    mismatch at [%d]: %.6f vs %.6f (diff=%.6f)\n",
+                    i, a[i], b[i], fabsf(a[i] - b[i]));
+            return false;
+        }
+    }
+    return true;
+}
+
+// Compute grouped conv1d on CPU naively for reference
+// kernel (F16): [K, IC_G, OC], input (F32): [L, IC, N], output: [OL, OC, N]
+static void conv1d_grouped_ref(
+        const ggml_fp16_t * kernel, const float * input, float * output,
+        int K, int IC, int OC, int L, int N, int groups, int stride, int padding) {
+    int IC_G = IC / groups;
+    int OC_G = OC / groups;
+    int OL = (L + 2 * padding - K) / stride + 1;
+
+    memset(output, 0, (size_t)OL * OC * N * sizeof(float));
+
+    for (int n = 0; n < N; n++) {
+        for (int g = 0; g < groups; g++) {
+            for (int oc = 0; oc < OC_G; oc++) {
+                int oc_global = g * OC_G + oc;
+                for (int ol = 0; ol < OL; ol++) {
+                    float sum = 0.0f;
+                    for (int ic = 0; ic < IC_G; ic++) {
+                        for (int k = 0; k < K; k++) {
+                            int il = ol * stride + k - padding;
+                            if (il >= 0 && il < L) {
+                                int ic_global = g * IC_G + ic;
+                                // kernel: [K, IC_G, OC] -> k + ic * K + oc_global * (IC_G * K)
+                                float w = ggml_fp16_to_fp32(kernel[k + ic * K + oc_global * (IC_G * K)]);
+                                // input: [L, IC, N] -> il + ic_global * L + n * (IC * L)
+                                float x = input[il + ic_global * L + n * (IC * L)];
+                                sum += w * x;
+                            }
+                        }
+                    }
+                    // output: [OL, OC, N] -> ol + oc_global * OL + n * (OC * OL)
+                    output[ol + oc_global * OL + n * (OC * OL)] = sum;
+                }
+            }
+        }
+    }
+}
+
+static bool run_test(const char * label, int IC, int OC, int K, int L, int groups, int stride, int padding,
+                     enum ggml_type kernel_type = GGML_TYPE_F16, int N = 1) {
+    printf("  TEST: %s (IC=%d OC=%d K=%d L=%d N=%d G=%d s=%d p=%d) kernel=%s\n",
+           label, IC, OC, K, L, N, groups, stride, padding,
+           ggml_type_name(kernel_type));
+
+    int IC_G = IC / groups;
+    int OL = (L + 2 * padding - K) / stride + 1;
+
+    size_t ctx_size = 256 * 1024 * 1024;
+    struct ggml_init_params params = {
+        /*.mem_size   =*/ ctx_size,
+        /*.mem_buffer =*/ NULL,
+        /*.no_alloc   =*/ false,
+    };
+    struct ggml_context * ctx = ggml_init(params);
+
+    // generate kernel data
+    std::vector<float> kernel_f32(K * IC_G * OC);
+    fill_random_f32(kernel_f32.data(), K * IC_G * OC);
+
+    // kernel for op: [K, IC_G, OC] and reference (F16)
+    struct ggml_tensor * a = ggml_new_tensor_3d(ctx, kernel_type, K, IC_G, OC);
+    std::vector<ggml_fp16_t> kernel_f16(K * IC_G * OC);
+    for (int i = 0; i < K * IC_G * OC; i++) {
+        if (kernel_type == GGML_TYPE_BF16) {
+            ggml_bf16_t b = ggml_fp32_to_bf16(kernel_f32[i]);
+            ((ggml_bf16_t *)a->data)[i] = b;
+            kernel_f16[i] = ggml_fp32_to_fp16(ggml_bf16_to_fp32(b));
+        } else {
+            kernel_f16[i] = ggml_fp32_to_fp16(kernel_f32[i]);
+        }
+    }
+    if (kernel_type != GGML_TYPE_BF16) {
+        memcpy(a->data, kernel_f16.data(), K * IC_G * OC * sizeof(ggml_fp16_t));
+    }
+
+    // generate reference input (F32)
+    std::vector<float> input_f32(L * IC * N);
+    fill_random_f32(input_f32.data(), L * IC * N);
+
+    // input for op: [L, IC, N]
+    struct ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, L, IC, N);
+    memcpy(b->data, input_f32.data(), L * IC * N * sizeof(float));
+
+    // reference
+    std::vector<float> ref(OL * OC * N);
+    conv1d_grouped_ref(kernel_f16.data(), input_f32.data(), ref.data(),
+                       K, IC, OC, L, N, groups, stride, padding);
+
+    // ggml
+    struct ggml_tensor * result = ggml_conv_1d_grouped(ctx, a, b, stride, padding, 1, groups);
+
+    struct ggml_cgraph * gf = ggml_new_graph(ctx);
+    ggml_build_forward_expand(gf, result);
+
+    ggml_backend_t backend = ggml_backend_cpu_init();
+    ggml_backend_graph_compute(backend, gf);
+
+    bool ok = true;
+
+    if (result->ne[0] != OL || result->ne[1] != OC || result->ne[2] != N) {
+        fprintf(stderr, "    FAIL: shape [%lld, %lld, %lld], expected [%d, %d, %d]\n",
+                (long long)result->ne[0], (long long)result->ne[1], (long long)result->ne[2], OL, OC, N);
+        ok = false;
+    }
+
+    if (ok) {
+        ok = all_close((float *)result->data, ref.data(), OL * OC * N);
+    }
+
+    printf("    %s\n", ok ? "PASS" : "FAIL");
+
+    ggml_backend_free(backend);
+    ggml_free(ctx);
+    return ok;
+}
+
+int main(void) {
+    srand(42);
+
+    printf("Testing ggml_conv_1d_grouped\n\n");
+
+    int n_pass = 0, n_fail = 0;
+
+    struct { const char * label; int IC, OC, K, L, G, s, p; int N; } scenarios[] = {
+        { "groups=1 (standard conv1d)", 128, 256, 3, 32,  1, 1, 0, 1 },
+        { "ZAYA1-8B exact params",      1280, 1280, 2, 16,  10, 1, 0, 1 },
+        { "small 2 groups",             4, 4, 2, 8,    2, 1, 0, 1 },
+        { "with padding",               8, 8, 2, 16,   4, 1, 1, 1 },
+        { "IC != OC",                   12, 6, 3, 10,  3, 1, 0, 1 },
+        { "stride=2",                   8, 8, 2, 16,   4, 2, 0, 1 },
+        { "longer sequence",            1280, 1280, 2, 128, 10, 1, 0, 1 },
+        { "depthwise (groups==IC==OC)", 8, 8, 3, 16,   8, 1, 0, 1 },
+    };
+
+    enum ggml_type kernel_types[] = { GGML_TYPE_F16, GGML_TYPE_BF16 };
+    for (auto kt : kernel_types) {
+        if (kt != GGML_TYPE_F16) {
+            printf("\n--- %s ---\n\n", ggml_type_name(kt));
+        }
+        for (auto &s : scenarios) {
+            if (run_test(s.label, s.IC, s.OC, s.K, s.L, s.G, s.s, s.p, kt, s.N)) { n_pass++; } else { n_fail++; }
+        }
+    }
+
+    printf("\nResult: %d passed, %d failed\n", n_pass, n_fail);
+    return n_fail > 0 ? 1 : 0;
+}