Optimize the custom SDPA kernel for BF16 - #21789

Open

JakeStevens wants to merge 13 commits into

pytorch:mainfrom

JakeStevens:bf16-sdpa-merged

Commits on Aug 11, 2026

[blas] Widen bf16 q@K.T to fp32 and let BLAS run it
JakeStevens
committed
03be6cf
[blas] Raise the bf16 attn@V gather threshold on aarch64
JakeStevens
committed
c1f2277
[blas] Block bf16 q@K.T by four output columns
JakeStevens
committed
95eb4c8
[llm] Keep attention weights in fp32 and widen V for large q blocks
JakeStevens
committed
35f06bc
[blas] Native AVX512-BF16 (vdpbf16ps) bf16 dot on x86
JakeStevens
committed
dcc1977
[blas] x86 counterpart of the four-column bf16 dot
JakeStevens
committed
a0f3e3c
[blas] Vectorize bf16 decode GEMV on x86 and ARM
JakeStevens
committed
defd0bd
[blas] Clarify ARM bf16 target macro
JakeStevens
committed
779b3b5

Commits on Aug 12, 2026

Merge branch 'pytorch:main' into bf16-sdpa-merged
JakeStevens
authored
0b95beb