Module fcd_ops 

FCD polish — hand-rolled forward/backward operators.

The training graph is FIXED (docs/RUST_FCD.md), so there is no autograd and no tape: every operator here is a (forward, backward) pair written out by hand, llm.c style. Ops are generic over the minimal Fp float trait so the SAME code that trains in f32 is gradchecked in f64 against central finite differences (tests/fcd_gradcheck.rs, rel err < 1e-3).

Conventions:

• all matrices are row-major; weight matrices use the runtime layout [out_dim, in_dim] (a matvec is out[o] = dot(w_row_o, x));
• every backward ACCUMULATES into its output gradient buffers (+=) — callers zero them once per step, residual branches then just add up naturally;
• the attention-weight ops (attn_head_*, nystrom_head_*) are meant to run in f64: the certified CPU probe computed the T×T weight matrices in f64, which also makes raw exp(±40) safe with no flash-shift bookkeeping.

Structs

GdnSeqCfg

GDN geometry + frozen elementwise weights for the sequence ops.

NysCfg

Nyström joint kernel geometry — matches nystrom::O1Cfg semantics.

Traits

Fp

Minimal float abstraction: just enough for the fixed graph. Not a general numeric tower — two impls (f32/f64), no external crates.

Functions

attn_head_bwd

Exact-attention backward (probs recomputed row by row — the trainer is layer-checkpointed, nothing is stored). Standard softmax chain: ds_j = p_j·(dp_j − Σ p·dp), dp_j = dout·v_j.

attn_head_fwd

Exact per-head causal attention: out[t] = softmax(q_t·Kᵀ/√d)·V over j ≤ t. q,k are [t,d], v is [t,dv], out is [t,dv].

ce_kl_position

One position of the polish loss (docs/RUST_FCD.md §2.4): L = (1−klw)·CE(student, target) + klw·KL(teacher‖student), both per-position; inv_n = 1/(B·T) folds the batch mean into dlogits. Returns (ce, kl) UNWEIGHTED for logging; dlogits gets the combined gradient (+= is NOT used — each position owns its slice).

gdn_conv_bwd

Conv+SiLU backward: dcq → draw (+=). Taps frozen (through-grad only).

gdn_conv_fwd

Depthwise causal conv + SiLU over the whole window: raw [t, c_dim] → (pre-activation [t, c_dim], cq = silu(pre)). Ring semantics of gdn_step from a fresh state: positions before 0 are zeros.

gdn_group_bwd

One GQA group backward: BPTT over the window given dout rows of this group’s v-heads. Accumulates (+=) into dcq (q/k channels of ko, v channels of its v-heads), dz, da, db. All weights frozen.

gdn_group_fwd

One GQA group (k-head ko, its rep = nv/nk v-heads) forward over the window. Writes out[t, nv·dv] slices of its v-heads only.

gdn_seq_bwd

Whole-layer GDN sequence backward: through-grads (+=) into the four projection streams.

gdn_seq_fwd

Whole-layer GDN sequence forward (serial over groups): raw projections → out [t, nv·dv]. The trainer parallelizes groups on the pool with the same gdn_group_* entry points.

gemm_dw

dW += dYᵀ · X — parallel over dW ROW ranges (each worker owns a disjoint slice of output neurons; X is shared read-only).

gemm_dx

dX += dY[n,m] · W[m,k] — parallel over row blocks (disjoint dX rows).

gemm_nt

y[n,m] = x[n,k] · w[m,k]ᵀ — parallel over row blocks; bit-identical to matmul_nt (disjoint rows, same dot kernel regrouped by NEON).

matmul_nt

y[n,m] = x[n,k] · w[m,k]ᵀ (serial reference; the f32 hot path is gemm_nt below — same math, blocked + pooled).

matmul_nt_dw

dW += dYᵀ[m,n] · X[n,k] — the weight-gradient half of matmul_nt.

matmul_nt_dx

dX += dY[n,m] · W[m,k] — the input-gradient half of matmul_nt.

nystrom_head_bwd

Nyström joint backward: dq/dk/dv (+=) with M constant. The whole graph is recomputed (layer checkpointing) — see docs/RUST_FCD.md §2.2 for the derivation. Chain summary: dw[t,j] = dout_t·(v_j − out_t)/den_t near: dlg = dw·w → dq, dk (±scale) far: da = dw·e^{−c}·[row kept] → dFu, dE through the two matmuls with M constant → dq, dk, dQ̃, dK̃ landmarks: segment-mean scatter back into the PREFIX rows of dq, dk (the seal only saw those). The far gate is per ROW (the aggregate guard), not per key.

nystrom_head_fwd

Nyström joint forward (teacher-forced matrix form). Falls back to exact attention for short windows (runtime guard: t ≤ W+sink+8).

rmsnorm_bwd

RMSNorm backward: through-grad into dx (+=) and, when the gain is trainable, gain grad into dw (+=). inv is the saved 1/rms.

rmsnorm_fwd

RMSNorm over n rows of width d = w.len(): Qwen style y = x̂·w, Gemma style y = x̂·(1+w), x̂ = x/√(mean x² + eps). Sum-of-squares accumulates in f64 (runtime discipline). inv_out stores the per-row 1/rms for the backward.

rope_bwd

RoPE through-grad: a rotation’s Jacobian is the rotation itself, so dL/dx = R(−θ)·dL/dy — the inverse rotation, in place.

rope_fwd

Rotate the first 2·inv_freq.len() dims of one head vector in place (half-split pairing, same convention as attention::rope_rotate).

seg_means

Contiguous segment means (Nyströmformer landmark recipe), the same integer split (i·t)/m as nystrom::seg_means and the torch probes.

seg_means_bwd

Segment-mean backward: the mean is linear, so the landmark grad scatters back uniformly over its segment (dx += dl/seg_len).

silu

silu_bwd

d silu / dx = σ(x)·(1 + x·(1−σ(x))).