本页学习说明
直接分步 attention 实现可能物化完整分数矩阵 S=QKᵀ,再写入 HBM;读出 S 做 softmax 得到概率 P,再写回;之后读取 P 与 V 相乘。蓝色箭头为读入片上、红色为写回。课堂路径代表会物化大中间表的一类直接实现,不表示所有现有库都如此。
T为序列长度。稠密关系按 query 与允许的 key 比较,分块不能随意漏掉本应参与的位置。片上工作区放不下整表时,可分块计算并保留可合并累计;难点是 softmax 仍需跨所有有效位置共同归一。
FlashAttention 的IO模型与范围;论文 §2、§3.1。
阅读完整讲义