FlashAttention for Scalable Vector Architectures

Transformer-based inference on CPUs is increasingly critical for edge deployment of Small Language Models. However, the attention module remains a severe performance bottleneck due to high memory bandwidth requirements and large intermediate matrices. While FlashAttention mitigates this on GPUs by fusing operations and improving data locality, existing CPU vector implementations tie the vector length directly to the head dimension. This leaves SIMD resources severely underutilized on architectures featuring long vector registers where vector lengths exceed the head dimension.

The authors propose FlashAttention-V, a blocked FlashAttention algorithm designed to efficiently scale across both short and very long vector architectures. The core innovation lies in exploiting parallelism across attention heads rather than restricting computation to a single head. By combining loop reordering, loop unrolling, and inter-head packing, FlashAttention-V maps multiple independent attention heads into a single wide vector register. Additionally, the design incorporates native support for Grouped-Query Attention and Multi-Head Attention, effectively eliminating redundant memory loads for shared key-value pairs and maximizing on-chip data reuse.

Implemented within the ggml and llama.cpp frameworks and evaluated on RISC-V and Arm SVE platforms using gem5 and physical hardware, FlashAttention-V demonstrates massive performance gains. At 512-bit vector lengths, it delivers 22x to 42x speedups over scalar FlashAttention in the prefill phase, and 8x to 11x speedups during token decoding. Simulation analyses up to 8192-bit vector lengths reveal that configurations with up to 64 vector lanes and 4096-bit vectors maintain strong scalability. Furthermore, the study identifies that Q8_0 quantization formats suffer from structural bottlenecks like heavy packing and masked reduction overheads, presenting a fundamental challenge to long-vector hardware utilization.