perf(backend-native-cpu): read B once per matmul in the BF16 kernel - #897
Merged
Conversation
i-p-j walks the whole of B once per row of A. For ffn_up 8B at m=16 that is 16 passes over 90 MiB, 1.4 GiB of traffic to do 1.4 GFLOP. Tile j instead and widen each B row once per tile into a small stack buffer, then multiply it into all m rows of C, so B is read once in total. Keep plain i-p-j at m == 1. There every B element is used exactly once either way, so tiling only trades sequential streaming for a column-block walk -- it cost the FP16 kernel 15% at m == 1, and m == 1 is the decode step of inference. Measured on i7-9750H / OpenJDK 21, median ms, fp32 column as the scale (the two runs differ by ~3% on the baseline): shape batch fp32 bf16 before bf16 after q_proj 1B 16 16.15 10.52 9.30 q_proj 8B 16 99.97 65.35 59.42 ffn_up 8B 16 270.27 178.94 145.41 ffn_down 8B 16 254.50 174.43 143.35 ffn_up 8B 1 108.53 56.78 56.94 9-19% at m=16 and unchanged at m=1. Worth noting that is far less than cutting memory traffic 16x would suggest: at 90 MiB the remaining B traffic is a few ms, so what is left is the FMA chain. This kernel is compute-bound at m=16, not bandwidth-bound, and the next real win there is a blocked microkernel or bfdot on ARMv8.6-A+, not more layout work. Accumulation into any given C element stays p ascending on both paths, so results are bit-identical to the previous formulation, not merely within tolerance. The new cross-path test asserts that on raw bits. Two coverage gaps closed while here: every existing parity shape was either m == 1 or n <= 256, so a tiled path would only ever have run as a single full tile with its boundary arithmetic never exercised -- n = 1100 adds two full tiles plus a 76-column remainder. Follows the same change to skainet_fp16_matmul in #896.
aharakal
approved these changes
Jul 29, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Follow-up to #896, which found this while fixing #887.
What
skainet_bf16_matmulused i-p-j order, which walks the whole of B once per row of A. Forffn_up 8Bat m=16 that is 16 passes over 90 MiB — 1.4 GiB of traffic to do 1.4 GFLOP.Tile
jinstead, widen each B row once per tile into a 512-float stack buffer, and multiply it into all m rows of C. B is then read once in total. No allocation enters the kernel.Keep plain i-p-j at m == 1. There every B element is used exactly once either way, so tiling only trades sequential streaming for a column-block walk. It measured 15% slower at m == 1 in the FP16 kernel, and m == 1 is the decode step of inference — the wrong place to lose 15%.
Measured
i7-9750H (AVX2), OpenJDK 21.0.11, median ms per call,
NarrowFloatMatmulBenchmarkin SKaiNET-transformers. The two runs differ by ~3% on the FP32 baseline, so that column is included as the scale:9-19% at m=16, unchanged at m=1.
Worth knowing before reviewing
That gain is far smaller than cutting memory traffic 16x would lead you to expect, and the reason is useful: with B traffic down to 90 MiB the remaining reads are a few ms, so what is left is the FMA chain. This kernel is compute-bound at m=16, not bandwidth-bound. The next real win is a blocked microkernel, or
bfdot/bfmmlaon ARMv8.6-A+ — not more layout work. I would not expect another restructuring of this shape to pay.Correctness
Accumulation into any given C element stays
pascending on both paths, so results are bit-identical to the previous formulation, not merely within tolerance.Two coverage gaps closed while here:
m == 1orn <= 256, so a tiled path would only ever have run as a single full tile with its boundary arithmetic never exercised.n = 1100adds two full tiles plus a 76-column remainder.Full
skainet-backend-native-cpusuite green on linux-x86_64.Platform caveat
Same as #896: I could only build and run the native library for linux-x86_64. This is plain C11 with no intrinsics, so it compiles wherever the existing kernel does, but the AArch64 and MSVC builds are unexercised locally and want CI confirmation.
Ordering note
This restores BF16 as the faster of the two narrow formats at m=16 (145 vs 161 ms against FP16 in #896), which is the expected ordering given BF16's dequant is a single shift. Before this change FP16 had overtaken it purely because FP16 got the amortization first.