KV 缓存卸载#

通过卸载到 CPU 内存或本地磁盘来扩展 KV 缓存容量,以获得更大的批量大小和减少 GPU 内存压力。

注意

Ray Serve 不提供开箱即用的 KV 缓存卸载,但可与 vLLM 解决方案无缝集成。本指南演示了一种此类集成:LMCache。

KV 缓存卸载的好处

  • 容量增加:通过使用 CPU RAM 或本地存储来存储更多 KV 缓存,而不是仅仅依赖 GPU 内存

  • 请求间的缓存重用:保存并重用先前计算的 KV 缓存,用于重复或相似的提示,减少预填充计算

  • 灵活的存储后端:从多种存储选项中选择,包括本地 CPU、磁盘或分布式系统

当您的应用程序有重复提示或多轮对话,并且您可以重用缓存的预填充时,请考虑 KV 缓存卸载。如果连续的对话查询不是立即发送,GPU 会将这些缓存驱逐出去,为其他并发请求腾出空间,导致缓存未命中。将 KV 缓存卸载到 CPU 内存或其他具有更大容量的存储后端,可以更长时间地保留它们。

使用 LMCache 部署#

LMCache 提供支持多种存储后端的 KV 缓存卸载。

先决条件#

安装 LMCache

uv pip install lmcache

基本部署#

以下示例显示了如何使用 LMCache 进行本地 CPU 卸载进行部署

from ray.serve.llm import LLMConfig, build_openai_app
import ray.serve as serve

llm_config = LLMConfig(
    model_loading_config={
        "model_id": "qwen-0.5b",
        "model_source": "Qwen/Qwen2-0.5B-Instruct"
    },
    engine_kwargs={
        "tensor_parallel_size": 1,
        "kv_transfer_config": {
            "kv_connector": "LMCacheConnectorV1",
            "kv_role": "kv_both",
        }
    },
    runtime_env={
        "env_vars": {
            "LMCACHE_LOCAL_CPU": "True",
            "LMCACHE_CHUNK_SIZE": "256",
            "LMCACHE_MAX_LOCAL_CPU_SIZE": "100",  # 100GB
        }
    }
)

app = build_openai_app({"llm_configs": [llm_config]})
serve.run(app)
applications:
  - name: llm-with-lmcache
    route_prefix: /
    import_path: ray.serve.llm:build_openai_app
    runtime_env:
      env_vars:
        LMCACHE_LOCAL_CPU: "True"
        LMCACHE_CHUNK_SIZE: "256"
        LMCACHE_MAX_LOCAL_CPU_SIZE: "100"
    args:
      llm_configs:
        - model_loading_config:
            model_id: qwen-0.5b
            model_source: Qwen/Qwen2-0.5B-Instruct
          engine_kwargs:
            tensor_parallel_size: 1
            kv_transfer_config:
              kv_connector: LMCacheConnectorV1
              kv_role: kv_both

使用以下命令部署

serve run config.yaml

使用 MultiConnector 组合多个 KV 传输后端#

您可以使用 MultiConnector 组合多个 KV 传输后端。当您希望在分离式部署中进行本地卸载和跨实例传输时,这很有用。

何时使用 MultiConnector#

当您使用预填充/解码分离并且希望同时进行跨实例传输 (NIXL) 和本地卸载时,请使用 MultiConnector 来组合多个后端。

以下示例演示了如何在预填充/解码部署中将 NIXL(用于跨实例传输)与 LMCache(用于本地卸载)结合使用

注意

连接器的顺序很重要。由于您希望优先通过 LMCache 进行本地 KV 缓存查找,因此它会出现在列表中的 NIXL 连接器之前。

from ray.serve.llm import LLMConfig, build_pd_openai_app
import ray.serve as serve

# Shared KV transfer config combining NIXL and LMCache
kv_config = {
    "kv_connector": "MultiConnector",
    "kv_role": "kv_both",
    "kv_connector_extra_config": {
        "connectors": [
            {
                "kv_connector": "LMCacheConnectorV1",
                "kv_role": "kv_both",
            },
            {
                "kv_connector": "NixlConnector",
                "kv_role": "kv_both",
                "backends": ["UCX"],
            }
        ]
    }
}

prefill_config = LLMConfig(
    model_loading_config={
        "model_id": "qwen-0.5b",
        "model_source": "Qwen/Qwen2-0.5B-Instruct"
    },
    engine_kwargs={
        "tensor_parallel_size": 1,
        "kv_transfer_config": kv_config,
    },
    runtime_env={
        "env_vars": {
            "LMCACHE_LOCAL_CPU": "True",
            "LMCACHE_CHUNK_SIZE": "256",
            "UCX_TLS": "all",
        }
    }
)

decode_config = LLMConfig(
    model_loading_config={
        "model_id": "qwen-0.5b",
        "model_source": "Qwen/Qwen2-0.5B-Instruct"
    },
    engine_kwargs={
        "tensor_parallel_size": 1,
        "kv_transfer_config": kv_config,
    },
    runtime_env={
        "env_vars": {
            "LMCACHE_LOCAL_CPU": "True",
            "LMCACHE_CHUNK_SIZE": "256",
            "UCX_TLS": "all",
        }
    }
)

pd_config = {
    "prefill_config": prefill_config,
    "decode_config": decode_config,
}

app = build_pd_openai_app(pd_config)
serve.run(app)
applications:
  - name: pd-multiconnector
    route_prefix: /
    import_path: ray.serve.llm:build_pd_openai_app
    runtime_env:
      env_vars:
        LMCACHE_LOCAL_CPU: "True"
        LMCACHE_CHUNK_SIZE: "256"
        UCX_TLS: "all"
    args:
      prefill_config:
        model_loading_config:
          model_id: qwen-0.5b
          model_source: Qwen/Qwen2-0.5B-Instruct
        engine_kwargs:
          tensor_parallel_size: 1
          kv_transfer_config:
            kv_connector: MultiConnector
            kv_role: kv_both
            kv_connector_extra_config:
              connectors:
                - kv_connector: LMCacheConnectorV1
                  kv_role: kv_both
                - kv_connector: NixlConnector
                  kv_role: kv_both
                  backends: ["UCX"]
      decode_config:
        model_loading_config:
          model_id: qwen-0.5b
          model_source: Qwen/Qwen2-0.5B-Instruct
        engine_kwargs:
          tensor_parallel_size: 1
          kv_transfer_config:
            kv_connector: MultiConnector
            kv_role: kv_both
            kv_connector_extra_config:
              connectors:
                - kv_connector: LMCacheConnectorV1
                  kv_role: kv_both
                - kv_connector: NixlConnector
                  kv_role: kv_both
                  backends: ["UCX"]

使用以下命令部署

serve run config.yaml

配置参数#

LMCache 环境变量#

  • LMCACHE_LOCAL_CPU:设置为 "True" 以启用本地 CPU 卸载

  • LMCACHE_CHUNK_SIZE:KV 缓存块的大小,以 token 为单位(默认值:256)

  • LMCACHE_MAX_LOCAL_CPU_SIZE:本地 CPU 最大存储大小(GB)

  • LMCACHE_PD_BUFFER_DEVICE:预填充/解码场景的缓冲区设备(默认值:“cpu”)

有关 LMCache 配置选项的完整列表,请参阅 LMCache 配置参考

MultiConnector 配置#

  • kv_connector:设置为 "MultiConnector" 以组合多个后端

  • kv_connector_extra_config.connectors:要组合的连接器配置列表。顺序很重要 — 列表中靠前的连接器具有更高的优先级。

  • 列表中的每个连接器都使用与独立连接器相同的配置格式

性能注意事项#

将 KV 缓存扩展到本地 GPU 内存之外会增加跨不同内存层次结构管理和查找缓存的开销。这会带来权衡:您会获得更大的缓存容量,但延迟可能会增加。请考虑这些因素

缓存未命中场景中的开销:当没有缓存命中时,卸载会增加适度的开销(约 10-15%),与纯 GPU 缓存相比,这是基于我们内部实验的。此开销来自额外的哈希、数据移动和管理操作。

缓存命中时的好处:当缓存可以重用时,卸载会显著减少预填充计算。例如,在多轮对话中,用户在几分钟的非活动时间后返回,LMCache 会从 CPU 检索对话历史记录,而不是重新计算它,从而显著减少后续请求的首次 token 时间。

网络传输成本:当将 MultiConnector 与跨实例传输(例如 NIXL)结合使用时,请确保分离的好处大于网络传输成本。

另请参阅#