Stanford CS336 | Assignment 2 - FlashAttention-v2 Pytorch & Triotn实现在Transformer架构的工程优化中,注意力机制的计算效率是核心瓶颈之一。标准的缩放点积注意力(Scaled Dot-Product Attention)存在 O(T²d) 的时间复杂度和内存占用问题——当序列长度T超过1k时,显存消耗会急剧增加,甚至导致训练中断。为解决这一问题,FlashAttention-v2通过分块计算和LogSumExp数值优化,在保持精度的前提下,将显存占用降低至O(Td),同时通过硬件感知优化提升计算速度。