Helix Parallelism: Rethinking Sharding Strategies for Interactive Multi-Million-Token LLM Decoding

作者/机构: Nidhi Bhatia, Ankit More, Ritika Borkar, Tiyasa Mitra, Ramon Matas, Ritchie Zhao, Max Golub, Dheevatsa Mudigere, Brian Pharris, Bita Rouhani (NVIDIA Corporation)

A1 主要贡献

本文旨在解决大型语言模型（LLM）在处理百万级Token上下文时，在严格的“词元到词元延迟”（Token-to-Token Latency, TTL）约束下进行实时自回归解码所面临的挑战。随着上下文历史变长，解码过程面临两大核心瓶颈：
1. 前馈网络（FFN）权重读取：生成每个新词元都需要从DRAM中加载巨大的FFN权重，在小批量下此开销无法被摊销，成为延迟主导因素。
2. 长KV缓存读取：在自注意力计算中，读取随上下文长度和批量大小线性增长的KV缓存，会迅速耗尽DRAM容量和带宽。

现有的张量并行（Tensor Parallelism, TP）策略虽然能缓解FFN权重读取的成本，但在注意力计算上扩展性不佳。一旦TP宽度超过KV头的数量（TP > K），就会导致低效的KV缓存重复，限制了并行度和批量大小。为了解决KV缓存扩展问题，近期工作如Medha【7，Medha: Efficiently Serving Multi-Million Context Length LLM Inference Requests Without Approximations, 2025】采用KV并行（KV Parallelism, KVP）来分片KV缓存，但它们在后续的FFN计算中没有复用这些GPU，导致FFN权重加载仍然是瓶颈。

为同时解决上述两个瓶颈，本文提出Helix Parallelism，一种混合执行策略，其核心创新点如下：

1. 解耦注意力和FFN的并行策略：Helix引入了一个时间流水线，在同一组GPU上针对解码过程的不同阶段采用最适合的并行策略。
   • 注意力阶段：应用KVP沿序列维度将KV缓存分片到多个GPU上，并结合在KV头维度上的TP（当TP ≤ K时），从而在不复制缓存的情况下，利用N = KVP × TP个GPU共同处理注意力计算。
   • FFN阶段：立即将同一组N个GPU重新配置，用于FFN权重的并行计算。对于密集型LLM，使用纯TP；对于混合专家（MoE）模型，则使用TP与专家并行（Expert Parallelism, EP）的组合。这种重构允许FFN的TP宽度超过KV头的数量，而不会在注意力阶段引入缓存复制。
2. 引入Helix HOP-B优化通信：为了最小化Helix引入的通信开销，本文提出了Helix HOP-B（Helix Overlap Pipeline – Batch-wise），一种通过批量维度的计算-通信重叠来有效隐藏通信延迟的优化方法，从而在提高GPU效率的同时保持较低的TTL。

通过将并行方案与每个计算阶段的独特需求对齐，Helix能够在固定批量大小下将TTL降低高达1.5倍，并在相同延迟预算下支持高达32倍的批量大小，从而显著提升了吞吐量-延迟的帕累托前沿，使超长序列的实时推理变得可行。

A3 关键观察与设计原则

Helix Parallelism的核心洞见在于：要实现具有数百万词元KV历史的实时解码，必须解耦注意力（attention）和前馈网络（FFN）的映射。图1从宏观上解释了为何这种分离至关重要。Helix在每个Transformer层内引入了一个时间流水线（temporal pipeline），允许同一组GPU在注意力和FFN计算之间复用，并为每个阶段应用不同的并行策略。

图1: KV缓存和线性权重读取的Roofline分析，假设一个密集型LLM，批量B=8，查询头Q=128，KV头K=8，头大小Hsz=128，FFN维度F=65536，运行在GB200 NVL72上。权重和KV缓存均以FP4格式存储和获取。图中未包含TP和KVP的通信开销；这些图仅显示GPU DRAM读取延迟随TP宽度和KVP宽度的变化。详细信息见附录A。（左）DRAM读取延迟与TP宽度的关系。当T P > K时，由于完整的KV复制，收益趋于平缓，凸显了Helix中进行KV分片的必要性。（中）DRAM读取时间与KV长度S的关系。自注意力成本随S线性增长，最终主导延迟。（右）DRAM读取时间与KVP宽度的关系。Helix在注意力计算中应用KVP以减少每GPU的内存流量并实现次线性扩展，从而支持数百万词元的推理。然后，相同的GPU被重新配置用于FFN中的TP或EP，以最小化延迟。

A2 方法细节

2.1 注意力划分

2.1.1 KV划分（KVP）

图2: 不同注意力分片策略的概述。此处以GQA为例，查询头Q=4，KV头K=2。每个块处理一个上下文长度为S的词元。（左）无TP：所有Q和KV头位于单个GPU上；无重复。（中左）TP=2：查询头在2个GPU上分割；由于T P ≤ K，KV头仍能被干净地划分。（中右）TP=4：分片数多于KV头数；GPU必须复制KV缓存以服务其分配的查询头，重新引入了DRAM容量和带宽的低效问题。（右）Helix (TP=2, KVP=2)：Helix使用KVP沿序列长度（S）分片KV缓存，因此每个KVP rank只持有一个切片（S/2）。通过将TP限制在K以内，并将剩余的GPU分配给KVP，Helix避免了重复，并形成了一个二维布局：TP分割头，KVP分割序列。

GPU配置与KV缓存分片。在注意力阶段，Helix将所有可用的GPU配置成一个大小为N = KVP × TPA（其中TPA ≤ K）的GPU池，然后沿着序列维度将KV缓存分片到KVP个GPU上，从而消除了完整的缓存复制，并降低了DRAM占用和带宽需求。为了避免在注意力计算前跨KVP个GPU进行昂贵的查询（queries）All-Gather操作，Helix让每个KVP GPU独立计算完整的QKV投影。具体来说，每个GPU接收完整的输入批次[B, H]，并将其与QKV权重矩阵WQ ∈ RH×(H/TPA)、WK ∈ RH×(⌈K/TPA⌉·Hsz)和WV ∈ RH×(⌈K/TPA⌉·Hsz)相乘，以产生完整的QKV投影。这意味着KVP中的每个GPU都持有自己完整的一套查询头以及相应的键/值投影，因此它可以在其KV分片上独立运行FlashAttention【9，Jay Shah等人，Flashattention-3: Fast and accurate attention with asynchrony and low-precision，2024】。每个GPU会生成一个部分的注意力输出和一个log-sum-exp标量；随后通过一次沿查询头轴的All-to-All操作交换这些片段，每个GPU再对接收到的数据进行重新缩放和求和，从而在一个通信回合内重构出精确的softmax归一化注意力结果，无需额外的同步或归一化步骤【10，Tri Dao等人，Flash-Decoding for long-context inference，2023】。这次All-to-All交换也为后续基于TP的执行重新对齐了数据分区。图2从高层次概述了不同的注意力分片方案及其对应的数据布局。

2.1.2 优化的all-to-all通信

通信开销与序列长度无关。在Helix并行中，通信量与KV序列长度S无关，仅随批次中的查询词元数B和隐藏维度H扩展。这种恒定的单位词元开销使得Helix并行具有高度的可扩展性，即使处理数百万词元的KV缓存，也能实现高效的解码时注意力计算。

2.1.3 批处理维度的计算-通信重叠

Helix HOP-B流水线策略。为了最小化暴露的通信时间，我们引入并部署了Helix HOP-B（Helix Overlap Pipeline – Batch-wise），这是一种细粒度的流水线策略，它将All-to-All通信与正在进行的跨批次维度的注意力计算进行重叠。如图3所示，一旦第一个查询词元的注意力输出计算完成，Helix会立即为该词元启动All-to-All通信，同时并行处理下一个词元的注意力计算。这种重叠保持了硬件的高利用率，并有效地隐藏了通信延迟，进一步降低了实时推理解码的TTL。

图3: KVP All-to-All的暴露时间：（上图，不使用HOP-B）所有8个请求同步执行，每个请求在启动9.6个单位的通信之前消耗16个单位的注意力计算时间，总耗时25.6个单位，无重叠；（下图，使用HOP-B）请求被流水线化，因此在一个请求的通信进行时，下一个请求立即开始其注意力计算。这里每个请求的计算（2个单位）和通信（1.2个单位）按比例绘制为离散块。一条虚线灰线和“TTL Saving”箭头突显了启用HOP-B后，总TTL从基线的25.6个单位减少到17个单位。

2.2 FFN划分

GPU复用与重配。在自注意力计算之后，Helix会复用最初为KV划分和TP配置的同一个N = KVP × TPA的GPU池，来执行FFN计算，其布局会根据模型类型（密集型或MoE）进行优化。图4详细展示了Helix在一个Transformer层中的端到端工作流程。为简化起见，图中省略了层归一化（layer-normalization）和残差连接（residual-addition）。

B:批量大小 S: KV序列长度 N: GPU总数 TPA: 注意力TP宽度 KVP: 注意力KVP宽度 TPF: FFN TP宽度 EP: FFN EP宽度 Hsz: 注意力头大小 Q: 查询头 K: KV头 H: 隐藏维度=Q*Hsz. KT: ceil(K/TPA) QT:Q/TPA F: 密集型FFN维度 M: 专家FFN维度 E: 专家数量 EPT: 每词元专家数 EB: 每专家批量 = B*EPT/E 假设E=EP, 每GPU专家数=1

图4: Helix在注意力和FFN阶段中每个GPU的工作流布局。Helix通过在每层基础上将N个GPU从N = KV P × T P A (T P A ≤ K)配置切换到N = T P F × EP来复用它们。（上）在注意力阶段，每个KV P GPU独立地将完整的批次[B, H]投影到QKV，在其KV分片上运行FlashAttention【9】以产生部分输出和log-sum-exp标量，然后参与一次沿查询头轴的All-to-All；每个GPU用接收到的片段重新缩放和求和，得到精确的softmax归一化张量，随后进行一次TP All-Reduce，形成最终的[B, H]注意力输出。（下）对于FFN，Helix根据模型类型采用两种模式：对于密集型FFN（EP = 1），它将所有N个GPU保留在张量并行模式T P F = N下，计算[B, H] → [B, F/N ] → [B, H]，然后进行TP All-Reduce；对于MoE FFN（EP > 1），它将N个GPU重新划分为一个T P F × EP网格，将词元路由到相应的专家，在每个专家组内对全连接层应用TP，执行专家内All-Reduce，然后是专家间All-Gather，最后通过局部归约得到[B, H]。Helix在这些配置之间无缝切换，实现了零停机时间的流水线并提高了GPU利用率。

后注意力线性投影的数据流。在局部注意力计算结束时，每个KVP GPU已经计算出了其KV分片上、对应查询头以及整个批次B的部分注意力输出。这些输出通过一次沿查询头维度的All-to-All进行交换：每个GPU将其B × H / (KVP × TPA)的部分结果发送给KVP域中的其他所有GPU，并接收相应的数据切片。在与每个分片的log-sum-exp统计数据进行重新缩放和求和后，每个GPU持有其分配到的H / (KVP × TPA)隐藏维度切片的完全归一化注意力输出，但覆盖了整个批次。这实际上形成了一个大小为TPA × KVP = N的TP组。这些归一化的输出被送入后注意力线性投影层，该层以张量并行的方式在相同的N个GPU上运行。每个TP rank持有一份形状为H/N × H的投影权重矩阵分片，在其完整的批次输入（形状为B × H/N）上执行局部矩阵乘法，然后参与N个GPU间的All-Reduce操作。这个All-Reduce将B × H的部分投影聚合为完整的B × H输出，然后通过层归一化进入FFN块。

FFN的动态重配。后注意力阶段从KVP到TP的重配过程，在密集型和MoE模型中是相同的。在后注意力线性投影及其输出在N个GPU上的All-Reduce之后，整个N个GPU池会为FFN重新进行配置（见图4）：
* 密集型FFN (EP = 1)：保持TPF = N，所有GPU协作以分摊权重读取成本，并在小批量上保持低延迟。
* MoE FFN (EP > 1)：重新划分为一个最优的TPF × EP网格，通过选择合适的TPF与EP来最好地匹配专家的规模和数量。

无缝流水线。通过将注意力阶段的KVP × TPA、后注意力线性投影的TP = N、以及灵活的TPF × EP FFN布局串联起来，Helix确保所有N个GPU在一个零停机时间的流水线中保持完全利用，从而在注意力、投影和FFN阶段最大化了可扩展性和吞吐量。

2.3 分布式KV拼接策略

轮询式KV缓存更新。在解码过程中，每个新生成的词元都会广播到所有的KVP GPU，以便每个设备都能访问到当前的查询。然而，Helix在KVP rank之间错开KV缓存的更新，以保持内存的均衡增长，并确保每个GPU均匀地追加KV条目。具体来说，系统会将固定数量的解码步骤（例如16个词元）的KV对追加到KVP Rank 0上的分片，然后切换到KVP Rank 1处理接下来的16个词元，以此类推，以轮询（round-robin）的方式循环遍历所有N个KVP rank。这种分阶段的KV拼接策略保证了无论批量大小或序列长度如何，所有KVP GPU都能对KV存储做出贡献，避免了热点问题，并均匀地分散了KV缓存的增长。

A4 实验环境

本次评估使用NVIDIA最新的GB200 NVL72硬件平台【8，NVIDIA blackwell architecture technical brief, 2024】，并采用FP4精度【11，Bita Darvish Rouhani等人，Microscaling data formats for deep learning, 2023】。

• 硬件配置：
  • GPU：GB200 NVL72 平台（评估1-64个GPU的配置，在一个GB200节点内）。
  • 互联：NVLink。
• 软件配置：
  • 模拟器：使用一个内部开发的高保真模拟器，该模拟器能精确建模GB200硬件的计算和通信成本，包括GPU间的NVLink传输延迟、DRAM带宽限制和FLOP吞吐量。
  • 精度：所有模型权重、KV状态和算术运算均假设使用FP4。
• 模型架构：
  • Llama-405B【12，The llama 3 herd of models, 2024】：一个405B参数的密集型模型，采用分组查询注意力（GQA），拥有128个查询头和8个分组KV头。
  • DeepSeek-R1【13，Deepseek-v3 technical report, 2025】：一个671B参数的混合专家（MoE）模型，采用多头潜在注意力（MLA），所有128个查询头共享一个等效的KV头。
• 实验设置：
  • 上下文长度：模拟解码时推理，KV缓存序列长度达到一百万词元及以上。
  • 基线方法：覆盖了当前最优的划分策略，包括张量并行（TP）、流水线并行（PP）、专家并行（EP）和普通的KV划分（KVP），并对批量大小进行了全面扫描。其中，EP指数据并行注意力与专家并行FFN的组合；普通KVP指Medha风格的分片方法【7】。
  • 评估指标：
    • GPU吞吐量 (tokens/sec/gpu)：每秒每GPU生成的总词元数，反映系统效率。
    • 用户交互性 (tokens/sec/user)：解码TTL的倒数，表示为单个用户生成新词元的速率。
    • 批量可扩展性：在固定的TTL预算下，系统能维持的最大并发用户请求数。

A4 实验结果

通过对超过100,000种配置（包括不同的模型划分策略、批量大小、GPU数量和LLM架构）进行详尽模拟，我们绘制了吞吐量与交互性之间的帕累托前沿。如图5和图6所示，Helix通过将帕累托前沿向外推，显著改善了吞吐量-延迟的权衡，实现了更高的系统吞吐量和更低的单用户延迟。

图5: 在GB200上服务具有100万上下文长度的DeepSeek-R1的帕累托前沿

• DeepSeek-R1上的结果：
  • 性能提升：Helix将用户交互性提高了最多1.5倍，同时支持高达32倍的并发用户（即实现了32倍的GPU吞吐量）。
  • 原因分析：这是因为Helix能够将KV缓存和FFN权重分片到所有可用设备上，减轻了DRAM压力并提高了计算效率。对于采用MLA注意力的DeepSeek-R1，传统的TP策略在TP>1时就会导致KV缓存复制，因此Medha那种将FFN和注意力的TP绑定的方法不适用。

图6: 在GB200上服务具有100万上下文长度的Llama-405B的帕累托前沿

• Llama-405B上的结果：
  • 性能提升：与TP分片相比，Helix将最大可实现交互性提高了1.13倍，并将吞吐量和批容量提高了4倍。
  • 原因分析：收益来源于两点：（1）通过KVP解除了TP因KV复制而产生的扩展瓶颈；（2）在不引入缓存复制的情况下进一步增加了FFN的并行度。与Medha相比，Helix解除了注意力和FFN之间TP宽度的绑定约束，从而获得了优势。此外，Med-ha暴露了所有通信开销，而Helix的HOP-B优化则能有效隐藏这部分开销。

消融研究：HOP-B的影响

我们通过在注意力阶段关闭Helix的批处理维度重叠策略（HOP-B）来分析其影响。

图7: 在GB200上，开启HOP-B与关闭HOP-B在100万上下文长度下的帕累托前沿对比。

• 实验模式：在“HOP-B OFF”模式下，通信和计算严格串行执行，导致GPU空闲，从而降低了交互性。
• 实验结果：
  • 对于Llama-405B（右图），关闭HOP-B导致在固定的GPU吞吐量下，用户交互性（Tokens/s/User）下降了约12%。
  • 对于DeepSeek-R1（左图），关闭HOP-B仅导致约1%的性能下降，因为其all-to-all通信仅占端到端解码延迟的约1%。
• 结论分析：这一鲜明对比表明，当通信开销在总TTL中占比较大时，计算-通信重叠变得至关重要。

A7 补充细节

4 相关工作

与序列并行技术的区别。大量先前工作集中在序列并行技术上，主要应用于训练和预填充（prefill）阶段【14，Loongserve: Efficiently serving long-context large language models with elastic sequence parallelism, 2024】【15，Usp: A unified sequence parallelism approach for long context generative ai, 2024】【16，Distflashattn: Distributed memory-efficient attention for long-context llms training, 2024】【17，Context parallelism for scalable million-token inference, 2025】。这些方法主要通过划分序列维度来提高模型执行的非自回归阶段的内存效率和并行性。虽然这些策略对预填充和训练有效，但它们并未针对解码阶段进行优化，因为在解码阶段，严格的TTL约束以及对不断增长的KV缓存的因果访问需求，引入了独特的系统级瓶颈，尤其是在处理数百万词元的KV历史时。

与耦合式KV并行方案的区别。另一类相关工作探索了KV并行，并将注意力和FFN层的TP绑定在一起以简化执行（例如【7，Medha: Efficiently Serving Multi-Million Context Length LLM Inference Requests Without Approximations, 2025】）。然而，这种紧密耦合限制了可扩展性，尤其是在新的模型架构上，如具有MLA和MoE FFN的DeepSeek-R1。

Helix的独创性。据我们所知，Helix是第一个专门为解决现代LLM架构在长上下文解码瓶颈方面设计的并行框架。通过解耦注意力和FFN的分片策略并引入时间执行流水线，Helix更好地将GPU利用率与每个阶段的计算特性对齐，从而在数百万词元的上下文长度下实现低延迟解码。Helix与包括GQA和MLA在内的多种注意力机制兼容，并且与Blackwell最新的功能（如其大型NVLink域）协同设计。

5 讨论

Helix的通用性。尽管本文主要讨论了通过Helix Parallelism进行的长上下文推理，但需要指出的是，Helix是一种通用的推理优化技术，其应用范围超出了长序列场景。其核心原则——分离注意力和FFN计算并独立分配它们——在各种上下文长度下都适用。

在短上下文场景下的应用。在短上下文场景（例如，上下文长度<4K）中，Helix简化为数据并行注意力和张量并行FFN，这是一种已在现代LLM推理服务框架【18，Nvidia tensorrt-llm】【19，Data parallelism attention for deepseek models】中广泛使用的模式。Helix捕捉到了处理不同上下文长度推理的增长趋势——为传统和新兴的长上下文工作负载提供了一个统一的抽象。

6 未来工作

对稀疏注意力的支持。本文介绍并评估了Helix并行，并展示了其在各种最新的密集注意力架构和MoE模型中优化百万级词元KV历史解码性能的有效性。一个自然的扩展是支持稀疏注意力机制，如原生稀疏注意力（Natively Sparse Attention, NSA）【20，Native sparse attention: Hardware-aligned and natively trainable sparse attention, 2025】，这些机制可以减少KV读取带宽，但不能减少总内存容量需求。鉴于Helix的模块化设计，我们预计其原则能够自然地应用于这些新兴架构。

A5 结论

Helix Parallelism代表了在严格延迟约束下运行的超长上下文LLM解码效率的范式转变。通过时间流水线解耦注意力和FFN层的并行及扩展策略，Helix克服了长上下文解码中的关键瓶颈，提升了系统的交互性和效率。Helix与现代LLM架构完全兼容，包括多样化的注意力机制（如GQA、MLA）和稀疏混合专家模型。其设计也与新兴的GPU平台（如Blackwell系统）无缝对接，确保为未来的硬件和模型趋势做好了准备。

A6 附录

A Roofline分析

本节展示了用于推导图1中所示的KV缓存和FFN权重读取时间的公式。

• B : 批量大小
• Q : Q头数量
• K : KV头数量
• Hsz : 注意力头大小
• S : KV序列长度
• H : 隐藏维度 = Q × Hsz
• F : 中间隐藏维度
• TPA : 注意力TP宽度
• TPF : FFN TP宽度
• KVP : KVP宽度
• bytesparam : 每参数字节数
• MemBW : GPU内存带宽

每个LLM层读取KV缓存的时间由以下公式给出。它假设独立的K和V头来代表像GQA这样的注意力变体。

每个LLM层读取权重的时间由以下公式给出。它假设FFN块中使用SwiGLU激活函数。

图1假设MemBW = 8000 GB/s。