llama.cpp
3f56636c - ggml-metal: FA tensor kernel: register-light softmax in the per-d-block path

Commit
4 days ago
ggml-metal: FA tensor kernel: register-light softmax in the per-d-block path The NBLK != 2 branch (NBLK == 1 for dv <= 128, NBLK == 4 for dv == 512) still used the shared-memory + 4 simdgroup barriers softmax per chunk, with the O rescale decoded the query from each accumulator element. Apply the same treatment as the NBLK == 2 branch: - online-softmax state (M, S[QPSG], alpha) in per-thread registers, reduced with simd_max/simd_sum: no shared memory and no barriers in the chunk loop; - one global running max (scalar) so the O rescale is a uniform scalar multiply; - skip the O rescale when the max did not change (it does not overlap the tensor core: it serializes between QK^T and PV); - sinks correction with the final global max; - drop the now-unused shared memory layout (the threadgroup buffer stays in the ABI; the host still allocates it). Measured on Apple M5 Max (llama-bench, Llama-2-7B Q8_0, pp2048, -fa 1): d0: tensor 1750, vec 1701 tok/s (1.03x) d4096: tensor 1014, vec 1025 tok/s (0.99x) d8192: tensor 715, vec 717 tok/s (1.00x) (128/128 FA in isolation: 31.3 vs 30.9 ms, was ~0.92x before) test-backend-ops -o FLASH_ATTN_EXT: 4809/4809 pass with the tensor path enabled and disabled.
Author
Parents
Loading