利用 Amazon SageMaker Sticky Session 实现大语言模型推理加速

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

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

SageMaker Sticky Session 技术原理与优势

技术原理

SageMaker Sticky Session 通过会话标识符（Session ID）实现智能路由机制。与传统的负载均衡策略不同，该机制确保来自同一会话的所有请求都被定向到特定的计算实例。这种“会话亲和性”（Session Affinity）设计使得 GPU 内存中的 KV Cache、中间结果等状态信息得以持续保存和复用，从而避免了重复的计算开销。

核心优势

启用粘性会话后，同一会话的所有请求都会路由到同一实例，这样您的 AI 应用程序就能重复使用先前处理过的信息，避免了不必要的重复计算，从而减少延迟并改善用户体验。当客户想要使用大型数据负载或需要无缝的交互体验时，这项功能尤其有用。通过利用以前的推理请求，客户现在可以借助这项功能在 SageMaker 上构建创新的状态感知 AI 应用程序。

工作流程

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

在第一个请求中，当调用 Boto3 SageMaker runtime 的 invoke_endpoint 时，在请求头中设置 session-id=NEW_SESSION，并在请求负载中指明这是一个打开会话类型的请求。然后 SageMaker 会创建一个新的会话并存储会话 ID。SageMaker 会向推理服务器发起一个打开会话的请求（这个 API 是由客户端定义的，也可以使用其他名称，比如 start_session），并返回 200 状态码，同时附带会话 ID 和生存时间（TTL），这些信息会被发送回客户端。
后续请求中，客户端在 invoke_endpoint 调用中传递相同的会话 ID。SageMaker 路由服务根据会话 ID 识别对应的推理实例，确保所有请求都路由到同一台服务器。这样，在在 LLM 推理的多轮对话场景中，GPU 内存中的 KV Cache 得以保持和累积，显著减少了计算开销，从而显著降低推理延迟。
当对话结束时，客户端可以发送关闭会话请求，您可以使用 invoke_endpoint 传入一个包含会话 ID 的关闭请求。SageMaker 路由器首先会检查该会话是否存在。如果存在，路由器会向推理服务器发起关闭会话的调用，服务器会返回一个成功的 200 响应以及会话 ID，这些信息随后会被发送回客户端。在会话 ID 不存在的情况下，路由器会返回一个 400 响应。

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

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

实践案例: 构建支持 Sticky Session 的 LLM 推理服务

本案例将以基于 SGLang 推理引擎的模型部署为例，展示如何在 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")
        if not 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})

构建客户端调用封装

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

import boto3, json

class StatefulSMEDPBuilder:
    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")

# 生成唯一会话 ID
session_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 的有效复用，显著提升多轮对话的性能表现。

完整的部署和使用代码可参考项目仓库：sample-sagemaker-sticky-session

总结

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

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

参考链接

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

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

亚马逊AWS官方博客