从 Tensor Shape 重新理解 Attention
沿着 batch、head、sequence 和 dimension 四个轴,追踪 Attention 的计算与存储。
公式告诉我们 Attention 做什么,Tensor Shape 则告诉工程系统需要搬运多少数据。
四个核心轴
设输入为 [batch, sequence, hidden],投影后通常重排为 [batch, head, sequence, head_dim]。
计算路径
QK 乘法生成序列间关系,softmax 后再与 V 相乘。序列长度的平方项解释了长上下文的主要压力。
从算法到系统
当 shape 明确后,FlashAttention、分块和融合算子的设计目标就更容易理解。