基于 Amazon SageMaker 有状态路由优化大规模推理集群下的 KV Cache 复用方案

摘要

随着大语言模型（LLMs）在各行各业的应用日益广泛，高效部署这些计算密集型模型已成为一项关键挑战。特别是在生成式 AI 服务中，用户对模型响应速度的期望与模型推理的计算成本控制之间存在明显的矛盾。为了解决这一挑战，研究界和工业界已开发出多种优化技术，其中 KV Cache（键值缓存）作为一种基础优化手段，通过存储和重用中间计算结果，显著提升了单次推理的效率。

随后，SGLang、vLLM 等前沿推理框架通过引入 Prompt Cache 机制，实现了跨请求的 KV Cache 复用，使相似提示结构的请求能够共享同一计算资源，在单机环境下显著降低了计算开销并提升了吞吐量。

然而，当我们将服务扩展至大规模推理集群时，负载均衡机制的挑战浮现出来。在分布式推理系统中，为了保证系统稳定性和资源利用率，请求通常被随机分发至集群中的不同节点。这种流量的随机分发机制，虽然在资源调度层面行之有效，却无意中破坏了 KV Cache 复用的基础条件——相似请求的连续性和局部性。当相似的请求被分散到不同的计算节点时，即使它们在逻辑上高度相关且可共享计算资源，物理隔离也使得缓存复用变得不可能。

具体而言，当相似或相关的推理请求被分散到不同计算节点时，即使这些请求理论上可以共享 KV Cache，由于计算实例间的区隔，复用机会也被大大削弱。这导致了计算资源的重复消耗，推理延迟的增加，以及整体系统效率的下降。

面对这一挑战，Amazon SageMaker 提供了一种创新的技术方案，它能够使得用户通过一系列机制，在基于集群均衡必须要求的同时，引导相关请求流向可复用 KV Cache 的计算节点，在大规模推理集群环境下构建了 KV Cache 复用的可行性。

推理引擎中的 KV Cache 复用

NVIDIA TensorRT-LLM KV Cache Early Reuse

NVIDIA 的 TensorRT-LLM 引入了先进的 KV Cache （缓存）重用技术，显著提升大语言模型推理效率。该技术主要包括两项创新：

早期 KV 缓存重用：允许在系统提示生成过程中实时共享缓存，避免了传统方法要求完成整个 KV 缓存计算后才能重用，在多用户环境下将首个 Token 生成速度提升数倍；
灵活的 KV 缓存 Block 大小控制：将支持将缓存 Block 从 64 个 Token 精细调整至最低 2 个 Token，优化内存使用并提高缓存命中率，在运行如 Llama 70B 模型时能将推理速度提升约 7%。这些优化策略通过减少冗余计算，有效解决了大规模部署中的性能瓶颈问题；
高效的 KV 缓存驱逐协议：引入了智能 KV 缓存驱逐协议，通过识别缓存 Block 之间的依赖关系树结构，优先驱逐叶节点（Dependent Nodes）而非分支节点（Source Blocks），即使叶节点有更高的使用频率。这种方法确保了在内存管理过程中不会错误地移除关键的 Source Blocks，从而避免了连锁反应式的缓存重建，显著减少了被驱逐的块数量，提高了缓存重用率，并降低了新用户提示的首个 Token 生成时间（TTFT）。

SGLang with RadixAttention

SGLang 在开发过程中发现了 LLM 在复杂推理中对性能优化至关重要的 KV 缓存复用（KV Cache Reuse）这一关键优化机会。 KV 缓存复用允许具有相同前缀的不同提示共享中间计算结果，从而减少内存占用和计算冗余。

下图展示了四种常见类型的 KV 缓存复用模式，包括少样本学习示例、自洽性问题、多轮对话历史以及思维树搜索历史等场景：

（引用: SGLang: Efficient Execution of Structured Language Model Programs）

图 2. KV 缓存共享示例。蓝色框是可共享的提示部分，绿色框是不可共享部分，黄色框是不可共享的模型输出

为系统化利用这些复用机会，SGLang 引入了 RadixAttention 技术，该技术能在运行时自动实现 KV 缓存复用。与传统方法不同，RadixAttention 不会在生成请求完成后丢弃 KV 缓存，而是将提示和生成结果的缓存保留在基数树（Radix Tree）中结构中。

基数树作为字典树的空间优化版本，其边可标记不同长度的元素序列，而非仅单个元素，从而提高效率。在 SGLang 系统中，基数树用于管理 Token 序列（键）与对应KV缓存张量（值）之间的映射。这些缓存张量以分页布局存储在 GPU 上，每页大小等同于一个 Token 。考虑到 GPU 内存限制，系统实现了递归淘汰叶节点的最近最少使用（LRU）策略，并辅以缓存感知调度，以提高命中率。

RadixAttention 与连续批处理和分页注意力等现有技术兼容，且可轻松扩展以支持多模态模型中图像 Token 的处理。在实际运行中，前端向 Runtime 发送完整提示，而运行时自动执行前缀匹配、复用和缓存操作。整个树结构存储在 CPU 上，维护开销极小，确保了系统的高效运行。

下图展示了在处理多个传入请求时如何维护基数树。前端始终向 Runtime 发送完整提示，而 Runtime 将自动进行前缀匹配、复用和缓存。每条树边都带有一个标签，表示一个子字符串或令牌序列。节点采用不同颜色编码以反映不同状态：绿色表示新添加的节点，蓝色表示在该时间点访问的缓存节点，红色表示已被淘汰的节点。

（引用：Fast and Expressive LLM Inference with RadixAttention and SGLang）

图 3. 采用 LRU 淘汰策略的 RadixAttention 操作示例，分九个步骤展示

集群环境下的挑战及解决方案

技术挑战

以 SGLang、TensorRT-LLM 等为代表的前沿推理框架，通过引入不同形式的 Prompt Cache 等创新机制，实现了跨请求的 KV 缓存复用，为具有相似提示词结构的请求带来了显著的计算节省，可在单机环境下实现高效的 KV 缓存重用机制，显著降低了计算开销并提升了吞吐量。

然而，将服务扩展到大规模推理集群时，由于系统通常采用随机分发流量的策略，将请求平均分配给集群内的各个节点。尽管在资源调度方面不可或缺，但却破坏了 KV Cache 复用的核心前提——相似请求的连续处理和本地聚集特性。当逻辑上高度相关且本可共享计算资源的相似请求被分散到不同计算节点处理时，由于物理上的分离，原本可能的缓存复用机会也随之消失。

具体而言，当相似或相关的推理请求被分散到不同计算节点时，即使这些请求理论上可以共享 KV Cache，由于计算实例间的显存区隔，复用概率也随着集群规模的增大而降低。这导致了计算资源的重复消耗，推理延迟的增加，以及整体系统效率的下降。

利用 Amazon SageMaker Stateful Session 实现集群 KV Cache 复用

Amazon SageMaker 在 Inference 中推出了有状态会话路由，该路由策略帮助推理集群利用以前处理过的信息改善其生成式人工智能应用程序的性能和用户体验。借助 Amazon SageMaker，可以更轻松地部署 ML 模型，包括各类基础模型（FM），从而以最佳性价比对任何使用案例发出推理请求。

启用有状态会话路由后，同一会话的所有请求都会路由到同一实例，因此模型服务器上的基础模型推理得以重复使用先前处理过的信息，从而减少延迟并改善用户体验，同时也可以利用该能力在 SageMaker 上构建创新的状态感知 AI 应用。启动时，需要在第一个请求中创建一个会话 ID，SageMaker 推理节点（Endpoint）完成首次会话的负载均衡路由后，将会使用该会话 ID 来指示推理节点应将所有后续请求路由到同一计算节点。会话完成后，也可以将其删除，为新会话预留更多资源。通过 SageMaker Endpoint 的自定义镜像的方式，使用有状态路由集群的简要方式如下：

1. 客户端请求模型时，通过 Amazon Boto3 SDK 创建 SageMaker Runtime 后，在 API 中指定 SessionId="NEW_SESSION"，同时将客户端创建的 <Session_i_Unique_ID> 作为请求 Payload 传入。

payload = {
    "request_body": ...,
    "inference_parameters": ...,
    "unique_ession_id": <Session_i_Unique_ID>
}

boto3_sagemaker_client.invoke_endpoint(
                EndpointName=<YOUR_ENDPOINT_NAME>,
                Body=json.dumps(payload),
                ContentType="application/json",
                SessionId="NEW_SESSION"
        )

2. 在部署于 SageMaker Endpoint 的模型服务器中，我们自定义实现如下判断逻辑：当收到首个 SessionId="NEW_SESSION" 的请求后，需在 Response 的 Header 中增加特定的 KV Pair 如下，并返回。

response.headers["X-Amzn-SageMaker-Session-Id"] = f'{<Session_i_Unique_ID>}; Expires="%Y-%m-%dT%H:%M:%SZ"'

当一个请求完成 1、2 两个步骤完成后，客户端获取了完整的模型返回。同时，SageMaker 的推理集群路由层也同时记录了该 Session_i_Unique_ID 的目标推理节点。

3. 该 Session 的后续请求，可以继续通过 API 中的 SessionId 参数，传入会话创建时在请求 Payload 中所使用的 Session_i_Unique_ID ，如：

boto3_sagemaker_client.invoke_endpoint(
                ...
                SessionId="<Session_i_Unique_ID>"
        )

使得 SageMaker 推理集群路由层将请求路由至同一目标节点。

4. 当某一 Session 需要终止，则在模型 WebServer 返回之前，在 Header 中增加如下的 KV Pair：

response.headers["X-Amzn-SageMaker-Closed-Session-Id"] = '<Session_i_Unique_ID>'

SageMaker 有状态路由的实现原理及其基于 SGLang 引擎的方案明细参考 – 使用 Sticky Session 加速 LLM 推理。

基于多轮会话场景的性能评估

为验证 KV Cache 复用所带来的业务收益，我们设计了如下实验：

多轮数据集选取：使用 lmsys/lmsys-chat-1m，抽取其中会话轮次在 4 轮以上的数据，避免由于单轮会话的存在摊低有状态会话路由所带来的性能收益
推理集群设置：使用 SageMaker Inference Endpoint 及 8 台 g5.2xlarge (Nvidia A10G GPU) 推理实例，避免集群规模过小而无法有效观察路由层策略所引入的效果
模型及推理引擎：选取 Meta-Llama-3.1-8B-Instruct 及 SGLang 作为模型推理引擎
负载模式设计：通过启动多个 Thread 来模拟多个客户端用户的请求。
- 对于普通推理集群（基线），采取完全随机的策略来进行每个请求（任一轮会话及其历史）的负载均衡
- 对于 SageMaker 有状态推理集群，当每个新的会话启动时，生成全局唯一的 Session ID 并将其通过上文提到的 invoke_endpoint() API 传入。同时该会话后续轮次的请求，采用相同的 Session ID 供 SageMaker Endpoint 识别并将其路由至固定的推理实例
选取相同的数据集 Shuffle Seed，保证不同集群上的请求长度对齐

考察当客户端并发量为 64 时，两个不同类型的集群负载均衡策略，其不同轮次首 Token 延时（TTFT）对比如下图：

图 4. 多并发下集群的不同路由策略在多轮会话下的首 Token 延时（TTFT）对比

其中，数据集中会话轮次的分布如下：

图 5. 不同轮次会话数量（样本数量）分布

可以看出，基于 SageMaker 的有状态推理集群，在多轮会话场景的不同轮次中，较随机路由策略的普通推理集群，TTFT 均有较明显的性能收益。

同时，由于数据集选取时，对于 4 轮以下的样本进行了过滤，因此轮次 1-4 的样本数量一致。其在使用进行随机负载均衡策略的普通集群中，中值 TTFT 随着轮次的增加（会话历史长度递增），而呈现出递增的趋势。在使用有状态路由策略的集群中，首轮会话（全量Prefill）的开销与随机负载集群基本一致，而第 2 轮会话开始有了显著下降，随后 4 轮内呈现递增的趋势，符合该实验设置的预期。

其他场景及优化

SageMaker Stateful Session 在其他场景的应用

如在 Build ultra-low latency multimodal generative AI applications 中的多模态推理场景，可以通过 SageMaker Stateful Session 来降低 VQA 场景中的图片传输开销。

提示结构优化

除了前文介绍的在集群基础设施层的有状态路由策略外，我们仍可以如业务层的优化来进一步提升缓存复用程度。例如，我们可以在保证效果的前提下，尽量按照该原则这样构建提示：先放置静态内容（如系统提示，对话历史），然后再放置动态内容（如用户输入）。这样做可以增加静态系统提示在多个请求之间被缓存的可能性。比如，可以将提示：

Here is a question + {question}. Please act as a math expert and solve the given question..<Long Instructions>..

修改为：

Please act as a math expert and solve the given question..<Long Instructions>.. Here is a question + {question}.

总结

本文探讨了在大规模 LLM 推理集群环境中优化KV 缓存复用的技术方案。在分布式推理集群中，传统的随机负载均衡机制破坏了KV 缓存复用所需的请求连续性，导致即使逻辑上相关的请求也无法共享计算资源。针对这一挑战，本文基于 Amazon SageMaker 的 Stateful Session（有状态会话路由）机制，通过唯一会话 ID 确保同一会话的所有请求路由到同一实例，从而在分布式环境下实现有效的 KV 缓存复用。实验结果表明，这种有状态路由策略显著降低了多轮会话场景中的首 Token 生成时间（TTFT），为生产场景大规模集群环境下的推理优化提供了实用解决方案。这一方法不仅适用于多轮会话场景，也可在多模态推理等其他应用场景获得性能收益。同时，可结合提示词结构优化等技术，来进一步提升缓存命中率和系统整体性能。通过这种集群级的优化策略，我们能够利用 SageMaker Endpoint 推理集群的弹性和可靠性的同时，充分发挥领先 LLM 推理引擎中 KV 缓存复用机制的潜力，为模型推理服务提供更快速、更经济的体验。

*前述特定亚马逊云科技生成式人工智能相关的服务仅在亚马逊云科技海外区域可用，亚马逊云科技中国仅为帮助您了解行业前沿技术和发展海外业务选择推介该服务。

亚马逊AWS官方博客

基于 Amazon SageMaker 有状态路由优化大规模推理集群下的 KV Cache 复用方案

摘要

相关技术背景

自回归生成的基本过程

KV Cache 的工作机制

KV Cache 到 Prefix/Prompt Cache