AIBrix KVCache 卸载框架

AIBrix KVCache 卸载框架#

对大型语言模型日益增长的需求，使得对高效内存管理和缓存的需求也随之增加，以优化推理性能并降低成本。在聊天机器人和基于代理的系统等多轮用例中，重叠的 token 序列导致预填充阶段的冗余计算，浪费资源并限制吞吐量。

许多推理引擎，如 vLLM，使用内置的 KV 缓存来缓解这个问题，利用空闲的 HBM 和 DRAM。然而，单节点 KV 缓存面临关键限制：受限的内存容量、特定于引擎的存储（阻止跨实例共享）以及难以支持 KV 迁移和预填充-解码分离等场景。

在 AIBrix v0.3.0 中，我们引入了一个生产就绪的 KVCache 卸载框架，它能够实现高效的内存分层和低开销的跨引擎复用。默认情况下，该框架利用基于 L1 DRAM 的缓存，通过减轻 GPU 内存压力同时不带来高延迟，已经提供了显著的性能改进。对于需要多节点共享或更大规模复用的场景，AIBrix 允许用户选择性地启用 L2 远程缓存，从而释放分布式 KV 缓存层的优势。

aibrix-kvcache-offloading-arch-overview

图 1. AIBrix KVCache 卸载框架

如图 1 所示，在数据平面上，它通过 AIBrix 卸载连接器与推理引擎（例如 vLLM）紧密集成，该连接器采用优化的 CUDA 内核显著加速 GPU 和 CPU 之间的数据移动。为了提高内存可伸缩性，其多层缓存管理器动态地平衡存储层之间的工作负载，缓解 GPU 内存容量限制，同时最大程度地减少延迟损失。该框架支持可插拔的逐出策略（例如 LRU，S3FIFO）和多样化的后端存储选项（例如 InfiniStore），从而能够选择性地卸载 KV 缓存以减少网络和 PCIe 争用。至关重要的是，其缓存放置模块可以与集中式分布式 KV 缓存集群管理器协调，以最大限度地提高全局 KV 缓存利用率。这使得跨引擎 KV 复用成为可能，并确保集群范围内的资源效率，将孤立的 KV 缓存实例转变为可扩展的共享 KV 缓存基础设施。

L1 引擎 DRAM 缓存管理#

现代模型日益增长的需求和 LLM 推理中上下文长度的增加，导致 KV 缓存消耗越来越多的 GPU 内存，甚至超出了最先进 GPU 的硬件限制。最近的系统，如 Dynamo、LMCache 和 MoonCake，已经开发出将 KV 缓存卸载到外部内存层次结构（从 CPU 内存到 SSD）的解决方案。KVCache 卸载 也支持通过仅启用其 DRAM 支持的 L1Cache 将 KV 缓存卸载到 CPU 内存。虽然这种方法不能实现多个引擎之间的 KV 缓存共享，但它消除了分布式 KV 缓存设置和配置的复杂性。更重要的是，通过利用 CPU 内存显著更大的容量，这种方法提供了实质性的性能提升——使其成为优先考虑可伸缩 KV 缓存容量而非跨引擎 KV 复用的用例的理想解决方案。

L2 分布式 KVCache 和跨引擎 KV 复用#

大型语言模型日益增长的需求显著增加了对大容量 KV 缓存的需求。虽然 CPU 内存卸载有效地解决了适度的扩展需求，但处理大规模、动态工作负载的生产环境需要更大的可扩展性——尤其是在内存需求超出单节点容量时。为了解决这个问题，AIBrix 启用分布式 KV 缓存服务作为其 L2Cache 后端，可以跨多个节点水平扩展以满足容量需求。

与此同时，随着 LLM 部署在集群中跨多个引擎扩展，引擎之间 KV 缓存的冗余会带来巨大的低效率。重复计算常见提示前缀会浪费 GPU 周期和 HBM 带宽。AIBrix 通过高性能、共享的分布式 KV 缓存实现高效的跨引擎 KV 复用来解决这一挑战，从而大规模优化资源利用率。

aibrix-infinistore-arch-overview

添加新的 KVCache 后端#

可以通过实现 Connector 接口轻松添加新的 KVCache 后端

  1@dataclass
  2class ConnectorFeature:
  3    """The features of the kv cache connector.
  4    Args:
  5        mput_mget: Whether the kv cache connector supports mput/mget
  6        prefetch: Whether the kv cache connector supports prefetch.
  7        rdma: Whether the kv cache connector supports RDMA.
  8        gdr_put: Whether the kv cache connector supports GDR put.
  9        gdr_get: Whether the kv cache connector supports GDR get.
 10    """
 11
 12    mput_mget: bool = False
 13    prefetch: bool = False
 14    rdma: bool = False
 15    gdr_put: bool = False
 16    gdr_get: bool = False
 17
 18
 19@dataclass
 20class ConnectorConfig:
 21    """The config of the kv cache connector."""
 22
 23    backend_name: str
 24    namespace: str
 25    partition_id: str
 26    executor: Executor
 27    block_spec_signature: str = ""
 28    key_builder_signature: str = ""
 29    layout_signature: str = ""
 30
 31
 32@dataclass
 33class ConnectorRegisterDescriptor:
 34    """The register descriptor"""
 35
 36    pass
 37
 38
 39class Connector(Generic[K, V]):
 40    """Connector interface."""
 41
 42    @classmethod
 43    @abstractmethod
 44    def from_envs(cls, conn_id: str, executor: Executor, **kwargs):
 45        """Create a connector from environment variables."""
 46        raise NotImplementedError
 47
 48    @property
 49    @abstractmethod
 50    def name(self) -> str:
 51        raise NotImplementedError
 52
 53    @property
 54    @abstractmethod
 55    def feature(self) -> ConnectorFeature:
 56        """Get the feature of the connector.
 57        Returns:
 58            The feature of the kv cache service.
 59        """
 60        raise NotImplementedError
 61
 62    @abstractmethod
 63    def open(self) -> Status:
 64        """Open a connection."""
 65        raise NotImplementedError
 66
 67    @abstractmethod
 68    def close(self) -> Status:
 69        """Close a connection."""
 70        raise NotImplementedError
 71
 72    async def prefetch(self, keys: Sequence[K]) -> None:
 73        """Prefetch a list of keys.
 74        Args:
 75            keys: The keys of the kv tensors.
 76        """
 77        pass
 78
 79    @abstractmethod
 80    async def exists(self, key: K) -> Status:
 81        """Check if key is in the store."""
 82        raise NotImplementedError
 83
 84    @abstractmethod
 85    async def get(
 86        self, key: K, mr: MemoryRegion | Sequence[MemoryRegion]
 87    ) -> Status:
 88        """Get a value.
 89        Args:
 90            key: The key of the kv tensor.
 91            mr: The memory region or MR list to place the fetched kv
 92                tensor. It is an MR list only if using GDR.
 93        Returns:
 94            The status of the get operation.
 95        """
 96        raise NotImplementedError
 97
 98    @abstractmethod
 99    async def put(
100        self, key: K, mr: MemoryRegion | Sequence[MemoryRegion]
101    ) -> Status:
102        """Put a key value pair.
103        Args:
104            key: The key of the kv cache.
105            mr: The memory region or MR list holding the kv tensors. It is an
106                MR list only if using GDR.
107        Returns:
108            The status of the put operation.
109        """
110        raise NotImplementedError
111
112    def register_slabs(self, slabs: List[torch.Tensor]) -> Status:
113        """Register slabs with backend-specific register function.
114        Args:
115            slabs: slabs to be registered.
116        Returns:
117            Status of the register operation.
118        """
119        raise NotImplementedError
120
121    def get_batches(
122        self,
123        keys: Sequence[Any],
124        mrs: Sequence[MemoryRegion | Sequence[MemoryRegion]],
125        batch_size: int,
126    ) -> Sequence[Sequence[Tuple[K, MemoryRegion | Sequence[MemoryRegion]]]]:
127        """Get a list of key MR batches that is used for mput and mget
128        operations.
129
130        Args:
131            keys: The keys of the kv tensors.
132            mrs: Memory regions or lists of MRs holding the kv tensors.
133            batch_size: The maximum number of key MR pairs in a batch.
134        Returns:
135            List of key MR/MR List batches.
136        """
137        raise NotImplementedError
138
139    async def mget(
140        self,
141        keys: Sequence[K],
142        mrs: Sequence[MemoryRegion | Sequence[MemoryRegion]],
143    ) -> Sequence[Status]:
144        """MGet a list of values. This function is optional and only connectors
145        have mput_mget feature enabled can implement this function.
146        Args:
147            keys: The keys of the kv tensors.
148            mrs: Memory regions or lists of MRs to hold the fetched kv
149                 tensors. It is an MR list only if using GDR.
150        Returns:
151            List of statuses.
152        """
153        raise NotImplementedError
154
155    async def mput(
156        self,
157        keys: Sequence[K],
158        mrs: Sequence[MemoryRegion | Sequence[MemoryRegion]],
159    ) -> Sequence[Status]:
160        """MPut a list of key value pairs. This function is optional and only
161        connectors have mput_mget feature enabled can implement this function.
162        Args:
163            keys: The keys of the kv tensors.
164            mrs: Memory regions or lists of MRs holding the kv tensors. It is
165                 an MR list only if using GDR.
166        Returns:
167            List of statuses.
168        """
169        raise NotImplementedError
170
171    @abstractmethod
172    async def delete(self, key: K) -> Status:
173        """Delete a key.
174        Args:
175            key: The key of the kv cache.
176        Returns:
177            The status of the delete operation.
178        """
179        raise NotImplementedError

请参考现有连接器了解更多详情。

环境变量参考#

本节描述了 AIBrix KVCache 卸载框架所有可用的环境变量。

核心配置#

变量

默认

描述

AIBRIX_KV_CACHE_OL_CHUNK_SIZE

“512”

操作的块大小。

AIBRIX_KV_CACHE_OL_BLOCK_SIZE

“-1”

kvcache 块中的 token 数量（最细的 IO 粒度）。默认为 -1，表示使用引擎的块大小。

AIBRIX_KV_CACHE_OL_MAX_SEQ_LEN

“-1”

最大序列长度。默认为 -1，表示没有限制。如果设置，超过此长度的 token 将被忽略。

AIBRIX_KV_CACHE_OL_TIME_MEASUREMENT_ENABLED

“1”

启用时间测量。

AIBRIX_KV_CACHE_OL_BREAKDOWN_MEASUREMENT_ENABLED

“1”

启用分解测量。

AIBRIX_KV_CACHE_OL_DOUBLE_GET_THRESHOLD

“4,0.1”

控制何时向 L2 缓存发出第二次获取请求。第一个值是最小缺失块数，第二个是比例阈值。

AIBRIX_KV_CACHE_OL_TOKEN_VALIDATION_ENABLED

“0”

是否验证 L2 缓存中的 token。禁用则使用更紧凑的内存布局。

AIBRIX_KV_CACHE_OL_TRANSPORT_RDMA_ADDR_RANGE

“::/0”

有效的 GID 范围（CIDR 格式）。类似于 NVSHMEM 的 NVSHMEM_IB_ADDR_RANGE。

AIBRIX_KV_CACHE_OL_PROFILING_ENABLED

“0”

启用性能分析。

AIBRIX_KV_CACHE_OL_PROFILING_SERVER_ADDRESS

“http://0.0.0.0:4040”

性能分析服务器地址。性能分析服务器负责收集性能分析数据并在 Web UI 中显示。

L1 缓存配置#

变量

默认

描述

AIBRIX_KV_CACHE_OL_L1_CACHE_ENABLED

“1”

启用 L1 缓存。

AIBRIX_KV_CACHE_OL_L1_CACHE_EVICTION_POLICY

“S3FIFO”

L1 缓存的逐出策略（“S3FIFO”、“LRU”或“FIFO”）

AIBRIX_KV_CACHE_OL_L1_CACHE_CAPACITY_GB

“10”

L1 缓存容量（GB）。

AIBRIX_KV_CACHE_OL_DEVICE

“cpu”

用于缓存操作的设备（“cpu”或“cuda”）

AIBRIX_KV_CACHE_OL_S3FIFO_SMALL_TO_MAIN_PROMO_THRESHOLD

“1”

S3FIFO 逐出策略：从小队列到主队列的提升阈值。

AIBRIX_KV_CACHE_OL_S3FIFO_SMALL_FIFO_CAPACITY_RATIO

0.3

S3FIFO 逐出策略：小 FIFO 的容量比。

L2 缓存配置#

变量

默认

描述

AIBRIX_KV_CACHE_OL_L2_CACHE_BACKEND

“”

L2 缓存的后端。

AIBRIX_KV_CACHE_OL_L2_CACHE_NAMESPACE

“aibrix”

L2 缓存的命名空间。

AIBRIX_KV_CACHE_OL_L2_CACHE_OP_BATCH

“32”

操作批次大小。

AIBRIX_KV_CACHE_OL_L2_CACHE_PER_TOKEN_TIMEOUT_MS

“20”

每 token 超时时间（毫秒）。

AIBRIX_KV_CACHE_OL_L2_CACHE_KEY_BUILDER

“ROLLING_HASH”

L2 缓存的键生成器（“RAW”、“ROLLING_HASH”或“SIMPLE_HASH”）

AIBRIX_KV_CACHE_OL_L2_CACHE_INGESTION_TYPE

“HOT”

摄取类型（“ALL”、“HOT”或“EVICTED”）。

AIBRIX_KV_CACHE_OL_L2_CACHE_INGESTION_MAX_INFLIGHT_TOKENS

“0”

最大未完成写入数量（0 表示同步）。

AIBRIX_KV_CACHE_OL_L2_CACHE_NUM_ASYNC_WORKERS

“8”

异步工作器数量。

AIBRIX_KV_CACHE_OL_L2_CACHE_PLACEMENT_POLICY

“SIMPLE”

放置策略（仅在使用元数据服务时适用）。

元数据服务配置#

变量

默认

描述

AIBRIX_KV_CACHE_OL_META_SERVICE_BACKEND

“”

元数据服务的后端。如果未设置元数据服务后端，L2 缓存后端将使用直接模式访问给定的缓存服务器。否则，我们将从元数据服务获取成员信息并构建 L2 缓存集群。

AIBRIX_KV_CACHE_OL_META_SERVICE_REFRESH_INTERVAL_S

“30”

刷新间隔（秒）。

AIBRIX_KV_CACHE_OL_META_SERVICE_URL

“”

元数据服务的 URL。

AIBRIX_KV_CACHE_OL_META_SERVICE_CLUSTER_META_KEY

“”

集群元数据键。

连接器配置#

InfiniStore 连接器配置#

变量

默认

描述

AIBRIX_KV_CACHE_OL_INFINISTORE_HOST_ADDR

“127.0.0.1”

主机地址。

AIBRIX_KV_CACHE_OL_INFINISTORE_SERVICE_PORT

“12345”

服务端口。

AIBRIX_KV_CACHE_OL_INFINISTORE_CONNECTION_TYPE

“RDMA”

连接类型。

AIBRIX_KV_CACHE_OL_INFINISTORE_IB_PORT

“1”

IB 端口。

AIBRIX_KV_CACHE_OL_INFINISTORE_LINK_TYPE

“Ethernet”

链路类型。

AIBRIX_KV_CACHE_OL_INFINISTORE_VISIBLE_DEV_LIST

“”

可见设备列表。自 0.2.42 版本起，InfiniStore 支持客户端配置中的 RDMA GID 索引，用户可以按此格式指定每个设备的 GID 索引：“mlx5_0:gid0,mlx5_1:gid1,mlx5_2:gid2”

HPKV 连接器配置#

变量

默认

描述

AIBRIX_KV_CACHE_OL_HPKV_REMOTE_ADDR

“127.0.0.1”

远程地址。

AIBRIX_KV_CACHE_OL_HPKV_REMOTE_PORT

“12346”

远程端口。

AIBRIX_KV_CACHE_OL_HPKV_LOCAL_ADDR

“127.0.0.1”

本地地址。

AIBRIX_KV_CACHE_OL_HPKV_LOCAL_PORT

“12345”

本地端口。

PrisKV 连接器配置#

变量

默认

描述

AIBRIX_KV_CACHE_OL_PRISKV_REMOTE_ADDR

“127.0.0.1”

远程地址。

AIBRIX_KV_CACHE_OL_PRISKV_REMOTE_PORT

“6379”

远程端口。

AIBRIX_KV_CACHE_OL_PRISKV_USE_MPUT_MGET

“0”

启用 MPUT/MGET。

AIBRIX_KV_CACHE_OL_PRISKV_PASSWORD

“”

密码。

EIC 连接器配置#

变量

默认

描述

AIBRIX_KV_CACHE_OL_EIC_CONFIG_FILE

“”

EIC 配置文件。

Mock 连接器配置#

Mock 连接器用于测试和性能分析。

变量

默认

描述

AIBRIX_KV_CACHE_OL_MOCK_USE_RDMA

“0”

在模拟连接器中使用 RDMA。

AIBRIX_KV_CACHE_OL_MOCK_USE_MPUT_MGET

“0”

在模拟连接器中使用 MPUT/MGET。

AIBRIX_KV_CACHE_OL_MOCK_USE_NOOP

“0”

所有操作都使用 NOOP。有助于分析框架开销。

RocksDB 连接器配置#

RocksDB 连接器用于测试目的。

变量

默认

描述

AIBRIX_KV_CACHE_OL_ROCKSDB_ROOT

“~/.kv_cache_ol/rocksdb”

RocksDB 的根目录。

AIBRIX_KV_CACHE_OL_ROCKSDB_TTL_S

“600”

TTL（秒）。

AIBRIX_KV_CACHE_OL_ROCKSDB_WRITE_BUFFER_SIZE

“67108864”

写入缓冲区大小。默认为 64MB。

AIBRIX_KV_CACHE_OL_ROCKSDB_TARGET_FILE_SIZE_BASE

“67108864”

目标文件大小基数。默认为 64MB。

AIBRIX_KV_CACHE_OL_ROCKSDB_MAX_WRITE_BUFFER_NUMBER

“3”

最大写入缓冲区数量。

AIBRIX_KV_CACHE_OL_ROCKSDB_MAX_TOTAL_WAL_SIZE

“134217728”

最大 WAL 总大小。默认为 128MB。

AIBRIX_KV_CACHE_OL_ROCKSDB_MAX_BACKGROUND_JOBS

“8”

最大后台作业数量。