From 6769a7aca9bb7a43d47f348fa3f011fa3ec800fc Mon Sep 17 00:00:00 2001 From: Mikyx-1 Date: Sat, 22 Aug 2026 19:43:56 +0700 Subject: [PATCH] Store pre_att_rms_out as BF16 pre_att_rms_out is only ever the A operand of the QKV MatMuls, and there is no f32 MatMul kernel: MaybeDecompressA in ops/matmul-inl.h 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" roadmap item from #164 to one more buffer. Same change proposed in #560, rebased onto the current templated activations, where every consumer (RMSNormBatched, LayerNormBatched, CallMatMul) is already generic over the element type. This adds no hardware bf16 requirement. MaybeDecompressA returns StridedViewBF in both branches, so the kernel sees the same type and runs the same instructions either way; targets without HWY_NATIVE_DOT_BF16 widen to f32 via PromoteEvenTo exactly as they did before. What goes away is one full pass over M x K, on every target. Measured on Apple M2, gemma2-2b-sfp-pt, 926-token prefill + 64 generated, profiler build: MM.DecompressA: 50,930 calls / 11.9 ms -> 14,816 calls / 3.5 ms The binary also shrinks 656 KB because the f32-A MatMul instantiations for these call sites are no longer emitted. End to end the win is real but humble: ~8 ms of a ~34 s run, below run-to-run noise here, because the conversion is O(M*K) in front of an O(M*K*N) matmul. Output is unchanged. 96 generated tokens are byte-identical before and after, gemma_test's cross entropy is identical to every printed digit (1.108557), and the attention_test goldens pass unchanged on NEON_BF16, NEON_WITHOUT_AES and EMU128. FillRandom in attention_test.cc is templated on the element type so it can still fill the buffer. --- gemma/activations.h | 6 ++++-- gemma/attention_test.cc | 7 ++++--- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/gemma/activations.h b/gemma/activations.h index 724931b9..3c22d64f 100644 --- a/gemma/activations.h +++ b/gemma/activations.h @@ -197,7 +197,9 @@ struct AttentionActivations { MatStorageT vit_K_T; MatStorageT vit_V_T; - MatStorageT pre_att_rms_out; + // BF16 because this is only ever the A operand of the QKV MatMuls, which + // would otherwise decompress it to BF16 on every call; see `pre_ffw_rms_out`. + MatStorageT pre_att_rms_out; MatStorageT att_out; // attention output MatStorageT att_out_reps; // attention output for each thread. MatStorageT softmax_max; // see OnlineSoftmaxState @@ -308,7 +310,7 @@ struct AttentionActivationsPtrs { MatPtrT vit_V_T; // Output of RMSNorm before attention, size batch_size x model_dim. - MatPtrT pre_att_rms_out; + MatPtrT pre_att_rms_out; // Attention output computed from att * V, size batch_size x (q_heads * // qkv_dim). MatPtrT att_out; diff --git a/gemma/attention_test.cc b/gemma/attention_test.cc index af0ae7c6..c7c45f4c 100644 --- a/gemma/attention_test.cc +++ b/gemma/attention_test.cc @@ -43,12 +43,13 @@ HWY_BEFORE_NAMESPACE(); namespace gcpp { namespace HWY_NAMESPACE { -void FillRandom(MatPtrT& mat, uint64_t seed) { +template +void FillRandom(MatPtrT& mat, uint64_t seed) { hwy::RandomState rng(seed); for (size_t r = 0; r < mat.Rows(); ++r) { - float* row = mat.Row(r); + T* row = mat.Row(r); for (size_t c = 0; c < mat.Cols(); ++c) { - row[c] = static_cast(RandomGaussian(rng)); + row[c] = hwy::ConvertScalarTo(RandomGaussian(rng)); } } }