以下文章来源于蚂蚁开源,作者罗泽轩
4 月 26 日,蚂蚁开源技术沙龙首站在 2050 大会新生论坛圆满落幕,业界资深技术专家与开源社区贡献者围绕 Cloud Native AI 的技术前沿展开了深入探讨。蚂蚁研发工程师、MOSN 社区核心成员罗泽轩分享了蚂蚁 AI Gateway 如何应对大模型推理的资源利用率不稳定、计算量大、对网关要求更高等挑战。
本文整理自罗泽轩的演讲,围绕蚂蚁 AI Gateway 在智能路由、cache-aware 调度、网关架构优化等方面的实践,来显著提升资源利用率与推理效率。
背景介绍
在进入主题之前,先与大家分享一些关键数据。
首先是通义大模型的下载量增长:去年 5 月是 700 万次,8 月突破 2000 万次,至今年 2 月已超 2 亿次,几乎是每三个月翻三倍。一年增长几十倍的业务,可以说是相当快迅猛的增长了。这些主要是去年的数据,那么今年会不会还能这么飞速且持续地增长下去呢?
对此,在今年英伟达的 GTC 大会上,英伟达 CEO 黄仁勋表示推理型模型(如 OpenAI O1 和 DeepSeek R1)的普及,未来推理服务需求将迎来 100 倍量级增长。可以说,AI 推理服务已突破传统线性增长模式,形成“多级火箭”式发展态势,每到不同的节点就会催生一些新场景,推动业务持续跃迁。面对日新月异的推理场景,和蒸蒸日上的推理需求,沿用传统的网关的思路无法很好地服务用户的需求,我们必须另辟蹊径,在推理场景里引入新的网关架构。
大模型推理场景特征:
1. 流式响应:基于 Token 逐步生成内容,相信大家都见过 AI 应用是如何一点点吐出推理结论的;
2. 耗时长:单次推理常持续数分钟;
3. 请求数低:一个万卡集群,对外提供 DeepSeek R1 的服务,QPS 只有几百个;
4. 成本高:一张 GPU 卡的租金往往需要几千元每月,而万卡集群只能提供几百 QPS 的服务。折算下来,单次请求成本高昂。
由此衍生的技术痛点:
1. 低流量场景下的资源利用率波动较大;
2. 推理请求计算量大:需要结合大模型推理的一些概念(Prefill、Decode、KV Cache 等)做优化,尽可能提升 QPS;
3. 对网关提出更高的要求:网关需要实现配置热更新与流式响应支持。
技术目标:
作为 AI 服务的核心入口,蚂蚁 AI Gateway 承担流量调度与算力优化的核心角色,聚焦如下三个技术目标:
优化资源利用率
加速推理
提升接入体验
下文将围绕这三大目标,详解技术实践中的关键突破。
核心挑战和应对策略
在低 QPS、高成本的大模型推理场景下,AI 网关面临双重技术挑战:我们既要让引擎持续充分地忙碌起来,又需避免过多请求在引擎侧堆积。因此,AI Gateway 需要尽可能好地主动平衡不同的请求。
目前蚂蚁 AI Gateway 通过监控引擎的运行队列和等待队列两个指标,来计算当前哪个引擎最适合处理请求。我们会根据运行任务数和等待任务数来计算得分。得分越低,就越可能被选中。其中等待任务数会乘上一个惩罚系数,放大它对总得分的影响。
此外,还有两个小调整:
1. 每次选中一个节点后,立刻将它的运行任务数 +1,因为引擎指标的更新不是实时的,通过先在网关侧 +1,我们可以有更短的反馈链路,等到引擎指标同步过来之后,再使用最新的引擎指标。
2. 因为网关是多实例的,有可能每个网关认为的最优节点都是同一个节点,导致多个请求同时打到同一个节点。所以我们并非直接取分数最低的节点,而是取最好的一个范围里的节点,从中随机选一个。
为什么推理引擎并发数很低?因为推理过程的计算量很大。以 Prefill 阶段生成的 KV Cache 为例,4K Input Token,精度为 FP8,在 MHA Attention 算法下会生成超过 16 GB 的 KV Cache。优化后的 MLA 算法可以大幅降低 KV Cache,但也需要接近 300 MB,好在推理引擎支持在请求间复用 KV Cache。在基于文本的推理场景下,不同请求往往共享同一个 System Prompt 抑或它们是多轮会话的一部分,所以我们开启了 Prefix Cache,同样前缀的请求共享一组 KV Cache。为了保证同样前缀的请求尽可能地落到同样的引擎,我们在全局维护一个近似的前缀树,能够根据用户输入的 Prompt 计算出每个引擎能够有百分之多少的复用率,然后把这个复用率加到打分的公式里。这样,在负载相近的前提下,复用率越高的节点越有可能被选中。
前面讲到的都是负载均衡相关的,让我们把视角放得更高一点,看看网关在这之外需要具备的能力。因为推理请求每个请求都是长连接,所以目前我们日常变更的配置都是动态加载,只有升级网关时才会重启。而且网关是流式处理响应,每生成一次 Token,都会流式响应给客户端。
下图为蚂蚁 AI Gateway 架构图,可以从全局上认识 AI Gateway。
蚂蚁 AI Gateway 架构图
首先在控制路径方面,目前是多个推理集群共享一个 Console 服务,每个集群里有多台 Istio 和更多的 Envoy。我们修改了 Istio,在里面加入自己的 CRD 来描述模型映射关系、限流、灰度等配置。在推理集群里,上游信息是通过蚂蚁的 AntVIP 服务获取的,因此我们增加了 AntVIP 到 CDS(Cluster Discovery Service)的转换功能,而不是像标准的 Istio 那样由 Istio 提供上游配置。
在数据链路方面,可以看到 AI gateway 是整体推理服务的入口,是用户和各式推理引擎之间的桥梁。另外,考虑到 Metrics 采集之类的能力是开发在独立的 Meta Center 上的。如果是由每个网关直接来采集 Metrics,随着服务规模增大,会导致 Metrics 采集量以 O(n²) 增大。假设整个集群扩大 10 倍,那么网关也需要扩大 10 倍,则单个引擎上的 Metrics 采集请求也会扩大 10 倍,哪怕采集间隔没有变化。从总体上看,Metrics 采集量会是集群扩大倍数 * 单引擎的采集请求扩大倍数,也就是一个 O(n²) 的量。而单独拆分出一个采集服务,则采集请求数只和集群规模有关,每个引擎也不会有额外的压力。集群扩大规模后,同时也要求更短的采集间隔。举个例子,一个引擎的 Metrics 变化通常在几十 ms 到 100 ms 之间(取决于具体一次 Step 花费的时间)。如果只有一个引擎,那么 100 ms 采集一次,几乎能感知到每一次变化。如果有 100 个引擎,则有可能在 0~100ms 之中每个 ms 都有一个引擎发生了变化,如果要达到之前那么高的灵敏度,需要把采集间隔调小。但是这种调小在大部分时候都是无用功,因为从单个引擎的角度看,Metrics 变化的周期还是 100ms 里发生一次。事实上,采用轮询方式获取引擎的指标是种浪费行为。因此,我们计划修改引擎代码,让引擎主动把各种指标推送给 Meta Center。
当前实践效果
目前蚂蚁内部已经有若干万卡集群部署了 AI Gateway。相比基线的优化程度,我们衡量网关效能的指标是利用率,计算方式为:在满足目标 TTFT(Time To First Token)和 TPOT(Time Per Output Token)的前提下,经过网关后的 QPS 除以后端的个数,再除以每个后端的基线 QPS,得到的百分比就是利用率。相比于没有任何智能加持的 RoundRobin 负载均衡,我们的利用率提升了 50% 以上。
业界动态
目前业界内也有许多发力推理场景下的 AI Gateway。严格来说,下面列出的 AIBrix 和 Dynamo 并不是纯粹的 AI Gateway,而是整套的分布式推理系统,AI Gateway 只是它们内部的一个组件。
SGLang-Router:
SGLang-Router 机制相对简单,Rust 写的 HTTP 服务器。它是业界较早做 Cache Aware 的负载均衡的实现,维护了一个近似的前缀树来判断 Cache 复用的情况。对于引擎执行任务数和缓存复用率之间的平衡,SGLang-Router 实现了两种负载均衡算法,当任务数均衡时它基于缓存复用率来调度;当任务数不均衡时,它基于任务数来调度。
SGLang-Router 实现了两种负载均衡算法
Gateway-API-Inference-Extension
包含两个组件:Envoy 和 Endpoint Picker。Endpoint Picker 是用 Go 实现的一个数据面的 Sidecar 组件,负责选择 Endpoint。请求打到 Envoy 之后和 Sidecar 通信,拿到目标 Endpoint。前面提到的智能调度就是实现在 Sidecar 中。目前通信的实现比较粗糙,虽然 Sidecar 在选 Endpoint 时只需要知道 Model 名称,但在通信时会将整个请求体都发过去(内含十几 KB 的 Prompt)。AI Gateway 则是在 Envoy 进程里面执行 Go 代码,所以几乎没有通信的开销。
AIBrix
AIBrix 中的 AI gateway 架构和 Gateway-API-Inference-Extension 是相同的。由于它们的 Metrics 采集是直接做在网关上,而且还是每个网关定期轮询所有推理引擎,所以在大规模场景下,还有比较大的优化空间。
AIBrix 中的 AI gateway 架构示意图
Dynamo
Dynamo 拓展性强,因为它把 Proxy、Endpoint Picker 和 Metrics 采集都放到独立模块里,任何一个模块都能根据实际需要部署更多实例。和我们的设计一样,Dynamo 也能有效避免随着服务规模增大,而导致 Metrics 采集量以 O(n²) 增大。Dynamo 并没有采用一个中心化的组件来采集 Metrics。它会在推理引擎上运行名为 Worker 的组件,这个组件可看作管理引擎的 Agent。Dynamo 修改了 vLLM 的代码,让它把最新的 Metrics 推送到 Worker 上面,Worker 再把 Metrics 提交到 NATS 上。Dynamo 的 Endpoint Picker 逻辑是实现在名为 KV Router 的组件上。KV Router 会定期去从 NATS 查询节点的 Metrics。这么一来,既能保证 Metrics 的实时性,又能保证 Metrics 采集的负载尽可能小。当然,Dynamo 的开发还处于早期阶段,比如 Metrics Publish 的能力,目前只在 vLLM 上支持。随着时间推移,它的一些实现细节很有可能会有大的变化。
未来计划与前沿探索
AI Gateway 相关功能会在 MOSN 社区开源,希望能够帮助到有同样需求的开发者。同时,我们也正在与多个推理领域的开源社区合作,推动 AI Gateway 标准化,求同存异,实现互惠共赢。
前面聊到的各种同类项目,它们都有不止一种负载均衡算法的实现,而且每种实现里都有许多可供调整的参数,包括蚂蚁自己的 AI Gateway 也是一样的。从某种意义上说,目前推理请求的负载均衡还是有一些调参的成分,往往是拿一些数据制作一个仿真场景,然后调控参数,看看效果如何。既然都是调参,那么能不能引入 AI 来做预测与决策,实现更智能的自适应调度?这是个值得思考的方向。
另外,以上所有的负载均衡算法的实现,基本上都依赖引擎提供的 Metrics。是否有一种方式,可以基于网络特征,实现引擎指标无关的负载均衡?如果这种方式真的存在,那么无疑是对现有各种基于 Metrics 的调参算法的一大降维打击。
👆扫描上方二维码,观看演讲视频《智算时代的流量枢纽:蚂蚁 AI Gateway 如何提升大模型推理效能》