S1-27: OUT_PROD_ID_GRP CPU kernel — the weight half of MUL_MAT_ID's backward

ggml/src/ggml-cpu/ggml-cpu.c

@@ -1845,6 +1845,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_out_prod(params, tensor);
             } break;
+        case GGML_OP_OUT_PROD_ID_GRP:
+            {
+                ggml_compute_forward_out_prod_id_grp(params, tensor);
+            } break;
         case GGML_OP_SCALE:
             {
                 ggml_compute_forward_scale(params, tensor);
@@ -2321,6 +2325,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_MUL_MAT:
         case GGML_OP_MUL_MAT_ID:
         case GGML_OP_OUT_PROD:
+        case GGML_OP_OUT_PROD_ID_GRP:
             {
                 n_tasks = n_threads;
             } break;

ggml/src/ggml-cpu/ggml-cpu.cpp

@@ -467,10 +467,16 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st
                     ((ggml_is_quantized(src0->type) || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16) &&
                      src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) &&
                 src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32;
-        case GGML_OP_OUT_PROD_ID:
         case GGML_OP_OUT_PROD_ID_GRP:
-            // learning-llamas (S1-25): declared, not yet implemented. The kernels land in S1-26 and
-            // S1-27, which flip this to a real check.
+            // learning-llamas (S1-27): d(as). All three operands are F32 on the training path --
+            // b is activations, grad is a gradient, and the expert stack it feeds is the F32 LoRA
+            // A/B (base experts are frozen). A quantized `as` would need a dequantizing variant,
+            // and there is no caller for one.
+            return src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 &&
+                   op->src[2]->type == GGML_TYPE_I32 && op->type == GGML_TYPE_F32;
+        case GGML_OP_OUT_PROD_ID:
+            // learning-llamas (S1-25): declared, not yet implemented. The kernel lands in S1-26,
+            // which flips this to a real check.
             //
             // This case is NOT redundant, and leaving it out is the trap. The default below returns
             // TRUE -- so a brand-new op with no dispatch case is reported *supported* by the CPU

ggml/src/ggml-cpu/ops.cpp

@@ -4512,6 +4512,101 @@ void ggml_compute_forward_out_prod(
     }
 }
 
+// ggml_compute_forward_out_prod_id_grp  (learning-llamas, S1-27)
+//
+// The WEIGHT half of MUL_MAT_ID's backward: d(as).
+//
+//   forward   dst[j,i,t] = sum_k as[k,j, ids[i,t]] * b[k, i % ne_b1, t]
+//   this op   d_as[k,j,e] = sum_{(i,t) : ids[i,t] == e}  b[k, i % ne_b1, t] * grad[j,i,t]
+//
+//   src0 = b     [n,  ne_b1, n_tok]   activations
+//   src1 = grad  [m,  n_ids, n_tok]   incoming gradient
+//   src2 = ids   [n_ids, n_tok]  I32  which expert each (slot, token) routed to
+//   dst          [n,  m,     n_expert]
+//
+// "grp" is for GROUPED: it reduces over every (slot, token) that routed to a given expert. That
+// reduction is the whole difficulty. An expert chosen by many tokens accumulates all of them, and
+// a token may route to several experts -- so this is a scatter-add, not a permutation.
+//
+// It is threaded BY EXPERT, and that is deliberate rather than incidental: each expert's output
+// slice is written by exactly one thread, so the accumulation needs no atomics and no reduction
+// barrier, and the summation order within a slice is fixed by the (t, i) loop rather than by which
+// thread got there first. Determinism is a hard requirement on a gradient path (ADR-0002) -- a
+// float sum whose order depends on thread scheduling makes a training run unreproducible, and it
+// does so silently.
+//
+// Experts with no tokens routed to them are not skipped: their slice is still zeroed, because it
+// is a gradient and the optimizer will read it whether or not anything routed there this step.
+static void ggml_compute_forward_out_prod_id_grp_f32(
+        const ggml_compute_params * params,
+              ggml_tensor * dst) {
+
+    const ggml_tensor * src0 = dst->src[0];  // b
+    const ggml_tensor * src1 = dst->src[1];  // grad
+    const ggml_tensor * src2 = dst->src[2];  // ids
+
+    GGML_ASSERT(dst->type  == GGML_TYPE_F32);
+    GGML_ASSERT(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT(src1->type == GGML_TYPE_F32);
+    GGML_ASSERT(src2->type == GGML_TYPE_I32);
+
+    GGML_ASSERT(dst->nb[0] == sizeof(float));
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    const int64_t n        = src0->ne[0];  // b's row length  == dst->ne[0]
+    const int64_t ne_b1    = src0->ne[1];  // b's broadcast dim
+    const int64_t n_tok    = src0->ne[2];
+    const int64_t m        = src1->ne[0];  // grad's row length == dst->ne[1]
+    const int64_t n_ids    = src1->ne[1];
+    const int64_t n_expert = dst->ne[2];
+
+    GGML_ASSERT(dst->ne[0] == n);
+    GGML_ASSERT(dst->ne[1] == m);
+    GGML_ASSERT(src1->ne[2] == n_tok);
+    GGML_ASSERT(src2->ne[0] == n_ids);
+    GGML_ASSERT(src2->ne[1] == n_tok);
+    GGML_ASSERT(n_ids % ne_b1 == 0);
+
+    for (int64_t e = ith; e < n_expert; e += nth) {
+        float * d_e = (float *) ((char *) dst->data + e*dst->nb[2]);
+
+        // Zero this expert's whole slice first -- including the experts nothing routes to.
+        for (int64_t j = 0; j < m; ++j) {
+            ggml_vec_set_f32(n, (float *) ((char *) d_e + j*dst->nb[1]), 0.0f);
+        }
+
+        for (int64_t t = 0; t < n_tok; ++t) {
+            const int32_t * ids_t = (const int32_t *) ((const char *) src2->data + t*src2->nb[1]);
+
+            for (int64_t i = 0; i < n_ids; ++i) {
+                if (ids_t[i] != (int32_t) e) {
+                    continue;
+                }
+
+                // The forward broadcasts b's columns across slots when ne_b1 < n_ids, so slot i
+                // reads column i % ne_b1 -- and the gradient must gather from the same column.
+                const float * b_col = (const float *) ((const char *) src0->data
+                                        + (i % ne_b1)*src0->nb[1] + t*src0->nb[2]);
+                const float * g_col = (const float *) ((const char *) src1->data
+                                        + i*src1->nb[1] + t*src1->nb[2]);
+
+                // d_as[:, j, e] += g_col[j] * b_col[:]
+                for (int64_t j = 0; j < m; ++j) {
+                    ggml_vec_mad_f32(n, (float *) ((char *) d_e + j*dst->nb[1]), b_col, g_col[j]);
+                }
+            }
+        }
+    }
+}
+
+void ggml_compute_forward_out_prod_id_grp(
+        const ggml_compute_params * params,
+              ggml_tensor * dst) {
+    ggml_compute_forward_out_prod_id_grp_f32(params, dst);
+}
+
 // ggml_compute_forward_scale
 
 static void ggml_compute_forward_scale_f32(

ggml/src/ggml-cpu/ops.h

@@ -49,6 +49,7 @@ void ggml_compute_forward_rms_norm_back(const struct ggml_compute_params * param
 void ggml_compute_forward_group_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_l2_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_out_prod(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_out_prod_id_grp(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_scale(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_set(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_cpy(const struct ggml_compute_params * params, struct ggml_tensor * dst);

tests/test-backend-ops.cpp

@@ -4313,9 +4313,12 @@ struct test_mul_mat_id : public test_case {
     const int64_t m;
     const int64_t n;
     const int64_t k;
+    // learning-llamas: bit 1 = ask for d(as) [OUT_PROD_ID_GRP], bit 2 = ask for d(b) [OUT_PROD_ID].
+    // 0 keeps the case eval-only, which is what every pre-existing instantiation wants.
+    const int grad_param;
 
     std::string vars() override {
-        return VARS_TO_STR8(type_a, type_b, n_mats, n_used, b, m, n, k);
+        return VARS_TO_STR9(type_a, type_b, n_mats, n_used, b, m, n, k, grad_param);
     }
 
     double max_nmse_err() override {
@@ -4337,29 +4340,51 @@ struct test_mul_mat_id : public test_case {
 
     test_mul_mat_id(ggml_type type_a = GGML_TYPE_F32, ggml_type type_b = GGML_TYPE_F32,
             int n_mats = 8, int n_used = 2, bool b = false,
-            int64_t m = 32, int64_t n = 32, int64_t k = 32)
+            int64_t m = 32, int64_t n = 32, int64_t k = 32, int grad_param = 0)
         : type_a(type_a), type_b(type_b), n_mats(n_mats), n_used(n_used), b(b),
-            m(m), n(n), k(k) {
+            m(m), n(n), k(k), grad_param(grad_param) {
             GGML_ASSERT(n_used <= n_mats);
         }
 
+    // MODE_GRAD's default objective is sum(out), and under it THIS OP'S WEIGHT GRADIENT IS
+    // VACUOUS. With an all-ones incoming gradient,
+    //
+    //     d_as[i,j,e] = sum_{(s,t): ids=e} b[i,s'] * 1
+    //
+    // is independent of j -- the result is constant along the entire output axis, so a kernel that
+    // ignored `grad` completely and just summed b-columns per expert would pass every case. A
+    // weighted sum makes the objective depend on grad's actual (j, slot, token) structure, which is
+    // the only thing that can catch a transposed or mis-strided read of it.
+    //
+    // Same trap, same fix as SOFT_MAX (S1-34): an op can be invisible to sum(out) for structural
+    // reasons, and then its grad test checks nothing while reporting OK.
+    ggml_tensor * grad_loss(ggml_context * ctx, ggml_tensor * out) override {
+        ggml_tensor * w = ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, out->ne);
+        ggml_set_name(w, "grad_loss_weights");
+
+        return ggml_sum(ctx, ggml_mul(ctx, out, w));
+    }
+
     ggml_tensor * build_graph(ggml_context * ctx) override {
         // C^T = A * B^T: (k, m) * (k, n) => (m, n)
         ggml_tensor * as = ggml_new_tensor_3d(ctx, type_a, k, m, n_mats);
         ggml_set_name(as, "as");
 
-        // learning-llamas (S1-25): ask for BOTH gradients.
+        // learning-llamas (S1-25/S1-27): WHICH gradients are asked for is a test parameter, and it
+        // has to be, because the two halves of MUL_MAT_ID's backward are separate ops landing in
+        // separate tickets:
         //
-        // `as` is not an exotic param. build_lora_mm_id computes
-        // mul_mat_id(B, mul_mat_id(A, cur, ids), ids), so the trainable LoRA A/B tensors ARE the
-        // 3D expert operand -- LoRA-only MoE training needs the weight-grad half, and an
-        // "activations only" backward would silently train nothing at all.
+        //   grad_param & 1  ->  as needs grads  ->  emits OUT_PROD_ID_GRP  (S1-27, implemented)
+        //   grad_param & 2  ->  b  needs grads  ->  emits OUT_PROD_ID      (S1-26, still pending)
         //
-        // These cases build a backward graph containing OUT_PROD_ID / OUT_PROD_ID_GRP, which no
-        // backend supports yet (the kernels are S1-26 / S1-27). They therefore register and report
-        // not-supported rather than executing -- which is the point: the wiring is exercised now,
-        // and the day a kernel lands these turn on with no test change.
-        if (type_a == GGML_TYPE_F32) {
+        // An `as`-only case therefore exercises S1-27's kernel ALONE. Asking for both would make
+        // every case depend on the op that does not exist yet, and S1-27 would have no green test
+        // to stand on.
+        //
+        // `as` is not an exotic param: build_lora_mm_id computes
+        // mul_mat_id(B, mul_mat_id(A, cur, ids), ids), so the trainable LoRA A/B tensors ARE the
+        // 3D expert operand. LoRA-only MoE training needs the weight-grad half.
+        if ((grad_param & 1) && type_a == GGML_TYPE_F32) {
             ggml_set_param(as);
         }
 
@@ -4372,7 +4397,7 @@ struct test_mul_mat_id : public test_case {
 
         ggml_tensor * b = ggml_new_tensor_3d(ctx, type_b, k, this->b ? 1 : n_used, n);
         ggml_set_name(b, "b");
-        if (type_b == GGML_TYPE_F32) {
+        if ((grad_param & 2) && type_b == GGML_TYPE_F32) {
             ggml_set_param(b);
         }
 
@@ -9154,6 +9179,31 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
                                                   256, 16, 16, {1, 1}, {nr2, 1}));
     }
 
+    // learning-llamas (S1-25/S1-27): MUL_MAT_ID gradient cases.
+    //
+    // Deliberately TINY. grad_nmax() is 10000 parameter elements, and every pre-existing
+    // test_mul_mat_id shape is far above it -- so they would be SILENTLY SKIPPED under MODE_GRAD
+    // and print OK. (That is trap #2 of ADR-0002, and it is why these shapes are 4x6x8 rather than
+    // anything realistic: 4*6*3 = 72 elements for `as`.)
+    //
+    // grad_param 1 = d(as) only -> exercises OUT_PROD_ID_GRP alone (S1-27, implemented).
+    // grad_param 2 = d(b)  only -> exercises OUT_PROD_ID alone      (S1-26, still pending, so these
+    //                              register and report not-supported rather than aborting).
+    // grad_param 3 = both.
+    //
+    // `b` toggles the forward's BROADCAST: false gives ne_b1 == n_used, true gives ne_b1 == 1.
+    // Both are live in the real LoRA MoE graph -- the inner mul_mat_id broadcasts, the outer does
+    // not -- and they are different index arithmetic in the kernel, so both are covered.
+    for (int grad_param : {1, 2, 3}) {
+        for (bool bcast : {false, true}) {
+            for (int n_used : {1, 2}) {
+                test_cases.emplace_back(new test_mul_mat_id(
+                    GGML_TYPE_F32, GGML_TYPE_F32, /*n_mats =*/ 3, n_used, bcast,
+                    /*m =*/ 6, /*n =*/ 5, /*k =*/ 4, grad_param));
+            }
+        }
+    }
+
     // add_id
     for (ggml_type type_a : {GGML_TYPE_F32}) {
         for (ggml_type type_b : {GGML_TYPE_F32}) {