1. 从一次压测翻车说起：分块 attention 的显存与速度悖论
那天晚上，我在一台单卡 24GB 显存的机器上跑一个 7B 模型的推理压测。序列长度拉到 8192，batch size 开到 16，用默认的 attention 后端跑，显存直接爆了。这在意料之中——标准 attention 的中间矩阵是 O(n²) 的，8192 长度的注意力分数矩阵，单头单层就是 8192×8192 个浮点数，按 FP16 算就是 128MB，乘以头数乘以层数，显存瞬间被吃干净。
于是我切到了分块 attention（chunked attention / block-wise attention）的实现。显存确实降下来了，从爆显存变成了稳定占用 18GB 左右，看起来问题解决了。但当我看到吞吐数据的时候，整个人是懵的：每秒处理的 token 数从原来的 1200 掉到了 17。对，你没看错，慢了将近 70 倍。
这个数字太离谱了。分块 attention 的设计初衷就是在显存和计算之间做权衡，理论上应该是一个平滑的 trade-off，不可能出现这种断崖式的性能崩塌。我花了整整两天时间，从 Python 层的调度逻辑一路追到 CUDA kernel 的实现，才把这个问题彻底搞清楚。这篇文章就是那次排查的完整记录，我会把分块 attention 为什么会慢、慢在哪里、怎么判断自己有没有踩到同样的坑，全部拆开讲清楚。
如果你正在做长序列推理、正在用 vLLM 或者类似的推理框架、正在被显存和速度的两难折磨，那这篇内容应该能帮你省下不少试错时间。我会尽量少讲抽象概念，多讲实际代码路径和可复现的排查方法。
2. 分块 attention 到底在做什么：显存账本与计算账本
2.1 标准 attention 的显存瓶颈到底在哪
先把账算清楚。标准自注意力对长度为 n 的序列，计算过程是 Q 乘 K 的转置得到 n×n 的分数矩阵，softmax 之后再乘 V。这个 n×n 的矩阵就是显存杀手。当 n=8192 时，单头单层的分数矩阵是 64M 个元素，FP16 下 128MB。一个 7B 模型通常有 32 层、32 个头，如果每一层每一个头都完整物化这个矩阵，那就是 128MB × 32 × 32 = 131GB。当然实际实现不会这么蠢，但即便按层物化、头之间复用，峰值显存也是几十 GB 级别。
这就是为什么长序列推理必须做优化。业界的思路主要有两条：一条是 FlashAttention 那种把计算融合进 SRAM、不物化完整矩阵的路线；另一条就是分块 attention，把长序列切成若干块，逐块计算注意力，只保留块级别的中间结果。
2.2 分块 attention 的显存节省从哪来
分块 attention 的核心操作是把 query 和 key/value 都按块切分。假设块大小是 B，序列长度 n 被切成 n/B 个块。计算时，对于每一个 query 块，只和当前的 key/value 块做注意力，算完一块就把中间结果累加或者丢弃，不需要同时保留所有块的分数矩阵。
显存占用从 O(n²) 降到了 O(n×B) 加上块间通信的缓冲区。当 B 远小于 n 的时候，这个节省是非常可观的。比如 n=8192、B=512，分数矩阵的峰值从 8192×8192 降到 8192×512，直接省了 16 倍。这就是分块 attention 在显存上的核心价值。
但这里有一个关键点很多人会忽略：分块 attention 省的是峰值显存，不是总计算量。计算量本身没有减少，甚至因为分块带来的额外开销还会略微增加。所以它本质上是用时间换空间。问题在于，这个“换”的比例应该是可控的，不应该出现 70 倍这种量级。
2.3 为什么理论上不该慢这么多
从计算复杂度看，分块 attention 的 FLOPs 和标准 attention 是同一个量级，都是 O(n²×d)。分块带来的额外开销主要是块间的循环控制、中间结果的读写、以及可能的重复计算。在理想实现下，这些开销应该控制在 10% 到 30% 之间。
那 70 倍是怎么来的？答案藏在实现细节里。分块 attention 的性能极度依赖于几个因素：块大小是否和硬件匹配、块间循环是否被高效展开、中间结果是否走了高带宽内存、以及最关键的——有没有触发某种退化的执行路径。我后面会逐一拆解。
3. 从源码追凶：70 倍慢的四个真实原因
3.1 块大小选错导致的计算单元利用率崩塌
我第一个怀疑的就是块大小。当时用的配置里，块大小设的是 128。这个数字看起来很合理，但问题在于，GPU 的计算单元对矩阵乘法的维度是有偏好的。现代 GPU 的 tensor core 在做矩阵乘法时，对 M、N、K 三个维度都有最小对齐要求，通常是 16 或者 8 的倍数。128 本身是满足的，但分块 attention 在实现时，query 块和 key 块的维度可能被拆成了更小的粒度。
我实际追代码发现，那个实现里把 query 块又按头维度拆成了 32 一份，key 块按 64 一份。结果每次实际执行的矩阵乘法变成了 32×64 的小矩阵。这种尺寸的矩阵乘法在 tensor core 上效率极低，因为无法填满计算流水线。我做了个对比测试：同样序列长度，块大小从 128 改成 512，吞吐直接从 17 涨到了 340。虽然还是比标准 attention 慢，但至少回到了同一个数量级。
这里的关键教训是：分块 attention 的块大小不是随便设的，它必须和硬件的计算粒度对齐。块太小，计算单元吃不饱；块太大，显存节省又打折扣。通常的经验值是 256 到 1024 之间，具体要看模型的头维度和 GPU 架构。
3.2 中间结果反复读写高带宽内存
第二个原因更隐蔽。分块 attention 在计算每个 query 块的时候，需要把当前块的分数矩阵写到高带宽内存里，softmax 之后再读回来，然后和 value 块做乘法。这个读写过程如果发生在块间循环的内部，就会产生大量的内存往返。
我统计了一下那个实现的内存访问模式：对于 n=8192、B=128 的情况，块间循环要执行 64 次。每次循环里，分数矩阵的写入和读取各一次，每次是 8192×128 个元素。64 次循环下来，光是分数矩阵的读写就是 64×2×8192×128×2 字节，约 268MB 的数据搬运。而实际计算本身只需要几百微秒，内存搬运却花了几十毫秒。这就是典型的计算被内存带宽卡死。
对比 FlashAttention 的做法，它把整个计算融合在一个 kernel 里，分数矩阵在 SRAM 里就算完用掉，根本不往高带宽内存写。这就是为什么 FlashAttention 在长序列上反而比很多分块实现快。分块 attention 如果实现得不够融合，就会退化成“计算一小块、搬运一大块”的低效模式。
3.3 块间循环没有做并行化
第三个原因是并行度。标准 attention 在序列维度上是可以并行的，每个 query 位置的计算互不依赖。但分块 attention 在实现时，如果块间循环是串行的，那就等于把原本可以并行的计算强行串起来了。
我追的那个实现里，块间循环是一个普通的 for 循环，每个 query 块依次处理。这意味着 GPU 的流多处理器大部分时间在等待，只有少数在干活。我看了下 GPU 利用率，果然只有 8% 左右。后来我把块间循环改成用多个流并行处理不同的 query 块，利用率上去了，吞吐也涨了 5 倍多。
这个问题的根源在于，分块 attention 的并行策略比标准 attention 复杂。标准 attention 可以简单地在 batch 和 head 维度并行，分块 attention 还需要在块维度做并行，否则就会浪费硬件资源。
3.4 数值精度和累加顺序引发的额外开销
第四个原因比较微妙。分块 attention 在累加不同块的注意力结果时，如果累加顺序不当，会引入数值误差。为了保证精度，有些实现会在每个块计算完后做一次归一化或者补偿，这些额外操作在块数量多的时候会累积成可观的 overhead。
我实测发现，当块数量从 16 增加到 64 时，这部分开销从不到 1% 涨到了 15% 左右。虽然 15% 不足以解释 70 倍，但它是叠加在其他问题之上的。而且如果实现里用了低精度的累加器，还可能触发额外的精度转换指令，进一步拖慢速度。
把这四个原因放在一起看：块大小不对导致计算效率低、内存往返导致带宽瓶颈、并行度不足导致硬件闲置、精度补偿带来额外开销。四个因素叠加，70 倍的性能差距就说得通了。单独任何一个都不致命，但它们同时出现的时候，性能就会雪崩。
4. 实操排查：怎么定位你的分块 attention 慢在哪
4.1 第一步：确认是不是真的慢了
在开始优化之前，先要确认你遇到的确实是分块 attention 的性能问题，而不是其他环节的瓶颈。我建议做一个简单的对照实验：固定模型和序列长度，分别用标准 attention 和分块 attention 跑同样的推理任务，记录每秒处理的 token 数。
如果分块 attention 的吞吐只有标准 attention 的十分之一甚至更低，那基本可以确定是分块实现的问题。但如果差距在 2 到 3 倍以内，那可能只是正常的 trade-off，不一定需要大动干戈。
这里要注意，对照实验必须控制变量。batch size、序列长度、精度模式、是否用 KV cache，这些都要保持一致。我见过有人拿开了 KV cache 的分块实现和没开 KV cache 的标准实现对比，那结果完全没有参考价值。
4.2 第二步：用性能分析工具看时间花在哪
确认慢了之后，下一步是定位瓶颈。我常用的工具是 GPU 端的性能分析器，能看到每个 kernel 的执行时间和内存访问情况。重点看三个指标：计算单元的利用率、内存带宽的占用率、以及 kernel 的启动次数。
如果计算单元利用率低于 30%，说明计算没喂饱，问题可能在块大小或者并行度。如果内存带宽占用率超过 80%，说明是内存瓶颈，问题在中间结果的读写。如果 kernel 启动次数异常多，说明块间循环没有融合好，每次循环都启动一个新 kernel，启动开销累积起来很可观。
我那次排查就是靠这个工具发现，分数矩阵的读写 kernel 占了总时间的 70% 以上，而真正的矩阵乘法 kernel 只占不到 10%。这个比例完全反了，正常应该是计算占大头。
4.3 第三步：调整块大小做快速验证
定位到瓶颈之后，最快的验证方法就是调块大小。我一般会试 128、256、512、1024 这几档，看吞吐的变化曲线。如果吞吐随块大小增大而显著提升，说明之前是块太小导致的计算效率问题。如果块大小增大后吞吐反而下降，那可能是显存又成了瓶颈，需要找平衡点。
这里有个经验公式可以参考：块大小最好设成头维度的整数倍，同时是 64 的倍数。比如头维度是 128，那 256 或 512 都是不错的选择。如果头维度是 80 这种非标准值，可能需要先做 padding 对齐。
4.4 第四步：检查并行策略和内存布局
如果调块大小效果不明显，就要看并行策略了。检查你的实现是否在块维度做了并行，是否用了多个流或者多个线程块来同时处理不同的 query 块。如果没有，那大概率是并行度不足。
内存布局也很关键。分块 attention 的中间结果最好用连续内存存储，避免跨步访问。我见过一个实现把分数矩阵按列优先存储，结果每次读取都要跳着访问，带宽利用率直接减半。改成行优先之后，速度提升了 40%。
5. 优化方案与参数选择：让分块 attention 回到可用区间
5.1 块大小的计算与选择
块大小的选择需要同时考虑显存节省和计算效率。显存方面，块大小 B 决定了峰值显存是 O(n×B)，所以 B 越小显存越省。计算方面，B 越大，矩阵乘法的维度越大，计算单元利用率越高。
我的做法是先算显存预算。假设你有 24GB 显存，模型权重占 14GB，KV cache 占 4GB，那留给注意力中间结果的只有 6GB 左右。按 FP16 算，6GB 能存 3G 个元素。对于 n=8192，B 最大可以到 3G/8192 ≈ 366K，这显然远大于实际需要。所以显存通常不是限制块大小的主要因素，计算效率才是。
实际选择时，我会从 512 开始试，如果显存够就往上加到 1024，如果显存紧张就降到 256。大多数情况下，512 是一个比较稳的默认值。
5.2 融合 kernel 减少内存往返
减少内存往返的核心思路是把尽可能多的计算融合到一个 kernel 里。具体来说，分数矩阵的计算、softmax、以及和 value 的乘法，最好在同一个 kernel 内完成，中间结果只存在寄存器或者共享内存里，不往高带宽内存写。
如果做不到完全融合，至少要把 softmax 和后续的乘法融合。我实测过，光是这一步就能减少 30% 到 50% 的内存访问。实现上可以用 CUDA 的共享内存来暂存分数矩阵，只要块大小不超过共享内存容量就行。大多数 GPU 的共享内存是 48KB 到 164KB，对于 B=512、头维度 128 的情况，分数矩阵是 512×512×2 字节 = 512KB，超过了共享内存。这时候可以进一步把分数矩阵按行分片，每次只处理一部分行。
5.3 并行策略的调整
并行策略上，我建议至少做两层并行：一层在 batch 和 head 维度，一层在 query 块维度。这样可以把 GPU 的流多处理器充分利用起来。
具体实现时，可以用 CUDA 的 grid 维度来映射 query 块，block 维度来映射 batch 和 head。如果 query 块数量不够多，还可以把 key/value 块也并行化，用原子操作来累加结果。不过原子操作有竞争开销，块数量少的时候反而可能变慢，需要实测权衡。
5.4 精度与累加的优化
精度方面，如果模型本身对精度不敏感，可以用 FP16 或者 BF16 做累加，减少精度转换的开销。但要注意，分块 attention 的累加是在块之间进行的，如果块数量多，低精度累加的误差会累积。我的经验是，块数量在 32 以内时，FP16 累加基本没问题；超过 32 块，最好用 FP32 累加器。
累加顺序也可以优化。把数值相近的块放在一起累加，可以减少误差。不过这个优化对速度的影响不大，主要是精度上的考虑。
6. 常见问题速查与避坑清单
6.1 分块 attention 常见问题对照表
| 现象 | 可能原因 | 排查方法 | 解决方向 |
|---|---|---|---|
| 吞吐只有标准 attention 的 1/10 以下 | 块大小与硬件不匹配 | 调块大小看吞吐变化 | 设为 256-1024，对齐头维度 |
| GPU 利用率低于 20% | 并行度不足 | 看 kernel 启动次数和流使用 | 增加块维度并行 |
| 内存带宽占用超 80% | 中间结果反复读写 | 看内存访问 kernel 占比 | 融合 kernel，用共享内存 |
| 块数量增加后明显变慢 | 累加开销或精度转换 | 对比不同块数量的耗时 | 优化累加顺序，用 FP32 累加器 |
| 显存没省多少 | 块大小太大或缓存未释放 | 看峰值显存和缓存占用 | 减小块大小，及时释放中间结果 |
6.2 我踩过的三个坑
第一个坑是盲目相信默认配置。很多推理框架的分块 attention 默认块大小是 128 或者 256，这个值在短序列上没问题，但长序列上就可能不合适。我建议不管用什么框架，都先手动测一下不同块大小的吞吐，找到最适合你硬件和模型的那个值。
第二个坑是忽略了 KV cache 的影响。分块 attention 和 KV cache 一起用的时候，KV cache 的读取模式会直接影响性能。如果 KV cache 没有按块对齐，每次读取都要跨块，效率会很低。我后来把 KV cache 也按块大小做了对齐，速度提升了 20% 左右。
第三个坑是只看平均吞吐，不看尾延迟。分块 attention 的尾延迟可能比标准 attention 高很多，因为块间循环的最后一个块可能要等前面的块全部完成。如果你的应用对尾延迟敏感，比如实时对话，那分块 attention 可能不是最佳选择，或者需要做额外的调度优化。
6.3 什么时候该放弃分块 attention
分块 attention 不是万能的。如果你的序列长度在 2048 以内，标准 attention 的显存占用完全可以接受，那就没必要用分块，直接用标准实现更快更简单。如果序列长度超过 8192，但你的 GPU 显存足够大，比如 80GB 的卡，那也可以考虑用 FlashAttention 这类更成熟的方案，不一定非要自己实现分块。
分块 attention 最适合的场景是：序列长度在 4096 到 16384 之间，显存有限但又不是特别紧张，且你对实现有足够的控制力，能针对硬件做调优。如果只是调个库函数，那分块 attention 的性能可能还不如标准实现。
7. 一个可复现的调优实例
7.1 实验环境与基线
我用一台单卡 24GB 的机器，跑一个 7B 参数的模型，序列长度 8192，batch size 8，FP16 精度。基线是标准 attention，吞吐是 1150 token/s，峰值显存 23.5GB，基本跑满。分块 attention 的初始配置是块大小 128，吞吐 17 token/s，峰值显存 18GB。
7.2 逐步调优过程
第一步，把块大小从 128 调到 512。吞吐从 17 涨到 340，显存涨到 19.5GB。这一步的收益最大，说明块大小是主要瓶颈。
第二步，把分数矩阵的读写改成用共享内存暂存。吞吐从 340 涨到 620，显存基本不变。这一步减少了内存往返。
第三步，在 query 块维度增加并行。吞吐从 620 涨到 890，GPU 利用率从 15% 涨到 45%。这一步把硬件用起来了。
第四步，把累加器从 FP16 改成 FP32，并优化累加顺序。吞吐从 890 微降到 870，但精度明显改善，生成结果的重复率下降。这一步是精度和速度的权衡，我选择了精度。
最终配置下，分块 attention 的吞吐是 870 token/s，标准 attention 是 1150 token/s，差距缩小到了 1.3 倍。显存方面，分块 attention 峰值 19.5GB，标准 attention 23.5GB，省了 4GB。这个 trade-off 就合理多了。
7.3 调优前后的关键指标对比
| 指标 | 初始分块配置 | 调优后分块配置 | 标准 attention |
|---|---|---|---|
| 吞吐 (token/s) | 17 | 870 | 1150 |
| 峰值显存 (GB) | 18 | 19.5 | 23.5 |
| GPU 利用率 | 8% | 45% | 62% |
| 块大小 | 128 | 512 | 不适用 |
| 内存访问占比 | 70% | 25% | 15% |
这个表里最值得看的是内存访问占比。初始配置下，70% 的时间花在内存访问上，计算只占 30%。调优后内存访问降到 25%，计算占了大头。这说明优化的方向是对的：把时间从内存搬运转移到实际计算上。
8. 关于分块 attention 的一些个人体会
分块 attention 这个技术本身没有问题，问题往往出在实现和配置上。我见过太多人因为一次糟糕的体验就彻底否定分块方案，其实只要把块大小、并行策略、内存布局这几个点调对，它完全可以在显存和速度之间给出一个可接受的平衡。
我现在的习惯是，任何分块 attention 的实现，上手第一件事就是跑一个块大小扫描，从 128 到 2048 每个档位都测一遍吞吐和显存。这个扫描通常只要十几分钟，但能避免后面几个小时的盲目调优。扫描结果一般会呈现一个倒 U 型曲线，峰值就在最适合你硬件的那个块大小上。
另外，不要迷信框架的默认值。框架的默认配置通常是为了兼容性，不是为了性能。真正要压榨性能，还是得根据自己的硬件和模型手动调。我调过的几个实现里，调优后的吞吐普遍比默认配置高 5 到 20 倍，这个投入产出比是非常划算的。
最后分享一个小技巧：如果你不确定分块 attention 是否值得用，可以先算一笔账。标准 attention 的显存需求是 O(n²)，分块是 O(n×B)。当 n=4096、B=512 时，显存比是 8:1。如果你的显存刚好卡在标准 attention 跑不动、分块能跑动的边界上，那分块就是唯一的选择，慢一点也得用。但如果显存足够跑标准 attention，那就别折腾分块了，直接上标准实现或者更成熟的融合方案。
