[ethanhq]

Overview

Follow-up to #29570 (#29570 (review)): add a BF16 variant of the tensor API FLASH_ATTN_EXT kernel, used with -ctk bf16 -ctv bf16 under the same conditions and thresholds as the F16 one.

The kernel is templated on the Q/K/V/P type. Scores and output are still accumulated in float: with a bfloat P*V destination (BF16 += BF16 x BF16) the speed is the same and the KLD at 32K is ~100x higher (last row below). The tensor kernel and its F16/BF16 instances are now in kernels/fa_tensor.metal, which is only built into ggml-tensor.metallib, so the BF16 instances need no GGML_METAL_HAS_BF16 guard.

Additional information

M5 Pro (20-core GPU), master vs PR:

• test-backend-ops -o FLASH_ATTN_EXT: 4978/4978, with new BF16 cases for head sizes 128, 192/128, 512, 576/512 and large logits (only 64 and 72 were covered).
• FLASH_ATTN_EXT with BF16 K/V, nb=512, kv=4096 and 32768: 1.7-2.5x for all six head sizes. Same speed and same crossover as the F16 variant, F16 is unchanged.
• With EMBED=OFF, the BF16 cases that take the tensor kernel pass; the others still crash as on master (default.metallib has no BF16 FA kernels).

llama-bench, Qwen3.8-27B Q8_0, -fa 1	master, BF16 KV	PR, BF16 KV	PR, F16 KV
pp512 @ d8192	323.8	355.3	354.2
pp512 @ d32768	184.2	265.6	268.5
tg32 @ d8192	9.09	9.25	9.32

KLD vs -fa off -ctk f32 -ctv f32	mean @ 8K	mean @ 32K	same top p @ 32K
F16 KV	1e-6	1e-6	99.957%
BF16 KV, master (output accumulated in half)	1.7e-5	4.05e-4	99.2%
BF16 KV, PR	3e-6	5e-6	99.921%
BF16 KV, PR with bfloat P*V destination	2.1e-5	6.18e-4	99.036%

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES, Codex and Claude Code assisted with writing the kernel variant and running the benchmarks.