S1-24: chunked attention — the kernel-free long-context path#26
Open
dillon-blake wants to merge 17 commits into
Open
S1-24: chunked attention — the kernel-free long-context path#26dillon-blake wants to merge 17 commits into
dillon-blake wants to merge 17 commits into
Conversation
MoE training is blocked by one missing backward case. MUL_MAT_ID has no case in
ggml_compute_backward and falls into the op-level default GGML_ABORT, so any MoE
graph dies the moment a gradient is requested.
Two new ops at the enum TAIL (rebase hygiene -- inserting renumbers every op after
it and conflicts across the whole backend matrix):
OUT_PROD_ID d(b) -- gather the expert matrix each (slot, token) used and
push the gradient back through it
OUT_PROD_ID_GRP d(as) -- outer product scattered into the expert slice each
(slot, token) selected, accumulated over all of them
Both are gather/scatter WITH accumulation, which is why neither is expressible in
existing ops: the expert axis is a selection, not a broadcast. Constructors and
shape contracts only -- the CPU kernels are S1-26 and S1-27.
The weight-grad half is not optional. 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 -- an "activations only" backward would silently train
nothing.
Two things the ticket got wrong, both found by building it:
- ggml_out_prod_id(as, grad, ids) is UNDER-DETERMINED. d(b) has b's shape, and b's
middle dim is not recoverable from as/grad/ids -- the forward broadcasts b's
columns across slots whenever ids->ne[0] is a multiple of it. Both modes are live
in the LoRA MoE graph: the inner mul_mat_id has ne_b1 == 1, the outer has
ne_b1 == n_used. So ne_b1 is passed, and the forward's broadcast rule asserted.
- CPU supports_op ends in `default: return true`. A new op with no dispatch case is
therefore reported SUPPORTED, gets scheduled, and hits ggml_compute_forward's
`default: GGML_ABORT` -- it would look implemented right up until it killed the
process. Both new ops now return false explicitly; S1-26/S1-27 flip them.
test-backend-ops: test_add_id and test_mul_mat_id never called ggml_set_param, so
`grad -o ADD_ID` and `grad -o MUL_MAT_ID` requested no gradients and printed OK
while having no backward at all. Both now ask for them.
ADD_ID 16 grad cases now EXECUTE and pass. Verified non-vacuous: scaling the
VJP by 2 turns 26 cases red.
MUL_MAT_ID 516 grad cases register and report not-supported rather than
aborting -- the wiring is exercised now, and they turn on with no test
change the day a kernel lands.
…backward
d_as[k,j,e] = sum_{(i,t) : ids[i,t] == e} b[k, i % ne_b1, t] * grad[j,i,t]
This is the op that makes MoE LoRA possible at all: build_lora_mm_id computes
mul_mat_id(B, mul_mat_id(A, cur, ids), ids), so the trainable A/B tensors ARE the 3D
expert operand, and without d(as) a LoRA MoE run trains nothing.
Threaded BY EXPERT, deliberately rather than incidentally. Each expert's output slice
is written by exactly one thread, so the scatter-add needs no atomics and no reduction
barrier, and the summation order inside a slice is fixed by the (token, slot) loop
rather than by which thread arrives first. Changing the thread count changes WHICH
thread owns an expert, never the order of the float sum within it. Verified
bit-identical at 1/2/4/8 threads on data chosen to make float addition order-sensitive
(magnitudes mixed 1e6 against 1e-6, where a reordered sum would show).
Experts nothing routes to still get their slice zeroed -- it is a gradient, and the
optimizer reads it whether or not anything routed there this step.
THE TEST WOULD HAVE PASSED ON A KERNEL THAT IGNORED THE GRADIENT ENTIRELY.
MODE_GRAD's default objective is sum(out), which makes the incoming gradient all-ones,
which makes
d_as[k,j,e] = sum_{(i,t): ids=e} b[k,i,t] * 1
independent of j -- constant along the whole output axis. So a kernel that never reads
`grad` and just sums b-columns per expert satisfies it. Measured, not argued:
poisoned kernel (grad ignored) + default sum(out) -> OK, passes
poisoned kernel (grad ignored) + weighted grad_loss -> FAIL, MAA 6.56 vs 1e-4
test_mul_mat_id therefore overrides grad_loss with a weighted sum -- the same hook, for
the same structural reason, that S1-34 had to add for SOFT_MAX.
test_mul_mat_id also takes a grad_param selector now, because the two halves of
MUL_MAT_ID's backward are separate ops in separate tickets: `as`-only cases exercise
this kernel ALONE, so S1-27 has a green test to stand on while OUT_PROD_ID (S1-26) is
still unimplemented and its cases report not-supported. Shapes are deliberately tiny
(4x6x8) -- grad_nmax() is 10000 elements and every pre-existing test_mul_mat_id shape
is far above it, so they would have been silently skipped.
Both forward broadcast modes are covered (ne_b1 == 1 and ne_b1 == n_used); both are
live in the real LoRA MoE graph and they are different index arithmetic in the kernel.
…backward
d_b[k,l,t] = sum_{i : i % ne_b1 == l} sum_j as[k,j, ids[i,t]] * grad[j,i,t]
`as` may be QUANTIZED, and that is the whole difficulty rather than a corner case. In a
LoRA MoE graph the base expert stacks are frozen Q4_K. They take 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. Same dequantize-a-row-at-a-time
path as OUT_PROD, accumulating in F32 (ADR-0002).
Threaded by token: one writer per output column, no atomics, no barrier.
MODE_GRAD CANNOT CHECK THE QUANTIZED PATH, and it took a while to work out why.
ggml's quantized matmul quantizes the ACTIVATIONS on the fly to vec_dot_type before the
dot product. So the forward is a STAIRCASE in b: its finite difference at the FD step
scale measures quantization edges, not a derivative. The analytic gradient is the
gradient of the intended smooth function; the harness evaluates the actual quantized one.
They disagree by construction (measured MAA 0.028-0.071 -- and the kernel is exact).
Upstream reached the same conclusion and encoded it without a word of comment:
// test_mul_mat
if (!ggml_is_quantized(type_a)) {
ggml_set_param(b); // b is a param ONLY when the weight is not quantized
}
So the quantized cases are deliberately absent from MODE_GRAD, and the dequantize path is
verified a stronger way instead: OUT_PROD_ID with a quantized `as` against OUT_PROD_ID with
that same `as` dequantized to F32. BIT-EXACT on q8_0, q4_K and q4_0.
TWO TEST FIXES, both found by measuring rather than assuming.
1. The token count is 32, not 5. mean_abs_asymm divides by (gn + ga), NOT by
(|gn| + |ga|), so any output element whose true gradient is near zero sends the ratio
to infinity. This op manufactures such elements: d_as[k,j,e] sums only over the slots
that routed to expert e, so with 5 tokens and 3 experts it is a sum of ONE OR TWO
random products. One element with asymm ~700 drags the mean of 90 to ~8.
5 tokens -> 3/10 runs fail, MAA up to 7.9 (kernel exact throughout)
32 tokens -> 0/30 runs fail
Condition the test; do not widen the tolerance until the flapping stops.
2. max_maa_err = 5e-3, measured on both sides:
worst FD noise, 40 runs 4.5e-4
ignore grad entirely 9.21
ignore the ne_b1 broadcast 3.39 / 1.15
read grad row 0 always 4.41
route every slot everywhere 5.30
10x above the noise floor, 200-1800x below every real defect. Five mutations
introduced, five caught.
Both kernels verified against a naive DOUBLE reference across 15 configurations covering
both broadcast modes: relative error ~1e-7, i.e. float32 epsilon. That reference is what
proved the FD failures were the harness and not the arithmetic -- though note it shares
the kernels' assumption about the broadcast rule, so it cannot check THAT. The rule was
confirmed by reading the forward: row_mapping.i1 is the SLOT index (the in-code comment
calling it "selected expert index" is wrong), so the forward's b column is slot % ne_b1.
…byte stride
Found by adversarial review, with a working reproduction. Two real bugs, both silent,
both in the kernels I had just declared verified.
1. GRAD (src1) was indexed as g_col[j] -- a hard-coded 4-byte dim-0 stride.
ggml's autodiff hands out-prod-shaped ops TRANSPOSED grads as a matter of course:
the MUL_MAT backward passes ggml_transpose(grad) straight to ggml_out_prod, which is
exactly why ggml_compute_forward_out_prod_f32 -- the kernel mine is modelled on --
reads src1 through `i1*nb10` rather than indexing a float*. I dropped that.
A TRANSPOSE node reaching here has nb[0] == 8. Reproduced: 24 of 36 output elements
wrong, max abs err 9.2. No assert fires. The run just trains on a wrong gradient.
2. IDS (src2) was indexed as ids_t[i], likewise hard-coded.
The FORWARD reads ids through nb[0]
(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 and the backward misread -- crediting the
wrong experts with each other's gradients. Reproduced: forward bit-identical across
two ids layouts, backward disagreeing by 1.03.
Both are now read through their nb[0]. b keeps a contiguity assert, because it is the
VECTOR operand of ggml_vec_mad_f32 -- a strided read there is not merely wrong, it is
unrepresentable.
WHY NOTHING CAUGHT THIS. test-backend-ops only ever builds a CONTIGUOUS grad, so no test
in the suite could see it -- and neither could my double-precision reference, which built
its own contiguous tensors. Every check I had said the kernels were exact, and they were,
on the only inputs anyone was feeding them.
So the guard is a new test case, not a comment. grad_transposed routes the objective
through cont(transpose(out)), which makes ggml_build_backward_expand hand MUL_MAT_ID's
backward a grad whose op is TRANSPOSE. Reintroducing the exact bug:
contiguous-grad cases (the entire old suite) -> PASS. invisible.
transposed-grad cases (the new guard) -> FAIL, MAA 21.0 / 2.57 / 2.55
Re-verified after the fix: double-precision reference across 15 configs (rel ~1e-7),
bit-identical across 1/2/4/8 threads, quantized-vs-dequantized bit-exact on q8_0/q4_K/q4_0,
strided ids identical to contiguous ids, MUL_MAT_ID grad 0/20 runs failing, the grad
allowlist and the MUL_MAT_ID forward unregressed.
ggml could differentiate exactly ONE member of the GLU family: split SwiGLU, via a
SILU_BACK composite. Fused SwiGLU tripped `GGML_ASSERT(src1 && "only split swiglu")`,
and REGLU / GEGLU / GEGLU_ERF / GEGLU_QUICK / SWIGLU_OAI hit
`GGML_ABORT("unsupported glu op for backward pass")`. Every Gemma (GEGLU) and gpt-oss
(SWIGLU_OAI) model was untrainable.
New op GGML_OP_GLU_BACK at the enum tail. One kernel, one row pass, whichever way the
caller packed its operands: dst carries d_a's shape, and a split forward views half 0 as
d_a and half 1 as d_b. Split SwiGLU keeps its existing SILU_BACK composite untouched --
it is the FFN of every dense llama, it is covered today, and routing it through the new
op would change nothing mathematically while risking a regression in the one place this
project cannot afford one.
Derivatives are taken against ggml's OWN scalar forwards in vec.h, not against a paper:
if ggml's gelu uses the tanh approximation, the correct derivative is the derivative OF
THAT, and a finite difference will insist on it.
An odd-width fused `a` was the first thing to bite. ggml_glu computes its output width as
a->ne[0]/2 with INTEGER division, so the sweep's ne_a[0]=5 leaves a trailing element the
forward never reads. It gets a zero gradient -- and it has to be WRITTEN zero, not left as
whatever the allocator handed us, or d_a is not a->ne[0] wide and compute_backward's
same-shape assert fires.
THE TICKET'S HEADLINE BLOCKER IS REFUTED BY MEASUREMENT.
It says the GGML_GELU_FP16 lookup table makes GEGLU's MODE_GRAD unpassable, because the
F32 forward reads an F16 table and the FD would differentiate a step function. Tested by
making the F32 GEGLU forward exact: the error got WORSE (0.024 vs 0.0049), and the
table-free variants moved around just as much. The table is not the cause. That saves a
risky change to inference numerics on every Gemma.
MODE_GRAD IS A WIRING CHECK FOR THIS OP, AND THE COMMENT SAYS SO.
mean_abs_asymm divides by (gn + ga), not (|gn| + |ga|), so a near-zero gradient element
sends the ratio to infinity -- and every GLU gradient is a PRODUCT (dx = dy*g*act'(x)), so
near-zero elements are ordinary, not exotic. Measured: FD noise reaches 0.80 while a
genuinely broken kernel measures 0.18-0.60. They OVERLAP. No tolerance separates them, so
none is pretended: the bound asserts the backward builds, schedules, has the right shapes
and does not abort. That much was impossible before this commit.
Correcting the metric to |gn| + |ga| drops the noise floor to 0.022 and catches every
defect with 3.6-12x margin. It is written and measured and it is NOT here, because it also
removes a NaN-based free pass that TANH, SIGMOID and CROSS_ENTROPY_LOSS have relied on
since S1-19 (0/0 = NaN; `NaN > tol` is false; they pass unconditionally). That is ticket
S1-37.
So the NUMERICS get a real oracle: tests/test-glu-back.cpp, a new ctest binary that
differentiates ggml's own scalar forwards in float64. All six variants match to ~1e-7. It
catches what MODE_GRAD cannot -- perturbing GEGLU's tanh argument by 5% is invisible to
MODE_GRAD and fails here instantly. (It also caught a stale library mid-session, which is
how I know it works.)
test_glu now calls ggml_set_param -- it never did, so the FUSED backward was completely
unexercised while `grad -o REGLU` still printed OK. The param is the BASE tensor, never the
view: ggml_set_param asserts op == GGML_OP_NONE.
Two bugs in MODE_GRAD's error metric, and they compound.
1. THE DENOMINATOR WAS THE SIGNED SUM.
const float asymm = (a[i] - b[i]) / (a[i] + b[i]);
So the ratio goes to infinity whenever the two gradients nearly cancel, and in
particular whenever the TRUE gradient is near zero and both are rounding noise. That
is not exotic: any op whose output is a product has such elements for ordinary inputs
(every GLU: dx = dy * g * act'(x)), and any op summing over a selected subset can
manufacture them (MUL_MAT_ID's expert routing).
Measured on GLU_BACK, with the kernel verified exact against a float64 reference the
whole time: the "noise" reached MAA 1.80 while a genuinely broken kernel measures
0.18-0.60. The noise OVERLAPPED the defects, and because the metric is unbounded, NO
tolerance was safe. S1-28 had to ship its grad cases with a 0.9 bound -- i.e. a test
that asserted almost nothing -- and it STILL flaked.
Now |a| + |b|. That is what a symmetric relative error is, it is bounded in [-1, 1],
and since |a|+|b| >= |a+b| it can only ever make MAA smaller. GLU's noise floor drops
from 1.80 to 0.022, and its bound is now 5e-2: 2.3x above the noise, 3.6-90x below
every one of six injected defects, all six caught.
2. AN EXACTLY-ZERO GRADIENT WAS AN UNCONDITIONAL PASS.
a == b == 0 gives 0/0 = NaN. `NaN > max_maa_err()` is FALSE. The case passes.
tanh(150) and sigmoid(150) are 1.0 to float precision -- derivative exactly zero -- and
test_unary initializes in [-150, 150]. TANH AND SIGMOID HAVE BEEN ON THE VENDOR-BUMP
ALLOWLIST SINCE S1-19 ON THE STRENGTH OF A NaN. So has CROSS_ENTROPY_LOSS, whose
softmax is one-hot at [-100, 100]. Their gradients had never once been compared.
With the metric fixed they are checked for the first time and fail: MAA 0.14, 0.37 and
0.60. The kernels are fine -- a grad_eps of 15 against a range of 150 measures nothing
a derivative would recognise. Conditioned (range +-3, eps 0.02 for the saturating
unaries; range +-3 for CE) the residual drops to 5.5e-4, 1.1e-2 and 3.8e-3, and each
gets a measured bound. The EVAL sweep keeps its wide range: it is hunting NaNs in the
tails, which is the opposite requirement.
Also pinned, both measured, both previously flaking about 1 run in 15 against the default
1e-4: ADD_ID (identity VJP, noise 1.2e-4 -> 1e-3) and CONCAT (routing VJP, noise 2e-4 ->
1e-3).
Full 16-op grad sweep: 0/15 runs failing.
Found by training one. Neither could have been found any other way, and that is the
point of the end-to-end test.
1. DIV'S BACKWARD DID NOT HANDLE A BROADCAST src1. MUL's always has.
d/d(b) of (a/b) is -a/b^2, which has a's shape. When b is broadcast against a, that
gradient must be REDUCED back onto b's shape with repeat_back -- exactly as MUL does
six lines above it. Without that, ggml_build_backward_expand aborts on its own
same-shape assert.
build_moe_ffn normalizes the top-k router weights with
weights = div(weights[n_used, n_tok], sum_rows(weights)[1, n_tok])
so EVERY MoE model with norm_w -- which is every Mixtral -- hits it the moment a
gradient reaches the router. MUL_MAT_ID having a backward was necessary and nowhere
near sufficient.
2. ggml_get_rows_back REQUIRED A MATRIX GRAD AND A VECTOR INDEX TENSOR.
GGML_ASSERT(ggml_is_matrix(a) && ggml_is_vector(b) && ...)
while ggml_get_rows itself has always been fully general:
out[i0, i10, i11, i12] = a[i0, b[i10,i11,i12], i11, i12]
So any get_rows with a 3D source and a 2D index tensor had a forward and no backward.
build_moe_ffn gathers each token's top-k router probabilities with exactly that:
weights = get_rows(probs[1, n_expert, n_tok], selected_experts[n_used, n_tok])
Generalized to the real semantics; the old 2-D loop is the i11 == i12 == 0 slice of the
new one. Its flat indexing of src0 was also only correct by accident -- it ignored the
structure entirely, which is why test_get_rows_back could build a grad of [n, r, b]
against a [r/2, b] index tensor and still pass. Those shapes are now consistent.
WHY NO TEST COULD SEE EITHER OF THESE.
test_bin_bcast gated MODE_GRAD on
// The backward pass supports broadcasting only for GGML_ADD:
const bool grad_supported = op == ggml_add && ggml_are_same_shape(a, b[0]) && ...
Two consequences. MUL and DIV were never grad-tested AT ALL, at any shape. And no
BROADCASTING case was ever grad-tested, for any op -- the same-shape clause excluded them.
So DIV's broadcast backward was structurally invisible.
The class was already tuned for it: grad_eps, grad_precise and max_maa_err all branch on
`op == ggml_div` / `op == ggml_mul`. All of that was dead code behind the gate.
Now: ADD 59 grad cases, MUL 59 (previously ZERO), DIV 57, GET_ROWS 111, GET_ROWS_BACK 55.
DIV's bound is 5e-3, measured -- its gradient is second order in b, so its finite
difference is intrinsically noisier than ADD's or MUL's.
Full 21-op grad sweep: 0/12 runs failing. Forwards unregressed.
S1-29 is marked done and delivered ONLY the CONCAT VJP. The SSM_CONV_BACK / SSM_SCAN_BACK enums, constructors and backward-switch cases its own ticket also promised were never written -- SSM_CONV_BACK appeared zero times in ggml.h. So S1-30 and S1-31 were both written against a dependency that did not exist. Two ops at the enum tail, their constructors, and the two backward cases that emit them. No kernels: CPU supports_op returns FALSE explicitly for both, which is not redundant -- the default returns TRUE, so a new op with no dispatch case is reported *supported*, gets scheduled, and hits ggml_compute_forward's `default: GGML_ABORT`. It would look implemented right up until it killed the process. Same trap S1-25 hit. SSM_CONV's backward is the convolution run backwards; only sx takes a gradient, since the conv weight is a frozen base weight on the LoRA path. SSM_SCAN has SEVEN sources and five of them take a gradient, so ggml_ssm_scan_back returns a PACKED 1-D tensor -- [d_s | d_x | d_dt | d_B | d_C] -- exactly as ggml_ssm_scan itself packs y with the final states, and the backward case views each region back onto its source. `grad` is the WHOLE packed gradient of the forward's dst, state region included: MODE_GRAD's objective sums over it, so it is NOT zero there, and a kernel that assumed otherwise would disagree with the finite difference and be wrong to. That region seeds the reverse recurrence at t = n_t. Decided here so S1-31 cannot get it wrong. THE TESTS CALLED ggml_set_param ZERO TIMES. So `grad -o SSM_CONV` and `grad -o SSM_SCAN` requested no gradients, compared nothing, and printed `Backend CPU: OK` -- while neither op had a backward case at all and both would have aborted the moment one was asked for. Both now ask. And every existing shape is far above grad_nmax() (10000), so they would have been SILENTLY SKIPPED even with a param. Tiny shapes added for both, and for SSM_SCAN in both the Mamba-1 (head_dim == 1, per-state A) and Mamba-2 (head_dim > 1, scalar A) branches -- they are different code in the kernel. The xbc_overlap variants make x, B and C views of one tensor, and ggml_set_param asserts op == GGML_OP_NONE, so those stay eval-only. Stated in the code rather than hidden: it is a real llama.cpp layout and its backward aliasing is exactly what goes wrong quietly. S1-31 owns it. Backward graphs now BUILD for both ops and report not-supported instead of aborting. Full 21-op grad sweep 0/8 failing; SSM forwards unregressed.
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.
The selective scan's VJP: a REVERSE recurrence, walking t from n_t-1 down to 0 and
carrying the state gradient backwards. Five gradients out of one op (d_s, d_x, d_dt, d_B,
d_C), packed exactly as ggml_ssm_scan itself packs y with the final states.
Both branches: Mamba-2 (one scalar decay per head) and Mamba-1 (one per state).
TWO THINGS THAT ARE EASY TO GET WRONG, AND BOTH ARE SILENT.
1. `grad` is the WHOLE packed gradient of the forward's dst -- y AND the final states --
and the state region is NOT zero. MODE_GRAD's objective sums over the packed dst, so
d(sum)/d(s_final) is 1, and a kernel that seeded the reverse recurrence with zeros would
disagree with the finite difference and be WRONG TO. It seeds ds at t = n_t.
That is not a hypothetical: seeding ds with zeros is one of the five mutations below,
and it measures MAA 0.59. In training the region genuinely IS zero -- the cached state
feeds nothing downstream of the loss -- so honouring it costs nothing there, and it makes
cross-ubatch BPTT nearly free later.
2. The forward OVERWRITES its state in place, so s_{t-1} is gone by the time the backward
needs it -- and it does need it, for d(dt) via the dA path. The states are recomputed and
STORED, all n_t + 1 of them. Store-all: correct, O(n_t) memory, and the honest starting
point for a CPU oracle. Checkpoint-every-K is the optimization and must be pinned
bit-for-bit against this.
Threaded by SEQUENCE, not by head. dB and dC accumulate over every head in a group, so a
head-partitioned kernel would have several threads writing the same (i0, g, t) and would
need atomics -- neither deterministic nor free. One thread per sequence owns every output
it touches. Parallelism is n_seqs, which is small; correctness and determinism (ADR-0002)
come first in the kernel the GPU ports get measured against.
THE ORACLE IS A FLOAT64 FINITE DIFFERENCE OF GGML'S OWN FORWARD, and it has to be.
A reverse recurrence has many ways to be subtly wrong -- a dropped dA path, a state read
one token late, a missing seed -- and a hand-written reference would share my derivation's
bugs. So each of the five gradients is checked by perturbing every input element and
re-running ggml_ssm_scan itself, under a NON-UNIFORM objective that includes the packed
state region. All five exact to ~1e-5 (float32 FD precision), in both branches.
max_maa_err 5e-2, measured: noise floor 2.1e-2 (the recurrence is exponential in dt*A, so a
finite difference of it amplifies its own rounding), five injected mutations at 0.38-0.59,
all five caught. 2.4x above the noise, 7.5-12x below every defect.
The tests called ggml_set_param zero times before S1-29b, and every shape was above
grad_nmax, so `grad -o SSM_SCAN` reported OK while the op had no backward at all.
The attention matrix is [n_kv, n_q, n_head]: QUADRATIC in context, and the reason
long-context training runs out of memory. Split the QUERY axis, softmax each chunk on
its own, and only [n_kv, chunk_q, n_head] is live at a time.
Chunking QUERIES is exact. A softmax row depends on its own row of scores and nothing
else, so partitioning the query axis is not an approximation -- the losses come out
BIT-IDENTICAL to the naive path and the gradients differ only by float32 reduction
order (~1e-6). Chunking KEYS is the other problem: a partial softmax needs rescaling
when a later chunk raises the running maximum. That is flash attention, and it is why
FA needs a kernel. This needs none: MUL_MAT, SOFT_MAX, CONCAT, OUT_PROD, all already
on every backend. It is the standing fallback until FA's backward kernels land.
CHUNKING ALONE DOES NOT SHRINK THE BACKWARD, and this is the part that fails silently.
SOFT_MAX_BACK reads the softmax's own output, so without recompute every chunk's P
stays live from the forward until the backward consumes it -- the memory is all still
there, just spread over more tensors. Under ggml_build_backward_expand_checkpointed
(S1-17) each P is a segment-interior node, rebuilt immediately ahead of the backward
node that reads it and dead again straight after. S1-17 removes the DEPTH factor;
S1-24 removes the n_ctx^2 factor WITHIN a layer. They multiply.
Measured (tiny fixture, peak compute buffer, one training step):
n_ctx naive chunk only ckpt only ckpt+chunk
1024 79 MiB 66 MiB 47 MiB 44 MiB 1.80x
2048 254 MiB 196 MiB 166 MiB 128 MiB 1.98x
4096 904 MiB 652 MiB 620 MiB 416 MiB 2.17x
The `chunk only` column IS the warning above, measured.
TWO THINGS I GOT WRONG FIRST, both worth the comments they now carry:
1. Slice the token axis BEFORE the permute. build_attn_mha permutes q to
[head_dim, n_tokens, n_head, n_stream]; slicing THAT takes a view of a
non-contiguous tensor, and ggml's autodiff does not survive it -- the gradient
returns through PERMUTE as a non-contiguous view and dies in ggml_scale's
GGML_ASSERT(ggml_is_padded_1d(a)) inside the LoRA scale's backward. An abort three
ops from the mistake, naming a function you never called. q_bhd keeps the
pre-permute tensor (token axis = ne[2], rows contiguous) and the backward chain is
then the same shape the naive path already proves works.
2. Reassemble with a BALANCED CONCAT TREE, not a left fold. A fold builds accumulators
of 1/C, 2/C, ... C/C of the output and hands the saved memory straight back: peak
bottomed out at 4 chunks and got WORSE from there (151.5 MiB at chunk_q=64 vs 128.6
at 512) -- the exact opposite of the point. A balanced tree has log2(C) levels and
only ~2 are ever live. CONCAT's VJP landed in S1-29, which is what makes any of this
differentiable with no new kernel.
Training only (inference keeps no attention matrix alive). Fixed for the life of a run:
chunking changes the graph's node COUNT and ggml-opt keys optimizer state by node index.
Softcap (gemma2) and ALiBi branches are written but have NO FIXTURE -- attn_soft_cap is
set in code by models/gemma2.cpp, not by a GGUF key, so it cannot be switched on for a
llama-arch fixture. Both are chunk-invariant by construction (ALiBi's slope is a
function of the HEAD index; softcap is elementwise) and mirror the naive path line for
line. That is an argument, not a test, and it is recorded as one.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…r trained on no gradient ggml-gallocr may place an op's output on top of a source buffer when the op is on ggml_op_can_inplace's list, and SOFT_MAX_BACK is. It prefers src[0] (the incoming grad dy) but falls through to src[1] -- the softmax OUTPUT y -- when src[0] is not reusable. Attention never hits it: its softmax output is also read by the V matmul backward, so it is never the lone consumer. A Mixtral router does hit it: the softmax output feeds only argsort (no gradient) and get_rows (which reads the ids, not the values), so at backward time SOFT_MAX_BACK is y's sole consumer and gallocr aliases the result onto y. The CPU kernel computed dx = scale*y*(dy - <y,dy>) as `cpy(dx,dy); dx-=<y,dy>; dx*=y`. That is safe when dx aliases dy but WRONG when dx aliases y: the cpy overwrites the whole y buffer before the final `dx*=y` reads it, so it returns ~(dy - <y,dy>)^2. On the router this drove d_logits to nearly zero -- the gate weights got no gradient, and every LoRA upstream of an MoE block trained on a dh that had silently dropped its router term. Forward, loss, and expert-weight grads all stayed correct, so nothing else caught it, and test-backend-ops cannot (it validates CPU vs CPU and both alias identically). Compute dx element-wise, reading y[i] and dy[i] before writing dx[i] with <y,dy> already reduced. Correct whether dst aliases y, dy, or neither -- and one pass instead of four. test-soft-max-back-inplace.cpp forces the alias and checks the kernel against a float64 reference; it fails ~0.22 on the old kernel and passes on this one. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…state cache ggml_ssm_scan's initial-state input is a view of the recurrent-state cache, and the forward overwrites that cache in place with the final state (mamba-base.cpp). ssm_scan_back recomputes the intermediate states from that initial state -- the forward destroys them in place -- so by backward time it starts from the FINAL state and every scan gradient (dx, ddt, dB, dC) is wrong. It presents as layer-0-wrong, last-layer-right: only the last layer's state buffer survives to the backward. Copy the initial state into its own buffer (ggml_cont) before the scan, in both the Mamba-1 and Mamba-2 paths, so the in-place update cannot reach the backward's input. The kernel is untouched; test-backend-ops grad still passes. Two more fixes the Mamba training path needed, neither of which had ever run: - VIEW backward: ggml_acc asserts nb[0] == sizeof(float), but the Mamba conv splits ssm_in's output with a view and feeds it through ggml_transpose, whose backward hands the view backward a non-contiguous gradient. Every Mamba training graph aborted at backward. ggml_cont the grad before ggml_acc_or_set, matching the RESHAPE backward's existing guard. - SSM_SCAN backward now asserts n_group == 1. The kernel's per-group dB/dC routing (g = h/(nh/ng)) has no finite-difference oracle -- every grad-enabled test_ssm_scan case is n_group == 1, and the n_group > 1 shapes exceed grad_nmax and are skipped. Mamba-1 always builds n_group == 1. Mamba-2 / Falcon-H1 with n_group > 1 now abort loudly rather than training on an unverified gradient. The composed Mamba-1 gradient is now verified per tensor against a self-audited float64 reference (learning-llamas tests/test_ssm_training.py): every LoRA gradient matches at ~1e-6, where before the layer-0 tensors were off by ~1e-2 and a finite difference of the real ggml forward sided with the reference.
The scan backward's group routing (g = h/(nh/ng), folding every head of a group into one
dB/dC slab, threaded by sequence) was written at S1-31 but never grad-checked: every
grad-enabled test_ssm_scan case was n_group == 1, where g collapses to 0 and the fold never
runs, and the n_group > 1 shapes exceed grad_nmax so MODE_GRAD skipped them. S1-47 asserted
n_group == 1 in ggml.c's backward switch rather than train on an unverified gradient.
The kernel arithmetic was already correct; this proves it and removes the assert.
- test-backend-ops: four TINY n_group > 1 test_ssm_scan cases (n_group in {2,4}, both the
per-state-A (head_dim == 1) and scalar-A (head_dim > 1) branches), sized a few hundred
elements so they land under grad_nmax and are actually COMPARED. n_head/n_group == 2, so a
group's dB/dC genuinely folds two heads. grad -o SSM_SCAN now compares 6 gradients (was 2).
- ggml.c: the ssm_B->ne[1] == 1 refusal is replaced by the only structural check the fold
needs, n_head % n_group == 0.
Verified end-to-end on the learning-llamas side against a self-audited float64 Mamba-2 oracle
(reference_mamba2.py + a gen_tiny_mamba2 fixture, arch mamba2, ssm.group_count = 2): the real
n_group=2 training gradient matches at ~1e-6, and is bitwise-identical across thread counts.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015uJLvFUB4aphCBCtx7tSDZ
…lid guard + seed Three fork-side pieces of the minor-findings sweep. - ggml-opt: the global gradient norm was computed inside the backward graph for the grad-clip factor but never surfaced. Capture it: the pre-clip norm is the existing `grad_norm` node; the post-clip norm is a NEW reduction over the CLIPPED gradients the optimizer steps on (measured, not inferred as pre*factor). Both are pinned as outputs and read back like the loss in ggml_opt_eval, into last_grad_norm_pre/post, exposed by ggml_opt_grad_norm_pre/post. NaN unless the eval ran the OPT graph with grad_clip > 0, so an unclipped run's graph is unchanged and reports nothing rather than a stale value. This is what the shim's ll_grad_norms getter (S1-10 item 3) reads. - test-backend-ops mean_abs_asymm: when the grad_expect filter discards every element, nvalid stays 0 and sum/nvalid was 0.0/0 = NaN, and NaN > max_maa_err() is false, so the case PASSED having compared nothing (the same free-pass class S1-37 closed per-element, reachable for CLAMP). Return +inf when nvalid == 0 so it FAILS loudly instead. CLAMP still passes 16884/16884 -- no registered case currently triggers it; the guard is protective. - test-backend-ops seeding: GGML_TEST_SEED makes the three shared float-init primitives (init_tensor_uniform, init_tensor_kq_mask, init_tensor_tril) deterministic run to run, so a MODE_GRAD case near tolerance no longer flips pass/fail between runs. Unset (the default), every init keeps its std::random_device behavior, byte-for-byte unchanged. Verified: two seeded `grad` runs of a project op produce byte-identical output. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015uJLvFUB4aphCBCtx7tSDZ
…platform margin MSVC's exp() produced a draw at MAA 0.0558 against the glibc-measured 2.4x margin over noise. 1e-1 keeps 3.8x below the smallest injected defect (0.38) while covering the observed cross-libm spread; the sharp per-element claim stays with the float64 e2e oracle, as the existing comment already argued. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ness audit - SSM_SCAN_BACK work buffer: MIN(ns, n_tasks) slabs, plus the matching `ith >= ns` early return that makes the bound sound. Sizing a slab per thread asked 12.9 GB to do 0.8 GB of work at realistic Mamba dims (ns == 1, 16 threads) -- a plan-time OOM for a step that fits the algorithm, invisible on the tiny fixtures. The two sides have to be read together: without the guard the excess threads form slab pointers past the end of the buffer. - mamba-base: gate the S1-47 ggml_cont on cparams.training. `ssm` spans the whole recurrent-state cache (the reshape is over get_size(), not this ubatch's n_seqs), so unconditionally it cost a cache-sized memcpy and a cache-sized compute-buffer slab per layer per eval -- on the inference path, which builds no backward and has no live reader of the pre-update value. - ggml-opt: build the pre/post-clip norm scalars in ctx_results rather than ctx_compute, and count them in tensors_const. On the static-graph path the graph that runs is a dup_graph copy inheriting its data pointer from the original, so the original must own a buffer; left in ctx_compute nothing allocated them and the read-back aborted in ggml_backend_tensor_get on a NULL buffer. No change on the dynamic path, where ctx_results is ctx_compute. - ops.cpp: GLU backward row addressing decomposed into i1/i2/i3 -- folding three indices through nb[1] is wrong for any non-contiguous 3-D grad (latent: no llama graph builds one today). OUT_PROD_ID_GRP now range-checks expert ids like OUT_PROD_ID already did, instead of silently dropping the slot's contribution to d(as) while its sibling aborts on the same input. - test-backend-ops: GGML_TEST_SEED argument validation (malformed input was silently seed 0); a GGML_TEST_MAA_REPORT=1 stream so a bound can be measured without hand-patching the source; test_glu_split / test_swiglu_oai 0.9 -> 5e-2, measured, worst 0.0184 across the three GLU classes over 450/300/240 comparisons; and ggml_set_name(a, "view_of_b") -> (b, ...) in both split classes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UqTpucYHeqfkA9eNMc7ABj
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The attention matrix is
[n_kv, n_q, n_head]— quadratic in context, and the reason long-context training runs out of memory. Split the query axis, softmax each chunk on its own, and only[n_kv, chunk_q, n_head]is live at a time.Chunking queries is exact; chunking keys is not
A softmax row depends on its own row of scores and nothing else, so partitioning the query axis is not an approximation — the losses come out bit-identical to the naive path, and the gradients differ only by float32 reduction order (~1e-6).
Chunking keys is the other problem: a partial softmax needs rescaling when a later chunk raises the running maximum. That is flash attention, and it is why FA needs a kernel. This needs none —
MUL_MAT,SOFT_MAX,CONCAT,OUT_PROD, all already on every backend. It is the standing fallback until FA's backward kernels land.Chunking alone does not shrink the backward
SOFT_MAX_BACKreads the softmax's own output, so without recompute every chunk'sPstays live from the forward until the backward consumes it — the memory is all still there, just spread over more tensors. Underggml_build_backward_expand_checkpointed(S1-17) eachPis a segment-interior node, rebuilt immediately ahead of the backward node that reads it and dead again straight after.S1-17 removes the depth factor; S1-24 removes the n_ctx^2 factor within a layer. They multiply.
The
chunk onlycolumn is the warning above, measured.Two things I got wrong first
1. Slice the token axis BEFORE the permute.
build_attn_mhapermutesqto[head_dim, n_tokens, n_head, n_stream]; slicing that takes a view of a non-contiguous tensor, and ggml's autodiff does not survive it — the gradient returns throughPERMUTEas a non-contiguous view and dies inggml_scale'sGGML_ASSERT(ggml_is_padded_1d(a))inside the LoRA scale's backward. An abort three ops from the mistake, naming a function you never called.q_bhdkeeps the pre-permute tensor (token axis =ne[2], rows contiguous).2. A balanced concat tree, not a left fold. A fold builds accumulators of
1/C, 2/C, ... C/Cof the output and hands the saved memory straight back: peak bottomed out at 4 chunks and got worse from there (151.5 MiB atchunk_q=64vs 128.6 at 512) — the exact opposite of the point. A balanced tree haslog2(C)levels and only ~2 are ever live. CONCAT's VJP landed in S1-29, which is what makes any of this differentiable with no new kernel.Scope
Training only (inference keeps no attention matrix alive). Fixed for the life of a run — chunking changes the graph's node count, and ggml-opt keys optimizer state by node index.
Softcap (gemma2) and ALiBi are written but have no fixture.
attn_soft_capis set in code bymodels/gemma2.cpp, not by a GGUF key, so it cannot be switched on for a llama-arch fixture. Both are chunk-invariant by construction (ALiBi's slope is a function of the head index; softcap is elementwise on the scores — neither mixes query rows) and mirror the naive path line for line. That is an argument, not a test, and it is recorded as one.Upstreaming disposition: fork-local for now.
🤖 Generated with Claude Code