The llama.cpp project has introduced sparse flash attention (FA) support for the Metal backend, specifically targeting the prefill phase of model inference. This update includes a new kernel that compacts finite mask entries into per-row index lists and extends the vector FA kernel with optional sparse index gathering.

  • The implementation adds a host-side gate to enable the sparse path when n_kv_max is greater than 0 and other conditions are met.
  • A fix was applied to correct row addressing bugs in the sparse flash attention kernels that previously caused failures in multi-row cases.
  • Optimization techniques were added to perform single-pass sparse index compaction, avoiding redundant mask reads by keeping finite positions in per-thread registers.
  • Test cases and performance benchmarks were added to validate head sizes, quant types, and various configuration parameters.

This change improves efficiency for Metal users by reducing memory traffic during the prefill stage when sparse attention hints are utilized.