现代 AI 应用越来越依赖于将庞大的参数数量与数百万个令牌的上下文窗口相结合的模型。无论是经过数月对话的 AI 智能体、通过像整个百科全书一样大的案例法进行推理的法律助理，还是在不断扩展的资源库中对 copilot 进行编码，保持远程上下文对于相关性和一致性至关重要。最重要的是，用户期望快速的交互式响应。
对如此庞大的数据进行解码以及让多个 GPU 快速扩展并相互通信的需求不断增长，这凸显了 FP4 计算和 NVIDIA Blackwell 系统提供的高带宽大型 NVLink 域的重要性。本博客中介绍的 Helix Parallelism 是与 Blackwell 共同设计的。与使用超长上下文进行实时解码的最知名的先前并行方法相比，它在给定延迟下的并发用户数量增加了 32 倍。
换言之，它让 AI 智能体和虚拟助理能够以前所未有的速度为更多人提供服务。
(注意：此博客中的上下文是指之前生成的令牌序列，其中间键和值表示形式存储为 KV 缓存，并在每个解码步骤中进行访问。)
解码瓶颈：KV 缓存和 FFN 权重读取
为支持大规模实时解码，系统必须克服解码 (即生成) 阶段的两大瓶颈：
- 键值 (KV) 缓存流式传输：在处理数百万个令牌的上下文时，每个 GPU 必须从每个样本的 DRAM 中读取大量过去令牌 ( KV 缓存) 的历史记录。这种恒定的流式传输反过来会使 DRAM 带宽饱和，增加 token 到 token 的延迟 (TTL) ，并随着上下文长度的增长迅速成为主要的瓶颈。
- 前馈网络 (FFN) 权重加载：在自回归解码期间，生成每个新 token 都需要从 DRAM 加载大型前馈网络 (FFN) 权重。在小批量的低延迟场景中，这种内存访问成本无法很好地分摊，这使得 FFN 权重读取成为主要的延迟来源。
使用传统的并行策略，很难同时优化 KV 缓存流和 FFN 权重加载这两个瓶颈。
我们以张量并行 (TP) 为例：增加 TP 可以通过在多个 GPU 之间分配权重负载和改进 TTL 来帮助减少 FFN 停顿，但只能达到一定程度。在 Llama 模型中使用的 Grouped Query Attention (GQA) 或 DeepSeek 模型中的 Multi-Latent Attention (MLA) 等注意力方案中，多个查询头共享有限数量的 KV 头。如图 2 (c) 所示，当 TP 超过 KV 头的数量时，系统最终会跨 GPU 复制每个样本的数百万个令牌 KV 缓存，以实现自注意力。因此，即使 TP 增加，KV 读取量仍然很高，再次使 DRAM 带宽饱和，并限制了可扩展性。就 MLA 而言，TP 的上限仅为 1，以避免 KV 缓存重复。
那么，开发者如何在不牺牲实时交互性的情况下扩展模型大小和上下文长度呢？Helix Parallelism 提供了一条前进之路。
Helix 执行流程
Helix 是一种混合分片策略，可在时间工作流中分解注意力和 FFN 的并行策略，有效解决数百万个令牌解码期间的 KV 缓存和 FFN 权重读取瓶颈。
图 1 (下图) 显示了 Helix 如何在单个 Transformer 层中编排注意力和 FFN 的执行。受 DNA 螺旋结构的启发，Helix 将多个并行维度 ( KV、张量和专家) 交织成一个统一的执行循环。通过将用于注意力和 FFN 的并行策略解，Helix 允许每个阶段在针对其自身瓶颈进行调整的配置中运行，同时重复使用相同的 GPU 池。Helix 重用方法可保持 GPU 在各个阶段的高效利用，消除计算流经模型时的空闲时间。
注意力阶段
Helix 通过沿 KVP GPU 序列维度对数百万个令牌的 KV 缓存进行分片来应用 KV 并行 (KVP) ，同时在注意力头 (TPA) 中应用张量并行 (TPA) ，其中 TPA 是每个 QKV 投影分割的 GPU 数量，并保持小于等于 KV 头数量以避免重复。
图 2 (d) 中还通过简化的玩具示例说明了这种分片策略。这将导致总共 N* 个 KVPxTPA GPU 协同进行注意力计算，而不会在 GPU 之间重复 KV 缓存。此处的 N 表示用于端到端推理执行的 GPU 总数。同一组 N 个 GPU 将在 FFN 阶段重复使用。
为避免注意力前的全集，Helix 可确保每个 KVP GPU 拥有与其本地 KV 头相关的所有查询头，并冗余计算 QKV 投影。这将在每个 KV 分片上启用完全本地的 FlashAttention。
在本地 FlashAttention 之后，沿着 KVP GPU 中的查询头维度的单个多对多，交换部分注意力输出和对数和表达式标量。重要的是，此通信成本随批量大小和隐藏维度而扩展，并且独立于 KV 缓存长度，即使上下文长度扩展为数百万个令牌，也能高效运行。每个 GPU 在本地重建 SoftMax 标准化的精确输出。
这种多对多还会触发将 KVP GPU 重新分配为 TP 组 (TP = N+ KVP x TPA) ，用于注意力输出线性计算。至关重要的是，这种多对多阶段受益于 NVLink/ NVL72 的高带宽互连，可实现大型 GPU 数量之间的快速集合通信。图 2 概要介绍了关注的不同分片方案及其相应布局。
为进一步降低 TTL，我们引入 Helix 重叠工作流 – 批量 (HOP-B) ，这是一种细粒度工作流技术，可重叠整个批量的通信和计算，如图 3 所示。计算一个 token 的注意力输出后，Helix 会启动该 token 的多对多交换，同时计算下一个 token 的注意力。这种紧密重叠隐藏了有用工作背后的通信延迟，保持了较高的 GPU 利用率，并进一步加速了实时解码。
在图 3 (上图) 中，顶部显示了在没有 HOP-B 的情况下，八个请求以锁步方式执行注意力计算。然后是顺序多对多通信。图表底部显示，使用 HOP-B 时，一个请求的通信与下一个请求的计算重叠，通过细粒度流水线减少 TTL。
FFN 阶段
注意力集中后，系统会重新调配相同的 N* KVPxTPA GPU 池，而不会在空闲时间内执行 FFN 块。多对多步骤的输出已按隐藏维度在 N 个 GPU 中进行划分，因此注意力后线性投影可在 TP 模式 (TP = N) 下立即运行。每个 GPU 使用其权重分片执行局部矩阵乘法运算，并参与跨 TP = N GPU 的全局归约，以构建正确的输出。
在注意力后线性投影之后，Helix 使用密集模型中的 1D TP (N+ TPF) 布局或 MoE 模型中的 2D TP x 专家并行 (N+ TPFxEP) 网格重新配置相同的 N 个 GPU 以进行 FFN 计算。
分布式 KV 连接
在解码期间，每个新 token 都会广播到所有 KVP GPU 进行查询计算。为了防止出现 DRAM 热点，Helix 以循环方式在 KVP 等级中错开 KV 缓存更新，例如，令牌 1 到 16 转到 KVP 0，17 到 32 转到 KVP 1，以此类推。这可确保统一的 KV 增长、平衡 GPU 之间的显存使用情况，并保持一致的吞吐量，而不管序列长度或批量大小如何。
Blackwell 上的模拟结果
Helix 为长上下文 LLM 解码设置了新的性能基准。图 4 (下图) 显示了在使用 (假设) 100 万个令牌的上下文进行解码时，DeepSeek-R1 671B 的标准化吞吐量延迟帕累托边界。通过对数千种配置进行全面模拟、系统化地改变模型分区策略 ( TP、EP、PP 和 KVP) 和批量大小，以找到吞吐量 – 延迟的最佳权衡，Pareto 前沿得到了推导。具体而言：
- 对于固定的延迟预算，Helix 可以将并发用户的数量增加多达 32 倍 (即，使 token/s/ GPU 增加 32 倍) 。
- 对于低并发设置，Helix 可将用户交互性提升高达 1.5 倍 (即，将可实现的最小 TTL 降低高达 1.5 倍) 。
通过在所有可用设备中对 KV 缓存和 FFN 权重进行分片，显著降低 DRAM 压力并提高计算效率，实现了这些提升。Helix 推动了吞吐量延迟方面的发展，即使在较低的延迟下也能实现更高的吞吐量。如需了解更多详情，请参阅此处的论文。
敬请关注
与 Blackwell 的最新功能共同设计的 Helix Parallelism 为在不影响交互性的情况下大规模提供数百万个令牌模型提供了蓝图。请持续关注我们将此优化引入推理框架。
