FlashAttention 是怎么省显存的？分块计算

大模型推理时最吃资源的环节之一是注意力计算，FlashAttention 这个名字几乎人人都听过，但它到底怎么省显存，多数人只停留在「更快」的印象上。这篇就把分块计算拆成三步，看每一块的进出账。

标准注意力的显存账怎么算

标准注意力要先算出一张 N×N 的分数矩阵：查询和键两两打分，N 是序列长度。这张中间矩阵必须完整落地，序列一长，显存占用和读写带宽同时吃紧。进来的是查询、键、值三组向量，发生的是全量打分加 softmax，出去的是中间矩阵加最终输出，中间那张大表就是瓶颈所在。

省显存的第一反应是「少算一点」，可注意力语义要求每个位置都能照看到所有位置，直接砍掉会改变结果。所以问题不在算不算，而在中间量怎么放：放法不改，显存账就永远跟着序列长度平方级上涨。

换句话说，只要中间矩阵要完整落地，显存就永远追不上序列增长，这是结构性问题，不是调参能解决的。举个直观的数：序列从一千翻到四千，中间矩阵的规模就翻十六倍，带宽压力同步翻倍。这也是长上下文产品早年迟迟推不动的根本原因，不是模型不想读长文，是账本扛不住。

分块计算的三步走

第一步切块。查询按行切成小块，键和值按列切成小块，块的大小跟着硬件缓存容量选，保证每一块能放进片上存储。进来的是长序列，发生的是按块重排，出去的是一批小矩阵。

第二步块内计算加在线融合。每个查询块对每个键值块算局部注意力，用滚动更新的方式合并 softmax 的归一化统计量，整张大表自始至终不完整出现。进来的是小块对小块，发生的是局部打分与增量归一化，出去的是逐步收敛的部分结果。

第三步输出。所有块扫完之后做最终合并，得到与标准注意力在数值上等价的结果。进来的是累加的归一化项，发生的是收尾合并，出去的是注意力输出，全程没有为整张 N×N 矩阵付过显存。

三步的进出账并排看，省显存发生在哪一步一目了然。

三步合起来，注意力计算从「先算完整张大表」变成「边扫边攒」，中间量的存亡方式被彻底改写。这也是它能在不改变结果的前提下把长序列跑起来的原因，省显存与保语义在这套流程里不再打架。

重计算换来了什么

块内数据要重复读，反向传播时也不存中间矩阵、用到时再重算，这叫重计算。它拿算力换显存和带宽，而 GPU 上往往最缺带宽，这笔交换通常划算。

另一个容易被忽略的好处是合并访存：块内数据一次性读进缓存，算完再走，避免了中间矩阵在显存和缓存之间的反复搬运。访存次数降下来，等显存的时间也跟着降。

代价也不是没有：实现更复杂，对块大小、数值精度都更敏感，换硬件时可能要重新调。你在 DeepSeek、Kimi 或豆包里提交长文档、长对话时，服务端能稳稳接住长上下文，背后就有这类分块实现的功劳。它不改变模型学到的知识，只改变注意力计算的排布方式。

• 速查：显存大账在 N×N 中间矩阵；分块让大表不进显存；在线融合保证结果等价；重计算用算力换带宽；块大小跟着硬件缓存走。

想看上下文优化的完整版图，可以顺带读位置插值是怎么拉伸上下文的？从原理到落地、PagedAttention 是怎么管缓存的？像内存分页、稀疏注意力的「稀疏」是怎么决定的？和语义缓存是怎么判断「相似问题」的？。

关于FlashAttention的三个常问问题

• 它主要省的是显存还是速度？：省的是显存和带宽，速度变快是顺带的好处，因为少了中间大表的来回搬运。
• 用了它输出会变吗？：设计目标是数值等价，具体实现可能有微小浮点差异，工程上按等价处理。
• 块大小怎么定？：跟着硬件缓存容量和头维度走，具体数值以官方技术报告为准。