当前位置：首页 > news >正文

告别KV Cache无法复用！Amazon SageMaker加速LLM推理

news 2026/3/26 15:59:35

随着大语言模型（LLM）的广泛应用，如何优化推理性能已成为企业级应用部署中的核心挑战，传统的无状态推理架构在处理LLM应用时面临诸多性能瓶颈：每次请求被随机路由到不同的计算实例，导致KV Cache无法有效复用、多轮对话上下文频繁重建、系统提示词重复处理等问题，严重影响了用户体验和系统效率。

针对这些痛点，Amazon SageMaker推理集群的粘性会话路由（Sticky Session Routing）功能，通过会话绑定机制，确保同一用户会话的所有请求都路由到相同的推理集群实例，从而实现缓存复用和状态保持，显著提升LLM推理性能和交互体验。

技术原理与优势

Amazon SageMaker Sticky Session

技术原理

Amazon SageMaker Sticky Session通过会话标识符（Session ID）实现智能路由机制。与传统的负载均衡策略不同，该机制确保来自同一会话的所有请求都被定向到特定的计算实例。

这种“会话亲和性”（Session Affinity）设计使得GPU内存中的KV Cache、中间结果等状态信息得以持续保存和复用，从而避免了重复的计算开销。

核心优势

启用粘性会话后，同一会话的所有请求都会路由到同一实例，这样您的AI应用程序就能重复使用先前处理过的信息，避免了不必要的重复计算，从而减少延迟并改善用户体验。

当客户想要使用大型数据负载或需要无缝的交互体验时，这项功能尤其有用。通过利用以前的推理请求，客户现在可以借助这项功能在Amazon SageMaker上构建创新的状态感知AI应用程序。

工作流程

一起看看客户端在Amazon SageMaker上启动粘性会话时的事件序列，Sticky Session的完整工作流程包含三个核心阶段：会话创建、会话维持和会话关闭:

1.在第一个请求中，当调用Boto3 SageMaker runtime的invoke_endpoint时，在请求头中设置session-id=NEW_SESSION，并在请求负载中指明这是一个打开会话类型的请求。然后Amazon SageMaker会创建一个新的会话并存储会话ID。Amazon SageMaker会向推理服务器发起一个打开会话的请求（这个API是由客户端定义的，也可以使用其他名称，比如start_session），并返回200状态码，同时附带会话ID和生存时间（TTL），这些信息会被发送回客户端。

2.后续请求中，客户端在invoke_endpoint调用中传递相同的会话ID。Amazon SageMaker路由服务根据会话ID识别对应的推理实例，确保所有请求都路由到同一台服务器。这样，在LLM推理的多轮对话场景中，GPU内存中的KV Cache得以保持和累积，显著减少了计算开销，从而显著降低推理延迟。

3.当对话结束时，客户端可以发送关闭会话请求，您可以使用invoke_endpoint传入一个包含会话ID的关闭请求。Amazon SageMaker路由器首先会检查该会话是否存在。如果存在，路由器会向推理服务器发起关闭会话的调用，服务器会返回一个成功的200响应以及会话ID，这些信息随后会被发送回客户端。在会话ID不存在的情况下，路由器会返回一个400响应。

这种设计带来的性能提升主要体现在：

首次响应时间优化：避免重复处理系统提示词和上下文信息
后续对话加速：利用已缓存的键值对，减少大量重复计算
内存使用效率：合理利用GPU内存，避免频繁的内存分配和释放

实践案例

构建支持Sticky Session的LLM推理服务

本案例将以基于SGLang推理引擎的模型部署为例，展示如何在Amazon SageMaker上构建支持Sticky Session的高性能LLM推理服务。整个实现涉及推理容器构建、模型部署和客户端调用三个核心部分。

构建支持会话管理的推理容器

在推理服务器端，使用Sanic Web框架构建了支持会话管理的推理服务。核心实现包括三种请求类型的处理逻辑：

@app.route("/invocations", methods=["POST"])async def generate(request): reqType = request.json.get("requestType") extSessID = request.json.get("extSessionID") # 处理新建Session请求 if'NEW_SESSION' == reqType: current_time = datetime.now(dt.timezone.utc) future_time = current_time + timedelta(minutes = int(os.environ['SES_TTL_MIN'])) formatted_time = future_time.strftime("%Y-%m-%dT%H:%M:%SZ") response = json({}) response.headers["X-Amzn-SageMaker-Session-Id"] = f'{extSessID}; Expires={formatted_time}' return response # 处理Session关闭请求 elif 'CLOSE_SESSION' == reqType: response = json({}) response.headers["X-Amzn-SageMaker-Closed-Session-Id"] = extSessID return response # 处理Session内推理请求 else: prompt = request.json.get("inputs") ifnot prompt: return json({"error": "inputs is required"}, status=400) inf_params = request.json.get("parameters") # 使用部署至 SGLang 推理引擎的模型产生图里结果 result = await engine.async_generate(prompt=prompt, sampling_params=inf_params) return json({"generation": result})

左右滑动查看完整示意

构建客户端调用封装

为了简化Amazon SageMaker Sticky Session的使用，本例封装了一个客户端工具类，提供了会话管理的完整功能：

import boto3, json classStatefulSMEDPBuilder: def __init__(self, endpoint_name): self.endpoint_name = endpoint_name ## Boto3 API self.sm_bt3_client = boto3.client("runtime.sagemaker") def start_session(self, extSessID): """创建新的粘性会话""" payload = { "extSessionID": extSessID, "requestType": 'NEW_SESSION' } response = self.sm_bt3_client.invoke_endpoint( EndpointName=self.endpoint_name, Body=json.dumps(payload), ContentType="application/json", SessionId="NEW_SESSION" ) return response def end_session(self, extSessID): """关闭指定会话""" payload = { "extSessionID": extSessID, "requestType": 'CLOSE_SESSION' } response = self.sm_bt3_client.invoke_endpoint( EndpointName=self.endpoint_name, Body=json.dumps(payload), ContentType="application/json", SessionId=extSessID ) return response def invoke(self, textPayload, sampling_params=None, extSessID=None): """在会话中进行推理调用""" if None == sampling_params: sampling_params = {"temperature":0.9, "max_new_token":128, "do_sample":True} payload = { "inputs": textPayload, "sampling_params": sampling_params, "extSessionID": extSessID, "requestType": 'SESSION' } response = self.sm_bt3_client.invoke_endpoint( EndpointName=self.endpoint_name, Body=json.dumps(payload), ContentType="application/json", SessionId=extSessID ) return response def invoke_stream(self, textPayload, sampling_params=None, extSessID=None): """流式推理接口""" pass

左右滑动查看完整示意

使用示例

下面展示如何使用封装好的客户端进行多轮对话：

import uuid # 初始化客户端client = StatefulSMEDPBuilder(endpoint_name="llama3-sticky-session-endpoint") # 生成唯一会话 IDsession_id = f"conversation-{uuid.uuid4().hex[:8]}" try: # 创建新会话 session_response = client.start_session(session_id) print(f"会话 {session_id} 创建成功") # 第一轮对话 response1 = client.invoke("你好，请介绍一下自己", extSessID=session_id) result1 = json.loads(response1['Body'].read().decode()) print(f"AI: {result1['generation']['text']}") # 第二轮对话（利用上下文） response2 = client.invoke("请详细说明你刚才提到的能力", extSessID=session_id) result2 = json.loads(response2['Body'].read().decode()) print(f"AI: {result2['generation']['text']}") # 第三轮对话（继续利用完整上下文） response3 = client.invoke("基于我们之前的对话，你觉得哪个能力最重要？", extSessID=session_id) result3 = json.loads(response3['Body'].read().decode()) print(f"AI: {result3['generation']['text']}") finally: # 会话结束后清理资源 client.end_session(session_id) print(f"会话 {session_id} 已关闭")

左右滑动查看完整示意

通过这种方式，同一会话的所有请求都会路由到相同的推理实例，确保KV Cache的有效复用，显著提升多轮对话的性能表现。

完整的部署和使用代码可参考项目仓库。

项目仓库：

https://github.com/aws-samples/sample-sagemaker-sticky-session

总结

本文深入介绍了Amazon SageMaker Sticky Session在LLM推理中的技术原理和实践应用。

通过会话亲和性设计，该技术有效解决了传统无状态推理架构在处理LLM应用时面临的性能瓶颈，这一功能特别适合需要维持对话上下文的应用以及要求低延迟的实时交互场景。后续博客将展示详细的性能对比测试结果。

参考链接

1.https://aws.amazon.com/cn/about-aws/whats-new/2024/09/sticky-session-routing-amazon-sagemaker-inference

2.https://aws.amazon.com/cn/blogs/machine-learning/build-ultra-low-latency-multimodal-generative-ai-applications-using-sticky-session-routing-in-amazon

本篇作者

胡凯

亚马逊云科技解决方案架构师，负责基于亚马逊云科技的云计算方案架构的咨询和设计，致力于亚马逊云科技云服务在企业的应用和推广。

郑昊

亚马逊云科技AI/ML解决方案架构师。主要专注于基础模型的训练、推理及性能优化；广告、排序算法等及其基于亚马逊云科技AI/ML技术栈的相关优化及方案构建，拥有多年排序、定价及竞价机制等算法研发经验。

新用户注册海外区域账户，可获得最高200美元服务抵扣金，覆盖Amazon Bedrock生成式AI相关服务。“免费计划”账户类型，确保零花费，安心试用。

星标不迷路，开发更极速！

关注后记得星标「亚马逊云开发者」

听说，点完下面4个按钮

就不会碰到bug了！

点击阅读原文查看博客！获得更详细内容！

查看全文