Optimized blas bf16 sdpa - #21553
Draft
JakeStevens wants to merge 3 commits into
Draft
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).
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. Both dispatch checks now go through cpuinfo_initialize(): cpuinfo_has_* reads a zeroed struct until it runs, so a caller that has not already initialized cpuinfo (via the threadpool, say) silently got the fallback. This was latent on the pre-existing ARM check too. 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).
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.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21553
Note: Links to docs will display an error until the docs builds have been completed. ❌ 78 New Failures, 1 Cancelled Job, 2 Unrelated FailuresAs of commit a591205 with merge base e3571c2 ( NEW FAILURES - The following jobs have failed:
CANCELLED JOB - The following job was cancelled. Please retry:
FLAKY - The following jobs failed but were likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
JakeStevens
marked this pull request as draft
August 4, 2026 20:04
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.
Custom SDPA runs its q@k.T and attn@v matmuls with bf16 inputs and fp32 accumulation, which instantiates gemm_transa_<BFloat16, float, float> and gemm_notrans_<BFloat16, float, float>. Neither had an optimized specialization, so both fell through to the generic scalar templates
This PR adds vectorized specializations for those two instantiations on top of the existing internal::bf16_dot_with_fp32_arith, which accumulates in fp32 and so is numerically equivalent to the scalar path it replaces It also makes bf16_dot_with_fp32_arith itself faster on x86 by adding an AVX512-BF16 path using mm512_dpbf16_ps, behind m__attribute_((target(...))) so the feature is enabled for that one function rather than the whole translation unit, and dispatched at runtime via cpuinfo_has_x86_avx512bf16.
Finally, it routes both dispatch checks through cpuinfo_initialize(). cpuinfo_has_* reads a zeroed struct until initialization runs, so a caller that hasn't already initialized cpuinfo will fall back to the scalar. This fixes this without relying on someone else (threadpool, in ET case currently) calling it first.