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.