机器学习平台技术栈之 Aibrix

随着大模型（LLM）的爆发式发展，各大企业和研究机构纷纷将业务核心向 AI 倾斜。随之而来的是对底层基础设施的巨大考验，尤其是如何高效、经济、稳定地部署和提供大模型推理服务。在这个背景下，云原生架构和机器学习平台的结合成为了必然趋势。

在众多的云原生 AI 推理解决方案中，Aibrix 逐渐崭露头角。作为一款专为大规模人工智能推理打造的云原生基础设施组件，Aibrix 致力于解决多模型部署、流量路由、资源调度，特别是 GPU 缓存（如 KV Cache）管理等核心痛点。

本文将深入探讨 Aibrix 的技术栈，从核心概念、架构设计到关键技术细节，为你全面解析这款强大的机器学习平台推理网关与调度引擎。

1. 什么是 Aibrix？

Aibrix（AI + Bricks）寓意为构建 AI 应用的基石。它本质上是一个部署在 Kubernetes 上的大模型推理网关与调度器。不同于传统的微服务网关（如 Nginx、Envoy 等），Aibrix 是专为 LLM 的 workload 定制的，能够感知 LLM 推理的上下文（如 Prompt、Prefix、KV Cache），并以此为基础进行智能路由和资源分配。

传统的无状态负载均衡在处理大模型推理时存在严重缺陷：LLM 推理是强状态相关的（Context Stateful）。如果一个包含了长前缀（Long Context）的请求被随机分配到一个没有缓存该前缀的 GPU 节点上，节点需要重新计算 Context（Prefill 阶段），这会消耗大量的 GPU 算力并带来极高的延迟老虎延迟。

Aibrix 的核心目标就是：让大模型推理的流量调度“懂”状态（State-aware），从而最大化复用 KV Cache，降低首字延迟（TTFT），提升整体系统的吞吐量。

2. 核心概念与关系

理解 Aibrix，首先需要理清以下几个核心概念及其相互作用。

2.1 Pod 与 Inference Engine (推理引擎)

在 Aibrix 中，底层的推理计算依然由成熟的推理引擎（如 vLLM, TensorRT-LLM, TGI 等）承担，这些引擎被打包进 Kubernetes 的 Pod 中。Aibrix 不负责具体的矩阵运算，而是负责统筹这些 Pod。

2.2 KV Cache (键值缓存)

在 Transformer 架构中，生成下一个 Token 需要依赖之前所有 Token 的计算结果（Key 和 Value 矩阵）。将这些结果缓存下来，就是 KV Cache。KV Cache 的命中率直接决定了 LLM 服务的性能。

2.3 Prefix Routing (前缀路由)

这是 Aibrix 最核心的概念之一。Aibrix 的网关会解析客户端发来的请求，提取公共前缀（如 System Prompt 或历史对话）。然后，它会查询集群中哪个 Pod 已经缓存了该前缀的最长 KV Cache，并将请求精确路由过去。

2.4 GPU 分时与 LoRA 适配器

通过调度，Aibrix 可以在同一组 GPU 资源上动态加载和卸载不同的 LoRA 权重，使得一个小规模的 GPU 集群能够看似同时提供成百上千个微调模型的服务。

概念之间关系图谱

graph TD
    User([User Request / Prompt]) --> Gateway[Aibrix Gateway / Router]
    Gateway -->|Parse Context & Routing| ControlPlane[Aibrix Control Plane]
    ControlPlane -->|State & Cache Info| ETCD[(State Store/Redis)]
    Gateway -->|Direct Request| Pod1[Pod 1 (vLLM) \n Cache: Prompt A]
    Gateway -->|Direct Request| Pod2[Pod 2 (vLLM) \n Cache: Prompt B]
    Gateway -->|Direct Request| Pod3[Pod 3 (vLLM) \n LoRA Weights: X]
    
    subgraph K8s Cluster
    Gateway
    ControlPlane
    Pod1
    Pod2
    Pod3
    end

3. 架构设计

Aibrix 的架构设计遵循控制面（Control Plane）与数据面（Data Plane）分离的微服务设计模式。

3.1 总体架构

Aibrix 的总体架构可以分为以下几个核心组件：

Aibrix Control Plane (控制器): 以 Kubernetes Operator 的形式运行，负责监听 CRD（Custom Resource Definition），管理 Model, Deployment 等资源的生命周期。它还负责收集各个推理 Pod 的状态和缓存信息。
Aibrix Gateway (数据面网关): 接收流量入口，通常是一个高性能的反向代理。它内置了智能路由逻辑，根据请求头或请求体进行解析。
KV Cache Metadata Tracker (状态追踪器): 组件会定期拉取或接收推理后端（如 vLLM）爆出的缓存状态，维护一个分布式的 Trie 树（前缀树），记录哪个 Pod 拥有哪个文本片段的 KV Cache。

3.2 智能路由架构细节

当一个用户的请求到达时，系统是如何处理的？

sequenceDiagram
    participant Client
    participant GW as Aibrix Gateway
    participant Router as KV-Aware Router
    participant Tracker as Metadata Tracker
    participant Pod1 as Inference Pod 1 
    participant Pod2 as Inference Pod 2

    Client->>GW: POST /v1/chat/completions (Prompt: "Sys: Act as expert. User:... ")
    GW->>Router: Forward Request Data
    Router->>Tracker: Lookup Longest Prefix Match
    Tracker-->>Router: Pod 2 has prefix "Sys: Act as expert." (Token length: 150)
    Router-->>GW: Route to Pod 2
    GW->>Pod2: Forward Request
    Pod2-->>GW: Stream Response Tokens
    GW-->>Client: Stream Response
    
    Note over Pod2,Tracker: Asynchronous KV Cache status update
    Pod2-)Tracker: Cache updated (New prefix generated)

4. 关键技术细节剖析

在这部分，我们将深入探讨支撑 Aibrix 达到高性能的几项关键技术。

4.1 基于 Prefix/Radix Tree 的状态感知路由 (State-aware Routing)

为什么说 Aibrix 比 Nginx 更适合 LLM？因为 Nginx 只能做针对 IP、轮询或是基于 Header 的 Hashing。但是在 LLM 中，真正的业务逻辑“主键”是 Token 序列。

Aibrix 在内存中维护了一棵分布式的 Radix Tree（基数树）。每一个节点代表一段 Token 序列，树的叶子或中间节点上附加了 Pod IP 列表。
当请求进入时：

Aibrix 使用内置的 Tokenizer 将 Prompt 转换为 Token IDs。
在 Radix Tree 中执行最长前缀匹配（Longest Prefix Match）。
找到拥有最长前缀的 Pod，将请求转发。

为了更清晰地说明这个过程，让我们看一个具体的例子。假设系统中有三个常用前缀：

前缀 A: “系统提示：你是一个专业的” （缓存于 Pod 1 和 Pod 2）
前缀 B: “系统提示：你是一个专业的程序员，擅长” （缓存于 Pod 1）
前缀 C: “帮我写一篇关于” （缓存于 Pod 3）

当一个新请求 “系统提示：你是一个专业的程序员，擅长 Python 语言” 到达网关时：

网关的 Rust Tokenizer 提取出 Token。
在分布式 Radix Tree 中匹配，发现前缀 B 匹配程度最高（Longest Prefix Match）。
这个前缀被存储在 Pod 1 中。
流量路由模块直接将请求打到 Pod 1。

Pod 1 在收到请求后，不需要为 “系统提示：你是一个专业的程序员，擅长” 这部分提示词重新计算 Attention 矩阵，直接复用已有的 KV Cache，仅需计算 “ Python 语言” 这几个 Token，极大降低了首字延迟（TTFT，Time To First Token）。

难点： Tokenization 过程本身需要算力消耗，且必须与后端的 LLM 使用的 Tokenizer 一致。Aibrix 目前通过轻量级的 Rust Tokenizer 实现，并在网关侧进行缓冲加速。此外，KV Cache 的状态是动态变化的，Pod 可能会因为 OOM 驱逐缓存，网关需要与后端保持低延迟的状态同步。

4.2 弹性 LoRA 管理与 Serverless 部署 (Serverless LoRA)

随着 PEFT（Parameter-Efficient Fine-Tuning）技术的普及，传统的全参微调被成本更低的 LoRA（Low-Rank Adaptation）取代。企业通常拥有一个基座大模型（如 Llama 3 8B）和数以千计的特定任务 LoRA 模型（如用于客服对话的 LoRA，用于代码生成的 LoRA，用于特定行业文档总结的 LoRA）。

如果在传统的架构中，为每一个 LoRA 部署一个独立的 Pod，不仅会造成巨大的显存浪费（即使没有流量也需要常驻显存），而且很难实现资源的动态伸缩。Aibrix 的 LoRA 调度方案完美解决了这个问题：

graph TD
    User([User Request with LoRA ID]) --> Gateway[Aibrix Gateway]
    Gateway -->|Check LoRA Cache| Scheduler[LoRA Scheduler]
    
    subgraph K8s Node (GPU)
        Engine[vLLM Engine (Base Model Loaded)]
        Cache[LoRA Weights Cache]
        Engine <--> Cache
    end
    
    Scheduler -->|If not locally cached| Storage[(S3 / OSS / PVC)]
    Scheduler -->|Download & Load into GPU| Engine
    Gateway -->|Forward Inference Request| Engine

Aibrix 支持通过一个统一的入口提供所有 LoRA 的推断。

统一入口多路复用：请求中携带 Header 或者是特定 API 参数指定 model: lora-id。
零容错自动加载：若目标 Pod 内没有加载该 LoRA 适配器，Aibrix 将控制流介入，触发 Backend（如 vLLM, SGLang）动态从对象存储（S3/OSS）下载并将其加载到 GPU 显存。
基于 LRU 的显存驱逐策略：考虑到单卡的显存上限，如果同时加载太多 LoRA 会导致 OOM。当 GPU VRAM 满时，Aibrix 调度器会根据 LRU（最近最少使用）或 LFU 算法，通知 Backend 卸载不常用的 LoRA 权重，腾出空间。这种情况使得即使只有很少的物理 GPU 资源，也能让成百上千个 LoRA 以 Serverless 的感觉提供服务。

4.3 调度层深度融合：K8s Scheduler Plugin 扩展

Aibrix 不仅仅是一个网络层的网关，它还是在调度层（Scheduling Layer）深度优化的引擎。它包含一个基于 Scheduler Framework 的自定义 K8s Scheduler Plugin。

在标准的 K8s 世界里，调度一个 Pod 只关乎资源本身：Node A 有多少空闲 CPU 和 Memory，Node B 有多少空闲 GPU。但在 LLM 推理场景中，**”引力”（Gravity）**发生了改变。缓存成为了最重要的资源引力。

Aibrix Scheduler 能够根据过去一段时间内的“前缀热点”和网络拓扑进行调度。
举个例子：

如果识别到“中文系统提示词”的请求极为频繁，Aibrix 会有意识地将请求路由引向一组特定的 Pod 集合。
如果需要 Scale-up（扩容），Aibrix 会倾向于将被扩容的 Pod 放置在与“富集了热门 KV Cache 的节点”网络距离近的地方，甚至通过某种机制在后台预热（Pre-warm）缓存数据。
这使得这部分 Pod 成为该前缀的专属高速缓存区，减少了整体机群的计算开销。

4.4 针对大模型特性的高可用与弹性缩放 (LLM-Native HPA)

在云原生领域，基于 CPU 或内存使用率进行水平扩展（Horizontal Pod Autoscaling, HPA）是标配。但针对 LLM 推理，这一策略常常失效：因为即使单卡推理一个简单的请求，GPU 利用率也可能瞬时飙升到 100%；同时，显存通常一启动就被预分配满了（如 vLLM 的 PagedAttention 机制）。

不同于依靠 CPU/Mem 缩放，Aibrix 暴露了基于 LLM 专属特征的 Custom Metrics，并通过与 Prometheus 及 KEDA（Kubernetes Event-driven Autoscaling）结合，实现了 LLM-Native 的弹性扩缩容：

**KVCache 碎片率和占用率 (KV Cache Usage)**：如果 KV Cache 开始拥挤，说明吞吐触顶，需要扩容。
**排队中的 Token 数量 (Pending Token Count / Queue Length)**：如果在排队的请求需要的总计算 Token 量开始攀升。
首字延迟 (TTFT) P99 指标：这是最直观的体验指标，TTFT 超过阈值直接触发扩容。
**每秒生成 Token 数 (TPOT, Time Per Output Token)**：用于判断生成的吞吐量是否受限。

借助 KEDA，Aibrix 能在 TTFT 开始飙升，或者等待队列变长时，快速拉起新的 Pod 副本应对突发流量；而在夜间流量低谷，又能快速回收 GPU 资源，极大地节省了算力成本。

5. 核心组件交互协议 (Protocol)

Aibrix 的高效运作离不开控制面与数据面、网关与后端推理引擎之间的紧密通信。这里简要介绍其内部使用的协议设计。

Gateway <-> Control Plane：主要使用 gRPC 流式接口。网关会定期（例如每秒）向控制平面报告当前的流量分布和路由命中统计，而控制平面则将最新的路由表（反映各 Pod 的存活状态和缓存分布）下发给网关。
**Control Plane <-> Inference Engine (Pod)**：控制平面需要获取每个 Pod 的 KV Cache 实际情况。通常，引擎（如 vLLM）会被植入一个 sidecar，或者直接支持特定的 metrics 接口暴露其 PagedAttention 的块分配信息和映射的前缀哈希。
Client <-> Gateway：对外提供标准的 OpenAI 兼容 API（RESTful JSON），使得现有的 LangChain, LlamaIndex 等生态应用可以无缝迁移，无需修改任何代码。

6. 与其他开源解决方案的对比

在了解了 Aibrix 之后，我们不可避免地会将其与市面上其他的解决方案进行对比。

特性	传统的 K8s Ingress (Nginx)	KServe / Triton	Aibrix
路由感知粒度	IP / Header 层级	Model 层级 (请求基于模型名路由)	Context / Token 层级 (感知 Prompt Prefix)
KV Cache 亲和性	无支持，随机分发导致 Cache Miss 高。	极弱，主要依赖底层引擎自己处理。	强，基于最长前缀匹配精确路由。
LoRA 多路复用	不支持。	需要配置多个后端或依赖 Triton 复杂的配置。	原生支持，基于请求头自动下载、加载和驱逐。
扩缩容指标	CPU / 内存 / QPS	Request 队列长度 / QPS	LLM-Native (TTFT, KV 容量, Pending Tokens)

可以看出，Aibrix 是真正意义上面向大规模、复杂场景的 LLM 云原生系统组件。它接管了传统微服务网关在 AI 时代力不从心的那一层上下文感知路由能力。

7. 总结与展望

Aibrix 代表了下一代机器学习平台基础设施的发展方向：从无状态的网络微服务架构走向强状态感知的 AI 原生架构。

在万物皆大模型的语境下，计算成本和响应延迟是决定产品生死的关键。通过深度解析请求上下文、协调分布式 KV Cache、动态调度轻量级 LoRA 权重，并在 Kubernetes 调度层进行精细干预，Aibrix 成功填补了云原生 LLM Serverless 治理领域的巨大空白。

展望未来，随着多模态大模型（Vision, Audio）的普及，缓存的内容将不再仅仅是文本 Token，还包含大量的图像特征和音频特征。Aibrix 这类架构如何高效地管理多模态特征的 Cache Affinity（缓存亲和性），以及如何应对更加复杂的跨节点推理（如 Tensor Parallelism 的多节点调度）挑战，将是这个技术栈演进的下一个里程碑。

对于所有底层 AI 平台开发工程师和 SRE 来说，深入理解并掌握“让调度器和网关理解大模型计算上下文”的思想，无疑是构建下一代高效智能基础设施的核心竞争力。