diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index a1521366835..6b60d51c3ec 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2024,6 +2024,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { ggml_compute_forward_ssm_conv(params, tensor); } break; + case GGML_OP_SSM_CONV_BACK: + { + ggml_compute_forward_ssm_conv_back(params, tensor); + } break; case GGML_OP_SSM_SCAN: { ggml_compute_forward_ssm_scan(params, tensor); @@ -2404,6 +2408,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_FLASH_ATTN_EXT: case GGML_OP_FLASH_ATTN_BACK: case GGML_OP_SSM_CONV: + case GGML_OP_SSM_CONV_BACK: case GGML_OP_SSM_SCAN: { n_tasks = n_threads; diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 2355b520200..aab67118102 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -475,9 +475,13 @@ 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_SSM_CONV_BACK: + // learning-llamas (S1-30): d(sx) for the depthwise causal convolution. F32 throughout -- + // the conv weight is F32 in every Mamba GGUF, and a gradient is F32 by policy. + return src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && + op->src[2]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; case GGML_OP_SSM_SCAN_BACK: - // learning-llamas (S1-29b): declared, not yet implemented. The kernels land in S1-30 - // and S1-31, which flip this to a real check. + // learning-llamas (S1-29b): declared, not yet implemented. The kernel lands in S1-31, + // which flips this to a real check. // // This case is NOT redundant. The default below returns TRUE, so a new op with no // dispatch case is reported *supported* by the CPU backend, gets scheduled, and then diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index e47a44fb374..842cac30103 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -9946,6 +9946,90 @@ void ggml_compute_forward_flash_attn_back( // ggml_compute_forward_ssm_conv +// ggml_compute_forward_ssm_conv_back (learning-llamas, S1-30) +// +// The forward is a depthwise causal convolution over a sliding window: +// +// y[i1, i2, i3] = sum_{i0 < d_conv} sx[i2 + i0, i1, i3] * c[i0, i1] +// +// so the gradient w.r.t. the window SCATTERS each output's gradient back across the d_conv inputs +// that produced it: +// +// d_sx[i2 + i0, i1, i3] += dy[i1, i2, i3] * c[i0, i1] +// +// Written as a scatter rather than the equivalent gather (`d_sx[j] = sum over t of dy[t]*c[j-t]`, +// with its two-sided bounds on t) because the scatter needs no boundary arithmetic at all: every +// (i2, i0) pair lands in range by construction. The leading d_conv - 1 columns of sx are the +// carried convolution state, and they receive a gradient like any other input -- dropping them +// would silently truncate the gradient at every sequence boundary. +// +// Threaded by (row, sequence): each (i1, i3) owns its own d_sx column, so there are no atomics, no +// barrier, and the accumulation order within a column is fixed by the i2/i0 loops rather than by +// thread arrival. Deterministic by construction (ADR-0002). +static void ggml_compute_forward_ssm_conv_back_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; // dy [d_inner, n_t, n_s] + const ggml_tensor * src1 = dst->src[1]; // sx -- shape only + const ggml_tensor * src2 = dst->src[2]; // c [d_conv, d_inner] + + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src2->type == GGML_TYPE_F32); + GGML_ASSERT(dst->nb[0] == sizeof(float)); + GGML_ASSERT(src2->nb[0] == sizeof(float)); + + const int ith = params->ith; + const int nth = params->nth; + + const int64_t nc = src2->ne[0]; // d_conv + const int64_t ncs = src1->ne[0]; // d_conv - 1 + n_t + const int64_t d_inner = src1->ne[1]; + const int64_t n_t = src0->ne[1]; + const int64_t n_s = src1->ne[2]; + + GGML_ASSERT(dst->ne[0] == ncs); + GGML_ASSERT(dst->ne[1] == d_inner); + GGML_ASSERT(dst->ne[2] == n_s); + GGML_ASSERT(src0->ne[0] == d_inner); + GGML_ASSERT(src0->ne[2] == n_s); + GGML_ASSERT(src2->ne[1] == d_inner); + GGML_ASSERT(ncs == nc - 1 + n_t); + + // rows per thread, over d_inner + const int64_t dr = (d_inner + nth - 1)/nth; + const int64_t ir0 = dr*ith; + const int64_t ir1 = MIN(ir0 + dr, d_inner); + + for (int64_t i3 = 0; i3 < n_s; ++i3) { + for (int64_t i1 = ir0; i1 < ir1; ++i1) { + float * d = (float *) ((char *) dst->data + i1*dst->nb[1] + i3*dst->nb[2]); + + // Zeroed here, not by a memset over the whole tensor: each thread owns exactly the + // columns it is about to write, so no barrier is needed between the clear and the + // accumulate. + ggml_vec_set_f32(ncs, d, 0.0f); + + const float * c = (const float *) ((const char *) src2->data + i1*src2->nb[1]); + + for (int64_t i2 = 0; i2 < n_t; ++i2) { + const float g = *(const float *) ((const char *) src0->data + + i1*src0->nb[0] + i2*src0->nb[1] + i3*src0->nb[2]); + + // d[i2 .. i2 + nc) += g * c[0 .. nc) + ggml_vec_mad_f32(nc, d + i2, c, g); + } + } + } +} + +void ggml_compute_forward_ssm_conv_back( + const ggml_compute_params * params, + ggml_tensor * dst) { + ggml_compute_forward_ssm_conv_back_f32(params, dst); +} + static void ggml_compute_forward_ssm_conv_f32( const ggml_compute_params * params, ggml_tensor * dst) { diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index d6f3bb1ff3b..96ea3da936a 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -96,6 +96,7 @@ void ggml_compute_forward_flash_attn_back( const bool masked, struct ggml_tensor * dst); void ggml_compute_forward_ssm_conv(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_ssm_conv_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_ssm_scan(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_win_part(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_win_unpart(const struct ggml_compute_params * params, struct ggml_tensor * dst); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 2f8fdc8abd3..8864172ac97 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4204,6 +4204,41 @@ struct test_ssm_conv : public test_case { std::array ne_b = {3, 3, 1, 1}) : type(type), ne_a(ne_a), ne_b(ne_b) {} + // Under MODE_GRAD's default sum(out) objective THIS OP'S GRADIENT TEST IS VACUOUS. + // + // sum(out) makes the incoming gradient all-ones, and the kernel's scatter is + // + // d_sx[i2 + i0] += dy[i1,i2,i3] * c[i0,i1] + // + // so with dy == 1 a kernel that IGNORES dy entirely and scatters c alone produces exactly the + // same answer. Measured: that mutation is not caught at all under sum(out), and is caught at + // MAA 0.47 with a weighted objective. Same trap, same fix, as SOFT_MAX (S1-34) and MUL_MAT_ID + // (S1-27) -- 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)); + } + + // MEASURED (S1-30), with numbers on both sides: + // + // worst FD noise, 25 runs 2.2e-3 (on the widest shapes, d_inner = 1024/2048) + // ignore the incoming gradient 0.84 -- and INVISIBLE without grad_loss, see above + // shift the scatter window by 1 0.52 + // drop the last conv tap 0.47 + // + // 1e-2 sits 4.5x above the noise and 47-84x below every real defect. Three mutations injected, + // three caught. + // + // The kernel is ALSO checked exactly: against a naive DOUBLE reference, and with a TRANSPOSED + // grad (nb[0] != 4) -- the stride trap that bit both MoE kernels and GLU_BACK. Exact to ~1e-7 + // in both layouts. + double max_maa_err() override { + return 1e-2; + } + ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne_a.data()); ggml_set_name(a, "sx");