kdafix — the two class-B fixes burn-kda's designed-red gates were waiting for

2026-10-02, branch wt/kdafix2 (off f4f7a0b). Owner-approved lane: align burn-kda to FLA’s own KDA references — the decay form (A/B queue arm 5, technology REPLACE) and the read scale — and close the three RED-ON-PURPOSE oracle tests tests/kda_oracle.rs carries since b485ad9 the honest way, code to the reference, no tolerance touched.

Prior art this lane replaces: 23f4b5f (wt/kdaci, NOT merged) made the three reds #[ignore]d COUNTED reds; that lane was right that the reds were [[test]]-target-blind CI, not platform drift (its own docs/reviews/kda-oracle-ci-2026-10-02.md proves byte-identical numbers everywhere, ratio exactly K**0.5) and wrong that nothing could fix them this quarter. The formula audit: docs/reviews/2026-09-30-kda-formula-audit.md §3.1 (softplus placement) and §3.2 (read scale). The oracle pin came first and survives byte-identical: tests/oracle/upstream/fla_ops_kda_gate.py / fla_ops_kda_naive.py (sha256 asserted by gen_kda_oracle.py before import), FLA at 9f38d24980c46d46bd38614e743cdacd21906578.

What changed

Fix 1 — decay form → FLA RELEASES (fix commit a14153e)

• was (src/lib.rs, the DecayFn::Softplus arm): g = -softplus(exp(A_h) · z) — exp(A_h) INSIDE.
• now: g = -exp(A_h) · softplus(z) — exp(A_h) OUTSIDE, evaluated the reference’s order (softplus first, then the -exp multiplier).
• Evidence: gate.py:50 (naive_kda_gate: g = -A_log.exp() * F.softplus(g + dt_bias)) and independently the triton twin gate.py:167 (b_yg = -exp(b_A) * softplus(b_g)). Two transcriptions, one original file; the gate is Moonshot’s own supported file (gate.py:8).
• The two are different FUNCTIONS: they agree only at A = 0; at the crate’s own init (A = −3, z = +1) upstream gives alpha = 0.5743, the pre-fix branch gave 0.5120, and the ORDER of the two answers flips with the sign of A.
• What we compute now: alpha = exp(−exp(A_h)·softplus(z)), g ≤ 0 elementwise — NaN-impossible by construction (exp of a non-positive), saturation = underflow to alpha 0. The a_log clamp [−10, 20] stays (the sigmoid arm’s overflow protection; on the softplus arm it only bounds the multiplier).
• What we still do NOT compute: the bounded/lower-bound variant of the softplus arm — upstream has no such variant (lower_bound in gate.py belongs to the K3 sigmoid twin, naive_kda_lowerbound_gate), and gen_kda_oracle.py’s case list keeps the §1 part of the audit (sweep, endpoints, A=0 control) alive for the fixed arm.

Fix 2 — read scale → FLA’s head_k_dim**-0.5 (same commit)

• was: scale = 1 threaded to chunk_wy_forward at both chunk call sites and hard-coded again inside fused.rs’s two dispatch arms, with the claim “no softmax scale in KDA” as the asserted reason.
• now: one published function, pub fn fla_read_scale(head_k_dim) = head_k_dim**-0.5 (reached as KdaModule::k_scale()), and EVERY path takes it: both chunk arms, the fused bare-CUDA arm, the fused autodiff op, forward_recurrent and the decode prefill read (folded into q, FLA’s own place — naive.py:57).
• Evidence: fla/ops/kda/chunk.py:474-475 and fused_recurrent.py:261-262 — both if scale is None: scale = K ** -0.5 — and fla/layers/kda.py:262-278 calls chunk_kda with NO scale argument, so the shipped layer runs at the default. burn-gdn2 itself does this (gdn2/src/module.rs:299) — burn-kda overriding to 1.0 was unique in the library.
• why no trained number moves: the factor enters only the read o = q·S (+ the intra-chunk scores) and never the state update; the next thing the module does is output()’s RMSNorm, which is invariant to a constant rescale — absorbed to O(eps / mean(o²)) ≈ 1e-5. The fixtures/pins move by K**0.5 (2.83× at K=8); the trained objective does not move beyond that eps floor. Arithmetic in place.

The tests, none ignored

read_scale_matches_fla_reference now DRIVES our code (per (batch, head) kda_step calls over the recur fixture, q pre-multiplied by burn_kda::fla_read_scale(k) — FLA’s fold) against FLA’s scaleK_o row. The former green twin chunked_wy_honours_the_read_scale_when_asked differed from the fixed red only in where the constant came from and was folded into the fixed chunked_wy_applies_the_fla_read_scale (both chunk arms) rather than kept as a duplicate. kimi_linear_softplus_decay_matches_fla_reference is the same body it always was — the code moved under it. RED-ON-PURPOSE markers removed from the header; falsify.sh’s B-direction runs in reverse now (see below).

Fix 3 — the REL_BAR widening in wt/kdaci is REJECTED (not merged here)

d53229d opened REL_BAR 5e-2 → 0.25 where the bar measured the runner’s host (OpenBLAS DYNAMIC_ARCH kernel selection per host CPU; the failing number was byte-identical on every red: loss 77.005417, autodiff -1.002479e-2, fd -1.068115e-2, rel 6.15e-2; OPENBLAS_NUM_THREADS=1 pins the scheduler, not the host kernel). Not merged. Instead:

• ops_batched_autodiff.rs carries a KNOWN-FLAKE header: what a red looks like (the bytes above), that a re-run is the first response, and when to believe the break (a repeated red with a DIFFERENT number, or one that survives a local re-run). The bar is untouched.
• The thread cap (efce4f7’s OMP/OPENBLAS_NUM_THREADS=1) is re-applied on the ndarray job — free, and it removes the scheduling axis.
• The three designed-red tests are green after fixes 1–2; no #[ignore] anywhere.

Smoke: does OUR init survive the FLA form?

Recipe: 100 steps, nano (d_model 512, 8 heads × 64, max_iter 4, rank 96, 3 experts, Engram 25k rows in-VRAM), real corpus (/mnt/e43497ab-0ff2-45b4-b45f-28de3339a53e/aria_data/pretrain/real/corpus.bin), no --detach (NaN firewall unconditional since 1949c39). The trainer’s running branch is K3 Sigmoid; the form fix lives in the Softplus arm, so each smoke BUILDS with decay_fn = Softplus wired in dormouse-core/src/attention.rs (a trial flip, reverted after the run — not committed: wiring the form is arm 5’s A/B, not a default change).

• Form + OUR init (a_log = −3, b_alpha = +1): RESULT-TBD — NaN episodes: N_A, loss dloss_a.
• Form + FLA’s init (exp(A_h) ∈ (1,16) here 8.5; dt_bias −4.6, FLA’s mean inv_dt): RESULT-TBD (only if (a) degrades).
• Decision rule (from the lane brief): (a) clean → keep our init; (a) degrades → try (b); both bad → revert the form, findings only.

Checkpoint compatibility

• The running model (K3 sigmoid) is byte-for-byte unchanged; every tree checkpoint and every benches/history.tsv row stays valid.
• The Softplus arm’s numbers change everywhere — any future Softplus-branch run starts from scratch (none exists in checkpoints/ today). The read scale absorbs through the RMSNorm to ~1e-5; recorded anchors/FPB are not re-taken for that reason.

Verdict/commits

• a14153e — fix(kda): decay form + read scale (the two fixes, tests green).
• <instrument> — falsify.sh reversed (A9 re-inversion mutant, A6 k_scale mutant), flaker header, ndarray-job thread cap.
• <docs> — this file.
• Pushed to origin/wt/kdafix2; fused-library CI watched (result noted in the report).