大模型（LLM）推理是当前人工智能领域的核心议题之一，然而，传统的推理系统往往采用单一的静态并行策略，无法充分适应 Prefill（预填充） 和 Decode（解码） 两个阶段不同的计算特性。近期，一篇发表于MLSys 2025 的论文 “SEESAW: High-Throughput LLM Inference through Dynamic Model Re-sharding” 提出了名为 Seesaw 的动态模型重分片技术，通过动态调整并行策略，显著提升了大模型推理的吞吐量。本文将深入解读 Seesaw 的核心原理、关键贡献以及实验结果，探讨其在大模型推理优化方面的巨大潜力。
并行性分析：Prefill与Decode的资源需求差异
Seesaw 的核心在于对 Prefill 和 Decode 阶段的并行性需求进行了深入的量化分析。Prefill 阶段需要并行处理大量的输入 Token，而 Decode 阶段则以自回归的方式逐个生成 Token。这种差异导致两种阶段对计算、通信和数据移动的需求截然不同。
论文中提出了两个关键的观察：
- 观察 1：张量并行（Tensor Parallelism）瓶颈 Prefill。 张量并行将模型的不同部分分配到不同的 GPU 上，需要通过 All-Reduce 操作同步激活值。然而，All-Reduce 的通信开销会随着并行度的增加而显著增长。这是因为，虽然模型权重被分片，但激活值是完全复制的，导致通信量不随并行度变化。此外，All-Reduce 带宽会随着 GPU 数量的增加而下降，进一步加剧了通信开销。在 Prefill 阶段，需要处理大量的 Token，通信开销变得尤为突出，导致张量并行在吞吐量方面表现不佳。举例来说，使用8个GPU进行张量并行时，如果网络带宽较低，频繁的All-Reduce操作会成为性能瓶颈，导致GPU利用率下降，整体吞吐量降低。
公式：Communication Overhead ≈ O(P^2 * A / B)
其中 P 是张量并行的程度，A 是激活的大小，B 是 All-Reduce 带宽。
- 观察 2：流水线并行（Pipeline Parallelism）瓶颈 Decode。 流水线并行将模型的不同层分配到不同的设备上，每个设备处理连续的层块。由于 LLM 解码的自回归特性，一个序列一次只能占用一个流水线阶段，迫使设备以互斥的方式处理批次。这导致每个步骤的有效批大小减少到全局批大小的 1/PP（PP 是流水线并行度），通常被称为微批大小。微批处理需要 PP 倍的前向传递才能处理相同的工作负载。在解码过程中，每个 Pass 生成一个 Token，这导致频繁、重复的权重从内存加载，显著增加了开销。因此，在解码过程中，最小化权重传输对于吞吐量至关重要，而流水线并行在这方面表现不佳。例如，当使用流水线并行时，为了生成一个长度为 100 的序列，需要进行 100 次权重加载，这显著增加了延迟。
与张量并行和流水线并行不同，数据并行在每个设备上复制模型并分配输入数据。虽然数据并行引入的通信开销最小，但它有两个关键缺点：增加的权重传输和更高的内存使用。权重移动量随着副本数量的增加而增加，超过了张量并行的权重移动量。模型复制减少了可用于 KV 缓存的可用 GPU 内存，从而限制了最大批大小并降低了吞吐量。
数据并行可以正交地应用于张量或流水线并行，以提高可扩展性。 然而，在这项工作中，数据并行保持静态，不会动态调整。
总之，张量并行非常适合解码，但在预填充中由于通信开销而受到影响。流水线并行在预填充中有效，但解码受到权重加载效率低下的影响。数据并行是正交的，但伴随着更高的内存使用和权重传输成本，限制了 KV 缓存空间和批大小。在 Seesaw 中，数据并行保持静态。
动态模型重分片：Seesaw的核心技术
Seesaw 引入了 动态模型重分片（Dynamic Model Re-sharding） 技术，允许在 Prefill 和 Decode 阶段采用不同的并行策略。 具体来说，Prefill 阶段采用 cp 并行策略，Decode 阶段采用 cd 并行策略，并自动在两者之间进行切换。这种解耦扩展了配置空间，并支持针对特定阶段的优化，从而提高吞吐量。
为了支持切换，模型权重需要重新从 CPU 内存加载，KV 缓存需要使用 OS 级别的共享内存重新分片，以实现跨 GPU 的协调访问。 动态模型重分片面临的主要挑战是转换开销，尤其是在连续批处理和预填充优先调度下，这会导致频繁的阶段切换。 虽然解码优先调度减少了转换，但会损害吞吐量。
与空间解耦（将 Prefill 和 Decode 分配到不同的硬件分区）不同，Seesaw 的 动态模型重分片 技术能够更加灵活地利用硬件资源。空间解耦可能导致吞吐量不平衡，并且配置空间有限。例如，一个 70B 的模型在 8 个 40GB 的 GPU 上只能被分成 4:4，导致资源利用率很差。
Seesaw 通过在 Prefill 和 Decode 阶段之间动态调整并行策略，克服了静态并行策略的局限性。 例如，在 Prefill 阶段，Seesaw 可以采用流水线并行来最大化吞吐量，而在 Decode 阶段，则可以切换到张量并行来最小化权重加载。这种动态调整使得 Seesaw 能够充分利用硬件资源，并适应不同的工作负载。
缓解转换开销：分层 KV 缓存和最小化转换调度
为了缓解动态模型重分片带来的转换开销，Seesaw 引入了两项关键的优化技术：分层 KV 缓存（Tiered KV Cache Buffering） 和 最小化转换调度（Transition-Minimizing Scheduling）。
分层 KV 缓存 将 CPU 内存用作辅助 KV 缓存，允许更大的 Prefill 批次。 Prefill 生成的 KV 条目被卸载到 CPU 内存，并在解码期间换回。 这使得 GPU KV 缓存能够保持充分利用，同时减少了转换。
最小化转换调度 仅在 CPU KV 缓存已满或为空时才会发生阶段转换：
- Prefill → Decode：当 CPU KV 缓存已满时触发。
- Decode → Prefill：当所有缓存条目都被消耗时触发。
通过这种方式，Seesaw 能够在保持 GPU KV 缓存充分利用的同时，最大限度地减少阶段转换的次数。 KV 缓存重分片在换出期间完成（每个 GPU 写入其 cp 分片），并在换入期间完成（每个 GPU 检索其 cd 分片）。这是使用共享内存实现的，从而实现了高效的 GPU 协调。
例如，假设一个用户输入了一个较长的 Prompt，导致 Prefill 阶段生成了大量的 KV Cache。如果 GPU 显存不足以容纳所有的 KV Cache，Seesaw 会将一部分 KV Cache 卸载到 CPU 内存中。当 Decode 阶段需要用到这些 KV Cache 时，Seesaw 会将它们从 CPU 内存换入到 GPU 显存中。
系统设计与实现：调度器-工作器架构和异步流水线
Seesaw 采用单调度器、多工作器架构：调度器管理请求批处理和任务调度。 每个工作器控制一个 GPU，处理来自其本地队列的任务。 在解码中，每个步骤仅处理 1/PP 的序列，从而优化流水线利用率。该架构支持异步执行和分层 KV 缓存。
虽然模型重新分片和分层 KV 缓存提高了吞吐量，但它们引入了与权重加载和 KV 缓存传输相关的新开销。 每个批次的模型权重重新加载是恒定的，并且可以在大批次上分摊；模型权重重新加载未优化。 KV 缓存交换开销随着批大小的增加而增加，使其更难隐藏。
为了缓解这个问题，Seesaw 实现了 异步流水线（Asynchronous Pipeline），将数据移动与计算重叠。Prefill 期间生成的 KV 缓存直到解码才需要，从而允许换出与正在进行的预填充计算重叠。 由于共享内存无法固定，因此数据传输分为两个阶段：GPU → 固定内存（与计算重叠）和固定内存 → 共享内存（在主机上同时运行）。 每个工作器都包含一个由调度器管理的后台预取线程，以执行非阻塞 KV 缓存交换。 当 GPU KV 内存有可用插槽时，调度器会分配预取任务。 完成交换后，会通知工作器以调度相应的解码任务。 只要输出序列足够长，交换就可以有效地与解码重叠。
具体来说，Seesaw 使用了 Pinned Memory（固定内存） 来加速数据传输。固定内存是操作系统无法移动或换出的主机（CPU）内存的一种类型。这种固定的物理位置允许 GPU 直接内存访问 (DMA)，与常规可分页内存相比，可以实现更快和异步的数据传输。在 GPU 加速系统中，固定内存对于重叠数据传输和计算至关重要，这可以提高整体性能，尤其是在大规模推理工作负载中。
但是，共享内存（例如用于进程间通信的 OS 级别的共享内存）无法固定。 这是因为共享内存由操作系统的虚拟内存系统管理，并且缺乏保证的固定物理地址。 此外，大多数 GPU API（如 CUDA）不支持固定通过共享内存机制（例如 mmap、shm_open）创建的内存区域。 作为一种解决方法，像 Seesaw 这样的系统首先将数据从 GPU 传输到固定内存（以利用与计算的重叠），然后在单独的 CPU 端操作中从固定内存传输到共享内存。
为了提高带宽效率，Seesaw 在 CPU 内存中对 KV 缓存使用 HND 布局（head、序列长度、head 维度）。 这种布局与张量并行对齐，张量并行沿着 head 维度分片，从而减少了非连续内存访问。 通常使用的 NHD 布局（序列长度、head、head 维度）由于在这种上下文中效率低下而被避免。
例如，当 GPU 需要将 KV Cache 卸载到 CPU 内存时，Seesaw 会首先将数据传输到固定内存，然后从固定内存传输到共享内存。 这样可以确保数据传输不会阻塞 GPU 的计算，从而提高了整体的推理性能。
实验评估：吞吐量提升与带宽敏感性分析
为了验证 Seesaw 的性能，论文作者在三种类型的 GPU 上进行了实验：NVIDIA A10、L4 和 A100。 这些 GPU 部署在 AWS 和 GCP 上。 所有 GPU 都使用 PCIe 4.0 (16 GiB/s)，A100 还受益于 NVLink (600 GiB/s)。 每个 GPU 分配 80 GiB 的 CPU 内存。 使用了三个 LLM 进行测试：LLaMA3–15B、CodeLLaMA-34B 和 LLaMA2–70B，均使用分组查询注意 (GQA) 和 float16 精度。
实验使用了两个数据集：ShareGPT（聊天风格，输入 ≈ 输出长度）和 arXiv 摘要（长输入，短输出），分别采样了 2,000 个和 500 个请求。 主要评估指标是端到端吞吐量。 将 Seesaw 与 vLLM 0.5.4（一种广泛使用的开源推理引擎）进行了比较。
结果表明，Seesaw 在所有测试场景中都优于 vLLM，平均吞吐量提升了 1.36 倍，最高提升了 1.78 倍。Seesaw 的性能优势归功于其 动态模型重分片 技术，该技术能够根据 Prefill 和 Decode 阶段的资源需求，动态调整并行策略。
论文还分析了输入输出长度比对 Seesaw 性能的影响。 实验结果表明，当 Prefill 和 Decode 的运行时长平衡时，模型重分片最为有效。 随着输出长度的增加，流水线并行的解码效率下降，导致其吞吐量下降，而张量重型配置则占据主导地位。 在所有情况下，Seesaw 都通过动态适应运行时平衡来实现最高的吞吐量。 预计它在具有可变输入输出长度的实际工作负载中将具有更大的优势。
此外，论文还评估了互连带宽对 Seesaw 性能的影响。 结果表明，在低带宽条件下，流水线并行表现最佳，而在高带宽条件下，张量并行表现最佳。 在从 0.1× 到 50× PCIe 带宽的广泛范围内，Seesaw 通过动态调整并行策略以匹配通信约束，始终实现更高的吞吐量。
结论与展望
Seesaw 通过 动态模型重分片 技术，显著提升了大模型推理的吞吐量。它深入分析了 Prefill 和 Decode 阶段的资源需求差异，并引入了 分层 KV 缓存 和 最小化转换调度 等优化技术，有效缓解了动态模型重分片带来的转换开销。 实验结果表明，Seesaw 在多种 GPU 平台和 LLM 模型上都表现出色，证明了其在大模型推理优化方面的巨大潜力。
未来，可以进一步探索 Seesaw 在更大规模的模型和更复杂的推理场景下的性能表现，并研究如何将其与其他优化技术（例如模型压缩、量化）相结合，以实现更高的推理效率。 Seesaw 的成功经验也为我们提供了新的思路，即通过动态地调整系统配置，以适应不同阶段的计算需求，从而实现更高的资源利用率和更优的性能。