优化变压器模型在CPU上的推理速度,特别针对短上下文场景。
FlashAttention for Scalable Vector Architectures

- 采用分块FlashAttention,通过并行头间打包提升向量寄存器利用率。
- 在512位向量长度下预填充阶段提速22倍至42倍,解码阶段提速8倍至11倍。
- 适合小型语言模型在低功耗设备上部署,尤其对长向量执行有显著加速。
在CPU上推理变压器模型的重要性日益增加,尤其是针对小型语言模型(SLMs),其向量架构正成为有前景的执行基础。注意力模块因高内存带宽需求成为主要瓶颈;FlashAttention通过融合操作改善数据局部性并减少中间内存流量来缓解此问题。本文提出FlashAttention-V,一种适用于可扩展向量架构的分块FlashAttention,通过跨注意力头的并行性、头间打包以高效利用超出头维度的向量长度,以及改进向量寄存器使用和内存访问局部性。我们将FlashAttention-V集成到ggml中,基于llama.cpp,在TinyLlama、Llama 3.2、Qwen2.5和Pythia-410M上进行评估,使用gem5模拟器和Banana Pi BPI-F3平台。在Banana Pi BPI-F3上,我们确认跨注意力头的循环重排序与展开是有效优化原则,性能随模型增大而提升,且在短上下文和解码阶段最为显著。仿真分析显示,当向量长度为512位时,FlashAttention-V在预填充阶段相比标量FlashAttention实现22x-42x加速,进一步扩展至64车道和4096位向量长度时额外获得2x-2.5x加速。在解码阶段,使用512位向量长度时比标量FlashAttention快8x-11x,但由于单标记、内存受限执行,性能对向量宽度和车道数的敏感度逐渐降低。我们还发现,量化线性层中的结构瓶颈限制了长向量执行下的算术摊销效果,该现象在RVV和Arm SVE中均一致,表明当前量化格式对长向量可扩展性构成根本挑战。
原文摘要 · Abstract (English)
Inference with transformer models on CPUs is increasingly important, especially for Small Language Models (SLMs), where vector architectures are emerging as a promising execution substrate. The attention module is a major bottleneck due to high memory bandwidth requirements; FlashAttention mitigates this by fusing operations to improve data locality and reduce intermediate memory traffic. In this paper, we present FlashAttention-V, a blocked FlashAttention for scalable vector architectures that adapts efficiently from short to very long vectors by exploiting parallelism across attention heads, inter-head packing to enable efficient utilization of vector lengths beyond the head dimension, and improving vector register utilization and memory access locality. We integrate FlashAttention-V into ggml within llama.cpp and evaluate it on TinyLlama, Llama 3.2, Qwen2.5, and Pythia-410M using gem5 and a Banana Pi BPI-F3. On the Banana Pi BPI-F3, we confirm that loop reordering and loop unrolling across attention heads are effective optimization principles, scaling performance gains with larger models and most pronounced with short contexts and during decoding. Simulation-based analysis shows that FlashAttention-V achieves 22x-42x speedup over scalar FlashAttention at 512-bit VL in prefill, with an additional 2x-2.5x gain scaling to 64 lanes and 4096-bit VL. During decode, FlashAttention-V achieves 8x-11x speedup using 512-bit vector lengths over scalar FlashAttention, with performance showing diminishing sensitivity to vector width and lane count due to single-token, memory-bound execution. We further identify structural bottlenecks in Q8_0 quantized linear layers that limit arithmetic amortization under long-vector execution, consistent across RVV and Arm SVE, indicating that current quantization formats pose a fundamental challenge to long-vector scalability.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。