Attention is a fundamental computational kernel that accounts for the majority of the workload in transformer and LLM computing. Optimizing dataflow is crucial for enhancing both performance and energy efficiency in attention computation. This optimization involves a range of decisions, such as tiling, computation ordering and buffer management, and can be applied at both intra-operator and inter-operator levels, resulting in a highly complex decision space. We propose a new approach to cross-operator dataflow optimization. Its centerpiece is an analytical performance model that spans a large decision space and enables matrix-based encoding of multiple candidate solutions. Built on this foundation, a vast number of solutions can be evaluated rapidly, and with the aid of an effective pruning technique, the optimal solution can be identified through exhaustive enumeration. We refer to our method as MMEE (Matrix Multiplication Encoded Enumeration). The ability to efficiently enumerate a large design space allows MMEE to deliver higher-quality solutions at a substantially faster speed compared to prior approaches. The MMEE approach is evaluated across various test cases for different accelerator configurations. For energy-driven optimization, MMEE reduces energy consumption by 48%-50% and latency by 31%-69%, compared to state-of-the-art methods. For latency-driven optimization, MMEE achieves simultaneous reductions of 40%-50% in energy consumption and 40%-69% in latency, respectively. Additionally, MMEE is $64\times$ to $343\times$ faster than previous works.

翻译：注意力是一种基础计算核，在Transformer和LLM计算中占据绝大部分工作负载。优化数据流对于提升注意力计算的性能和能效至关重要。该优化涉及一系列决策，如分块、计算顺序和缓冲区管理，并可应用于算子内和跨算子层面，从而形成高度复杂的决策空间。我们提出了一种新的跨算子数据流优化方法。其核心是一个分析性能模型，该模型覆盖了广阔的决策空间，并支持基于矩阵对多个候选解进行编码。基于此，可以快速评估大量解，并借助一种有效的剪枝技术，通过穷举枚举确定最优解。我们将此方法称为MMEE（矩阵乘法编码枚举）。相较于先前方法，MMEE能够高效枚举大型设计空间，从而以更快的速度提供更高质量的解决方案。该方法在不同加速器配置的多个测试案例中进行了评估。针对能效驱动优化，与最先进方法相比，MMEE将能耗降低了48%–50%，延迟降低了31%–69%。针对延迟驱动优化，MMEE分别实现了能耗减少40%–50%和延迟减少40%–69%。此外，MMEE的速度比先前工作快$64\times$至$343\times$。