arXiv:2411.17720cs.DCcs.AI2024-11中稿 · MLSys 2025,被引 3

针对边缘设备设计高效注意力计算,提升速度并降低能耗。

MAS-Attention: Memory-Aware Stream Processing for Attention Acceleration on Resource-Constrained Edge Devices

  • 通过分层分块调度向量与矩阵计算单元,实现异构资源并行处理。
  • 在真实边缘芯片上实现最高1.76倍加速,能耗降低54%。
  • 适合资源受限的边缘推理场景,尤其适用于移动和嵌入式设备。

基础模型的兴起推动了自然语言处理、计算机视觉等领域的突破,其中注意力机制因能有效捕捉序列关联而成为核心组件。然而,随着上下文长度增长,注意力带来二次方复杂度的内存与计算开销。尽管已有大量融合型精确加速算法在数据中心级GPU上应用,但在计算单元有限、片上缓存紧张的边缘神经加速器上仍面临挑战。本文提出一种面向内存受限边缘加速器的精确注意力推理加速方案,通过并行利用向量处理单元与矩阵处理单元,在多层级分块策略下将注意力计算划分为向量与矩阵两个数据流,同时遵守任务依赖关系。通过搜索最优分块因子以最大化并行度,并引入主动缓存覆盖策略避免无效缓存溢出。基于开源仿真框架的实验表明,在边缘场景下相较最先进的融合方法FLAT,最高提速2.75倍,能耗降低54%。真实边缘神经处理单元上的测试进一步验证,注意力运算最高提速1.76倍,且不影响模型输出精度。

原文摘要 · Abstract (English)

The advent of foundation models have revolutionized various fields, enabling unprecedented task accuracy and flexibility in computational linguistics, computer vision and other domains. Attention mechanism has become an essential component of foundation models, due to their superb capability of capturing correlations in a sequence. However, attention results in quadratic complexity in memory and compute as the context length grows. Although many fusion-based exact attention acceleration algorithms have been developed for datacenter-grade GPUs and accelerators leveraging multi-core parallelism and data locality, yet it remains a significant challenge to accelerate attention on resource-constrained edge neural accelerators with limited compute units and stringent on-chip caches. In this paper, we propose a scheme for exact attention inference acceleration on memory-constrained edge accelerators, by parallelizing the utilization of heterogeneous compute units, i.e., vector processing units and matrix processing units. Our method involves scheduling workloads onto these different compute units in a multi-tiered tiling scheme to process tiled vector workloads and matrix workloads in attention as two streams, respecting the workload dependencies. We search for tiling factors to maximize the parallelization of both compute units while considering I/O overhead, and propose a proactive cache overwrite strategy to avoid undesirable cache spills in reality. Extensive results based on open-sourced simulation frameworks show up to 2.75x speedup and 54% reduction in energy consumption as compared to the state-of-the-art attention fusion method (FLAT) in the edge computing scenario. Further experiments on a real-world edge neural processing unit demonstrate speedup of up to 1.76x for attention as compared to FLAT, without affecting model output accuracy.

注意力机制边缘计算加速优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。