锐评主流AI推理负载均衡开源实现

负载均衡一向是业务架构里不可或缺的一部分，AI 场景下也不例外。由于推理请求业务量大，而且具有鲜明的特色，所以开发者会部署专门针对推理请求的负载均衡。这里我们就来讲讲推理请求负载均衡的一些开源实现。由于推理请求负载均衡这个名字实在太长，以下请允许我简称它为 ILB。

目前该领域负载均衡的一些实现，既有作为独立的组件对外提供，也有以网关的身份整合到一个整体的推理平台里面。前者有 Gateway API Inference Extension （以下简称 GAIE），后者则囊括了 AIBrix、Kthena、Dynamo 等等。作为一个推理平台，也可以使用第三方的负载均衡独立组件，而不是用自己内置的。比如 AIBrix 里面有篇文档提到如何禁用自己的 gateway，改为采用 GAIE 来做 PD 分离。以下为了简单明了，当我们讨论某个平台时，实际上指的是这个平台的网关组件。

在讨论具体的实现之前，让我们看看一个 ILB 需要满足什么功能。

像传统业务的负载均衡一样，它必须有两个基本功能：负载均衡和限额。前者可以拆成两部分 - 负载 和 均衡。一个负载均衡首先需要懂得跑在自己上面的负载是怎么样的，其次还要能知道离均衡状态还有多远。这样一来，我们就得到了考察 ILB 实现的三个最基础的问题：

1. 如何 tokenize 某个请求来知道当前对应的负载是怎么样的
   2a. 如何获取跟均衡相关的指标
   2b. 如何根据指标来达成均衡性

至于限额，如果你已经对当前负载了如指掌，无法是跑一些会更新计数器的自定义策略而已。如果是一些面向底层网络的负载均衡可能还会搞点拥塞控制，但对于收到完整的 prompt 才开始推理的推理请求业务，限额触发的操作简单很多。

AIBrix

tokenizer

AIBrix 支持三种 tokenizer：

1. 简单地按照 byte 切分
2. 调用 tiktoken 库切分
3. 调用 remote tokenize API

这里简单说说我自己的观点。tiktoken 是 OpenAI 开源出来，用于解析 GPT 系列模型的 tokenizer。我之前有篇文章聊过它：tiktoken vs hf tokenizer：AI网关如何本地高效统计Token。考虑到 AIBrix 不太可能只用来服务 GPT 系列，采用专为 GPT 系列模型开发的 tiktoken 就显得不够专业了。何况 AIBrix 用 tiktoken 时选择的 tokenizer encoding 配置还是已经被 OpenAI SOTA 模型弃用的 cl100k_base。即使用 AIBrix 来管理 GPT OSS 系列模型，也应该用 GPT OSS 系列的 o200k_harmony encoding。选错 tokenizer 对英文内容可能影响不大，但是涉及没有分隔符的中文，差别就很显然了。怀疑 AIBrix 对 tokenizer 的使用并不多，以致于错误的 tokenizer 也不会带来多大影响。

Remote tokenize API 是由维护推理引擎的团队来提供专门的 tokenize 接口。我曾经一度认为它是正确的做法，但现在看来这是个受迫于部门墙的次优选择。因为既然模型是自己部署的，有对应的 tokenizer 配置，为什么不直接在网关上通过 tokenizer 库解析，而是要把请求流量复制一遍，发送到另外一个团队维护的额外服务上呢？如果不是为了更好地分饼，这条串联链路实际上是没必要的。

我个人推崇的想法是采用 hf tokenizer 直接在网关上完成 tokenization，后续和 token 粒度相关的操作都直接读取 tokenize 后的输出。这样在整个请求的生命周期，都不需要额外发送 tokenize 请求，而且只需要一次 tokenization。当然这受限于许多实际因素（比如团队职责划分、技术栈是否能简单嵌入 hf tokenizer 库等）

metric collecting

AIBrix 采集指标的机制有三套：

1. 在每个网关上启动 worker 协程（默认 10 个），以默认 50 ms（AIBRIX_POD_METRIC_REFRESH_INTERVAL_MS） 轮询推理引擎的 metrics 接口
2. 同时还以 200 ms （AIBRIX_PROMETHEUS_QUERY_INTERVAL_MS）的频率查询 prometheus
3. 从消息队列中消费推理引擎推送的 kv event（如新增或删除 kvcache block）

第二套机制我觉得是毫无意义的。因为 Prometheus 数据来源，不也是和第一套机制一样高频轮询推理引擎的 metrics 接口得来的嘛。何况你还保证不了 Prometheus 那一条路径的数据及时性。比如说我 Prometheus 上配置采集的时间间隔是 15s，你 AIBrix 每 200 ms 频率查询结果也拿不到多高频的数据。如果说 Prometheus 接口的好处是做了一些聚合，那网关上同样也可以聚合，而且拿了 Prometheus 的数据后，也需要和第一套的高频数据做聚合。

在大规模的集群里，由分布式网关直接去采集每个推理引擎，消耗会很大。因为集群规模越大，引擎越多，网关也需要越多，这是个 O(n^2) 复杂度的事情。假设一个万卡集群有 500 个引擎，20 个网关，按 50 ms 默认间隔采集，一秒下来一共有 500 20 20 = 200k 个请求，平均一个引擎接收 400 个请求，一个网关发送 10k 个请求。如果还是按默认 10 个协程，一个协程要达到 1k 每秒的发送速度，则平均一个请求需要在 1ms 内完成。显然是不可能的。当然用户可以调大协程数，但总消耗的资源不会变；也可以调大查询间隔，但间隔大了实时性就要降低。按照经验，负载的均衡性和实时性呈正相关 - 决策方式不变，决策时用的数据越及时，上游负载越均匀。由每个网关定时轮询每个引擎，在大规模场景下架构是不太行的。

routing

AIBrix 有很多套路由算法可供选择。其中 kvcache aware 的路由算法会用到前面提到的第三套指标采集机制，并结合 tokenizer 做前缀树匹配，来找到最能复用 kvcache 的推理引擎节点。这里简单说明下为什么要按 block 粒度的前缀来做 kvcache 复用。推理引擎会按照一个可配置的 block 粒度来共享同样的 kvcache。如果一个请求的部分前缀命中 kvcache，则不需要重新计算这部分内容来生成 kvcache。为什么要根据前缀来复用，而不是任意相同序列呢？因为在计算 kvcache 时，计算后面的内容会依赖前面的，所以如果中间出现不同的 token，结果会不一样。

Kthena

Kthena router 和 AIBrix Gateway 很像，架构上差不多。比起 AIBrix，我觉得有两点改进：

1. 支持多个路由算法按权重编排。AIBrix 不支持复用路由算法，复杂的算法需要自己用代码实现，比如 vtc-basic routing 就是直接用代码实现了一套多因子路由。
2. 只是一个单二进制的 Go 实现的 router。不像 AIBrix 那样需要 Envoy + Go sidecar 作为数据面。作为一个简单主义者，我更喜欢做一个核心的数据面，不够了再加。而不是先引入一个支持 retry 和 ratelimit 的数据面，然后把核心作为 sidecar 实现。一开始引入复杂的架构，后面腾挪的时候就不容易了。再次以 AIBrix 的 vtc-basic routing 为例，因为 response 路径不经过 sidecar，所以它的实现里对于 output token 只能 estimated，拿不到真实的 token 数量。当然 AIBrix 可以把 response 路径也走一遍 sidecar，但是这么一来另一个数据面存在的意义就又被削弱了。

不过 Kthena router 也有一些鹦鹉学舌的地方。比如它还是用 cl100k_base 作为 tiktoken 的 encoding。如果是 AIBrix 用 cl100k_base 还算是代码写好了一直没更新，作为 2025 年中提交的代码还在用 cl100k_base，Kthena 这种就属于不知其所以然了。

Gateway API Inference Extension

tokenizer

GAIE 只支持 “简单地按照 byte 切分” 这种 tokenizer。准确来说是 bytes / averageCharactersPerToken，其中 averageCharactersPerToken = 4。

metric collecting

GAIE 采集指标的机制有两套：

1. 在每个实例上轮询每个推理引擎
2. 从推理响应结尾的 token usage 字段获取 token 数量

我怀疑第二套机制的有效性。考虑到推理请求动辄几分钟的耗时，这个 token 数量是几分钟内的总和，基本上没有什么实时性。如果说不用于负载均衡，只用于可观测性，完全可以在数据面上采集。

和本文中其他项目不同，GAIE 并不是个数据面，而是一个中心化的“大脑”（在该项目内叫做 EPP， Endpoint picker）。其他数据面（如 Envoy AI Gateway、AgentGateway）通过实现对接 EPP 的协议，由 EPP 决定路由到哪个推理引擎节点。理论上 GAIE 的第一套指标采集机制应该不会有 AIBrix 那种 O(n^2) 复杂度的问题，因为你可以部署少量 EPP 来专门做采集和调度。但是！由于 EPP 要求请求和响应路径都要走它，所以哪怕数据面是用 C++ 实现或者 Rust 实现，瓶颈都会在它这一个 Go 实现的组件上。以前面万卡集群的例子，20 个网关可能也要对应有 20 个 EPP。这么算下来还不如 sidecar，至少后者网络路径更短。

routing

GAIE 有很多套路由算法可供选择。和 Kthena 一样，它支持多套路由算法通过权重进行编排。

注意 GAIE 没有消费推理引擎的 kv event，那么它怎么做 kvcache aware 的 routing 呢？只要路由了请求过去，它就会认为目标推理引擎创建了 kv block，而推理引擎删除 kv block 的事件是通过 LRU 模拟的。

Dynamo

tokenizer

支持三种 tokenizer，全部都是本地 tokenizer：

• HuggingFace Tokenizer
• TikToken Tokenizer （基于 OpenAI tiktoken 的 fork）
• Fast Tokenizer：（encode 时使用 https://github.com/Atero-ai/fastokens，decode 还是用 HuggingFace Tokenizer）

routing

和其他项目一样，Dynamo 也支持配置多种 routing 策略。其中核心的策略是使用成本函数来选择最优 worker。

成本 = overlap_score_weight × prefill_blocks + decode_blocks

• prefill_blocks: 需要从头计算的 token 块数 = (总输入 token 数 - 重叠块数 × 块大小) / 块大小
• decode_blocks: 基于 worker 当前活跃序列估算的解码块数，参考了 router 本身维护的当前路由历史和 kv event 里面的 active blocks 指标。
• overlap_score_weight: 平衡缓存命中和负载分布的权重（默认1.0）

Dynamo 的创新之举在于，它在评估各个推理引擎的当前负载，只使用了 kv event 里面的 active blocks 指标和 router 本身维护的路由历史记录。思路是这样的：无论是当前计算的 token 数还是 GPU util，都能通过当前在用的 blocks 体现出来；至于排队中的 token 数，router 本身又是知道的。所以它的 metric collecting 要比其他家简单很多，只有消费 kv event 这一套。由于不需要轮询后端推理引擎，且消息队列本身已经在推理引擎和网关之间完成了解耦，它的指标采集消耗是最少的。

读者会问，如果我部署了多个 router，那每个 router 自己看到的数据不就是片面的，还不是需要去全量轮询每个推理引擎？Dynamo 通过 --router-replica-sync 启用 router 副本间的实时状态同步。另外还有 --router-temperature 选项引入适度随机性，避免惊群。需要承认的是，即使有同步机制，由于网络延迟，毫秒级的状态不一致仍可能发生，但 Dynamo 的异步事件系统确保了最终一致性。