metal: GATED_DELTA_NET_BACK — the gated delta rule backward on Apple silicon (stacked on #5; card 61ab9fb8, S3b)

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

@@ -692,6 +692,25 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv(ggml_metal_
     return res;
 }
 
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net_back(ggml_metal_library_t lib, const ggml_tensor * op) {
+    char name[256];
+    snprintf(name, 256, "kernel_gated_delta_net_back_f32_%d", (int) op->src[2]->ne[0]);
+    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_gated_delta_net_back_zero(ggml_metal_library_t lib) {
+    const char * name = "kernel_gated_delta_net_back_zero";
+    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_gated_delta_net(ggml_metal_library_t lib, const ggml_tensor * op) {
     char base[256];
     char name[256];

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

@@ -135,6 +135,8 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan_ssd_mma  (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv              (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net   (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net_back(ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net_back_zero(ggml_metal_library_t lib);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_solve_tri         (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_ext        (ggml_metal_library_t lib, const struct ggml_tensor * op, int nsg, int nxpsg, int r1ptg);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm            (ggml_metal_library_t lib, const struct ggml_tensor * op);

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

@@ -1736,6 +1736,8 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
             return true;
         case GGML_OP_GATED_DELTA_NET:
             return has_simdgroup_reduction && op->src[2]->ne[0] % 32 == 0;
+        case GGML_OP_GATED_DELTA_NET_BACK:
+            return has_simdgroup_reduction && ggml_metal_op_gated_delta_net_back_supported(op);
         case GGML_OP_SOLVE_TRI:
         case GGML_OP_MUL_MAT:
         case GGML_OP_MUL_MAT_ID:

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

@@ -991,6 +991,25 @@ typedef struct {
     uint64_t nb3;
 } ggml_metal_kargs_gated_delta_net;
 
+// GATED_DELTA_NET_BACK: element strides (floats) of q/k/v, the shape, and the segment size
+typedef struct {
+    int64_t H;
+    int64_t T;
+    int64_t n_seqs;
+    int64_t neq1;
+    int64_t neq3;
+    int64_t nek1;
+    int64_t nek3;
+    int64_t sq1, sq2, sq3;
+    int64_t sk1, sk2, sk3;
+    int64_t sv1, sv2, sv3;
+    int64_t C;       // checkpoint segment length (~sqrt(T))
+    int64_t off_dk, off_dv, off_dg, off_db, off_ds; // float offsets of the packed output blocks
+    int32_t kda;
+    int32_t K;
+    float   scale;
+} ggml_metal_kargs_gated_delta_net_back;
+
 typedef struct {
     int32_t  ne00;
     int32_t  ne01;

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

@@ -369,6 +369,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) {
             {
                 n_fuse = ggml_metal_op_rwkv(ctx, idx);
             } break;
+        case GGML_OP_GATED_DELTA_NET_BACK:
+            {
+                n_fuse = ggml_metal_op_gated_delta_net_back(ctx, idx);
+            } break;
         case GGML_OP_GATED_DELTA_NET:
             {
                 n_fuse = ggml_metal_op_gated_delta_net(ctx, idx);
@@ -1870,6 +1874,115 @@ int ggml_metal_op_rwkv(ggml_metal_op_t ctx, int idx) {
     return 1;
 }
 
+// GATED_DELTA_NET_BACK: the checkpoint length and the scratch after dst (see the kernel)
+static int64_t ggml_metal_gdn_back_C(int64_t T) {
+    int64_t C = 1;
+    while (C*C < T) {
+        ++C;
+    }
+    return C;
+}
+
+bool ggml_metal_op_gated_delta_net_back_supported(const ggml_tensor * dst) {
+    for (int s = 0; s < 7; ++s) {
+        if (dst->src[s] == nullptr || dst->src[s]->type != GGML_TYPE_F32) {
+            return false;
+        }
+    }
+    const ggml_tensor * q = dst->src[0], * k = dst->src[1], * v = dst->src[2],
+                      * g = dst->src[3], * b = dst->src[4], * st = dst->src[5], * gr = dst->src[6];
+    const int64_t S_v = v->ne[0];
+    if (S_v != 32 && S_v != 64 && S_v != 128) {
+        return false;
+    }
+    if (q->ne[0] != S_v || k->ne[0] != S_v || !(g->ne[0] == 1 || g->ne[0] == S_v)) {
+        return false;
+    }
+    if (v->ne[3] % q->ne[3] != 0 || v->ne[3] % k->ne[3] != 0 || v->ne[1] % q->ne[1] != 0 || v->ne[1] % k->ne[1] != 0) {
+        return false;
+    }
+    return dst->type == GGML_TYPE_F32 && ggml_is_contiguous(dst) &&
+           ggml_is_contiguous_rows(q) && ggml_is_contiguous_rows(k) && ggml_is_contiguous_rows(v) &&
+           ggml_is_contiguous(g) && ggml_is_contiguous(b) && ggml_is_contiguous(st) && ggml_is_contiguous(gr);
+}
+
+size_t ggml_metal_op_gated_delta_net_back_extra(const ggml_tensor * dst) {
+    const ggml_tensor * v = dst->src[2];
+    const int64_t S_v = v->ne[0], H = v->ne[1], T = v->ne[2], n_seqs = v->ne[3];
+    const int64_t C = ggml_metal_gdn_back_C(T);
+    const int64_t n_ck = (T + C - 1)/C;
+    return (size_t) (n_seqs*H*n_ck*S_v*S_v + n_seqs*H*(C + 3)*S_v*S_v)*sizeof(float) + 256;
+}
+
+int ggml_metal_op_gated_delta_net_back(ggml_metal_op_t ctx, int idx) {
+    ggml_tensor * op = ctx->node(idx);
+    ggml_metal_library_t lib = ctx->lib;
+    ggml_metal_encoder_t enc = ctx->enc;
+
+    GGML_ASSERT(ggml_metal_op_gated_delta_net_back_supported(op));
+
+    const ggml_tensor * q = op->src[0], * k = op->src[1], * v = op->src[2], * g = op->src[3], * b = op->src[4];
+
+    const int64_t S_v = v->ne[0], H = v->ne[1], T = v->ne[2], n_seqs = v->ne[3];
+    const int64_t C    = ggml_metal_gdn_back_C(T);
+    const int64_t n_ck = (T + C - 1)/C;
+
+    ggml_metal_kargs_gated_delta_net_back args = {
+        /*.H      =*/ H,
+        /*.T      =*/ T,
+        /*.n_seqs =*/ n_seqs,
+        /*.neq1   =*/ q->ne[1],
+        /*.neq3   =*/ q->ne[3],
+        /*.nek1   =*/ k->ne[1],
+        /*.nek3   =*/ k->ne[3],
+        /*.sq1    =*/ (int64_t) (q->nb[1]/sizeof(float)), (int64_t) (q->nb[2]/sizeof(float)), (int64_t) (q->nb[3]/sizeof(float)),
+        /*.sk1    =*/ (int64_t) (k->nb[1]/sizeof(float)), (int64_t) (k->nb[2]/sizeof(float)), (int64_t) (k->nb[3]/sizeof(float)),
+        /*.sv1    =*/ (int64_t) (v->nb[1]/sizeof(float)), (int64_t) (v->nb[2]/sizeof(float)), (int64_t) (v->nb[3]/sizeof(float)),
+        /*.C      =*/ C,
+        /*.off_dk =*/ ggml_nelements(q),
+        /*.off_dv =*/ ggml_nelements(q) + ggml_nelements(k),
+        /*.off_dg =*/ ggml_nelements(q) + ggml_nelements(k) + ggml_nelements(v),
+        /*.off_db =*/ ggml_nelements(q) + ggml_nelements(k) + ggml_nelements(v) + ggml_nelements(g),
+        /*.off_ds =*/ ggml_nelements(q) + ggml_nelements(k) + ggml_nelements(v) + ggml_nelements(g) + ggml_nelements(b),
+        /*.kda    =*/ g->ne[0] == S_v ? 1 : 0,
+        /*.K      =*/ ggml_get_op_params_i32(op, 0),
+        /*.scale  =*/ 1.0f/sqrtf((float) S_v),
+    };
+
+    ggml_metal_buffer_id bid_dst  = ggml_metal_get_buffer_id(op);
+    ggml_metal_buffer_id bid_ckpt = bid_dst;
+    bid_ckpt.offs += GGML_PAD(ggml_nbytes(op), 16);
+    ggml_metal_buffer_id bid_seg  = bid_ckpt;
+    bid_seg.offs  += (size_t) n_seqs*H*n_ck*S_v*S_v*sizeof(float);
+
+    // dq/dk accumulate across broadcast heads: zero the packed output first
+    {
+        const uint64_t n = (uint64_t) ggml_nelements(op);
+        auto pz = ggml_metal_library_get_pipeline_gated_delta_net_back_zero(lib);
+        ggml_metal_encoder_set_pipeline(enc, pz);
+        ggml_metal_encoder_set_buffer  (enc, bid_dst, 0);
+        ggml_metal_encoder_set_bytes   (enc, (void *) &n, sizeof(n), 1);
+        const int nth = 256;
+        ggml_metal_encoder_dispatch_threadgroups(enc, (int) ((n + nth - 1)/nth), 1, 1, nth, 1, 1);
+        ggml_metal_op_concurrency_reset(ctx);
+    }
+
+    auto pipeline = ggml_metal_library_get_pipeline_gated_delta_net_back(lib, op);
+    int ida = 0;
+    ggml_metal_encoder_set_pipeline(enc, pipeline);
+    ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), ida++);
+    for (int s = 0; s < 7; ++s) {
+        ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[s]), ida++);
+    }
+    ggml_metal_encoder_set_buffer  (enc, bid_dst,  ida++);
+    ggml_metal_encoder_set_buffer  (enc, bid_ckpt, ida++);
+    ggml_metal_encoder_set_buffer  (enc, bid_seg,  ida++);
+    ggml_metal_encoder_dispatch_threadgroups(enc, (int) H, (int) n_seqs, 1, (int) S_v, 1, 1);
+    ggml_metal_op_concurrency_reset(ctx);
+
+    return 1;
+}
+
 int ggml_metal_op_gated_delta_net(ggml_metal_op_t ctx, int idx) {
     ggml_tensor * op = ctx->node(idx);
 

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

@@ -62,6 +62,9 @@ int ggml_metal_op_ssm_conv          (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_ssm_scan          (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_rwkv              (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_gated_delta_net   (ggml_metal_op_t ctx, int idx);
+int    ggml_metal_op_gated_delta_net_back(ggml_metal_op_t ctx, int idx);
+bool   ggml_metal_op_gated_delta_net_back_supported(const struct ggml_tensor * op);
+size_t ggml_metal_op_gated_delta_net_back_extra(const struct ggml_tensor * op);
 int ggml_metal_op_solve_tri         (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_set               (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_cpy               (ggml_metal_op_t ctx, int idx);

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

@@ -237,6 +237,10 @@ static size_t ggml_backend_metal_buffer_type_get_alloc_size(ggml_backend_buffer_
             {
                 res += ggml_metal_op_out_prod_extra(tensor);
             } break;
+        case GGML_OP_GATED_DELTA_NET_BACK:
+            {
+                res += ggml_metal_op_gated_delta_net_back_extra(tensor);
+            } break;
         case GGML_OP_CROSS_ENTROPY_LOSS:
             {
                 res += ggml_nrows(tensor->src[0])*sizeof(float); // the per-row losses before the sum