Skip to content

S1-30: SSM_CONV_BACK CPU kernel#24

Open
dillon-blake wants to merge 1 commit into
ticket/S1-29b-ssm-back-opsfrom
ticket/S1-30-ssm-conv-back
Open

S1-30: SSM_CONV_BACK CPU kernel#24
dillon-blake wants to merge 1 commit into
ticket/S1-29b-ssm-back-opsfrom
ticket/S1-30-ssm-conv-back

Conversation

@dillon-blake

@dillon-blake dillon-blake commented Jul 14, 2026

Copy link
Copy Markdown
Owner

Stacked on #23 (S1-29b).

d_sx[i2 + i0, i1, i3] += dy[i1, i2, i3] * c[i0, i1]

The forward is a depthwise causal convolution over a sliding window, so its gradient scatters each output's gradient back across the d_conv inputs that produced it. Written as a scatter rather than the equivalent gather (d_sx[j] = sum_t 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 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 — no atomics, no barrier, fixed accumulation order.

Under sum(out) this test was vacuous, and I only found out by mutating the kernel

sum(out) makes the incoming gradient all-ones, and the scatter is dy * c — so with dy == 1, a kernel that ignores dy entirely and scatters c alone produces exactly the same answer.

mutation + objective result
ignore the incoming gradient, default sum(out) not caught at all
ignore the incoming gradient, weighted grad_loss MAA 0.84

Third instance of the same structural trap (SOFT_MAX in S1-34, MUL_MAT_ID in S1-27). The weighted grad_loss hook is now the default assumption for any new VJP, not an afterthought.

Verification

max_maa_err = 1e-2, measured with numbers on both sides:

MAA
worst FD noise, 25 runs 2.2e-3 (on the widest shapes)
ignore the incoming gradient 0.84
shift the scatter window by 1 0.52
drop the last conv tap 0.47

Three mutations injected, three caught.

Also checked exactly — against a naive double-precision 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; this kernel reads dy through src0->nb[0] from the start.

54 grad cases now execute. They previously reported OK while requesting no gradients at all.

    d_sx[i2 + i0, i1, i3] += dy[i1, i2, i3] * c[i0, i1]

The forward is a depthwise causal convolution over a sliding window, so its gradient
scatters each output's gradient back across the d_conv inputs that produced it. 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 no atomics, no
barrier, and the accumulation order within a column is fixed by the loops rather than by
thread arrival.

UNDER sum(out) THIS TEST WAS VACUOUS, AND I ONLY FOUND OUT BY MUTATING THE KERNEL.

sum(out) makes the incoming gradient all-ones, and the scatter is dy * c -- so with dy == 1
a kernel that IGNORES dy entirely and scatters c alone produces exactly the same answer.
Measured:

    ignore the incoming gradient + sum(out)      NOT CAUGHT AT ALL
    ignore the incoming gradient + grad_loss     MAA 0.84

Third instance of the same structural trap (SOFT_MAX in S1-34, MUL_MAT_ID in S1-27). The
weighted grad_loss hook is now the default assumption for any new VJP, not an afterthought.

max_maa_err 1e-2, measured: noise floor 2.2e-3 on the widest shapes, three injected
mutations at 0.47-0.84, all three caught.

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; this kernel reads dy through src0->nb[0] from the start.

54 grad cases now EXECUTE. They previously reported OK while requesting no gradients at all.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant