Softmax codegen wins for the canonical fp8 prefill shape (b1, sq=sk=75600,
hq=hk=5, d128, non-causal), matching the hand-tuned ASM softmax instruction mix:
* Packed score shift (UA_FA4_PACKED_SHIFT, default on): each thread holds one
rowmax, so the shift addend (-scale_s*max) is uniform across the thread's score
elements. Broadcast it and emit v_pk_fma_f32 (2 f32/instr) instead of 64 scalar
v_fma_f32 -> 237 v_pk_fma in the ISA. Bit-identical. +4.5%.
* Packed alu1 o_acc partial rescale (UA_FA4_PACKED_ALU1_RESCALE, default on):
pack the 6-register *= o_acc_scale with v_pk_mul_f32, halving the asm-volatile
scheduling boundaries. Bit-identical. +4% (barrier-wait drops 2504->1544 cyc).
Combined +8% (1649 -> 1783 TF/s standalone; ~1774 TF/s contiguous prod) over the
prior baseline; full regression matrix (bf16/fp8 prefill/decode/splitkv/long/ps16
+ fixtures) PASS.
Also consolidates the fastest fp8 prefill config as the compile-time default:
kv128 tile (UA_PREFILL_D128_BLOCKSIZE=128) + cooperative K/V load + single-sp +
wide 32x32x64 MMA, all 0-spill.
Gated-off experiments kept with measured verdicts: UA_FA4_EXP2_APPROX (Schraudolph
2^x, null -- exp is hidden by MFMA overlap) and UA_FA4_PACKED_ROWSUM (-13%, serial
chain worse than the tree reduce).
Co-authored-by: Cursor <cursoragent@cursor.com>
29e0f75e19