Mnemosyne: Parallelization Strategies for Efficiently Serving Multi-Million Context Length LLM Inference Requests Without Approximations

文章标题：Mnemosyne：无需近似即可高效服务数百万上下文长度LLM推理请求的并行化策略
作者/机构：Amey Agrawal (佐治亚理工学院), Junda Chen (加州大学圣地亚哥分校), Íñigo Goiri (微软), Ramachandran Ramjee (微软), Chaojie Zhang (微软), Alexey Tumanov (佐治亚理工学院), Esha Choukse (微软)

A1 主要贡献

本文旨在解决大语言模型（LLMs）在处理百万级甚至千万级长上下文推理请求时面临的独特挑战。现有的长上下文训练技术（如环形注意力和条带注意力）无法很好地适应推理场景，因为推理过程包含具有不同延迟约束（首个令牌生成时间TTFT和令牌间生成时间TBT）的预填充（prefill）和解码（decode）阶段。此外，现有方案不支持对长上下文请求进行批处理以提高硬件利用率。

本文提出的系统名为Mnemosyne，通过三项关键创新来支持高效、无近似的交互式长上下文LLM推理：

  • 序列流水线并行（Sequence Pipeline Parallelism, SPP）：一种结合了预填充分块（chunking）和流水线并行的新策略，旨在处理数百万上下文长度的预填充阶段时，避免队头阻塞（head-of-line blocking）问题，同时不牺牲TTFT延迟。
  • 自适应分块（Adaptive Chunking）和KV缓存并行（KV Cache Parallelism, KVP）：自适应分块技术根据工作负载特性动态调整块大小，以平衡预填充和解码延迟。KVP技术则通过在多个工作节点间分散KV缓存来并行化和加速令牌生成，从而在批处理请求中动态权衡TTFT和TBT。
  • Mnemosyne 3D并行：将张量并行（TP）、序列流水线并行（SPP）和KV缓存并行（KVP）相结合。Mnemosyne是首个能够将LLM推理扩展到至少1000万令牌上下文长度的系统，同时通过混合批处理保持高效率，并满足生产级的TBT延迟服务等级目标（SLO，如30ms）。


图1：Mnemosyne在极长上下文上的性能。系统能够支持高达1000万令牌的请求，并通过在128个NVIDIA H100 GPU上高效扩展来保持低延迟，即使对于海量输入也能维持交互式速度。

下表对比了各种用于长上下文LLM推理的并行化技术。


表1：长上下文LLM推理并行化技术对比。*可抢占性（Preemptability）显示了该并行策略是否能与分块预填充结合以支持细粒度的抢占。

A3 背景知识与挑战

2.1 长上下文Transformer

自回归生成模型的两个推理阶段。自回归生成式Transformer模型在推理时分为两个阶段。预填充（prefill）阶段处理输入上下文，为所有输入令牌构建内部表示（即KV缓存）。随后的解码（decode）阶段则基于先前处理的上下文，逐一生成输出令牌。


图2：生成式Transformer模型中的两阶段推理过程。首先，在预填充阶段处理输入提示；随后，在解码阶段使用KV缓存顺序生成输出令牌。

长上下文模型的能力与挑战。近期的研究通过重新缩放位置嵌入【10, Enhancing LLM Context Length with RoPE Scaling, 2024; 29, World Model on Million-Length Video And Language With Blockwise RingAttention, 2024; 46, Scaling Rotational Embeddings for Long-Context Language Models】表明，大语言模型可以被微调以处理百万级令牌的上下文长度。这些长上下文Transformer解锁了新的能力，例如多模态处理和对多本书籍长度的文本数据进行推理。Google的Gemini 1.5模型【39, Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context, 2024】支持高达200万的上下文长度，是这些进步的例证。然而，扩展的上下文窗口带来了巨大的计算挑战，尤其是在预填充阶段，其计算复杂度随输入大小呈二次方增长。

资源需求的量级。以Llama-3 70B模型为例，处理一个包含100万令牌的长请求，仅存储KV缓存就需要320GB内存，而预填充计算则需要惊人的2.4 exaFLOPs。

2.2 性能指标

三个关键性能指标。Transformer模型的推理性能由三个关键指标来衡量：
1. 首个令牌生成时间 (Time to First Token, TTFT)：从提交输入到生成第一个输出令牌的延迟。TTFT由预填充阶段决定，对交互式应用至关重要。
2. 令牌间生成时间 (Time Between Tokens, TBT)：在解码阶段生成连续令牌之间的延迟。TBT影响模型输出的流畅感。
3. 吞吐量 (Throughput)：在线推理服务系统在满足请求延迟的服务等级目标（SLO）约束下能够持续处理的请求负载。吞吐量反映了系统的整体效率。

长上下文推理的现状与挑战。当前的生产服务【22, Announcing Supermaven 1.0; 39, Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context, 2024】支持对高达100万令牌的请求进行交互式推理，TTFT约为50秒，TBT约为20毫秒。然而，当上下文长度接近1000万令牌时，注意力操作的二次方复杂度使得交互式预填充变得极具挑战性。为解决此问题，运营商引入了按需上下文缓存【44, Prompt Caching (beta); 45, Context caching】，允许用户将长提示作为批处理作业提交，然后利用生成的KV缓存进行后续查询的交互式处理。

2.3 流行的并行策略

流水线并行 (Pipeline Parallelism, PP)。PP将模型的层划分到不同的阶段，每个阶段在独立的设备上运行。一个阶段处理完一批数据后，将其传递给下一阶段。PP通过将内存负载分散到多个设备，为KV缓存释放了更多内存，从而支持更大的批处理大小和更高的吞吐量。由于阶段间通信量很小，PP可以跨多个节点进行扩展。然而，PP的主要限制是它无法改善推理延迟，因为流水线阶段之间存在顺序依赖。

张量并行 (Tensor Parallelism, TP)。TP将单个模型层内的张量进行划分，将矩阵运算（如注意力机制中的运算）分布到多个设备上（如图3所示）。通过在层内粒度上并行化计算，TP可以加速大型操作的执行并减少内存瓶颈，因此能有效改善延迟和吞吐量。然而，TP中频繁且大量的通信操作需要低延迟和高带宽，这将其可扩展性限制在NVLINK域内（通常是单个计算节点），与PP的广泛可扩展性形成对比。实践中，通常结合使用PP和TP来优化资源利用、最小化延迟并最大化吞吐量。

长上下文推理的局限。虽然PP和TP的组合可以扩展推理部署以满足长上下文的内存需求，但只有TP有助于降低延迟。由于TP的扩展受限于单个节点，当处理百万级令牌的请求时，这些传统并行化方法无法满足延迟目标。

2.4 批处理策略

连续批处理与预填充-解码干扰。在推理系统中，批处理是提高利用率和吞吐量的常用方法。Orca【49, Orca: A Distributed Serving System for Transformer-Based Generative Models, 2022, OSDI】引入了连续批处理，允许请求在单个迭代的粒度上动态加入或退出批次。然而，简单的迭代级批处理会因长预填充请求的干扰而导致解码延迟的尾部延迟过高【7, Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve, 2024, OSDI】。

解决干扰的两种策略。为同时实现高吞吐和低延迟，出现了两种策略来解决预填充-解码干扰问题：
1. 分块预填充 (Chunked Prefills)【7, Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve, 2024, OSDI; 8, SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills, 2023; 19, DeepSpeed-FastGen: High-throughput Text Generation for LLMs via MII and DeepSpeed-Inference, 2024】：将输入上下文分割成小块，并将这些块的预填充计算与现有的解码迭代“捎带”在一起。这种混合批处理方法允许以较小的增量成本高效计算预填充，并保持可预测的解码延迟。图4展示了这种方法如何调度请求B的预填充块，从而避免阻塞请求A的解码阶段。
2. 预填充-解码分离 (Prefill-Decode Disaggregation)【35, Splitwise: Efficient generative LLM inference using phase splitting, 2024, ISCA; 51, DistServe: Disaggregating Prefill and Decoding for Goodputoptimized Large Language Model Serving, 2024】：将预填充和解码阶段的执行解耦，并将其计算分配到不同的设备上。专用的预填充和解码设备运行同构批次，延迟更可预测。


图4：分块预填充对连续批处理的影响。标准的连续批处理（上图）在长请求B到达时会出现队头阻塞，延迟了请求A的后续解码。分块预填充批处理（下图）将B的预填充块与A的解码捎带在一起，降低了两个请求的延迟，并提高了整体吞吐量。

长上下文推理的适用性问题。尽管分块预填充能有效降低解码延迟，但先前研究【8, SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills, 2023; 51, DistServe: Disaggregating Prefill and Decoding for Goodputoptimized Large Language Model Serving, 2024】认为，由于KV缓存的读取放大效应，分块预填充的计算效率低下，不适合长上下文推理。另一方面，分离技术需要在预填充和解码设备之间移动整个KV缓存，这减少了可用于KV缓存的内存，使其对KV缓存可能达数百GB的长上下文推理吸引力下降。


图3：Transformer注意力并行化策略对比。张量并行（TP）沿头维度分片计算，环形注意力（RA）沿序列维度分布计算并通过循环KV缓存传输，而KV并行（KVP）则将KV缓存分片到不同GPU上。箭头表示数据流，编号步骤显示计算顺序。

3.1 资源需求

预填充和解码阶段的资源特性。LLM推理的两个主要阶段——预填充和解码——表现出不同的资源需求【8, SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills, 2023; 35, Splitwise: Efficient generative LLM inference using phase splitting, 2024, ISCA】。预填充阶段是计算密集型的，因为需要以批处理方式并发处理多个提示令牌。相比之下，解码阶段是内存密集型的，因为每个解码令牌是顺序生成的。随着请求上下文长度的增加，这两个阶段的资源需求呈非对称增长。

表2：方程中使用的符号定义。

预填充的二次方计算增长。在预填充过程中，每个提示令牌需要关注序列中所有在它之前的令牌。因此，预填充注意力计算所需的算术操作呈二次方增长。对于一个有$n$个输入令牌的提示，计算FLOPs ($F_a(n)$) 如下：

解码的线性内存增长。在解码阶段，需要扫描整个提示的KV缓存，这导致内存读取量呈线性增长。KV缓存的内存容量需求 ($M_{kv}(n)$) 和解码阶段的内存读取量 ($R_a(n)$) 可以计算如下：



值得注意的是，预填充的计算FLOPs随输入令牌数呈二次方增长，而KV缓存的内存需求仅线性增长。

计算成为主要瓶颈。现有的在线推理服务（如Gemini【39, Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context, 2024】）为100万输入上下文提供亚分钟级的TTFT和每秒超过40个令牌的输出速率。图5a展示了在单个DGX-H100节点上，为Llama-3 8B模型满足30秒TTFT和20毫秒TBT的SLO时，理论上可支持的最大输入令牌数。计算资源在76.8万输入令牌时迅速成为瓶颈，而内存容量的扩展性最好。反之，图5b显示了随着输入令牌数增加，为满足此SLO所需的GPU数量。100万上下文长度需要20个GPU，200万输入上下文需要80个GPU。由于生产服务旨在提供交互式用户体验，预填充的SLO不能随提示长度呈二次方增长，即使计算需求是这样。这种矛盾导致资源需求出现严重倾斜，预填充阶段的二次方计算成为最显著的瓶颈。


(a) 8个H100 GPU上每种资源类型支持的最大令牌数。(b) 满足给定上下文长度所需各种资源的GPU数量。
图5：服务Llama-3 8B并满足30秒TTFT和20毫秒TBT SLO的理论资源需求，说明计算是交互式长上下文LLM推理工作负载的主要扩展瓶颈。

3.2 并行化长上下文计算

训练并行技术的局限性。近期的工作【11, Striped attention: Faster ring attention for causal transformers, 2023; 30, Ring Attention with Blockwise Transformers for Near-Infinite Context, 2023】提出了新的注意力并行技术来解决长上下文Transformer模型训练的计算需求。例如，环形注意力（Ring Attention）【30, Ring Attention with Blockwise Transformers for Near-Infinite Context, 2023】（如图3所示）通过在参与的工作节点间划分查询（query）来并行化序列维度的计算。然而，将这些先进的训练并行技术应用于推理时面临以下几个挑战：
1. C1. 队头阻塞（Head-of-line blocking）：环形和条带注意力不允许抢占长上下文的预填充阶段。对于数百万上下文长度，这可能导致持续数分钟的队头阻塞。
2. C2. 批处理支持：这些并行方法应用于推理时仅能帮助降低预填充延迟，但不支持批处理或混合批处理，导致系统利用率和吞吐量低下。
3. C3. 刚性并行与可变输入长度：现实世界的服务中输入序列长度各异。环形和条带注意力的令牌分配方式是为满足最长上下文的延迟SLO而优化的。这种分配是刚性的，对于较短的请求也必须复用。注意力操作的算术强度与令牌数成正比（见公式(4)）。在环形注意力中，输入令牌$n$被均匀分配给所有工作节点($p_{ra}$)，导致每个节点的算术强度相应降低。当每个工作节点的令牌数过少时，操作变为内存/网络密集型，KV缓存块的传输成为瓶颈。因此，环形注意力的最佳并行度取决于序列长度，这在延迟目标和硬件利用率之间造成了权衡。


4. C4. 缺乏对解码计算的支持：由于这些技术最初是为长上下文训练提出的，它们可以并行化长上下文推理的预填充阶段，但不能直接扩展到解码阶段，使得这些解决方案对于端到端的长上下文推理是不完整的。


图6：分块预填充中KV缓存读取放大的过程图。随着新块（Q1, Q2）被处理，它们必须再次读取所有先前计算的KV缓存，导致总读取量呈二次方增长。这一现象曾被认为是导致分块预填充对长上下文效率低下的原因。

3.3 高效长上下文服务系统的需求

三大核心需求。基于以上分析，构建一个高效的长上下文服务系统需要满足以下要求：
* R1: 满足长上下文交互式服务的严格TTFT和TBT SLO。
* R2: 提高硬件利用率，以增加每台设备的吞吐量，并降低单次推理的成本。
* R3: 在单个系统内同时高效处理各种上下文长度的请求。

A2 方法细节

4.1 重新审视长上下文的分块预填充

分块预填充的潜力与误解。如2.4节所述，将输入提示分块进行预填充可以实现更好的跨请求和跨预填充/解码阶段的调度，因为最大批处理时间可以与预填充请求的输入长度解耦。然而，之前的分析【8, SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills, 2023; 51, DistServe: Disaggregating Prefill and Decoding for Goodputoptimized Large Language Model Serving, 2024】因读取放大问题（read amplification）而认为分块预填充不适合处理长上下文。对于一个有$n$个输入令牌的请求，连续预填充阶段的GPU核心读取数据量为：


分块预填充将单个预填充阶段在时间上分布到多个模型前向传播中，总读取量随块数线性增加（如图6所示），公式如下。这导致总缓存读取量从$O(n)$增加到$O(n^2)$。

关键洞察：算术强度与块大小相关。本文挑战了上述观点，核心洞察是预填充块的算术强度仅取决于块大小（chunk size），而非序列长度。这是因为在分块预填充中，虽然处理每个块需要从内存中获取所有先前生成的KV缓存令牌，但我们也需要为每个KV缓存令牌执行$c$次算术运算（对应预填充块中的每个令牌）。因此，虽然更长的序列需要更多的KV缓存读取，但每次读取的算术运算次数保持不变，由块大小决定。现代LLM通过分组查询注意力（grouped-query attention）【9, GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, 2023】进一步放大了此效应（见公式(7)），使得在NVIDIA H100 GPU上，仅需约40个令牌的预填充块就足以饱和GPU计算。这一发现使我们能够通过将数百万令牌的预填充分解为数千个小的、可管理的块来实现有效的批处理和细粒度抢占。

先进注意力内核的作用。先前研究【8, SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills, 2023】对分块预填充的误解还有一个原因：传统的注意力内核通过在查询令牌上分配工作来并行化预填充计算，这在查询和KV令牌数相等时有效，但在分块预填充中因查询令牌数有限而受限。受FlashDecoding【20, FlashDecoding++: Faster Large Language Model Inference on GPUs, 2023】的启发，先进的注意力内核【14, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, 2023; 48, Accelerating Self-Attentions for LLM Serving with FlashInfer, 2024】现在可以同时在查询和KV令牌维度上并行化预填充计算。这种二维并行化策略使得即使对于非常长的上下文，分块预填充计算也能保持高效。图7证实了即使使用极小的块大小（32-128个令牌），对于100万令牌的预填充，注意力计算时间也只比大块大小有边际开销。


图7：对于使用8个H100的Llama-3 70B模型进行1M令牌预填充，块大小对注意力计算时间的影响。即使非常小的块大小（32-128令牌）与较大块大小相比也只产生边际开销。这表明使用小块的注意力计算仍然高效，与之前关于长上下文分块预填充的假设相悖。

4.2 自适应分块预填充

静态块大小的权衡。虽然小块（如32）在注意力计算上开销不大（图7），但由于线性层的计算效率低下和其他固定的CPU开销，使用小块会导致显著的端到端性能下降。例如，在8个NVIDIA H100上运行Llama-3 8B，对于100万令牌的请求，块大小为32的预填充延迟比块大小为4096的高1.75倍。另一方面，较大的块大小会导致与之一起批处理的其他请求的解码延迟更高（图8a），从而在预填充和解码延迟之间形成一个不良的权衡。

动态调整块大小的策略。为了避免这种权衡，我们采用了一种动态块大小调整策略。该策略基于一个观察：在预填充处理的后期，每块的延迟较高，此时注意力运行时间占主导地位，因此较小的块变得更高效。为了在不牺牲预填充效率的情况下保持低解码延迟，我们从一个大的块大小开始，并随着预填充的进行动态减小它。我们使用Vidur模拟器【6, Vidur: A LargeScale Simulation Framework For LLM Inference, 2024】中的运行时预测组件来确定在不违反解码延迟SLO的情况下可以使用的最大块大小。自适应分块可以获得明显更优的预填充-解码延迟权衡，如图8b所示。这种方法为未来结合更复杂的调度目标（如公平性【40, Fairness in Serving Large Language Models, 2023】或截止时间感知调度【5, Etalon: Holistic Performance Evaluation Framework for LLM Inference Systems, 2024】）奠定了基础。


图8：展示自适应分块在平衡混合批处理中预填充和解码延迟方面优越性能的帕累托前沿。(a) 静态块大小强制进行权衡：较大的块减少预填充延迟，但增加每批（解码）延迟。(b) 自适应块大小策略从较大的块开始，并逐渐减小块大小，以使每批的延迟保持相似。这允许在预填充效率和为混合工作负载保持低解码延迟之间实现更好的折衷。

4.3 序列流水线并行 (Sequence Pipeline Parallelism)

传统并行策略的局限性。预填充延迟随序列长度的二次方增长对用户体验构成重大挑战。张量并行虽然能降低延迟，但因通信开销大而难以扩展到单个计算节点之外。流水线并行虽能高效扩展，但主要提高吞吐量而非降低延迟。

SPP的核心思想。我们引入一个关键洞察：将分块预填充与流水线并行结合，通过优化的流水线调度可以显著降低预填充延迟。传统的PP推理系统【7, Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve, 2024, OSDI; 49, Orca: A Distributed Serving System for Transformer-Based Generative Models, 2022, OSDI】使用交错的微批次来保持流水线效率，即请求A的块$i+1$通常在块$i$完成所有流水线阶段后才被调度（图9a）。这种方法对于自回归解码是必要的，但对于预填充阶段却是次优的，因为各个预填充块的处理独立于前一个块的模型输出。我们的核心创新在于，在预填充期间，块$i+1$在块$i$完成第一个流水线阶段后立即调度（图9b）。我们称此方法为序列流水线并行（SPP）。这种密集的流水线调度高效地并行化了预填充处理，随着GPU数量的增加产生接近线性的加速比，如下式所示：


其中，$T_{sppp}(n, c)$是SPP预填充时间，$T_p(n, c)$是标准预填充时间，$p_{spp}$是SPP的并行度，$T_{ppcomm}(c)$是阶段间通信时间。对于大的$n$，通信开销项变得可以忽略不计，从而实现近线性扩展。


图9：预填充处理的流水线并行策略对比。(a) 标准流水线并行使用微批次来提高吞吐量，但不能为长序列减少延迟。(b) 序列流水线并行（SPP）在各阶段间重叠块的处理，显著降低了长上下文的预填充延迟，同时保持了高GPU利用率。

SPP的优势与相关工作。与环形注意力不同，SPP的有效性与输入序列长度的变化无关，并且SPP支持批处理和抢占，便于更高效的调度。我们的方法与TeraPipe【27, TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models, 2021】提出的技术有相似之处，后者在模型训练中利用序列块的流水线来提高效率。但两者目标不同：TeraPipe旨在通过最小化流水线气泡来优化训练吞吐量，而Mnemosyne将类似概念应用于推理阶段，专门针对降低在线服务长上下文请求的延迟。

4.4 KV并行 (KV Parallelism)

解决解码延迟的挑战。SPP能有效降低预填充延迟，但由于自回归解码中存在跨迭代依赖，无法用于优化解码延迟。为解决此挑战，我们提出了KV并行（KVP），一种通过并行化KV缓存读取来有效降低解码延迟的新技术。

KVP的工作原理。在KVP中，KV缓存沿序列维度被分片到多个GPU上。在每次迭代中，我们将查询令牌复制到所有GPU，并基于本地的KV缓存分片计算部分注意力输出。然后使用在线softmax【32, Online normalizer calculation for softmax, 2018】将这些部分输出合并。KVP相比环形注意力等技术的一个关键优势是，其通信成本（$T_{kvpcomm}$）与KV缓存长度无关，仅取决于查询令牌的数量，这使其在管理长上下文请求的解码延迟方面极为有效。KVP的性能提升可由下式建模：


其中$T_{kvpd}(n)$是使用KVP的解码时间，$T_{attnd}(n)$是注意力计算时间，$p_{kvp}$是KV并行度，$T_d(n)$是总解码时间。

KVP在混合批处理和预填充中的应用。实验表明，KVP在混合批处理场景中也能有效降低预填充对其他批处理请求解码延迟的影响。例如，处理一个400万上下文长度的请求时，即使使用128个令牌的小块，与之批处理的请求的P95解码延迟也会达到近100毫秒，而KVP可以缓解此问题。KVP的概念也自然地扩展到分块预填充。对于长序列，KVP的通信成本（$iT_{kvpcomm}(c)$）远小于注意力计算本身，如下式所示：

动态增长的KVP工作节点。为优化资源利用，我们采用动态增长策略，而非预先分配所有KV并行工作节点。我们为每个请求定义了单个KV并行工作节点能管理的最大KV缓存令牌数。初始时分配一个工作节点，当达到令牌数限制时，再加入一个新的工作节点。这种方法允许每个KV并行副本在协作处理长请求的同时，独立地批处理其他短请求，确保在不同工作负载下的资源高效利用。


图10：KV并行（KVP）随着处理的上下文长度增长动态扩展资源。最初，单个GPU存储KV缓存（t=0）。随着更多块被处理，缓存在该GPU内扩展（t=3）。当单个GPU的延迟容量达到极限时，KVP增加一个新的GPU（t=4）以适应进一步的增长。这种渐进式扩展使KVP即使对于极长的上下文长度也能限制每块的延迟。


图11：Mnemosyne中KV并行和张量并行的组合。此图展示了KV并行（水平分布）和张量并行（垂直分布）如何协同工作以高效处理长序列。KV并行将序列跨GPU 0-1和2-3分割，而张量并行则在每个KV分片内跨注意力头划分计算。

4.5 Mnemosyne 3D并行

结合三种并行策略。为了满足长上下文LLM推理中预填充和解码延迟的苛刻要求，Mnemosyne引入了一种新颖的3D并行策略。该方法结合了序列流水线并行（SPP）、KV并行（KVP）和张量并行（TP），以在数百个GPU上有效扩展性能。这些并行技术的协同作用解决了推理过程的不同方面：SPP加速预填充计算，KVP降低解码延迟，而TP通过模型级并行化增强预填充和解码两个阶段。

配置优化与扩展特性。确定此3D并行的最佳配置是一项复杂任务，需要考虑延迟SLO、请求长度混合以及各并行维度的扩展效率。我们的实验表明，SPP表现出极佳的可扩展性，扩展到16个服务器时仍保持超过80%的效率（图15），特别适合降低超长上下文的预填充延迟。KVP对满足严格的解码延迟目标至关重要，但在超过2-4个服务器后可扩展性有限，原因是其通信开销相对较高且并行化仅限于注意力计算，受阿姆达尔定律限制。鉴于优化这种3D并行策略的复杂性，可以利用现有的LLM推理模拟器【6, Vidur: A LargeScale Simulation Framework For LLM Inference, 2024; 12, Llmservingsim: A hw/sw cosimulation infrastructure for llm inference serving at scale, 2024】来根据具体的工作负载需求找到最佳并行化配置。


图12：Mnemosyne的3D并行架构：结合张量并行（TP）、序列流水线并行（SPP）和KV并行（KVP），实现可扩展的长上下文LLM推理。每个KVP单元包含一个具有多个流水线阶段的完整模型副本，这些阶段又在服务器内的GPU之间利用张量并行。这种分层设计通过加速预填充和解码计算来高效处理长序列，同时为混合工作负载提供灵活的调度。

5. 平台优化

系统级优化的必要性。高效的长上下文LLM推理不仅需要新颖的并行策略，还需要细致的平台级优化。Mnemosyne扩展了Sarathi-Serve框架【7, Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve, 2024, OSDI】，在三个关键领域解决了规模化带来的独特挑战：进程间通信、模型执行和页表管理。如图13b所示，这些优化相比现有系统带来了高达4倍的解码延迟降低。

具体的优化措施：
* 进程间通信：传统系统如vLLM【3, vLLM: Easy, fast, and cheap LLM serving for everyone】和Sarathi-Serve【7, Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve, 2024, OSDI】依赖中央调度器来分配内存和通信。随着序列长度增加，这种方法的开销很大。我们通过在调度器和所有GPU工作节点间复制序列状态来减少通信量，并用ZeroMQ【4, ZeroMQ】替换Ray【1, Apache Ray】进行调度器-工作节点通信，消除了扩展到数百个工作节点时的全局解释器锁（GIL）【16, Thread and process efficiency in Python, 2019】竞争。
* 模型执行：我们集成了Flashinfer【48, Accelerating Self-Attentions for LLM Serving with FlashInfer, 2024】内核，它能有效地在查询和KV令牌上分配工作。这一增强实现了高效的分块预填充计算。为确保即使使用小的预填充块也能满足严格的延迟目标，我们为混合批次实现了CUDA图。
* 页表管理：我们观察到每次迭代中将大型页表从CPU复制到GPU内存会引入相当大的开销。我们的解决方案是实现GPU侧页表。我们在请求初次加入时引导页表，之后只向GPU传输增量更新。这种方法显著减少了CPU和GPU之间的数据移动以及分页注意力对长上下文长度的延迟影响。


图13：vLLM【3】与Mnemosyne-1D TP在Llama-3 8B和8个H100 GPU（$p_{tp}=8$）上处理长上下文的延迟比较。解码延迟降低3.8倍，小块预填充延迟降低约6倍，反映了Mnemosyne中实现的各种CPU优化的效果。

A4 实验环境

  • 平台：Mnemosyne在Sarathi-Serve框架【7, Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve, 2024, OSDI】的基础上实现，并加入了第5节中描述的基线优化以及SPP、KVP和3D并行策略的设计变更。
  • 模型：使用Llama-3 8B和Llama-3 70B模型，通过RoPE【43, Roformer: Enhanced transformer with rotary position embedding, 2024】缩放支持高达1000万令牌的上下文长度。由于Mnemosyne是精确推理系统，对模型准确性没有影响，因此评估不依赖特定数据集或评分系统。
  • 硬件：使用多达16个通过InfiniBand连接的DGX-H100系统【31, ND-H100-v5 sizes series, 2024】。每个DGX-H100服务器包含8个NVIDIA H100 GPU【34, DGX H100: AI for Enterprise, 2024】，每个GPU拥有80GB高带宽内存，总计最多128个GPU。服务器内GPU通过NVLINK4.0连接（900GBps双向带宽），跨服务器GPU通过InfiniBand连接（每GPU对50GBps）。
  • 系统配置：
    • Mnemosyne 2D (SPP+TP)：$p_{tp}$设为8（服务器内），SPP用于跨服务器扩展以实现快速和可抢占的预填充。
    • Mnemosyne 2D (KVP+TP)：$p_{tp}$设为8，KVP用于跨服务器扩展以实现快速解码。
    • Mnemosyne 3D (KVP+SPP+TP)：$p_{tp}$设为8，同时使用SPP和KVP来扩展最终的Mnemosyne设计。

A4 实验结果

6.2 使用SPP实现快速且可抢占的预填充

  • 实验内容：对比Mnemosyne 2D (SPP+TP) 与当前最强的长上下文预填充基线——条带注意力（Striped Attention）【11, Striped attention: Faster ring attention for causal transformers, 2023】，评估100万令牌预填充的TTFT延迟和抢占能力（队头阻塞情况）。
  • 实验结果：
    • TTFT延迟：如图14a所示，在使用128个GPU（16台服务器）处理100万令牌时，Mnemosyne 2D SPP+TP比条带注意力快64%，TTFT延迟低于15秒。
    • 抢占能力：条带注意力不支持分块预填充，导致长上下文预填充完全阻塞。如图14b所示，这导致了高达120秒的队头阻塞延迟，而Mnemosyne 2D SPP+TP仅为62毫秒。
  • 分析结论：Mnemosyne 2D SPP+TP实现了比最强基线更快且更具可抢占性的预填充阶段。


图14：条带注意力与Mnemosyne 2D Parallel在Llama3 8B上进行1M令牌序列预填充的性能比较。Mnemosyne不仅在预填充计算上实现了更好的扩展效率，还通过支持细粒度抢占帮助消除了队头阻塞。

  • 实验内容：评估Mnemosyne 2D (SPP+TP) 从100万到1000万令牌的扩展性能。
  • 实验结果：如图15所示，随着SPP流水线深度的增加，Mnemosyne 2D SPP+TP在Llama-3 8B和70B上均实现了接近线性的扩展效率（超过80%）。对于Llama-3 8B，使用16台DGX服务器可以在200万上下文长度上满足30秒的TTFT目标。
  • 分析结论：SPP具有很强的扩展性，表明增加更多服务器可以为更长的上下文实现更短的TTFT。


图15：Mnemosyne 2D并行（SPP + TP）在长上下文预填充处理中的扩展效率。对于Llama-3 8B和70B，随着SPP并行度增加到128个H100 GPU，Mnemosyne 2D在预填充延迟上表现出近线性的降低（80%+的扩展效率）。红色叉表示因内存限制而不可行的配置。

6.3 使用KVP实现快速解码

  • 实验内容：评估SPP和KVP对解码阶段性能（TBT）的影响。基线为Mnemosyne 2D (SPP+TP)，目标SLO为30ms。
  • 实验结果：
    • SPP对TBT的影响：如图16所示，随着SPP并行度（$p_{spp}$）的增加，TBT会略有恶化，因为流水线并行的通信开销是固定的。
    • KVP对TBT的影响：如图17所示，在Mnemosyne 3D并行设置下，增加KVP并行度（$p_{kvp}$）可以显著降低TBT，帮助达到交互性目标。这种延迟效益不是线性的，但随着上下文长度的增加而变得更加明显。例如，对于Llama-3 8B的1000万上下文，将$p_{kvp}$从1增加到4，TBT降低了2.5倍。
  • 分析结论：KVP是降低长上下文解码延迟的关键。对于Llama-3 8B，直到400万上下文，以及Llama-3 70B直到200万上下文，$p_{kvp}=1$即可满足TBT目标，超出此范围则建议增加KVP并行度。


图16：SPP扩展对Mnemosyne 2D并行（SPP+TP, $p_{tp}=8$）中解码延迟的影响。即使流水线包含16个阶段，解码延迟也仅受到轻微影响。


图17：在Mnemosyne 3D并行的纯解码批次中使用KV并行的TBT降低效果。对于1000万上下文长度的Llama-3 8B解码，$p_{kvp}=2$导致延迟降低近40%，允许以约每秒30个令牌的速度解码。

6.4 Mnemosyne 3D的端到端推理

  • 实验内容：评估Mnemosyne 3D并行在混合批处理（预填充和解码）场景下的TTFT与TBT权衡。
  • 实验结果：如图18所示，对于给定的KVP并行度（$p_{kvp}$），增加块大小会减少TTFT但增加TBT。而对于给定的块大小，增加$p_{kvp}$通常可以同时降低TTFT和TBT，从而在权衡空间中达到更优的点。
  • 分析结论：Mnemosyne的并行策略组合在TBT/TTFT权衡空间中创造了更多选择，这对于在严格的延迟SLO下操作至关重要，增加了满足特定SLO对的可行性。


图18：使用Mnemosyne 3D并行（$p_{tp}=4, p_{spp}=4$）的Llama-3 8B在不同KVP并行度和块大小（32-256）下的TTFT与P95 TBT权衡。

  • 实验内容：展示KVP实例动态添加的时间线。
  • 实验结果：如图19所示，在处理200万令牌时，Mnemosyne 3D从32个GPU（$p_{kvp}=1$）开始，并根据已处理的序列长度逐步增加KVP实例，最终达到128个GPU。不断增长的上下文长度和不断增加的并行度形成了对抗力量，使得迭代执行时间保持一致。
  • 分析结论：动态资源分配能维持近乎恒定的每迭代延迟。


图19：使用Mnemosyne 3D（$p_{tp}=8, p_{kvp}=4, p_{spp}=4$）处理200万令牌Llama-3 8B时随时间变化的GPU数量。Mnemosyne从一个KVP工作组（4台服务器）开始，并随着序列长度增加逐步增加12台服务器。通过动态增长资源分配，Mnemosyne保持了近乎一致的每迭代延迟。

6.5 实现高吞吐量

  • 实验内容：通过模型FLOPs利用率（MFU）和模型带宽利用率（MBU）评估硬件资源利用率，并通过批处理对延迟的影响来展示高吞吐量潜力。
  • 实验结果：
    • MFU和MBU：如图20和21所示，Mnemosyne在计算密集型的预填充阶段实现了50-75%的MFU，在内存密集型的解码阶段MBU高达92%。这些指标表明硬件利用率非常高。
    • 批处理效率：如图22所示，将批处理大小增加到128个请求，批处理执行时间的增幅不超过5%。
  • 分析结论：Mnemosyne能够高效利用硬件资源，并支持大批量处理，从而实现高吞吐量。


图20：Mnemosyne 2D并行（TP + SPP）的模型FLOPS利用率【13, PaLM: Scaling Language Modeling with Pathways, 2023】（MFU）。系统在不同序列长度和并行度下持续达到50-60%的利用率。


图21：Mnemosyne 2D并行（TP + KVP）的模型带宽利用率（MBU）。


图22：在8个H100（TP）上使用Llama-3 8B进行长请求（1M）的分块预填充和多个解码请求（各1K）的P95混合批次执行延迟。结果显示了当我们改变预填充的块大小和批处理的解码阶段数量时批处理执行时间的变化。红线表示仅运行分块预填充而未批处理任何解码阶段的批处理执行时间。

A7 补充细节

7. 使用Mnemosyne设计系统

  • 调度策略：Mnemosyne支持预填充抢占和自适应分块，这需要动态调度策略。此外，KVP实例的独立调度为提高系统吞吐量开辟了新空间。由于每个KVP实例内有完整的模型副本，当它们不协同处理长上下文推理时，可以被独立调度用于处理较短的请求。
  • 寻找正确的并行策略：Mnemosyne引入的3D并行性在TTFT、TBT和吞吐量之间提供了丰富的权衡空间，需要通过使用代表性追踪数据进行性能剖析来找到最合适的并行配置。
  • 在线与离线推理：尽管SPP和KVP的扩展性很强，但要为极长上下文（如1000万）满足交互式TTFT延迟目标（如30秒）可能需要过多的GPU。对于更长的上下文，可能需要采用近似方法，或采用先前工作【35, Splitwise: Efficient generative LLM inference using phase splitting, 2024, ISCA; 51, DistServe: Disaggregating Prefill and Decoding for Goodputoptimized Large Language Model Serving, 2024】中提出的预填充和解码阶段分离的架构。这种分离式架构允许通过较慢的预填充阶段进行离线上下文构建，然后将结果发送到为解码优化的系统以满足交互式TBT延迟SLO。
  • 精确与近似：虽然本文将Mnemosyne作为一种无近似的长上下文解决方案，但这并不妨碍其并行策略与第8节中描述的近似技术结合使用。

8. 相关工作

  • 长上下文LLM：近期研究集中在有效训练和高效服务长上下文LLM模型。一些工作提出了新的注意力并行技术【11, Striped attention: Faster ring attention for causal transformers, 2023; 26, Sequence parallelism: Long sequence training from system perspective, 2021; 30, Ring Attention with Blockwise Transformers for Near-Infinite Context, 2023】。同时，计算内核优化【14, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, 2023; 20, FlashDecoding++: Faster Large Language Model Inference on GPUs, 2023; 48, Accelerating Self-Attentions for LLM Serving with FlashInfer, 2024】增强了GPU内部的并行性。此外，像LongRoPE【15, Longrope: Extending llm context window beyond 2 million tokens, 2024】这样的技术允许在无需额外微调的情况下扩展现有LLM的上下文窗口。一个与SPP类似但没有自适应分块的概念，称为“令牌并行”，曾在TeraPipe【27, TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models, 2021】中用于在训练期间并行化微批次以减少流水线气泡。
  • 近似替代方案：为了降低计算复杂性，状态空间模型（SSMs）【17, Mamba: Linear-time sequence modeling with selective state spaces, 2023; 18, Efficiently modeling long sequences with structured state spaces, 2021】提出了替代传统注意力架构的方案。其他方法，如局部敏感哈希【24, Reformer: The efficient transformer, 2020】、压缩注意力【33, Leave no context behind: Efficient infinite context transformers with infini-attention, 2024】以及提示或KV缓存压缩【23, Longllmlingua: Accelerating and enhancing llms in long context scenarios via prompt compression, 2023; 25, {InfiniGen}: Efficient generative inference of large language models with dynamic {KV} cache management, 2024, OSDI 24; 50, \$H_2O\$: Heavy-Hitter Oracle for Effi cient Generative Inference of Large Language Models, 2023】，旨在减少计算和内存占用。这些方法以牺牲准确性换取效率，而本文专注于保持完整上下文和准确性的Transformer模型。
  • LLM推理效率：为提高LLM推理服务效率，研究人员提出了多种技术，包括分块预填充【8, SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills, 2023】、预填充-解码分离【35, Splitwise: Efficient generative LLM inference using phase splitting, 2024, ISCA; 51, DistServe: Disaggregating Prefill and Decoding for Goodputoptimized Large Language Model Serving, 2024】、弹性并行【47, Loongserve: Efficiently serving long-context large language models with elastic sequence parallelism, 2024】和分布式KV缓存【28, Infinite-llm: Efficient llm service for long context with distattention and distributed kvcache, 2024】。百万级令牌规模的推理请求带来了新的挑战，需要精心设计以在规模上有效地整合这些方法。同期工作Mooncake【38, Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving, 2024】提出了与SPP类似的分块流水线并行，但仅关注高达128K令牌的预填充请求。

A5 结论

长上下文LLM推理的应用场景持续增长，不断推动着“长”上下文的边界。通过Mnemosyne，本文提出了一套新颖的并行策略，使得端到端推理服务能够扩展到至少1000万令牌的上下文长度，并具有快速的预填充和解码阶段。该系统能够满足高达200万令牌的交互式TTFT目标，以及高达1000万令牌的TBT目标。此外，Mnemosyne支持可变长度的输入请求，并通过自适应分块实现混合批处理。这使得基于截止时间的动态调度成为可能，从而在支持不同上下文长度输入的同时，以有延迟保证的方式提高系统吞吐量。