Store pre_att_rms_out as BF16

[Mikyx-1]

pre_att_rms_out is only ever the A operand of the QKV MatMuls, and there is no f32 MatMul kernel — MaybeDecompressA converts it to BF16 on every call, with no caching across calls. Storing it as BF16 lets RMSNormBatched write the final type directly and removes that pass.

Its siblings pre_ffw_rms_out, x_bf and att_sums are already BF16, so this applies the "change most activations to bf16" item from the #164 roadmap to one more buffer.

This is the same change proposed in #560 by Fabian Schuetze, rebased onto the current activations. That patch no longer applies (it predates MatStorageT, and had to hand-fix a const float* x); today every consumer — RMSNormBatched, LayerNormBatched, CallMatMul — is already generic over the element type, so it reduces to two type changes.

Honest framing: this is cleanup, not a speedup

The redundant work is measurably gone. Profiler build, Apple M2, gemma2-2b-sfp-pt, 926-token prefill + 64 generated:

	MM.DecompressA calls	time
before	50,930	11.9 ms
after	14,816	3.5 ms

The binary also shrinks 656 KB, since the f32-A MatMul instantiations for these call sites are no longer emitted, and the buffer's footprint halves (37.7 → 18.9 MB at prefill_tbatch 4096).

But end-to-end there is no measurable speedup, matching the null result reported in #560. Interleaved A/B, 4 reps each: prefill 32.96 → 33.28 tok/s (σ 1.20), total 34.80 → 34.08 s. That's inside one standard deviation.

It can't be otherwise: the conversion is O(M·K) sitting in front of an O(M·K·N) matmul, so it is inherently ~1/N of the work it precedes (N ≈ 2560 here). Worth noting the null result is uniform across architectures — #560 saw it on both a non-bf16 x86 laptop and a Cortex-X3 with native bf16, and this is an M2 with NEON_BF16. It isn't a bf16-hardware split.

No new hardware bf16 requirement

MaybeDecompressA returns StridedViewBF in both branches, so the kernel receives the same type and runs identical instructions either way. Targets without HWY_NATIVE_DOT_BF16 widen to f32 via PromoteEvenTo/PromoteOddTo exactly as they did before — that widening is a property of the kernel and the CPU, not of how this buffer is stored. What goes away is one pass over M×K, on every target.

Behaviour is unchanged

• 96 generated tokens are byte-identical before and after (temperature 0, 926-token prompt, gemma2-2b-pt).
• gemma_test cross entropy identical to every printed digit: 1.108557.
• attention_test goldens pass unchanged on NEON_BF16, NEON_WITHOUT_AES and EMU128 — no golden regeneration needed.

FillRandom in attention_test.cc is templated on the element type so it can still fill the buffer; without this the Bazel build breaks.

Testing limitations

• arm64 only; I have no x86 machine to hand. The change is type-only, but the attention_test tolerance is loose (KV_t = BF16 → 3e-2), so that pass is not a tight numerical proof — the byte-identical generation is the stronger evidence.
• Only the Gemma2 text path was exercised at runtime. pre_att_rms_out also feeds DeepSeek MLA, T5Gemma and ViT/Gemma4-ViT; those compile but were not run.
• tiled_attention_test could not be built to verify — it needs gmock+absl and doesn't build under CMake at all (pre-existing). It uses the already-templated FillMatPtrT, so it should be unaffected.

Happy to drop this if it isn't worth the review cycle — the value is genuinely marginal and I'd rather not misrepresent it as a perf win.

🤖 Generated with Claude Code

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.