perf: SIMD MLA-absorb score + value-mix reductions (#442, +4.7x)

[woolcoxm]

Addresses the attention half of #442 (cc @cdhdt — analysis and exact change-site breakdown are all from that issue; the build-flag audit, the reduction-site identification, and the correctness-gate framing made this tractable).

What

Two scalar f32 reductions on the MLA-absorb decode path that ran on every (s,h) of every layer:

score:  a = dot(qabs, Lt)          # colibri.c:2459, kvl=512 on GLM-5.2
vmix:   clat += sc[jj] * Lt         # colibri.c:2467

Both were scalar a += x[i]*y[i] loops that -O3 cannot reassociate (no -ffast-math, no #pragma omp simd anywhere in-tree). SIMD-ified under AVX2 (8-lane fmadd + hsum256, same shape as the int8 matmul_q kernel in quant.h:91) and NEON, with a scalar tail for the kvl%8 remainder.

• score: 8-wide fmadd accumulator + hsum256. Reassociation is accepted here — softmax downstream softens the ~1.8e-6 rounding flip (measured), which does not cross a near-tie threshold.
• vmix: 8-wide fmadd writing back per-lane (no horizontal reduction). Strictly bit-identical to scalar (measured diff 0.0 — there is no reduction to reassociate).

Correctness

glm_tiny TF oracle holds at its 30/32 baseline — same two mismatched positions (5, 25), same wrong tokens, no new mismatches, no position shifts. All in-tree tests pass (test_logit_nan, test_dsa_select, test_idot, test_topp, etc.).

A note on the 30/32 floor: it is a pre-existing forward-pass numeric discrepancy on dev (engine argmax is correct; the logits at 2 positions are a near-tie that breaks differently vs the Python oracle — e.g. pos 5: engine 197@1.754 vs oracle 34@1.721). Unrelated to this change. I investigated it as a side quest (it briefly looked like it might be the same bug class #442 is about — it is not); will file separately with the evidence.

Measured (microbench tests/bench_mla_simd)

GLM-5.2 dims (kvl=512, nt=256, H=128), x86-64-v3 AVX2+FMA, median of 3:

	per-head	speedup
scalar	~290 µs	—
SIMD	~62 µs	4.7× (79% faster)

The bench re-implements the old scalar reductions inline and A/Bs them against the new SIMD ones verbatim from colibri.c, mirroring the bench_idot / bench_topp pattern. Not a gate; build on demand: make tests/bench_mla_simd.

The end-to-end score-softmax-value PROFILE sub-bucket is where this lands on a real decode — that bucket grows with context length (nt in the inner loop).

Follow-up (the router half)

The other half of #442 — the f32 router matmul at quant.h:80-84 — is deliberately not in this PR. It is riskier: reassociating the router can flip the top-k expert selection, which is load-bearing for output quality (a wrong expert is not softmax-softened the way an attention score is). That one needs the stronger correctness gate (a full eval_glm.py quality A/B, not just the tiny-model oracle), so it stays as a separate follow-up once this attention half lands.

Credit

@cdhdt — the #442 analysis, the change-site identification, the build-flag audit, and the correctness-gate framing. Co-authored / reported-by.

[cdhdt]

This is really clean, @woolcoxm — thanks for picking it up and, especially, for holding it to the oracle gate. Splitting the softmax-softened attention reductions from the load-bearing router matmul is exactly the right call, and the score-vs-vmix distinction is spot on.

A few notes from re-reading it:

vmix bit-identity — agreed, with a footnote. It's bit-identical because each lane recomputes clat[i] = fma(a, Lt[i], clat[i]) independently (no reassociation), and because the scalar clat[i] += a*Lt[i] already contracts to a single FMA under the tree's -O3 (FP-contract is on by default). So the measured 0.0 is exact for the project's build; the only way it would diverge is a -ffp-contract=off build, which colibri never uses. Might be worth a half-line in the code comment so a future reader knows the "0.0" is contraction-dependent, not order-independent.

score — one free perf lever. The dot uses a single acc accumulator, so it's bound by FMA latency (the vfmadd dependency chain), not throughput — the same thing the merged #444 (idot independent accumulators) addressed on the int8 path. Four independent __m256 accumulators, summed at the end, give the out-of-order engine 4 chains to overlap and would likely push past 4.7× on the kvl=512 dot, for a couple more lines. Optional, and orthogonal to correctness.

Independent corroboration on the full model. I built your branch and its parent commit and ran an interleaved greedy 24-token decode on the real GLM-5.2 744B int4 (DIRECT=1, temp 0, same prompt, frozen pin state): the generated tokens come out byte-identical (sha-identical completion, patched == base). So the score reassociation holds on the 744B too — the ~1.8e-6 drift doesn't cross a decode near-tie here. One prompt / 24 tokens so it isn't exhaustive, but it's a real-model datapoint on top of your 30/32 tiny-oracle.

Router half — agreed it stays separate. The top-k flip risk is real and not softmax-softened; that one wants the full eval_glm.py quality A/B, not just the tiny oracle (the caveat from #442). Happy to run that gate on the 744B here when you have the router patch ready — I have the checkpoint local.

Nice work.

---

Reviewed with AI assistance (Claude); the 744B corroboration was an interleaved greedy A/B of your branch vs its parent commit.