FlashAttention-V Boosts Transformer Inference on Vector CPUs

Transformer model inference on CPUs is becoming increasingly important, particularly for Small Language Models (SLMs), with vector architectures emerging as a promising execution platform. The attention module is a major performance bottleneck due to its high memory bandwidth demands. FlashAttention addresses this by fusing operations to enhance data locality and reduce intermediate memory traffic. This research presents FlashAttention-V, a blocked FlashAttention specifically designed for scalable vector architectures. FlashAttention-V achieves efficient adaptation across varying vector lengths by leveraging parallelism across attention heads, employing inter-head packing for vector lengths exceeding head dimensions, and optimizing vector register utilization and memory access locality. Integrated into ggml within llama.cpp, evaluations on various SLMs show significant speedups. Simulation analysis indicates 22x-42x speedup over scalar FlashAttention in prefill with 512-bit vector length, scaling further with larger vector widths. During decoding, it yields 8x-11x speedup. The study also identifies quantization formats in linear layers as a bottleneck for long-vector scalability, consistent across different vector instruction sets.