S1-26: OUT_PROD_ID CPU kernel — the activation half of MUL_MAT_ID's backward

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

@@ -1849,6 +1849,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_out_prod_id_grp(params, tensor);
             } break;
+        case GGML_OP_OUT_PROD_ID:
+            {
+                ggml_compute_forward_out_prod_id(params, tensor);
+            } break;
         case GGML_OP_SCALE:
             {
                 ggml_compute_forward_scale(params, tensor);
@@ -2326,6 +2330,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_MUL_MAT_ID:
         case GGML_OP_OUT_PROD:
         case GGML_OP_OUT_PROD_ID_GRP:
+        case GGML_OP_OUT_PROD_ID:
             {
                 n_tasks = n_threads;
             } break;
@@ -2880,6 +2885,19 @@ struct ggml_cplan ggml_graph_plan(
                             cur = ggml_type_size(GGML_TYPE_F32) * node->src[0]->ne[0] * n_tasks;
                         }
                     } break;
+                case GGML_OP_OUT_PROD_ID:
+                    {
+                        // learning-llamas (S1-26): d(b) propagates back through the base expert
+                        // stack, which in a LoRA MoE graph is frozen and QUANTIZED. The kernel
+                        // dequantizes one row of `as` at a time into a per-thread F32 buffer.
+                        //
+                        // Sized unconditionally, not just for quantized `as`. The buffer costs one
+                        // row per thread, and making it conditional means a later change to the
+                        // kernel's fast path silently overruns the work buffer -- which surfaces as
+                        // a corrupted tensor several ops downstream, not as a crash here.
+                        cur = ggml_type_size(GGML_TYPE_F32) *
+                              (node->src[0]->ne[0] + CACHE_LINE_SIZE/sizeof(float)) * n_tasks;
+                    } break;
                 case GGML_OP_SOFT_MAX:
                 case GGML_OP_ROPE:
                 case GGML_OP_ROPE_BACK:

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

@@ -475,14 +475,14 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st
             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
-            // backend, gets scheduled, and then hits ggml_compute_forward's `default: GGML_ABORT`.
-            // The op would look implemented right up until it killed the process.
-            return false;
+            // learning-llamas (S1-26): d(b). `as` may be QUANTIZED -- in a LoRA MoE graph the base
+            // expert stacks are frozen Q4_K, and the activations flowing into them still carry a
+            // gradient because an earlier layer's LoRA is upstream. Same dequantize-a-row-at-a-time
+            // path as OUT_PROD, and the same set of types.
+            return (src0->type == GGML_TYPE_F32 || ggml_is_quantized(src0->type) ||
+                    src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16) &&
+                   src1->type == GGML_TYPE_F32 && op->src[2]->type == GGML_TYPE_I32 &&
+                   op->type == GGML_TYPE_F32;
         default:
             return true;
     }

ggml/src/ggml-cpu/ops.cpp

@@ -4552,6 +4552,12 @@ static void ggml_compute_forward_out_prod_id_grp_f32(
 
     GGML_ASSERT(dst->nb[0] == sizeof(float));
 
+    // b is the VECTOR operand of ggml_vec_mad_f32, so its dim-0 has to be contiguous -- a strided
+    // read is not merely wrong here, it is unrepresentable. src1 (grad) and src2 (ids) are read
+    // element-by-element through their strides instead, and must NOT be asserted contiguous: ggml
+    // hands this op transposed grads, and the forward accepts strided ids.
+    GGML_ASSERT(src0->nb[0] == sizeof(float));
+
     const int ith = params->ith;
     const int nth = params->nth;
 
@@ -4578,23 +4584,40 @@ static void ggml_compute_forward_out_prod_id_grp_f32(
         }
 
         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]);
+            const char * ids_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) {
+                // ids is read through nb[0], NOT as int32_t*[i]. The FORWARD reads it that way
+                // (ggml_compute_forward_mul_mat_id: `*(int32_t *)(ids->data + iid1*nb[1] + id*nb[0])`)
+                // and ggml_mul_mat_id puts no contiguity constraint on ids -- so a strided ids is a
+                // legal tensor that the forward computes CORRECTLY. A hard-coded 4-byte stride here
+                // would make the backward disagree with the forward about which expert each slot
+                // chose, and credit the wrong experts. Silently.
+                const int32_t e_i = *(const int32_t *) (ids_t + i*src2->nb[0]);
+                if (e_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]);
+                const char * g_col = (const char *) src1->data
+                                        + i*src1->nb[1] + t*src1->nb[2];
 
-                // d_as[:, j, e] += g_col[j] * b_col[:]
+                // d_as[:, j, e] += grad[j,i,t] * 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]);
+                    // grad is read through nb[0] too, and this is not paranoia: ggml's autodiff
+                    // hands transposed grads to out-prod-shaped ops as a matter of course. The
+                    // MUL_MAT backward passes ggml_transpose(grad) to ggml_out_prod, which is why
+                    // ggml_compute_forward_out_prod_f32 reads src1 through `i1*nb10` rather than
+                    // indexing a float*. This kernel is modelled on that one and had dropped it.
+                    //
+                    // A TRANSPOSE node reaching here has nb[0] == 8, and a float*[j] read of it
+                    // silently returns the wrong element. No assert fires; the run just trains on a
+                    // wrong weight gradient.
+                    const float gj = *(const float *) (g_col + j*src1->nb[0]);
+                    ggml_vec_mad_f32(n, (float *) ((char *) d_e + j*dst->nb[1]), b_col, gj);
                 }
             }
         }
@@ -4607,6 +4630,118 @@ void ggml_compute_forward_out_prod_id_grp(
     ggml_compute_forward_out_prod_id_grp_f32(params, dst);
 }
 
+// ggml_compute_forward_out_prod_id  (learning-llamas, S1-26)
+//
+// The ACTIVATION half of MUL_MAT_ID's backward: d(b).
+//
+//   forward   dst[j,i,t] = sum_k as[k,j, ids[i,t]] * b[k, i % ne_b1, t]
+//   this op   d_b[k,l,t] = sum_{i : i % ne_b1 == l} sum_j as[k,j, ids[i,t]] * grad[j,i,t]
+//
+//   src0 = as    [n, m, n_expert]     the expert stack -- and this one may be QUANTIZED
+//   src1 = grad  [m, n_ids, n_tok]
+//   src2 = ids   [n_ids, n_tok] I32
+//   dst          [n, ne_b1, n_tok]
+//
+// `as` being quantized is the whole difficulty, and it is not a corner case. In a LoRA MoE graph
+// the base expert stacks (ffn_up_exps and friends) are frozen Q4_K -- they receive no gradient
+// themselves, but the activations flowing INTO them do, because an earlier layer's LoRA is
+// upstream. So d(b) has to propagate back through a quantized weight, exactly as OUT_PROD does
+// for the dense case: dequantize one row of `as` at a time into a per-thread F32 buffer and
+// accumulate in F32. Accumulation stays F32 on a gradient path, per ADR-0002.
+//
+// Threaded BY TOKEN. Each token owns its own dst columns, so -- as in OUT_PROD_ID_GRP -- there is
+// exactly one writer per output element, no atomics, no barrier, and the summation order within a
+// token is fixed by the (slot, row) loops rather than by thread arrival. Note the slots of one
+// token can share a dst column (that is precisely what ne_b1 == 1 means, and it is what
+// build_moe_ffn produces), so this accumulation is real and its order matters.
+static void ggml_compute_forward_out_prod_id_f32(
+        const ggml_compute_params * params,
+              ggml_tensor * dst) {
+
+    const ggml_tensor * src0 = dst->src[0];  // as
+    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(src1->type == GGML_TYPE_F32);
+    GGML_ASSERT(src2->type == GGML_TYPE_I32);
+    GGML_ASSERT(dst->nb[0] == sizeof(float));
+
+    const ggml_type type = src0->type;
+    ggml_to_float_t const to_float = ggml_get_type_traits(type)->to_float;
+
+    const bool as_is_f32 = (type == GGML_TYPE_F32);
+    GGML_ASSERT(as_is_f32 || to_float != NULL);
+
+    // A row of `as` must be contiguous for a single to_float call to make sense.
+    GGML_ASSERT(src0->nb[0] == ggml_type_size(type));
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    const int64_t n        = src0->ne[0];  // == dst->ne[0]
+    const int64_t m        = src0->ne[1];  // == grad->ne[0]
+    const int64_t n_expert = src0->ne[2];
+    const int64_t ne_b1    = dst->ne[1];
+    const int64_t n_tok    = dst->ne[2];
+    const int64_t n_ids    = src2->ne[0];
+
+    GGML_ASSERT(src1->ne[0] == m);
+    GGML_ASSERT(src1->ne[1] == n_ids);
+    GGML_ASSERT(src1->ne[2] == n_tok);
+    GGML_ASSERT(src2->ne[1] == n_tok);
+    GGML_ASSERT(n_ids % ne_b1 == 0);
+
+    // One dequantized row per thread. Sized in ggml_graph_plan -- miss it and this writes past the
+    // end of the work buffer, which is the kind of bug that shows up as a corrupted tensor three
+    // ops downstream.
+    float * wdata = (float *) params->wdata + (n + CACHE_LINE_SIZE_F32) * ith;
+
+    for (int64_t t = ith; t < n_tok; t += nth) {
+        // Zero this token's columns. Slots that route nowhere still need a defined gradient.
+        for (int64_t l = 0; l < ne_b1; ++l) {
+            ggml_vec_set_f32(n, (float *) ((char *) dst->data + l*dst->nb[1] + t*dst->nb[2]), 0.0f);
+        }
+
+        const char * ids_t = (const char *) src2->data + t*src2->nb[1];
+
+        for (int64_t i = 0; i < n_ids; ++i) {
+            // ids and grad are both read through nb[0] -- see the long comment in
+            // ggml_compute_forward_out_prod_id_grp_f32. The forward reads ids strided, and ggml's
+            // autodiff routinely produces transposed grads; a hard-coded 4-byte stride on either
+            // is a silent wrong answer, not a crash.
+            const int32_t e = *(const int32_t *) (ids_t + i*src2->nb[0]);
+            GGML_ASSERT(e >= 0 && e < n_expert);
+
+            // The forward broadcasts b's column across slots when ne_b1 < n_ids, so several slots
+            // of this token accumulate into the SAME dst column. That is the ne_b1 == 1 case.
+            float * d_col = (float *) ((char *) dst->data + (i % ne_b1)*dst->nb[1] + t*dst->nb[2]);
+            const char * g_col = (const char *) src1->data + i*src1->nb[1] + t*src1->nb[2];
+
+            for (int64_t j = 0; j < m; ++j) {
+                const char * as_row = (const char *) src0->data + j*src0->nb[1] + e*src0->nb[2];
+
+                const float * a;
+                if (as_is_f32) {
+                    a = (const float *) as_row;
+                } else {
+                    to_float((const void *) as_row, wdata, n);
+                    a = wdata;
+                }
+
+                const float gj = *(const float *) (g_col + j*src1->nb[0]);
+                ggml_vec_mad_f32(n, d_col, a, gj);
+            }
+        }
+    }
+}
+
+void ggml_compute_forward_out_prod_id(
+        const ggml_compute_params * params,
+              ggml_tensor * dst) {
+    ggml_compute_forward_out_prod_id_f32(params, dst);
+}
+
 // ggml_compute_forward_scale
 
 static void ggml_compute_forward_scale_f32(

ggml/src/ggml-cpu/ops.h

@@ -50,6 +50,7 @@ void ggml_compute_forward_group_norm(const struct ggml_compute_params * params,
 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_out_prod_id(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);