Optimize the custom SDPA kernel for BF16 - #21789
Open
JakeStevens wants to merge 13 commits into
Open
Conversation
custom SDPA's q@k.T and attn@v matmuls (bf16 inputs, fp32 accumulation)
instantiated the generic scalar gemm_{transa,notrans}_<BFloat16,float,float>
templates -- no vectorization, so bf16 SDPA ran ~5.5x slower than fp32.
Add optimized specializations that use the existing vectorized
internal::bf16_dot_with_fp32_arith (fp32 accumulate, identical numerics to
the scalar path). notrans packs a's strided-k rows into a contiguous buffer
so the dot applies. Kept serial: custom SDPA already parallelizes its outer
head loop and executorch's threadpool does not support nested parallelism.
gemma-3-1b bf16 8da4w: custom_sdpa 668->372ms, Method::execute 977->669ms
(prefill @16c 974->670ms); argmax unchanged (107).
cpuinfo_has_* reads a zeroed struct until cpuinfo_initialize() has run, so bf16_dot_with_fp32_arith silently took the widening fallback for any caller that had not already initialized cpuinfo elsewhere -- the threadpool, say. Nothing failed; the fast dot just never ran. This is not an x86 concern despite arriving alongside one: the check it fixes is the aarch64 one, and an unnoticed fallback there costs the whole bfdot path. Split out of the AVX512-BF16 commit so architectures that do not want the x86 dot still get it. Authored with Claude Code.
Cover internal::bf16_dot_with_fp32_arith and the gemm_{transa,notrans}_
<BFloat16, float, float> specializations against a sequential fp32 reference.
Lengths straddle the vector-loop, cleanup-loop and scalar-tail boundaries of
both bfdot paths (128/32 bf16 per iteration on x86, 32/8 on ARM), the gemm
cases use padded leading dimensions and a range of alpha/beta, and a separate
case pins the beta == 0 overwrite semantics with a NaN-filled output.
Verified the tests bite: with the x86 scalar tail loop deleted they fail at
exactly the non-multiple-of-32 lengths, and they pass with dispatch forced to
the portable fallback. That mutation is also what caught the missing
cpuinfo_initialize() fixed in the previous commit -- before it, the tests
passed no matter what the AVX512 path computed.
Only the dot implementation the host dispatches to is covered by a given run.
gemm_notrans_<BFloat16, float, float> gathers each of a's rows into a contiguous buffer so bf16_dot_with_fp32_arith can consume it. The gather costs k strided loads and buys n dots, so at n == 1 it is one gather per multiply-add. Custom SDPA hits exactly that at decode, where the q block is a single row. a's m-dimension is already contiguous, so the alternative needs no gather: convert to fp32 and accumulate along it, with l outermost so a's column is loaded once and reused across c's columns. Per-tile at the decode shape, the gather-free form is 29x faster on Genoa (AVX512-BF16) and roughly 4x on an M4 Max. Above n == 1 the gather amortizes and is kept, unchanged. Authored with Claude Code.
BLAS has no bf16 entry point, so a bf16 gemm falls to our own kernels while the fp32 overload gets a packed, blocked BLAS. Custom SDPA calls both, so q@K.T -- a short reduction (headSize) feeding an m*n tile -- was a per-output dot competing against a packed BLAS, paying a cross-lane reduction per result. Widen the operands and use the same BLAS. Conversion is O(mk + nk) against O(mnk) of multiply, so it only pays for a short k with enough output columns; the thresholds are empirical and decode's single column is excluded. Done at the SDPA call site rather than inside cpublas so the scratch comes from ctx.allocate_temp alongside the buffers cpu_flash_attention already allocates, rather than a thread_local in shared kernel code that memory planning cannot see. cpublas exposes gemm_uses_blas() so the caller can tell whether widening is worth it; the build flag that answers that is only visible inside that library. BFloat16 only: the widened path reinterprets the operands, and Half was never measured through it. Authored with Claude Code.
On aarch64 the gather-free form keeps winning well past n == 1. Measured on an M4 Max (m=64, k=512, GFLOP/s, runtime-valued dimensions): the gathered dot runs 6 at n=1, 28 at n=8, 37 at n=32 and 47 at n=256, against a fairly flat 25-42 without the gather, so the crossover sits near n=32. x86 keeps the previous behaviour. vdpbf16ps is much stronger there relative to fma, a per-tile measurement on Genoa put the gathered dot ahead by 1.19x at n=256, and this host cannot measure that case. Note the dimensions must be runtime values to see this: with them constant-folded, clang unrolls the gather-free inner loop to ~60 GFLOP/s and the crossover disappears, which is not what the kernel sees. Authored with Claude Code.
gemm_transa_<BFloat16, float, float> issued one bf16_dot_with_fp32_arith per output element. Custom SDPA calls it with k = headSize -- 64 for a typical LLM -- while producing an m x n tile of 512 x 256, so each of the 131k dots per tile pays its own cross-lane reduction over a reduction length of only 64. At that length the horizontal add is a large fraction of the dot, and it is paid once per output. Add bf16_dot4_with_fp32_arith: four dots against a shared vec1, accumulating into four registers that collapse with a single vpaddq tree rather than four independent cross-lane reductions. gemm_transa_ consumes four columns at a time and falls back to the single dot for the remainder. Measured on an M4 Max at the shape and strides SDPA uses (m=512, n=256, k=64, lda=512, ldb=2048, runtime-valued dimensions): 22.9 -> 42.6 GFLOP/s, against 9.8 for the fp32 specialization. End to end on Llama 3.2 1B 8da4w at 2048 context, bf16 relative to fp32: decode 1.09x -> 1.15-1.19x, prefill 0.66x -> 0.70x, and TTFT from 1.51x to 1.42x worse. Only ratios are quoted because macOS cannot pin cores and the absolute numbers drifted 13% between runs; the ratios reproduced across two. Only the aarch64 bfdot path is blocked. Elsewhere the new entry point issues four sequential dots, which is exactly what the caller did before, so nothing changes there. The same reduction sharing should apply to vdpbf16ps, but there is no AVX512-BF16 host here to measure it on. libblas_test passes. Authored with Claude Code.
For a reduced-precision activation dtype, cpu_flash_attention casts the fp32 attention weights down to the activation dtype so attn@V can run as a reduced-in/float-out gemm. That cast only earns its keep if the bf16 attn@V kernel is the one that runs. Above a q block of kMinQBlockForWidenedAV it is faster to leave the weights alone, widen V instead, and let BLAS do the multiply -- which also deletes the cast, so the two effects compound. It is also more accurate, because the weights are no longer rounded to bf16. On a 1984-token prefill the logits move measurably toward the fp32 reference (mean |logit| 1.596078 native, 1.591862 widened, against fp32's 1.591681), argmax unchanged. That is one sample, not a numerics validation. Threshold 64, from the q block sizes that actually occur (qSplitSize is 32, 64 or 256 by sequence length, and 1 at decode). Back-to-back A/B on one build, prefill: 256 gains 7.5-8% (2891 -> 2676ms, 2960 -> 2725ms), 64 gains 2-5%, 32 is unresolvable either way, and decode's single output column costs 15% -- the same failure mode as every other conversion that cannot amortize. Gating at 64 keeps decode at 1.19-1.31x of fp32 while prefill reaches parity. Scratch is buf_qdq_ptr rather than a new allocation: it is already per-thread, ctx.allocate_temp'd, 64-byte aligned and sized kvSplitSize * headSize, which is exactly what a widened V needs, and only the quantized branch uses it. Restricted to BFloat16. Half takes the existing path -- the same argument probably applies, but it was not measured. Authored with Claude Code.
bf16_dot_with_fp32_arith had a native bf16-dot path only for ARM (vbfdotq_f32); x86 fell back to convert-bf16->fp32 + fp32 FMA. Add an x86 AVX512-BF16 path using _mm512_dpbf16_ps (native bf16 pairwise dot, fp32 accumulate), behind __attribute__((target(...avx512bf16))) so it compiles without TU-wide -mavx512bf16, dispatched at runtime via cpuinfo_has_x86_avx512bf16(). Identical fp32-accumulate numerics; falls back to the existing path on non-avx512bf16 hw. Used by custom SDPA's bf16 q@k.T / attn@v dots. gemma-3-1b bf16 8da4w: custom_sdpa 372->56ms, Method::execute 669->351ms (prefill @16c 670->357ms, now faster than the fp32 build's 401ms); argmax unchanged (107). Unchanged from the original commit except that the cpuinfo_initialize() fix it also carried now lands separately, before this, since it fixes the aarch64 dispatch too.
dot4_with_fp32_arith_bfdot gave gemm_transa_ a four-output block on aarch64. Add the AVX512-BF16 equivalent so x86 gets the same structure. The win is larger here than the aarch64 argument alone suggests. The x86 single-dot path already carries four accumulators for instruction-level parallelism, but custom SDPA calls it with len == headSize (64 for a typical LLM), which is below the 128-element main loop: it fills one accumulator and then reduces all four, so three of every four cross-lane reductions run over zeros. Giving each accumulator a distinct output makes the same four reductions do four times the work. Dispatch mirrors the single-dot path, cpuinfo_has_x86_avx512bf16 behind COMPILER_SUPPORTS_X86_BF16_TARGET, and the scalar tail keeps the numerics of the no_bfdot fallback.
Custom SDPA decode reaches both attention GEMMs with `n == 1`, where the existing general loops leave bf16 slower than fp32. Add register-blocked attn@V kernels that keep fp32 output accumulators live while streaming bf16 inputs: AVX512 conversion plus FMA on x86, and NEON `shll` plus `fmla` on AArch64. For Q@K on x86, tile four key rows through the native AVX512-BF16 dot path. Preserve portable fallbacks, `alpha` and `beta` behavior, and vector-width tail handling. Keep the new architecture-specific implementations and dispatch together in one ARM/x86/portable conditional block. On AMD Genoa, bf16 decode rises from 19.58 to 24.77 tok/s at s2048 and from 9.30 to 17.14 tok/s at s8192, beating fp32 by 8.5% and 45.1%. Output digests are unchanged. Validated with the production `portable_lib` build, the decode GEMV correctness harness, AArch64 cross-compilation and qemu execution, clang-format, and lintrunner. Authored with Codex.
Rename the ARM-only compiler capability macro to distinguish it from the corresponding x86 bf16 target check. Validated with clang-format and lintrunner. Authored with Codex.
JakeStevens
requested review from
larryliu0820,
manuelcandales and
mergennachin
as code owners
August 12, 2026 19:38
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21789
Note: Links to docs will display an error until the docs builds have been completed. ❌ 10 New Failures, 4 Unrelated FailuresAs of commit 0b95beb with merge base d92a619 ( NEW FAILURES - The following jobs have failed:
FLAKY - The following jobs failed but were likely due to flakiness present on trunk:
BROKEN TRUNK - The following jobs failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
JakeStevens
changed the title
This PR needs a
|
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
1 participant
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.
Optimizes end-to-end bf16 custom SDPA for both prefill and decode. It adds bf16-input/fp32-output kernels for Q@K and attn@V, four-output dot tiling, native AVX512-BF16 and Arm BFDOT dispatch, and dedicated n == 1 GEMV paths on x86 and NEON.
For larger blocks, operands are widened once and routed through optimized fp32 BLAS, amortizing conversion while keeping attention weights and accumulation in fp32. Decode avoids that materialization: Q@K uses native bf16 reductions, while attn@V widens lanes to fp32 before FMA because each output lane must remain independent. Portable fallbacks, feature detection, tails, and alpha/beta behavior are covered by tests.
Verified on a Genoa server and M4 mac with Llama 3.2 1B instruct