本地存储#

概述#

CPU RAM 和本地存储是将 KV Cache 卸载到同一台运行推理的机器的非 GPU 内存的两种方式。

配置 LMCache 磁盘卸载的两种方式:#

1. 环境变量:

# 256 Tokens per KV Chunk
export LMCACHE_CHUNK_SIZE=256
# None if disabled
# Otherwise, enable by setting the directory where LMCache will
# create files for each KV cache chunks
# (this directory does NOT need to exist beforehand)
export LMCACHE_LOCAL_DISK="file://$HOME/local/disk_test/local_disk/"
# 5GB of Disk
export LMCACHE_MAX_LOCAL_DISK_SIZE=5.0

# Disable page cache
# This should be turned on for better performance if most local CPU memory is used
export LMCACHE_EXTRA_CONFIG='{'use_odirect': True}'

2. 配置文件:

通过 LMCACHE_CONFIG_FILE=your-lmcache-config.yaml 传入

# 256 Tokens per KV Chunk
chunk_size: 256
# Enable Disk backend
local_disk: "file:///local/disk_test/local_disk/"
# 5GB of Disk memory
max_local_disk_size: 5.0

# Disable page cache
# This should be turned on for better performance if most local CPU memory is used
extra_config: {'use_odirect': True}

多路径(多设备)磁盘卸载#

如果您有 **多个 NVMe 设备**(或任何独立的挂载点),您可以为每个 GPU 分配自己的磁盘路径,以便每个设备写入专用驱动器。

local_disk 中指定一个 以逗号分隔的路径列表。每个路径可以选择性地使用 file:// 前缀。 local_disk_path_sharding 选项控制每个 GPU 工作线程如何选择其路径。目前仅支持 "by_gpu"``(默认),它根据设备索引(``device_id % num_paths)选择路径,因此来自给定 GPU 的所有 KV 缓存文件都会落在同一个 NVMe 上。当 GPU 和 NVMe 设备共享 PCIe 交换机或 NUMA 节点时,这尤其有用。

例如,使用两个 GPU 和两个路径:

  • cuda:0/mnt/nvme0/kvcache/

  • cuda:1/mnt/nvme1/kvcache/

max_local_disk_size 是所有路径共享的 总预算

环境变量示例:

export LMCACHE_LOCAL_DISK="file:///mnt/nvme0/kvcache/,file:///mnt/nvme1/kvcache/"
export LMCACHE_LOCAL_DISK_PATH_SHARDING="by_gpu"
export LMCACHE_MAX_LOCAL_DISK_SIZE=20.0   # combined budget (GB)

YAML 示例:

local_disk: "/mnt/nvme0/kvcache/,/mnt/nvme1/kvcache/"
local_disk_path_sharding: "by_gpu"
max_local_disk_size: 20.0

备注

每个 GPU 工作节点仅使用其分配的路径,因此 O_DIRECT 对齐由该路径的文件系统块大小决定。不同设备可能具有不同的块大小而不会出现问题。

小技巧

如果您能够使用内核级 RAID 0(例如 mdadm --level=0),您将获得真正的块级条带化(即使是单个大文件也可以同时使用两个设备的带宽)。当您无法或不想重新配置块设备时,多路径功能最为有用——例如,当它们已经包含其他数据时。

本地存储说明:#

与 CPU RAM 卸载不同,磁盘卸载默认是 禁用 的(local_disk 设置为 None),最大本地磁盘大小设置为 0GB,而不是像默认最大本地 CPU 大小那样设置为 5GB,因为磁盘空间对于 LMCache 的功能并不是严格必要的。

此外,与固定的 CPU RAM 不同,磁盘后端不会贪婪地提前分配最大空间,而是会在存储时为每个 KV 缓存块创建一个文件,如果超出容量则逐出(当前使用 LRU)。

磁盘和远程后端(参见 RedisMooncakeValkeyInfiniStore)具有异步 put() 操作,因此 IO 延迟不会在阻塞 get() 操作的同时减慢推理。 本地磁盘后端还具有 prefetch() 操作,该操作将主动将 KV 缓存从磁盘移动到 CPU 内存卸载存储(即应设置 LMCACHE_LOCAL_CPU=True,参见 CPU RAM),以便为指定的令牌预取(这些 KV 缓存仍然保留在磁盘中)。

架构概述#

下图展示了本地磁盘后端的整体架构:

        %%{init: {'theme': 'base', 'themeVariables': { 'fontSize': '18px', 'fontFamily': 'arial', 'primaryColor': '#e3f2fd', 'primaryTextColor': '#000', 'primaryBorderColor': '#1976d2', 'lineColor': '#424242', 'secondaryColor': '#f5f5f5', 'tertiaryColor': '#ffffff', 'background': '#ffffff', 'clusterBkg': '#f8f9fa', 'clusterBorder': '#495057' }}}%%
flowchart TB
    subgraph Engine["<b>LMCache Engine</b>"]
        E["<b>Request Save/Load Operations</b>"]
    end

    subgraph LDB["<b>LocalDiskBackend</b>"]
        subgraph Meta["<b>Metadata Dictionary</b>"]
            Dict["<b>self.dict: CacheEngineKey → DiskCacheMetadata</b>
            (path, size, shape, dtype, pinned, positions)"]
        end

        subgraph Policy["<b>Cache Policy</b>"]
            CP["<b>Configurable Policy</b>
            (LRU, LFU, FIFO, MRU)
            Decides what to evict"]
        end

        subgraph Worker["<b>LocalDiskWorker</b>"]
            PQ["<b>Priority Queue Executor (4 workers)</b>"]
            P0["<b>Priority 0: PREFETCH</b>"]
            P1["<b>Priority 1: DELETE</b>"]
            P2["<b>Priority 2: PUT</b>"]
        end

        CPU["<b>LocalCPUBackend</b>
        (memory allocator)"]
    end

    subgraph Disk["<b>Local Filesystem</b>"]
        Files["/cache/vllm@model@[email protected]
        /cache/vllm@model@[email protected]
        /cache/vllm@model@[email protected]"]
    end

    E --> Dict
    Dict --> CP
    CP --> PQ
    PQ --> P0
    PQ --> P1
    PQ --> P2
    Worker --> Files
    CPU -.-> Worker
    

关键组件:

  • 元数据字典:将每个 CacheEngineKey 映射到其磁盘元数据(文件路径、大小、形状、数据类型、固定状态)

  • 缓存策略:可配置的逐出策略(LRU、LFU、FIFO 或 MRU),跟踪访问模式并决定在需要空间时逐出哪些条目

  • LocalDiskWorker: 带有优先级队列的异步任务执行器 - 预取任务优先运行(优先级 0),然后是删除(优先级 1),最后是保存(优先级 2)

  • 本地磁盘:KV Cache 块作为单独的 .pt 文件存储的文件系统

保存流程 (PUT)#

        %%{init: {'theme': 'base', 'flowchart': {'useMaxWidth': false, 'htmlLabels': true, 'nodeSpacing': 30, 'rankSpacing': 30}, 'themeVariables': { 'fontSize': '18px', 'fontFamily': 'arial', 'primaryColor': '#e3f2fd', 'primaryTextColor': '#000', 'primaryBorderColor': '#1976d2', 'lineColor': '#424242', 'secondaryColor': '#f5f5f5', 'tertiaryColor': '#ffffff', 'background': '#ffffff', 'clusterBkg': '#f8f9fa', 'clusterBorder': '#495057' }}}%%
flowchart LR
    A["<b>MemoryObj</b><br/>(KV cache in CPU memory)"] --> B{<b>Disk space<br/>available?</b>}
    B -->|"No"| C["<b>Evict via policy</b><br/>Delete .pt files"]
    C --> B
    B -->|"Yes"| D["<b>Track in put_tasks</b><br/>Queue async write<br/>(Priority 2 - lowest)"]
    D --> E["<b>LocalDiskWorker</b><br/>write_file()"]
    E --> F[("<b>Disk</b><br/>.pt file")]
    F --> G["<b>Add to metadata dict</b>"]

    style A fill:#e1f5fe
    style F fill:#c8e6c9
    style C fill:#ffcdd2
    

加载流程 (GET)#

        %%{init: {'theme': 'base', 'flowchart': {'useMaxWidth': false, 'htmlLabels': true, 'nodeSpacing': 30, 'rankSpacing': 30}, 'themeVariables': { 'fontSize': '18px', 'fontFamily': 'arial', 'primaryColor': '#e3f2fd', 'primaryTextColor': '#000', 'primaryBorderColor': '#1976d2', 'lineColor': '#424242', 'secondaryColor': '#f5f5f5', 'tertiaryColor': '#ffffff', 'background': '#ffffff', 'clusterBkg': '#f8f9fa', 'clusterBorder': '#495057' }}}%%
flowchart LR
    A["<b>Request</b><br/>(CacheEngineKey)"] --> B{<b>Key exists<br/>in dict?</b>}
    B -->|"No"| C["<b>Return None</b><br/>(cache miss)"]
    B -->|"Yes"| D["<b>Update policy</b><br/>Mark as accessed"]
    D --> E["<b>Allocate buffer</b><br/>via LocalCPUBackend"]
    E --> F["<b>Read from disk</b><br/>read_file()"]
    F --> G[("<b>Disk</b><br/>.pt file")]
    G --> F
    F --> H["<b>MemoryObj</b><br/>(KV cache ready)"]

    style A fill:#e1f5fe
    style H fill:#c8e6c9
    style C fill:#ffcdd2
    

在线推理示例#

这个示例与 CPU RAM 示例几乎相同。

让我们感受一下 TTFT(首次令牌时间)差异!

前提条件:

  • 一台至少配备一块 GPU 的机器。根据您的显存和希望使用的长上下文调整 vllm 实例的最大模型长度。

  • 安装了 vllm 和 LMCache (安装指南)

  • Hugging Face 访问 meta-llama/Meta-Llama-3.1-8B-Instruct

export HF_TOKEN=your_hugging_face_token
  • 一些软件包:

pip install openai transformers

步骤 0. 为此示例设置一个目录:

mkdir lmcache-local-disk-example
cd lmcache-local-disk-example

步骤 1. 准备一个长上下文!

我们希望上下文足够长,以至于 vLLM 的前缀缓存无法将 KV 缓存保留在显存中,因此需要 LMCache 将 KV 缓存保留在非显存中:

# 382757 bytes
man bash > man-bash.txt

步骤 2. 启动一个启用磁盘卸载的 vLLM 服务器:

通常不推荐这样做,但我们将禁用 CPU 卸载,以便仅感受磁盘卸载的延迟。

创建一个名为 disk-offload.yaml 的 lmcache 配置文件

示例 config.yaml

chunk_size: 256
local_cpu: false
max_local_cpu_size: 5.0
local_disk: "file:///local/disk_test/local_disk/"
max_local_disk_size: 5.0

如果您不想使用配置文件,请取消注释前五个环境变量,然后注释掉下面的 LMCACHE_CONFIG_FILE

# LMCACHE_CHUNK_SIZE=256 \
# LMCACHE_LOCAL_CPU=False \
# LMCACHE_MAX_LOCAL_CPU_SIZE=5.0 \
# LMCACHE_LOCAL_DISK="file:///local/disk_test/local_disk/" \
# LMCACHE_MAX_LOCAL_DISK_SIZE=5.0 \
LMCACHE_CONFIG_FILE="disk-offload.yaml" \
vllm serve \
    meta-llama/Llama-3.1-8B-Instruct \
    --max-model-len 16384 \
    --kv-transfer-config \
    '{"kv_connector":"LMCacheConnectorV1", "kv_role":"kv_both"}'
  • --kv-transfer-config: 这是实际告诉 vLLM 使用 LMCache 进行 KV Cache 卸载的参数。
    • kv_connector: 指定 vLLM V1 的 LMCache 连接器

    • kv_role: 设置为 "kv_both" 以同时存储和加载 KV Cache(重要,因为我们将运行两个查询,第一个将生成/存储一个 KV Cache,而第二个将消费/加载该 KV Cache)

步骤 3. 使用 LMCache 查询 TTFT 改进:

一旦 Open AI 兼容的服务器在默认的 vllm 端口 8000 上运行,让我们用相同的长上下文查询两次!

创建一个名为 query-twice.py 的脚本,并粘贴以下代码:

import time
from openai import OpenAI
from transformers import AutoTokenizer

client = OpenAI(
    api_key="dummy-key",  # required by OpenAI client even for local servers
    base_url="http://localhost:8000/v1"
)

models = client.models.list()
model = models.data[0].id

# 119512 characters total
# 26054 tokens total
long_context = ""
with open("man-bash.txt", "r") as f:
    long_context = f.read()

# a truncation of the long context for the --max-model-len 16384
# if you increase the --max-model-len, you can decrease the truncation i.e.
# use more of the long context
long_context = long_context[:70000]

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B-Instruct")
question = "Summarize bash in 2 sentences."

prompt = f"{long_context}\n\n{question}"

print(f"Number of tokens in prompt: {len(tokenizer.encode(prompt))}")

def query_and_measure_ttft():
    start = time.perf_counter()
    ttft = None

    chat_completion = client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model=model,
        temperature=0.7,
        stream=True,
    )

    for chunk in chat_completion:
        chunk_message = chunk.choices[0].delta.content
        if chunk_message is not None:
            if ttft is None:
                ttft = time.perf_counter()
            print(chunk_message, end="", flush=True)

    print("\n")  # New line after streaming
    return ttft - start

print("Querying vLLM server with cold LMCache Disk Offload")
cold_ttft = query_and_measure_ttft()
print(f"Cold TTFT: {cold_ttft:.3f} seconds")

print("\nQuerying vLLM server with warm LMCache Disk Offload")
warm_ttft = query_and_measure_ttft()
print(f"Warm TTFT: {warm_ttft:.3f} seconds")

print(f"\nTTFT Improvement: {(cold_ttft - warm_ttft):.3f} seconds \
    ({(cold_ttft/warm_ttft):.1f}x faster)")

然后运行:

python query-twice.py

由于我们处于流式模式,您将能够实时感受到 TTFT 差异!

请注意,如果我们启用 LMCACHE_LOCAL_CPU=True,我们将仅使用来自 CPU RAM 的相同示例,因为 LMCache 在检查磁盘之前会先检查 CPU RAM。实际上,磁盘能够存储更多的 KV 缓存,因此 CPU RAM 的卸载只能存储磁盘 KV 缓存的一个子集。

示例输出:

Number of tokens in prompt: 15376
Querying vLLM server with cold LMCache Disk Offload
Bash is a Unix shell and command-line interpreter that reads and executes
commands from standard input or a file, incorporating features from the
Korn and C shells. It is a conformant implementation of the IEEE POSIX
specification and can be configure to be POSIX-conformant by default,
supporting a wide range of options, built-in commands,
and features for scripting, job control, and interactive use.

Cold TTFT: 6.314 seconds

Querying vLLM server with warm LMCache Disk Offload
Bash is a Unix shell and command-line interpreter that reads and
executes commands from the standard input or a file, and is designed
to be a conformant implementation of the IEEE POSIX specification. It
is a powerful tool for automating tasks, managing files and directories,
and interacting with other programs and services, with features such as
scripting, conditional statements, loops, and functions.

Warm TTFT: 0.148 seconds

TTFT 改进:6.166 秒(快 42.6 倍)

如果你查看 vLLM 服务器的日志,你应该会看到(日志已为整洁而截断):

# Cold LMCache Miss and then Store

LMCache INFO: Reqid: chatcmpl-8676f9b9ebf04c79a5d47b9ada7b65fd, Total tokens 15410,
LMCache hit tokens: 0, need to load: 0

# you should see 8 of these storing logs total
# 2048 tokens is a multiple of the chunk size
LMCache INFO: Storing KV cache for 2048 out of 12288 tokens for request
chatcmpl-8676f9b9ebf04c79a5d47b9ada7b65fd

LMCache INFO: Storing KV cache for 2048 out of 14336 tokens for request
chatcmpl-8676f9b9ebf04c79a5d47b9ada7b65fd

LMCache INFO: Storing KV cache for 1074 out of 15410 tokens for request
chatcmpl-8676f9b9ebf04c79a5d47b9ada7b65fd

# Warm LMCache Hit!!

LMCache INFO: Reqid: chatcmpl-136d9dac1ba94bd4b4ae85007e8ad437, Total tokens 15410,
LMCache hit tokens: 15409, need to load: 1

查看您 SSD 中的 KV Cache:

ls "$HOME/local/disk_test/local_disk/"

提示:#

  • 如果您想多次运行 query-twice.py 脚本,您需要重启 vLLM LMCache 服务器或更改您传入的上下文前缀,因为您已经预热了 LMCache。

  • 这里的最大模型长度是通过仅使用 23GB 显存的 L4 运行决定的。如果您有更多内存,可以增加最大模型长度并修改 query-twice.py 以使用更多的长上下文。随着上下文长度的增加,LMCache 的 TTFT 改进变得更加明显!