SageMaker推前缀路由：LLM首Token延迟降低77%

当你在大语言模型（LLM）上构建应用时，发送给模型的 prompt 通常由两部分组成：一部分是固定内容，用来建立上下文（指令、参考文档、对话历史），另一部分是可变内容，即用户的实际输入。以客服机器人为例，每个请求都以相同的文本块开头："你是 AnyCompany 的支持代理，以下是我们的政策……"，后面跟着客户输入的内容。顶部的指令可能是 3000 个 token，底部客户的问题可能只有 50 个 token。

这意味着，在数百或数千个请求中，你的模型在反复处理那相同的 3000 token 前缀。

vLLM 和 TensorRT-LLM 等 LLM 服务框架为此提供了一种解决方案。它们会缓存已见过 prompt 前缀的计算结果——键值（KV） pairs。当新的请求中出现相同的前缀时，模型重用缓存的计算结果，只处理末尾的新 token。这就是前缀缓存（prefix caching），它可以显著降低首 token 响应时间（TTFT）。

但当扩展到单个实例之外时，问题就出现了。如果你在端点后有一组机器，请求会分发到所有机器上。同样的 3000 token 前缀，对一个请求落在实例 A 上，对下一个请求落在实例 B 上，再下一个落在实例 C 上。每个实例都从头计算，因为没有哪个实例频繁看到它来建立可靠的缓存。前缀缓存功能是有的，但路由层把请求分散得太薄，使其无法发挥作用。

今天，Amazon SageMaker Inference 推出了前缀感知路由（prefix-aware routing）。这是一种新的路由策略，它会查看每个请求的开头，将具有相同开头的请求一致地发送到同一个实例。这样该实例上的 KV 缓存会真正积累起来并被重用。在我们对 Llama 3.1 70B 的基准测试中，这使 P50 TTFT 降低了最多 77%，吞吐量提高了最多 16%。KV 缓存命中率也从大约 25% 提升到了 80% 以上。

前缀感知路由的工作原理

当请求到达你的端点时，Amazon SageMaker 会查看载荷的开头，并用它来决定应该由哪个实例处理。相同的开头发送到同一个实例。不同的开头分散到不同的实例。如果 10 个请求共享一个前缀或开头，全部 10 个都会发送到同一台机器，而该机器的缓存在那个前缀上保持热状态。

你不需要为请求打标签或自己管理亲和性。端点根据请求的内容来处理。

它有两个内置的保护机制：

过载保护。如果某个前缀非常流行且目标实例已满载，端点会将请求路由到负载较低的实例。你配置并发限制，端点会遵守它。在那一个请求上你可能错过一次缓存命中，但避免了压垮单台机器。

扩缩容期间的稳定行为。当你添加或移除实例时，大多数请求继续发送到它们之前发送的同一实例。只有一小部分流量会转移以适应变化的集群。每次扩缩容时你的缓存不会失效。

性能基准测试

我们使用 Llama 3.1 70B Instruct 在 7 个 ml.p5.48xlarge 实例（vLLM，启用前缀缓存）上，将前缀感知路由与默认的随机路由基线进行了基准测试。我们运行了 16 种测试配置，涵盖单模型端点、推理组件端点、原生 Invoke API 和 OpenAI 兼容 API。所有测试均以 100% 成功率完成。

长上下文工作负载

8000 token 共享前缀，持续 1 小时：

• P90 TTFT：降低了 33–37%
• P50 TTFT：降低了 71–77%
• KV 缓存命中率：从大约 25% 提升到 82%
• 吞吐量：提高了 15–16%

短上下文工作负载

可变长度的 ShareGPT 风格对话，30 分钟：

• P90 TTFT：降低了 24–37%
• P50 TTFT：降低了 13–16%
• KV 缓存命中率：从大约 30% 提升到 80%
• 吞吐量：提高了 1.7–2.0%

共享前缀越长，收益越大。长上下文工作负载受益最多，因为每次缓存命中可以跳过的计算量更大。短上下文工作负载也有收益，但由于共享前缀较小，每请求的节省比例也较小。

前缀感知路由逻辑每个请求增加 1.3–1.9 毫秒。本测试中模型 TTFT 范围为 63–280 毫秒。路由开销可以忽略不计。

在所有场景中，流量分配保持均衡。7 个实例各自接收了 13.3–15.4% 的请求，偏差在理想均匀分配的 1% 以内。没有热点实例。

SageMaker Inference 的路由策略

此次发布后，Amazon SageMaker Inference 为实时端点提供了三种路由策略：

RANDOM（默认）：将请求均匀分发到各个实例。适用于通用工作负载、非 LLM 模型，或者请求可互换、发送到特定实例没有收益的场景。

LEAST_OUTSTANDING_REQUESTS：将每个请求发送到飞行中请求数最少的实例。适用于请求处理时间各异、希望所有实例保持同等忙碌的场景。有助于防止慢请求在某一台机器上堆积而其他机器空闲。

PREFIX_AWARE（新增）：将共享相同 prompt 前缀的请求发送到同一实例。适用于许多请求在开头共享共同文本且你的服务框架已启用前缀缓存的 LLM 工作负载。

你可以在端点配置中为每个生产变体设置策略。更新端点配置即可切换策略，无需重新部署模型。

何时使用前缀感知路由

当你的请求在开头共享文本时，此功能会发挥作用。以下是影响最大的场景：

检索增强生成（RAG）应用。你检索一个文档并将其前置到用户问题之前。当多个用户询问同一文档的问题时，他们都将该文档作为前缀共享。前缀感知路由将他们发送到同一实例，该文档的 KV 缓存已经处于热状态。

多轮对话。对话中的每一轮都包含之前轮次的完整历史。随着对话增长，共享历史成为更长、更昂贵的前缀。基于该前缀进行路由，使对话的缓存在各轮次中保持在同一实例上。

模板化机器人和助手。带有长结构化指令（策略、格式化规则、角色定义）的机器人每次请求都发送相同的指令。只有末尾的用户消息发生变化。前缀感知路由意味着那条昂贵的指令块只处理一次，而不是数千次。

代码补全。编码助手将文件内容作为上下文包含。当开发者在同一文件中工作时，每个补全请求都将该文件内容作为前缀共享。

配置参数

你可以在创建端点配置时配置前缀感知路由。两个参数控制行为：

PrefixLength（1024–65536）：用于路由的请求部分大小。对于原生 Amazon SageMaker Invoke API，这是请求体开头的字节数。对于 OpenAI 兼容 API，这是从提取的消息文本中的字符数。设置为覆盖你的共享前缀加上足够的唯一内容，以将不同工作负载分散到不同实例。

ConcurrencyThreshold（1–1024）：溢出触发前目标实例上的最大飞行中请求数。如果目标实例达到此限制，请求将转到负载较低的实例。

aws sagemaker create-endpoint-config \
    --endpoint-config-name example-llm-config \
    --production-variants '[{
        "VariantName": "AllTraffic",
        "ModelName": "example-llm-model",
        "InitialInstanceCount": 3,
        "InstanceType": "ml.p5.48xlarge",
        "RoutingConfig": {
            "RoutingStrategy": "PREFIX_AWARE",
            "PrefixAwareRoutingConfig": {
                "PrefixLength": 4096,
                "ConcurrencyThreshold": 10
            }
        }
    }]'

然后按常规方式创建端点：

aws sagemaker create-endpoint \
    --endpoint-name example-llm-endpoint \
    --endpoint-config-name example-llm-config

无需更改模型容器或服务框架。前缀感知路由完全在端点路由层操作。

调用端点

调用端点的方式没有变化。InvokeEndpoint 和 InvokeEndpointWithResponseStream API 与之前完全一样：

aws sagemaker-runtime invoke-endpoint \
    --endpoint-name example-llm-endpoint \
    --content-type application/json \
    --body fileb://request.json \
    output.json

OpenAI 兼容的 Chat Completion API 同样如此：

from openai import OpenAI
from sagemaker.core.token_generator import generate_token

client = OpenAI(
    base_url=f"https://runtime.sagemaker.us-west-2.amazonaws.com"
             f"/endpoints/example-llm-endpoint/openai/v1",
    api_key=generate_token(region="us-west-2")
)
response = client.chat.completions.create(
    model="example-model",
    messages=[
        {"role": "user", "content": "What is your return policy?"},
    ],
)

多租户前缀隔离

如果不同租户共享相同的 prompt 指令，但你希望将他们分开路由（保持缓存上下文独立），可以传递一个可选 ID：

• 原生 Invoke API：设置 X-Amzn-SageMaker-Prefix-Aware-Id 请求头（最多 64 个 ASCII 字符）
• OpenAI API：在请求体中包含 prompt_cache_key 字段

此 ID 与前缀组合，使具有相同前缀但不同 ID 的请求落在不同实例上。

推理组件和 LoRA 适配器

前缀感知路由与推理组件端点和动态低秩适应（LoRA）适配器配合使用。对于推理组件，其行为与单模型端点相同。对于 LoRA 适配器，它在适配器的粘性实例集中操作，在已加载该适配器的实例之间使用基于前缀的选择。

最佳实践

在服务框架中启用前缀缓存。前缀感知路由将重复的前缀送到同一实例，但你的容器需要启用前缀缓存才能实际存储和重用这些 KV pairs。在 vLLM 中，最近版本默认启用此功能。其他框架可能需要显式配置。

保持请求序列化一致。对于原生 Invoke API，PrefixLength 对原始字节操作。JSON 空白符、键顺序和格式化都会影响路由。如果你以不同方式序列化同一 prompt，它们可能会落在不同实例上。使用一致的序列化。

谨慎设置 PrefixLength。太短会导致所有具有相同短前缀的请求被导向一个实例，触发溢出。太长则会使小有效载荷差异（如 temperature 值）将本应在一起的请求分散开。从你的共享前缀长度加上适当中缓冲开始设置。

至少需要两个实例。只有一个实例时，无论策略如何，所有请求都发送到同一位置。

监控缓存命中率。启用 SageMaker 详细可观测性来跟踪模型级别的 KV 缓存命中率。这可以确认前缀感知路由对你的特定工作负载是否有效。

前缀感知路由现已可在 SageMaker 实时推理端点上使用。更新你的 AWS SDK 或 CLI 到最新版本以访问新的 RoutingStrategy 和 PrefixAwareRoutingConfig 参数。参考此 notebook 了解在端点创建期间如何启用它的示例。