Follow-up to #978: extend fa_helpers multi-arch coverage beyond AVX-512F

Follow-up to PR #978 (which resolves #975 — CPU flash-attention regression).

PR #978 adds AVX-512F variants of three flash-attention helpers (llamafile_fa_vec_dot_f16, llamafile_fa_fp16_to_fp32_row, llamafile_fa_simd_gemm), runtime-dispatched via sgemm.cpp's GemmFuncs. CPUs without AVX-512F currently see no benefit — they fall through to upstream's helpers (which is exactly what they had before, so it's not a regression, just no win).

For full CPU coverage we should add per-architecture variants the same way #978 did for AVX-512F.

What's universally useful (independent of CPU)

The first three #978 commits already help every CPU:

• 43d62a6 — -fa auto → off on CPU-only. Behavioral default, no codegen involved. Helps everyone.
• e459002 — f32 VKQ accumulator in _one_chunk when hardware lacks native f16 FMA. Helps every CPU except those with native f16 FMA (AVX-512-FP16 on Sapphire Rapids+, NEON FP16 FMA on Apple Silicon / Pi 5, RISC-V Zvfh).
• 8d606c2 — extends the V-tile dequant hook to the tiled FA function. Also helps every CPU via the same fallback path.

So the question is specifically about the fa_helpers and fa_simd_gemm AVX-512F kernels (commits 893120f and bf713ad) — what would AVX2, ARM82, etc. variants look like.

Per-architecture analysis

x86 AVX2 / AVX-VNNI (no AVX-512)

helper	AVX2 wrapper viable?	expected speedup
fa_vec_dot_f16	Marginal — our cosmocc baseline already gets AVX2 in ops.cpp/vec.cpp. An AVX2-tuned wrapper might match or marginally improve via better unrolling.	≤ 10 %
fa_fp16_to_fp32_row	Same — already AVX2 via _mm256_cvtph_ps from the F16C-baseline path.	≤ 5 %
fa_simd_gemm	Same — already gets GEMM_RM=6, GEMM_RN=2 AVX2 codegen. Marginal at best.	≤ 10 %

Verdict: low priority. The AVX2 baseline path our build already produces is what an explicit AVX2 wrapper would also produce. Maybe worth a tiny manual unroll, but not a structural win.

ARM aarch64 baseline (ARMv8.0, no FP16 FMA)

helper	NEON wrapper viable?	expected speedup
fa_vec_dot_f16	Yes — cosmocc baseline aarch64 may not enable NEON-with-FP16 conversion intrinsics.	1.5-2× if baseline is scalar
fa_fp16_to_fp32_row	Same.	1.5-2×
fa_simd_gemm	Marginal — simd-gemm.h already uses GEMM_RM=4, GEMM_RN=4 for __ARM_NEON__.	≤ 10 %

Verdict: worth confirming cosmocc baseline ARM target. If __ARM_NEON__ is defined (it usually is for ARMv8), the helpers may already be reasonable.

ARM aarch64 with FP16 FMA (Apple Silicon, Raspberry Pi 5 — ARMv8.2-a + fp16)

helper	ARM82 wrapper viable?	expected speedup
fa_vec_dot_f16	Yes — major win. ops.cpp baseline aarch64 likely doesn't enable __ARM_FEATURE_FP16_VECTOR_ARITHMETIC (cosmocc baseline armv8.0). An armv8.2+fp16 wrapper would use native NEON FP16 FMA — 2× SIMD lane count + no cvt round-trip.	2-3×
fa_fp16_to_fp32_row	Same — native NEON FP16 conversion.	2×
fa_simd_gemm	Already AVX-512-equivalent tile config on NEON; modest win from FP16.	≤ 20 %

Verdict: highest-value follow-up. Apple Silicon CPU-only and Pi 5 inference would benefit substantially. Build pattern already exists — llamafile/tinyblas_cpu_sgemm_arm82.o uses private TARGET_ARCH += -Xaarch64-march=armv8.2-a+dotprod+fp16.

x86 AVX-512-FP16 (Sapphire Rapids+)

The f16 helpers (fa_vec_dot_f16, fa_fp16_to_fp32_row) could use _mm512_fmadd_ph (native f16 FMA) instead of cvt-to-f32. But in practice SPR users almost certainly have GPUs and won't run CPU-only FA. The e459002 guard (__AVX512FP16__ keeps the f16 accumulator path) is already correct for this case.

Verdict: low priority — niche audience.

RISC-V Zvfh

Similar story to ARM FP16 FMA. RISC-V users on hardware with Zvfh would benefit from a wrapper. Tiny audience today.

Verdict: low priority; revisit when RISC-V Zvfh hardware is common.

Proposed approach

Order by impact / effort:

1. ARM82 fa_helpers variant (Apple Silicon + Pi 5). Highest impact. Pattern is well-established (tinyblas_cpu_sgemm_arm82.o). One new file plus an extra dispatch branch in sgemm.cpp.
2. ARM82 fa_simd_gemm variant. Smaller win but completes the ARM82 coverage.
3. Baseline NEON variants if the cosmocc default ends up scalar.
4. (Probably skip) AVX2 variants — marginal gain, not worth the maintenance.
5. (Probably skip) AVX-512-FP16 / RISC-V Zvfh — niche.

Each variant adds a numerical-equivalence test case to tests/fa_helpers_test.cpp (the test already auto-skips when the helper returns false, so coverage just expands as new variants come online).

Related

• Parent: #978 (AVX-512F helpers — the architectural template)
• Sibling: #980 (dense-model gap, different surface but same fix shape)
• Original: #975 (CPU flash-attn regression)