fix(memit): accumulate FFN covariance in f64; escalate Cholesky ridge before failing

[Ai-Eli-ML]

Problem

estimate_ffn_covariance (crates/larql-inference/src/forward/trace.rs) accumulates KᵀK in f32. The comment above it says "Float64 would be safer but Array2 suffices at our scales" — that stops being true on larger models.

Observed on Gemma 4 26B-A4B (ffn_dim = 2112, activation magnitudes ~1e3) while running MEMIT locally: each accumulator entry grows to ~1e6 × N over thousands of sampled token positions; with a 24-bit mantissa the low-order contributions are dropped, and the accumulated Gram matrix is no longer numerically PSD. The Cholesky in memit.rs then hits a negative pivot and MEMIT fails with Cholesky failed.

memit.rs already converts the covariance to f64 (cov_f64) before Cholesky, but that cannot recover precision lost during accumulation — the damage is upstream of that cast.

Change

1. trace.rs — accumulate KᵀK in f64, scale in f64, downcast to f32 once at the boundary when returning. Public signature unchanged (Option<(Array2<f32>, usize)>). Memory: one temporary ffn_dim² matrix at 8 bytes instead of 4 (≈36 MB vs 18 MB at 2112).
2. memit.rs — if Cholesky still fails on an architecture nobody has profiled, escalate ridge ×10 up to 1000× before returning the error, logging each escalation via eprintln!. Defensive only; (1) is the fix. Behaviour at the original ridge is unchanged when the first attempt succeeds.

Not included: the ridge value. My local fork hard-coded ridge = 10.0 for Gemma 4 in April; LARQL_MEMIT_RIDGE (now in compile/into_model.rs / into_vindex.rs) is the right design, so nothing to add there.

Verification

• cargo build -p larql-inference -p larql-lql on the pinned 1.98.0 toolchain: clean.
• cargo clippy -p larql-inference -- -D warnings: clean.
• rustfmt --check on both files: clean.
• cargo test -p larql-inference: exit 0 — but note the crate's tests are #[ignore] model-gated (0 ran, 4 ignored on Linux/CPU), so this is a compile/clippy-level check, not behavioural coverage of the change.
• I do not have a small, deterministic repro that fails in-tree without a Gemma-4-sized model; if you'd like one I can add a synthetic test that builds a (N, ffn_dim) matrix with ~1e3 magnitudes and asserts the f32 vs f64 Gram matrices differ beyond tolerance.

Context

This was fixed in a local fork on 2026-04-21/22 and sat there; a repo cleanup surfaced it and a check of upstream trace.rs:303 showed the f32 accumulator still present, so here it is. First PR to this repo — happy to reshape it to whatever conventions you prefer.