摘要

我们介绍了AIBrix,一个云原生的开源框架,旨在优化和简化云环境中的大规模LLM部署。与传统的云原生技术栈不同,AIBrix遵循协同设计理念,确保基础设施的每一层都为与vLLM等推理引擎的无缝集成而量身定制。 AIBrix引入了多项关键创新以降低推理成本并提升性能,包括用于动态适配器调度的高密度LoRA管理、LLM专属的自动扩缩容机制,以及前缀感知、负载感知的路由策略。为了进一步提升效率,AIBrix集成了分布式KV缓存,促进了跨节点的令牌重用,从而使吞吐量提升50%,推理延迟降低70%。AIBrix还支持统一的AI运行时,在保持供应商无关的引擎兼容性的同时,简化了模型管理。 对于大规模多节点推理,AIBrix采用混合编排策略——利用Kubernetes进行粗粒度调度,利用Ray进行细粒度执行——以平衡效率与灵活性。此外,一个SLO驱动的GPU优化器能够动态调整资源分配,优化异构服务,在维持服务保障的同时最大化成本效益。最后,AIBrix通过AI加速器诊断工具增强了系统可靠性,实现了自动故障检测和模拟测试,从而提升了容错能力。AIBrix可在 https://github.com/vllm-project/aibrix获取。

1 引言

大语言模型已经彻底改变了AI应用,推动了聊天机器人、自动内容生成和高级推荐引擎等领域的创新。虽然基于OpenAI和Anthropic等专有模型的API服务已被广泛采用,但许多企业出于数据安全顾虑、定制化需求或专有解决方案的成本考虑,常常寻求开源替代方案。托管LLaMA、Deepseek、Qwen和Mistral等开源模型并提供生产级API的需求日益增长,这带来了新的挑战——即如何在保持低推理延迟和成本效益的同时,大规模高效地部署它们。 在生产环境中部署LLM不仅需要优化的模型,还需要一个跨越多个层次的系统性方法:

  • 开源模型​:AI应用的基础,通过模型架构优化、多头潜在注意力(MLA)(Liu et al.[2024])、蒸馏和适配器微调等技术来提升性能和适应性。
  • 推理引擎​:如vLLM(Kwon et al.[2023])或TensorRT-LLM等推理引擎,通过KV缓存管理、模型并行、注意力优化等技术提高模型服务效率。
  • 系统级编排(AIBrix)​:这一关键但常被忽视的层次,通过管理资源调度、自动扩缩容、请求路由、异构性管理以及多集群或多区域资源优化,决定了实际的成本效益和可扩展性。

虽然模型和引擎优化至关重要,但系统级编排是实现真正成本效益的关键。没有设计良好的基础设施,即使是最先进的模型和推理引擎也难以应对实际部署中的挑战,如自动扩缩容、缓存感知路由、异构资源管理等。 为了解决这些挑战,我们推出了AIBrix,一个新颖的云原生框架,旨在简化和优化LLM推理基础设施,为用户提供一键部署体验,同时确保一流的性能和成本效益。我们的主要贡献包括:

  • 高密度LoRA管理​:AIBrix支持在vLLM中动态注册LoRA和谱系支持,简化了LoRA适配器管理,降低了管理微调模型的成本。
  • LLM专属自动扩缩容​:AIBrix支持多种场景驱动的LLM自动扩缩容策略。它还采用了滑动窗口指标聚合等优化,以减少实时指标的传播延迟。
  • 高级LLM网关和路由策略​:AIBrix引入了LLM感知的API网关,扩展了Envoy Gateway以优化实例路由并支持多种路由策略。与盲目分发请求的传统网关不同,AIBrix分析令牌模式、预填充缓存可用性和计算开销,以提升不同部署场景下的路由效率。
  • 统一的AI运行时与GPU流式加载器​:AIBrix作为统一的运行时层,管理推理引擎Pod与控制平面之间的交互。它自动化模型工件处理,配置推理引擎,并提供实时可观测性,确保供应商无关的兼容性。此外,AIBrix还具备GPU流式加载器功能,可绕过磁盘I/O瓶颈,加速模型加载和执行。
  • 分布式与解耦的KV缓存池​:AIBrix引入了分布式KV缓存,支持高容量、跨引擎的KV重用,同时优化网络和内存效率。关键创新包括抗扫描的淘汰策略、减少冗余数据传输、异步元数据更新和基于共享内存的数据交换,从而提升推理吞吐量和效率。
  • 混合粒度的多节点推理编排​:AIBrix引入了一种混合方法进行多节点推理,将Ray用于细粒度的应用编排,将Kubernetes用于粗粒度的资源管理。与推理引擎在分布式环境中的原生支持(如vLLM)强调并行性而非面向服务的需求相比,AIBrix平衡了分布式执行与生产级编排,实现了可扩展性、滚动升级和高效资源分配。
  • 成本高效且SLO驱动的异构服务​:AIBrix引入了一个GPU优化器,平衡成本效益与SLO遵从性,根据工作负载特性和硬件可用性动态选择最优的GPU配置,确保成本高效的异构GPU利用。
  • 加速器诊断与故障模拟工具​:AIBrix引入了一个诊断工具,利用AI加速器的内置功能来检测和诊断硬件故障。同时,一个故障模拟工具可以模拟硬件故障,实现严格的容错和恢复测试。

AIBrix是一个云原生的开源框架,它简化并优化了生产环境中的LLM部署,为工业界和学术界的AI从业者提供了一个灵活、可扩展且经济高效的服务解决方案。通过将模型和引擎优化与系统级编排深度集成,AIBrix弥合了效率与灵活性之间的差距,为大规模LLM推理工作负载设立了新标准。

2 相关工作

现有的云原生和机器学习服务框架为模型推理提供了基础架构,但缺乏大规模LLM推理所需的关键优化。我们将先前的工作分为两大类:基于微服务的无服务器框架和传统的ML模型服务框架,并重点说明它们在应用于LLM工作负载时的局限性。 ​基于微服务的系统​:如Knative和Istio等基于微服务的框架为管理无状态服务提供了强大的解决方案,强调高级流量控制机制,如请求速率限制、请求拦截、身份验证、授权以及基于QPS和并发指标的自动扩缩容。然而,这些解决方案并非为基于GPU的推理工作负载设计,未能解决LLM应用带来的根本性变化。与传统的微服务不同,LLM推理不需要复杂的服务网格或广泛的路由功能。相反,它引入了新的挑战,如基于令牌的速率限制、KV缓存感知的自动扩缩容和模型特定的调度约束。例如,Knative中基于熔断器的速率限制与LLM基于令牌的推理约束不兼容,而基于QPS的自动扩缩容无法准确捕捉GPU绑定的资源使用模式(如KV缓存内存压力)。此外,Istio服务网格的开销使其不适合LLM应用,在这种情况下,轻量级、推理专用的优化更为有效。 ​传统ML模型服务系统​:如KServe和RayServe等传统ML模型服务框架为模型部署、请求处理和自动扩缩容提供了解决方案,使其非常适合传统的深度学习推理。这些框架支持模型URI管理、动态扩缩容和模型版本控制等功能。然而,它们缺乏与LLM推理引擎的深度集成,未能解决LLM工作负载特有的关键挑战。LLM推理引入了高度可变的输入输出长度、巨大的模型大小和有状态执行(如KV缓存管理)等独特特性,这些都需要自定义路由、模型分发策略和GPU感知调度。虽然KServe和RayServe可以部署LLM模型,但它们并未提供针对高效批处理调度、KV缓存协调或异构GPU利用的专门优化。因此,它们无法充分利用现代LLM推理引擎(如vLLM和TensorRT-LLM)的性能潜力。

3 AIBrix:用于LLM服务的云原生基础设施

3.1 AIBrix架构

在本节中,我们将描述AIBrix的架构,如图1所示。AIBrix包含控制平面组件和数据平面组件。控制平面的组件管理模型元数据的注册、自动扩缩容、模型适配器注册,并强制执行各类策略。数据平面组件提供可配置的组件,用于分发、调度和服务推理请求,实现灵活且高性能的模型执行。 ​AIBrix控制平面​:AIBrix控制平面通过自动化模型管理、优化资源分配和实现智能扩缩容来简化LLM部署。目前,AIBrix并未为基础模型构建任何抽象,它仍然利用Kubernetes部署进行基本的生命周期操作。LoRA适配器控制器通过支持每个Pod部署多个LoRA来增强灵活性,提高了可扩展性和资源利用率。另一个名为RayClusterFleet的控制器管理由vLLM服务的大规模模型的多节点推理。 为了桥接推理引擎和控制平面,AIBrix配备了一个AI运行时,它作为一个轻量级的边车,卸载管理任务、强制执行策略,并为vLLM、SGLang(Zheng et al.[2024])和TensorRT-LLM等系统抽象引擎交互。与此相辅相成的是,冷启动管理器跟踪DRAM、本地存储和云存储中的模型工件,确保模型加载到最快的可用节点上,以最小化启动延迟。 同时,LLM专属的自动扩缩容器支持实时的毫秒级扩缩容,利用KV缓存利用率和推理感知指标来动态优化资源分配。这些组件共同形成了一个高度自适应的控制平面,将云基础设施和推理引擎协同设计相结合,提供可扩展、成本高效和高性能的LLM推理。

AIBrix数据平面​:AIBrix的数据平面负责处理用户请求、管理模型实例、执行推理以及优化缓存策略,以确保高效且可扩展的LLM服务。它被设计为同时具备应用感知和资源感知能力,将深度推理优化与云原生编排相结合。 在入口点,API网关作为中央请求分发器,执行公平策略、速率控制(TPM/RPM)和工作负载隔离,同时根据KV缓存位置、令牌吞吐量和GPU异构性动态优化流量。这种自适应路由机制确保了在不同硬件配置上高效执行推理。 模型执行由服务单元(Serving Unit)促进,该单元包括推理引擎和AI运行时边车。AIBrix还引入了分布式KV缓存运行时,它扩展了外部缓存服务,以管理推理过程中动态生成的KV缓存。KV缓存重用对于减少冗余计算和提高令牌生成效率至关重要。基于DRAM的分布式KV缓存运行时支持跨节点的可扩展、低延迟缓存访问,并支持高级优化,如前缀缓存扩展、未来的预填充/解码解耦远程池(Zhong et al.[2024])和请求迁移,进一步提升了内存受限环境下的性能,并在某些特性不兼容时(例如,在vLLM v0.7.1中,当为Deepseek-R1启用MLA时,必须禁用vLLM中的前缀缓存)发挥关键作用。 这些组件共同形成了一个高度优化的数据平面,确保AIBrix为LLM应用提供可扩展、成本高效和高性能的推理,同时平衡延迟、资源利用率和 workload 公平性。

3.2 AIBrix特性

AIBrix汇集了多项创新,以简化企业级LLM基础设施,增强可扩展性和效率。

3.2.1 高密度LoRA管理

扩展低秩自适应(LoRA)微调模型传统上受到刚性部署的限制,导致灵活性降低和成本增加。传统的服务基础设施将LoRA适配器视为基础模型的静态附件,使得动态扩展不切实际。缺乏集成的资源管理会导致分配效率低下、驱逐不可靠以及故障处理欠佳。 AIBrix引入了高密度LoRA管理(图2),支持动态加载和卸载适配器、智能调度和LoRA感知路由,以提高推理效率。通过动态注册LoRA适配器,AIBrix支持高密度部署,显著降低了推理成本——这对于长尾场景尤其有益。利用Kubernetes的Service和EndpointSlice机制,AIBrix优化了LoRA模型的发现和放置,最大限度地减少了干扰并提高了资源利用率。此外,对vLLM的增强加强了LoRA管理能力,降低了操作开销,并提高了混合工作负载下的推理性能。

3.2.2 高级LLM网关和路由策略

传统的API网关难以应对LLM推理,因为请求的复杂性多种多样,从简单查询到需要复杂令牌管理的多轮交互。通用路由会导致流量分配效率低下和延迟峰值。AIBrix通过提供LLM感知的网关来解决这个问题,该网关扩展了Envoy Gateway以支持实例路由、前缀缓存感知和基于最小GPU内存的策略,如图3所示。与盲目分发请求的传统系统不同,AIBrix分析令牌模式、预填充缓存可用性和计算开销来优化流量。这使得高级路由定制和用户定义的策略成为可能。对于每个待处理的请求,当前版本的AIBrix根据以下路由策略之一确定目标实例:

  • random​:随机选择一个可用实例。
  • throughput​:选择每秒令牌数吞吐量最低的实例。
  • least-request​:选择已接纳请求数量最少的实例。
  • least-kv-cache​:选择平均KV缓存使用率最低的实例。
  • least-latency​:选择平均请求延迟最低的实例。对于每个请求,这是由其排队延迟和服务延迟之和得出的。
  • prefix-cache-aware​:优先选择具有可重用前缀缓存的实例,且缓存命中率超过阈值。

通过选择合适的路由策略,AIBrix能够将平均延迟降低19.2%,P99延迟降低79%,确保大规模LLM推理的高效和公平。AIBrix团队也正与Google及其他贡献者合作,致力于gateway-api-inference-extension Google项目,以备未来采用。

3.2.3 统一的AI运行时与GPU流式加载器

推理引擎领域正在快速发展,某些引擎在性能方面迅速超越其他引擎。此外,一些引擎提供独特的功能,例如对特定模型优化或功能组合的专门支持。因此,用户通常倾向于根据性能优势或功能集来利用不同的引擎。 然而,由于这些引擎使用的协议多种多样,在控制平面中直接支持它们是不可扩展的。为了解决这个问题,需要一个抽象层来统一和简化与这些不同推理引擎的交互,从而实现无缝集成和更高效的管理。AIBrix引入了统一的AI运行时(图4),作为AIBrix控制平面和推理引擎Pod之间的桥梁。该运行时管理模型、配置引擎、提供可观测性,并实现供应商无关的支持。它确保了核心组件(如LoRA适配器控制器、自动扩缩容器和冷启动管理器)之间的无缝通信,促进了动态的、云原生的资源管理。

3.2.4 LLM专属自动扩缩容以实现性能优化

LLM推理的自动扩缩容面临独特挑战,原因包括DCGM指标的限制、非线性的扩缩容行为,以及QPS或并发数等传统指标的不适用性。请求复杂性和I/O大小差异很大,常常在自动扩缩容器做出反应之前就使系统不堪重负。此外,大型模型镜像和缓慢的模型分发会给新Pod带来2-3分钟的延迟,使得快速扩缩容效率低下。AIBrix通过支持可配置的LLM专属自动扩缩容策略来缓解这些挑战。它绕过了自定义指标路径,直接在自动扩缩容器中维护滑动窗口指标聚合,以实现实时负载报告。通过利用Knative Pod Autoscaler(KPA)和AIBrix Pod Autoscaler(APA)等高级自动扩缩容算法,与原生HPA相比,AIBrix能够将延迟降低11.5%,将令牌吞吐量提高11.4%,并将扩缩容振荡减少33%。未来的工作将探索基于令牌的主动扩缩容和SLO驱动的自动扩缩容,以进一步提高效率和响应能力。

3.2.5 分布式KV缓存池

对大型语言模型日益增长的需求需要高效的内存管理和缓存策略,以优化推理性能并降低成本。在聊天机器人和基于代理的系统等多轮应用中,重叠的令牌序列通常会在预填充阶段导致冗余计算,造成资源浪费和吞吐量限制。虽然像vLLM这样的推理引擎内置了KV缓存,但单节点缓存受限于内存约束、引擎特定的存储限制,并且缺乏对KV迁移和预填充-解码解耦的支持。 为了解决这些挑战,AIBrix引入了分布式KV缓存(图5),支持高容量、跨引擎的KV重用,同时优化网络和内存效率。该系统采用抗扫描的淘汰策略来选择性持久化热门的KV张量,减少不必要的数据传输。此外,异步元数据更新最小化了开销,而缓存-引擎共址通过共享内存加速了数据传输。 表1展示了我们在4个Nvidia A10 GPU上使用Bird-SQL基准测试对我们的分布式KV缓存进行的性能评估。我们的结果表明,即使与vLLM的内置前缀缓存相比,将分布式KV缓存与前缀缓存结合使用,也能将峰值吞吐量提高约50%,将平均和P99 TTFT分别降低约60%和70%,并将平均和P99 ITL分别降低约30%和70%。这些发现证明了显著的效率提升。

表1​:vLLM和具有不同配置的AIBrix分布式KV缓存的性能比较。(此处vLLM默认配置为禁用分块预填充和前缀缓存。)

3.2.6 混合粒度的多节点推理编排

Llama-3.1-405B和Deepseek-R1的发布显著增加了对多节点推理的需求。然而,vLLM等现有框架优先考虑并行性而非面向服务的需求(如扩缩容和滚动升级),因此需要外部编排。 虽然Kubernetes和Ray提供了编排能力,但它们各有取舍:Kubernetes操作符提供粗粒度的资源管理,但对于细粒度调度可能很复杂;而Ray擅长分布式通信,但缺乏整体的资源控制。 为了解决这些局限性,AIBrix引入了一种混合方法,将Ray用于细粒度的应用编排,将Kubernetes用于粗粒度的资源管理。这种方法简化了操作符设计,确保能够适应未来推理编排模式的转变。我们观察到,编排方式常常由于技术不足(如预填充和解码解耦)而发生变化。保持编排设计的灵活性至关重要。借鉴字节跳动在管理Kubernetes和Ray工作负载方面的丰富经验,AIBrix利用自适应编排策略来克服工作负载通信挑战并优化多节点推理执行。

3.2.7 成本高效且SLO驱动的异构服务

最近的研究,如Mélange和QLM表明,在特定SLO约束下,LLM的吞吐量取决于输入/输出令牌数量以及在异构GPU环境中的模型选择。 此外,即使对于相同的模型,请求的成本效益在不同GPU上也有所不同,因为不同的GPU在不同工作负载下表现出不同的性能特征。加剧这些挑战的是,生产用户经常面临GPU可用性限制,无法获得一致的GPU类型。 为了解决这些问题,AIBrix引入了一个GPU优化器——一个离线路径组件,旨在通过平衡成本效益和SLO遵从性来优化异构GPU服务。如图8所示,该架构包含三个关键组件:

  • 负载监控器​:跟踪部署变化,假设不同的模型部署使用不同的GPU,并分析AIBrix网关统计数据以识别主要的工作负载模式。
  • GPU优化器​:动态选择最优的GPU组合,以平衡成本效益和SLO遵从性。
  • Pod自动扩缩容器​:从GPU优化器读取外部MetricSource数据,以动态调整GPU分配。目前,GPU优化器支持一个受Mélange启发的基于ILP的解决方案,需要部署前进行性能剖析。AIBrix提供了用于工作负载基准测试和性能剖析的工具包。

在我们比较异构工作负载(使用A10和L20)与同构设置(仅使用L20)的实验中,我们评估了一个包含ShareGPT和内部Text2SQL工作负载的混合数据集。异构配置导致延迟增加高达20%,但仍保持在指定的SLO范围内。然而,与同构GPU部署相比,该设置实现了约10%的成本降低。

图7​:加速器类型的选择因所用工作负载而异。

3.2.8 AI加速器诊断与故障模拟工具

GPU故障和性能下降在大规模AI部署中构成重大挑战。静默错误、过热、内存泄漏和间歇性故障可能会降低模型性能、增加延迟,甚至导致系统崩溃。在异构环境中诊断GPU问题尤其困难,因为不同的GPU模型在不同工作负载下表现出不一致的行为。 为了应对这些挑战,AIBrix加速器工具引入了:

  • GPU诊断和问题识别​:AIBrix自动化故障检测,帮助用户在工作负载受到影响之前主动识别和解决GPU相关的性能问题(图9a)。
  • GPU故障模拟工具​:AIBrix模拟GPU故障,允许开发人员测试和构建有弹性的AI框架,能够从硬件故障中优雅地恢复。目前支持Nvidia GPU和Ascend 910B NPU,并计划在不久的将来扩展对更多加速器的兼容性(图9b)。

4 结论

总结​:我们介绍了AIBrix,一个旨在解决大规模LLM推理挑战的新颖框架。AIBrix利用无服务器特性,包括LLM专属的自动扩缩容、冷启动优化和高密度部署,在系统层面显著降低了推理成本。除了成本效益,AIBrix还引入了分布式和解耦的服务特性,实现了跨不同工作负载的可扩展且灵活的LLM推理。分布式KV缓存池、混合编排和异构服务优化器等创新进一步推动了性能优化的边界,在保持低运营成本的同时确保高效的资源利用。这些特性共同使AIBrix成为一个高度适应性强且经济高效的大规模LLM部署解决方案。通过将深度推理引擎优化与云原生基础设施相结合,AIBrix提供了一个可扩展、高性能的服务平台,弥合了AI推理工作负载中效率与灵活性之间的差距。 ​局限性与未来工作​:我们的一些实验未能全面评估路由策略、非理想工作负载下的异构服务,限制了将这些特性推广到不同工作负载特征的能力。此外,基于性能剖析的自动扩缩容和异构GPU调度目前依赖于离线模型剖析,这引入了一个额外的步骤,对于动态工作负载可能并不总是实用。为了缓解这一点,一个潜在的解决方案是采用屋顶线模型分析来简化剖析过程,这可以为剖析异构推理性能提供更结构化和轻量级的方法。此外,我们旨在扩展评估范围以涵盖多样化的实际工作负载场景,改进路由策略、GPU分配机制和分布式编排,以进一步增强AIBrix在各种LLM部署环境中的适应性。 展望未来,我们将继续探索推理引擎与系统架构之间更深层次的协同设计,确保AIBrix始终是一个高度优化、可扩展且适用于生产环境的大规模LLM推理解决方案。


Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐