Skip to content

perf(backend-native-cpu): read B once per matmul in the BF16 kernel - #897

Merged
michalharakal merged 1 commit into
developfrom
perf/bf16-matmul-amortize-decode
Jul 30, 2026
Merged

michalharakal merged 1 commit into
developfrom
perf/bf16-matmul-amortize-decode

Conversation

@michalharakal

Copy link
Copy Markdown
Contributor

Follow-up to #896, which found this while fixing #887.

What

skainet_bf16_matmul used i-p-j order, which 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, 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, NarrowFloatMatmulBenchmark in SKaiNET-transformers. The two runs differ by ~3% on the FP32 baseline, so that column is included as the scale:

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, 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/bfmmla on 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 p ascending on both paths, so results are bit-identical to the previous formulation, not merely within tolerance.

Two coverage gaps closed while here:

  • A cross-path test asserts row 0 of a multi-row call matches the same row computed alone, compared on raw bits rather than a tolerance — that is what pins the bit-identical claim.
  • 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.

Full skainet-backend-native-cpu suite 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.

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.
@michalharakal
michalharakal merged commit ebfbacb into develop Jul 30, 2026
14 checks passed
@michalharakal
michalharakal deleted the perf/bf16-matmul-amortize-decode branch July 30, 2026 04:23
michalharakal added a commit to MacOS/SKaiNET that referenced this pull request Aug 11, 2026
Bump version 0.37.0 -> 0.38.0 (gradle.properties, docs/antora.yml, README
quickstart). Promote CHANGELOG [Unreleased] to [0.38.0].

The [Unreleased] section only carried the dynamic-dimension work (SKaiNET-developers#891),
so this also writes up the narrow-float line that merged without changelog
entries: the shared BF16/FP16 codec and KEEP_NATIVE loading (SKaiNET-developers#886), the
zero-copy transpose for input-major narrow weights (SKaiNET-developers#888/SKaiNET-developers#895), the native
FFM FP16 kernel and the missing matmulFp16 accessor behind it (SKaiNET-developers#887/SKaiNET-developers#896),
and read-once weight tiling in both native narrow kernels (SKaiNET-developers#897). Records
the measured result -- both narrow formats now beat the FP32 SGEMM, where
FP16 used to be 2-18x slower -- and the finding that these kernels are
compute-bound on the FMA chain at batch 16 rather than bandwidth-bound, so
the next win is not more layout work.

Two behaviour notes worth a reader's attention are called out rather than
buried: Fp16Codec.decode now quiets NaN, which changes exactly the 1022
signaling patterns and nothing else, and a row-major narrow buffer still
takes the generic transpose path on purpose.

Refresh README "What's New" around the two headline features and add a
Contributors (0.38.0) section.

No tag yet.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants