当记忆成为瓶颈: 把序列拆成流水的时机与方法

为什么我们仍然在为“逐个生成”的问题焦虑

你有没有想过，生成一段文本的延迟并不只是模型大小的问题，而更像是一条单行的高速公路上堵车的体验? 在大型语言模型推理里，前一刻还能并行处理的输入，到了生成阶段却被迫一次只放行一个 Token。这种从并行到串行的突变，正是当前推理效率的核心张力。

很多优化聚焦于把模型参数和算力做得更高效，或者把数值精度降到 8 位甚至 4 位。這些都是重要的改进，但它们没有触及一个更基础的问题: 当生成进入自回归解码阶段时，计算性质从算力受限转为显存和带宽受限，也就是典型的Memory Bound。在这个环境里，传统的计算并行技巧会失去很多价值。

当“记忆比算力更稀缺”时，最聪明的工程不是再追求更大的算力，而是重新划分序列和依赖关系，让并行回到赛道上。

---

两相世界: Prefill 与 Decoding 的本质差异

要明白为什么要重新划分流水，我们需要先分清推理的两相行为。

1. Prefill 阶段: 输入 Tokens 可以并行完成一次前向，模型可以像批处理程序一样高效利用矩阵乘矩阵操作。此阶段延迟低，计算密集且善用矩阵乘矩阵优势。

2. Decoding 阶段: 从第一个输出 Token 开始，模型进入自回归模式，必须一次生成一个 Token。自注意力需要关注此前所有生成的 Token，于是产生 KV Cache 来避免重复计算。但是每一步的计算性质从矩阵乘矩阵退化为矩阵乘向量，这使得计算变为 Memory Bound。随着序列增长，KV Cache 占用显存增加，单步延迟也会上升。

这个两相世界带来一个简单但被忽视的结论: 针对 Prefill 优化的硬件和调度策略，并不能自动帮忙解决 Decoding 的瓶颈。你需要不同的工具箱。

---

重新审视流水并行: 什么时候把序列沿维度切开才有意义

在过去几年，有一种看似晦涩但潜力巨大的并行策略被提出: 沿序列维度把 Token 切分到不同设备上并行处理，也就是常说的按 Token 划分的流水并行。它之所以冷冷清清，是因为对通用自回归文本生成而言，负载均衡和依赖调度太难，使得收益并不明显。

但这里有一个关键洞见: 并行策略的好坏不是绝对的，而是和“问题的依赖结构”紧密相关。换句话说，应用的依赖图决定了并行的可行性。

举个直观比喻: 如果你要在一条单行道上运送货物，所有货物必须按严格顺序上下车，任何并行尝试都会撞车。相反，如果你在一座分叉的城市，货物可以分配到不同街区再汇总，那么分工就能显著提升吞吐量。对于模型推理，Diffusion 类型的图像合成属于后者，它天然适合按 Token 或 Patch 做流水并行；而纯文本自回归可能更像单行道。

因此，判断按序列切分是否有意义，可以用两个维度来量化:

1. 依赖范围: 每个输出是否需要关注全量历史，还是只依赖局部上下文。
2. 序列长度与 KV 成本比: 序列越长，KV 占用越大；如果每个 Token 带来的显存成本显著高于单步计算量，说明 Memory Bound 严重。

当依赖范围小且序列很长时，按序列切分更可能产生实际收益。

---

两个实用框架: 计算-内存谱系 和 依赖粒度矩阵

為了把讨论落到实操层面，我提出两个可以立刻使用的决策框架。

框架一: 计算-内存谱系

把任务在谱系上定位: 横轴是每 Token 的计算强度, 纵轴是每 Token 带来的显存成本。谱系的四象限告诉你优先级:

• 高计算低内存: 典型的矩阵乘矩阵密集步骤, 适合数据并行或模型并行。
• 高计算高内存: 同时需要算力和带宽, 适合混合并行策略。
• 低计算高内存: Decoding 的常见位置, 说明应优先考虑减少 KV 或分散 KV 的存储负担。
• 低计算低内存: 不是瓶颈, 用最简单的调度即可。

如果你的任务落在“低计算高内存”区域, 那么按序列切分是有吸引力的方向。

框架二: 依赖粒度矩阵

构建一个二维矩阵, 横轴是依赖范围: 从局部到全局; 纵轴是可切分性: 从天然可分到难以分。每个具体任务或模型把自己标在矩阵上, 位置越靠近“局部+可分”，按 Token 划分并流水并行的收益越高。

例如, 一个图像扩散模型中, 每个时间步的噪声预测可以局部并行; 而纯文本自回归通常落在“全局+难以分”。

---

把抽象变成操作: 一套可执行的工程清单

当你判断按序列切分值得尝试時, 以下是具体可执行的步骤和考虑。

1. 量化 KV 成本: 计算每增加一个 Token 会占用多少显存, 以及单步计算变慢的幅度。
2. 评估依赖图: 把模型的注意力模式、层间依赖画成图, 看是否存在局部化或可拆分的区域。
3. 选择拓扑: 如果你有 NVLink 或类似高带宽互联, 跨卡通信的成本会低很多; 没有时就要更谨慎。
4. 利用混合精度与稀疏性: 在支持的硬件上把权重以 4 位存储, 但在计算时转换到 FP16 或 FP8, 并利用稀疏矩阵乘法来减少通信与存储压力。
5. 考虑 PipeFusion 思路: 在流水并行中把多个阶段的通信与计算合并, 减少小包通信开销, 提升设备利用率。

具体的工程示例: 如果你的生成序列平均长度超过几百 Token, 且模型的自注意力有办法局部化注意窗口, 那么把序列沿时间轴切成若干段, 在多卡上并行生成并周期性合并 KV，可以显著降低每卡的显存占用并提高吞吐。

---

一个直观的类比: 面包店的流水线

想像一个面包店: Prefill 相当于把大量面团一次性放进烤箱, 批量高效。Decoding 就像是手工装饰面包, 每个面包必须完成前一个的某些步骤才能继续。传统做法是让一个师傅从头到尾做完。但如果你把装饰过程拆成独立可并行的小片段, 并把不同面包按阶段在多名师傅间传递, 就像按序列切分的流水并行, 整体吞吐可以翻倍。

但如果每个面包的装饰都高度依赖前一个面包的特殊用料或者风味，分工反而会增加等待，效率会更差。这正是依赖范围决定并行收益的直观体现。

---

Key Takeaways

1. 识别 Memory Bound 场景: 当单步计算变为矩阵乘向量, 并且 KV Cache 成为显存瓶颈時, 按序列切分可能带来比再扩张算力更大的收益。
2. 用依赖粒度做决策: 只有当模型的依赖可以局部化或弱化全局依赖時, 才值得大规模尝试按 Token 划分的流水并行。
3. 优先优化通信拓扑: 高带宽互联可以把跨设备 KV 交换的成本降低好几个数量级。
4. 结合混合精度和稀疏计算: 存储上量化, 计算上使用 FP16 或 FP8, 并启用稀疏化以降低通信与内存负担。
5. 把调度视为艺术: PipeFusion 式的阶段合并和小包通信聚合常常比单纯增加设备更能提升实际上吞吐。

---

结语: 以问题为中心再设计并行策略

我们常犯的一个工程直觉错误是把“更大的模型”和“更多的算力”当作万能钥匙。然而，当生成从并行走向串行時，钥匙不再是算力，而是如何重新划分依赖和记忆的策略。按 Token 划分的流水并行并不是对所有场景的银弹，但把它纳入你的工具箱，并用依赖粒度与计算-内存谱系来判断何时启用，你會发现很多被算力浪费覆盖的问题，其实是可以通过更合适的划分来解决的。

最重要的工程能力不是把更多资源堆到问题上，而是找到问题的结构，然后让系统的形状匹配这个结构。