添加新的设备后端#

本指南解释了如何在 多进程 (MP) 模式 下向 LMCache 添加一个 新加速器。

For basic users: read Part 1 only — adding a DeviceSpec and a DeviceOps subclass is all you need to get your device working with the built-in torch baseline ops.

对于高级用户： 继续阅读 Part 2 以了解原生操作和高级传输模式。

第 1 部分 — 基本功能启用#

For the majority of devices, a DeviceSpec class plus a minimal DeviceOps subclass are sufficient. LMCache ships with a complete torch baseline ops layer (lmcache/v1/platform/torch_ops.py) that works on any device supporting standard PyTorch tensor operations — no custom kernels required.

前提条件#

您的 PyTorch 后端应支持：

设备发现与状态：

  • torch.<device>.is_available() → bool

  • torch.<device>.device_count() → int

设备上下文与同步：

  • torch.<device>.set_device(device) → None

  • torch.<device>.current_device() → int

  • torch.<device>.synchronize() → None

数据移动:

  • tensor.to(device) / tensor.cpu() （主机↔设备传输）

Step 1: Add a FooDeviceSpec#

创建 lmcache/v1/platform/foo/__init__.py：

# SPDX-License-Identifier: Apache-2.0
"""Foo-specific platform primitives."""

from __future__ import annotations

from typing import TYPE_CHECKING

from lmcache.v1.platform.base.device_spec import DeviceSpec

if TYPE_CHECKING:
    from lmcache.v1.platform.base.device_ops import DeviceOps


class FooDeviceSpec(DeviceSpec):
    """Foo device specification for LMCache registry discovery."""

    @property
    def device_type(self) -> str:
        return "foo"

    @property
    def torch_module_name(self) -> str:
        return "foo"

    @property
    def ops_cls(self) -> type[DeviceOps]:
        from lmcache.v1.platform.foo.device_ops import FooDeviceOps

        return FooDeviceOps

    def is_available(self) -> bool:
        """Check backend availability without importing lmcache.__init__."""
        try:
            import torch

            return hasattr(torch, "foo") and torch.foo.is_available()
        except Exception:
            return False

Step 2: Add a FooDeviceOps#

Create lmcache/v1/platform/foo/device_ops.py:

# SPDX-License-Identifier: Apache-2.0
"""Foo ops backend."""

from __future__ import annotations

from typing import ClassVar

from lmcache.logging import init_logger
from lmcache.v1.platform.base.device_ops import DeviceOps

logger = init_logger(__name__)


class FooDeviceOps(DeviceOps):
    device_type: ClassVar[str] = "foo"

    def ensure_native(self) -> None:
        if self._native_bound:
            return
        self._native_bound = True
        try:
            import foo_device.ops as native
        except ImportError:
            logger.warning(
                "foo native ops not found; FooDeviceOps stays on "
                "the torch baseline for all ops."
            )
            return
        self.bind_native(native)

备注

If you have no native extension yet, simply leave ensure_native as a no-op (just pass). The torch baseline handles everything.

关键属性：

Property / Method

必需的

目的

DeviceSpec.device_type

是

Device type string (e.g. "cuda", "musa", "xpu")

DeviceSpec.torch_module_name

是

torch 包上的属性（例如 "cuda" → torch.cuda）

DeviceSpec.ops_cls

不

Returns the DeviceOps subclass for this device. Base returns DeviceOps itself (pure torch baseline).

DeviceSpec.is_available()

是

当设备可用时返回 True

DeviceOps.ensure_native()

不

Called once on first use; override to bind native ops. Base is a no-op (e.g. HpuDeviceOps inherits it unchanged).

备注

上面显示的 hasattr(torch, "foo") 保护仅在使用树外 PyTorch 扩展时需要（例如，当作为插件安装的 torch.musa、torch.xpu）。对于 PyTorch 本身提供的加速器（如 torch.cuda），只需使用简单的 torch.foo.is_available() 即可。

That's it. Defining these two classes is enough for auto-discovery — no global list or manual registration call is required. All ops automatically route through the torch baseline in torch_ops.py, which is not performant but is functionally applicable to any device that supports the standard PyTorch tensor surface.

验证#

启动 LMCache 服务器：

lmcache server --l1-size-gb 10 --eviction-policy LRU --port 5555

使用 MP 连接器运行 vLLM。如果主机上可见多个加速器，请设置 DEVICE_TYPE 强制 LMCache 选择新的后端，而不是自动检测：

export DEVICE_TYPE=foo            # optional; only when auto-detection picks the wrong device

vllm serve <your-model> \
    --kv-transfer-config '{
        "kv_connector": "LMCacheMPConnector",
        "kv_connector_module_path": "lmcache.integration.vllm.lmcache_mp_connector",
        "kv_role": "kv_both",
        "kv_connector_extra_config": {
            "lmcache.mp.host": "tcp://localhost",
            "lmcache.mp.port": "5555",
            "lmcache.mp.mp_transfer_mode": "engine_driven"
        }
    }' \
    --no-enable-prefix-caching \
    --port 8000

备注

默认传输模式为 auto``（CUDA → LMCache 驱动，其他设备 → 引擎驱动）。上述示例明确设置为 ``engine_driven，以便新的非 CUDA 设备在没有额外能力检查的情况下正常工作。有关 lmcache_driven 模式（IPC 零拷贝），请参见 Advanced transfer mode。

Check the LMCache logs — without a native extension you should see:

torch_dev=..., torch_device_type=foo
foo native ops not found; FooDeviceOps stays on the torch baseline for all ops.

This confirms your DeviceSpec was discovered and the torch baseline is active. Once you provide a native extension and ensure_native succeeds, the warning disappears and ops call through your compiled kernels.

调试检查清单：

  • [ ] torch.foo.is_available() 返回 True。

  • [ ] 设置 DEVICE_TYPE=foo 以强制选择，如果未自动选择。

  • [ ] Log shows either the "stays on torch baseline" warning (no native ops) or silent success (native ops bound).

  • [ ] 引擎驱动的传输端到端工作正常（检查 LMCache 日志以确认选择了 SHM 还是 Pickle 子路径 — 两者都应该成功）。

  • [ ] 存储/检索的正确性已验证。

  • [ ] TP>1 / 多工作进程行为已验证。

第二部分 — 性能优化#

一旦基本功能得到验证，添加设备特定的优化。

Device-specific ops via bind_native#

The DeviceOps base class delegates every op to the torch baseline in lmcache/v1/platform/torch_ops.py. Each vendor may replace any subset of these functions with a device-specific implementation.

How it works. When ensure_native() calls self.bind_native(native_module), the method walks the module's public symbols and rebinds them as instance attributes — overriding the base-class methods that delegate to torch_ops:

callers  →  lmcache.c_ops (shim)  →  DeviceOps instance
                                       │
                                 bind_native() overlay:
                                   ├── native.multi_layer_kv_transfer  ← vendor CUDA/SYCL kernel
                                   ├── native.calculate_cdf            ← vendor kernel
                                   └── (everything else)               ← torch_ops baseline

集成合同#

无论您如何构建操作模块，以下契约必须成立：

  • Same symbol names. Every function you override must be exposed under the exact name used in torch_ops.py (e.g. multi_layer_block_kv_transfer).

  • Same call signature. Positional/keyword arguments, argument order and semantics must match the baseline; callers invoke the DeviceOps instance without knowing which backend answered.

  • Importable Python module. Your native module must be importable via import (how it gets there — a pure-Python file, a pybind11 extension, a ctypes wrapper, a Rust PyO3 module, etc. — is your choice).

  • Partial override is allowed. You do not have to reimplement every function. Anything you leave out keeps using the torch baseline, so incremental optimization is supported.

  • Types from the native module are also bound. bind_native also binds types (classes) — this is how native plan types like StagingCopy, KernelGroupSpec, etc. overlay the stubs in ops_types.py.

实现说明#

  • multi_layer_block_kv_transfer and lmcache_memcpy_async are the hot entry points for both engine-driven and LMCache-driven transfer; other functions in torch_ops.py can be overridden as needed.

  • If you fall back to the generic path from inside a device-specific wrapper (e.g. when inputs are unsupported), call the corresponding lmcache.v1.platform.torch_ops function directly to preserve semantics.

  • ensure_native() is called once when DeviceSpec.get_ops() first creates the singleton. It is safe to fail soft (log a warning and return without binding).

For concrete reference implementations, see:

  • lmcache/v1/platform/cuda/device_ops.py — binds the compiled lmcache.c_ops pybind11 extension.

  • lmcache/v1/platform/xpu/device_ops.py — binds the SYCL lmcache.xpu_ops extension.

  • lmcache/v1/platform/musa/device_ops.py — method overrides without a separate native module.

高级传输模式#

默认情况下，传输模式为 AUTO：路由器严格按照 device_type 分发 — device_type == \"cuda\" 的设备使用 LMCacheDrivenTransferContext``（IPC 零拷贝），其他设备则使用 ``EngineDrivenTransferContext。支持 IPC 句柄传输的非 CUDA 设备仍然可以显式选择 LMCache 驱动（如下所述）。

备注

ROCm 不是一个单独的后端：在 PyTorch 下，ROCm GPU 报告 device_type == \"cuda\" 并重用 CudaIPCWrapper``（CUDA IPC）用于 LMCache 驱动的路径，因此它在 AUTO 模式下工作，无需额外设置。目前没有专门的 ``platform/rocm 包。

When the caller (or LMCACHE_MP_TRANSFER_MODE) explicitly requests lmcache_driven, _build_lmcache_driven_context performs two hard checks — both must succeed, otherwise the factory raises ValueError (no silent degradation):

  1. Your DeviceSpec subclass must bind a DeviceIPCWrapper subclass (exposing a wrap classmethod) via ipc_wrapper_cls. resolve_kv_wrapper_factory() reads that binding off the registered spec — no separate registry / auto-scan.

  2. DeviceSpec.is_handle_transfer_available() 必须返回 True``（基类默认值；仅当您的设备缺少 IPC 句柄传输时才重写为 ``False）。

Separately, the LMCache-driven server module also requires a BaseCacheContext subclass under lmcache/v1/platform/foo/cache_context.py and a matching DeviceSpec.create_cache_context override that lazy-imports and instantiates it. The platform-agnostic factory lmcache.v1.platform.cache_context.create_cache_context dispatches by device_type through the DeviceSpec registry and invokes that hook; the default DeviceSpec.create_cache_context raises NotImplementedError so a missing override surfaces loudly instead of silently falling back. The cache context itself manages the KV cache layout and pointers used for IPC transfer.

通过 pin_memory_backend 在主机端固定是 可选 的，仅影响暂存缓冲区的性能；启用 LMCache 驱动模式并不需要它。

Event IPC capability#

The LMCache-driven multiprocess handle path also requires a platform event-IPC backend. The capability is declared by DeviceSpec.event_ipc_backend and is intentionally separate from DeviceOps and DeviceIPCWrapper:

  • The base DeviceSpec returns None. A concrete device must explicitly opt in.

  • CUDA-style event APIs can use DefaultEventIPCBackend(event_module=..., device_type=...).

  • Devices with a different event ABI should implement an EventIPCBackend in their own lmcache/v1/platform/foo/ package.

  • Concrete device specs should cache the backend because request futures may query this property repeatedly.

The backend contract covers event creation, handle export/import, event recording, stream wait, query, and synchronization. check_event_support(device) must raise RuntimeError when those operations are unavailable for the requested device. The default backend checks for a CUDA-style Event type that accepts interprocess=True and provides from_ipc_handle. A custom backend should instead validate the equivalent prerequisites for its own event ABI.

For a CUDA-style device, bind and cache the default backend as follows:

from typing import TYPE_CHECKING

from lmcache.v1.platform.base.device_spec import DeviceSpec

if TYPE_CHECKING:
    from lmcache.v1.platform.base.event_ipc import EventIPCBackend

class FooDeviceSpec(DeviceSpec):
    _event_backend_cache: "EventIPCBackend | None" = None

    @property
    def event_ipc_backend(self) -> "EventIPCBackend":
        backend = self._event_backend_cache
        if backend is None:
            import torch

            from lmcache.v1.platform.base.event_ipc import (
                DefaultEventIPCBackend,
            )

            backend = DefaultEventIPCBackend(
                event_module=torch.foo,
                device_type=self.device_type,
            )
            self._event_backend_cache = backend
        return backend

Event IPC operations must preserve producer/consumer stream ordering without adding a device-wide synchronization. query_event must remain non-blocking so request futures can poll completion safely.

备注

If the device does not support Event IPC, leave the base event_ipc_backend implementation unchanged so it returns None. Engine-driven mode does not require this capability. LMCache-driven mode raises an explicit error rather than falling back to CUDA or to the accelerator active in the process.

The capability is checked during worker/server registration and before constructing a device-aware completion future. This keeps unsupported platforms from entering an asynchronous transfer path that cannot order KV-cache memory safely. The STORE and RETRIEVE message wire format is unchanged: the existing event handle bytes are still carried in the request and response payloads.

在你的 DeviceSpec 中重写这些方法：

class FooDeviceSpec(DeviceSpec):
    @property
    def ipc_wrapper_cls(self):
        """Bind the DeviceIPCWrapper subclass for this device.

        Lazy import so the accelerator-specific module is only
        pulled in when the LMCache-driven path is actually used.
        """
        from lmcache.v1.platform.foo.ipc_wrapper import FooIPCWrapper

        return FooIPCWrapper

    def is_handle_transfer_available(self) -> bool:
        """Return True if your device supports IPC handle transfer."""
        return True  # base-class default; override to False if unsupported

    @property
    def pin_memory_backend(self):
        """Return a PinMemoryBackend subclass, or None.

        Optional; only affects host staging performance.
        """
        return None  # default

    def create_cache_context(self, *args, **kwargs):
        """Lazy-import and instantiate the BaseCacheContext for this device.

        Required for LMCache-driven mode; the base-class default
        raises ``NotImplementedError``.
        """
        from lmcache.v1.platform.foo.cache_context import FooCacheContext

        return FooCacheContext(*args, **kwargs)

通过在 Part 1 中显示的 vLLM kv_connector_extra_config 中将 lmcache.mp.mp_transfer_mode 设置为 lmcache_driven，或通过导出 LMCACHE_MP_TRANSFER_MODE=lmcache_driven 来选择 LMCache 驱动模式。如果任何硬检查失败，工厂将引发 ValueError 并拒绝构造上下文——请切换回 engine_driven 或 auto。

参考文献#

主题

路径

设备规格基础

lmcache/v1/platform/base/device_spec.py

Device ops base

lmcache/v1/platform/base/device_ops.py

Torch ops baseline

lmcache/v1/platform/torch_ops.py

Ops types and enums

lmcache/v1/platform/ops_types.py

Event IPC base

lmcache/v1/platform/base/event_ipc.py

Backend loading (resolve_device_ops / _detect_device)

lmcache/v1/platform/__init__.py

c_ops shim install (_install_c_ops_shim)

lmcache/__init__.py

缓存上下文基类

lmcache/v1/platform/base/cache_context.py

缓存上下文工厂

lmcache/v1/platform/cache_context.py

Reference DeviceSpec (CUDA)

lmcache/v1/platform/cuda/__init__.py

Reference DeviceOps (CUDA, bind_native)

lmcache/v1/platform/cuda/device_ops.py

Reference DeviceOps (XPU, bind_native)

lmcache/v1/platform/xpu/device_ops.py

Reference DeviceOps (MUSA, method overrides)

lmcache/v1/platform/musa/device_ops.py

引擎驱动的调用站点

lmcache/v1/multiprocess/transfer_context/worker_transfer.py (EngineDrivenTransferContext, create_transfer_context)