在 LMCache 多进程模式下接入新加速器DeviceSpec 与 DeviceOps 插件化开发完整指南【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCache本指南以 LMCache 官方文档 adding_a_new_device_backend.rst 为主体骨架讲解如何在MultiprocessMP模式下为 LMCache 接入新的加速器Device Backend。读完本文你将掌握两条完整的技术路径只用一个DeviceSpec 最小DeviceOps子类即可让设备跑通 LMCache 的 torch 基线算子Part 1以及在功能验证通过后通过bind_native、Event IPC 与DeviceIPCWrapper等机制把传输链路升级为高性能的 LMCache-driven 零拷贝模式Part 2。为什么需要设备后端抽象LMCache 的 KV Cache 层需要把张量从推理引擎搬进搬出缓存。不同加速器CUDA、XPU、MUSA、Neuron 乃至未来的新设备的显存管理、IPC 句柄、事件同步机制完全不同如果把这些差异写进核心逻辑核心代码将被if device ...的分支淹没。LMCache 的解法是在 lmcache/v1/platform/ 下建立一套跨平台抽象层DeviceSpec描述这台设备是什么、如何检测、该用哪套算子是设备在 LMCache 注册表中的身份证DeviceOps统一每设备算子入口每个方法默认委托给 torch_ops.py 的纯 torch 基线实现平台无关的工厂与分发逻辑缓存上下文、IPC 包装、事件后端全部基于DeviceSpec注册表数据驱动新增设备不需要改任何分发代码。从源码看平台注册表的构建逻辑位于 _device_detect.py内置子包通过pkgutil/子类发现机制扫描外部插件通过 Python entry-point 组lmcache.device_plugins加载二者统一汇入_build_backend_registry()再按device_type分组建出_build_device_registry()。运行时检测_detect_device()同文件 L229返回(torch_device_module, device_type, backend_name)最终在 lmcache/init.py 中解析出包级单例lmcache.device_ops供全局调用。两种维护模型in-tree 集成 与 外部 wheel 插件LMCache 为设备厂商提供两种所有权模型二者实现的是同一套DeviceSpec/DeviceOps接口、走同一条运行时检测与分发路径模型代码位置适合场景In-tree 集成lmcache/v1/platform/device/目录下后端需要随 LMCache 一起发布、测试与发版外部 wheel 插件厂商独立仓库 lmcache.device_pluginsentry-point 组需要独立发布节奏或要自带原生包分发关键点两种模型都不需要维护任何全局设备名列表。所有算子默认路由到torch_ops.py基线直到该后端主动覆盖。方案 A集成进 LMCache 仓库在 LMCache 仓库内直接创建平台子包lmcache/v1/platform/foo/ ├── __init__.py └── device_ops.pyFooDeviceSpec定义在__init__.py中。LMCache 会自动扫描内置平台子包因此无需 entry point 或注册列表。之后把代码、测试与设备文档一起提交 LMCache PR合并后后端即跟随 LMCache 的发布生命周期。方案 B维护外部 wheel使用src布局让 wheel 只持有厂商自己的命名空间lmcache-foo-device/ ├── pyproject.toml └── src/ └── lmcache_foo/ ├── __init__.py ├── device.py └── device_ops.pypyproject.toml必须声明一个位于lmcache.device_plugins组的 entry point其name 是DeviceSpec.backend_name的小写形式value 指向DeviceSpec类本身# pyproject.toml [build-system] requires [setuptools77, wheel] build-backend setuptools.build_meta [project] name lmcache-foo-device version 0.1.0 dependencies [lmcache] [project.entry-points.lmcache.device_plugins] foo lmcache_foo.device:FooDeviceSpec [tool.setuptools.packages.find] where [src]发布前应把lmcache依赖固定到插件实际测试过的版本范围。LMCache 将DeviceSpec/DeviceOps视为插件接口接口不兼容的改动会被该依赖范围拦截。从 _device_detect.py 的_load_external_device_specs()可以看出加载器的校验逻辑entry point 必须解析为DeviceSpec的子类不能是DeviceSpec本身、构造出的实例必须校验device_type/backend_name为非空小写字符串、且entry-point 的 name 必须与backend_name严格一致任何导入失败、类型错误或构造异常都只会记 warning 后跳过不影响其他设备使用。Part 1基础功能使能最小可运行实现对绝大多数设备而言一个DeviceSpec类加一个最小DeviceOps子类就足够了。LMCache 内置了完整的 torch 基线算子层torch_ops.py只要设备支持标准 PyTorch 张量运算就无需任何自定义 kernel。前置条件PyTorch 后端需要支持的能力能力类别必须支持的 API设备发现与状态torch.device.is_available()→booltorch.device.device_count()→int设备上下文与同步torch.device.set_device(device)→Nonetorch.device.current_device()→inttorch.device.synchronize()→None数据搬运tensor.to(device)/tensor.cpu()host↔device 传输Step 1选择所有权模型见上节Step 2实现FooDeviceSpecIn-tree 位置是lmcache/v1/platform/foo/__init__.py外部 wheel 位置是src/lmcache_foo/device.py。两种布局的实现完全相同# SPDX-License-Identifier: Apache-2.0 LMCache platform registration for Foo devices. 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 backend_name(self) - str: return foo property def torch_module_name(self) - str: return foo property def ops_cls(self) - type[DeviceOps]: from .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三个必须实现的属性含义如下DeviceSpec 基类 中的默认值分别是空字符串和False裸DeviceSpec()只会被当作 CPU/未知设备的安全兜底device_type面向 torch 的设备类别字符串如cuda、musa、xpu。多个后端可以合法共享同一个device_typebackend_nameLMCache 专属的选择器唯一标识一个具体后端实现必须是唯一、非空的小写字符串。基类默认复用device_type共享设备类型的专门后端应覆盖为不同的名字torch_module_nametorch包上的属性名如cuda对应torch.cuda。补充约束外部 wheel 的 entry-point 目标必须是类本身不能是实例或工厂函数且类必须有无参构造函数——LMCache 每个进程只实例化一次并缓存见 device_spec.py 的get_ops()首次访问时实例化、调用ensure_native()、存入_ops_cache单例。关于hasattr(torch, foo)守卫它只对树外 PyTorch 扩展如以插件形式安装的torch.musa、torch.xpu是必要的对 PyTorch 内置的加速器如torch.cuda直接写torch.foo.is_available()即可。可对照参考 CudaDeviceSpec.is_available() 的实现它还会额外检查torch.version.cuda是否存在。Step 3实现FooDeviceOps创建device_ops.pyin-tree 放在 spec 包旁边外部 wheel 放在src/lmcache_foo/下# SPDX-License-Identifier: Apache-2.0 Foo ops backend. from __future__ import annotations from typing import ClassVar from lmcache.v1.platform.base.device_ops import DeviceOps class FooDeviceOps(DeviceOps): device_type: ClassVar[str] foo def ensure_native(self) - None: Keep the torch baseline until native ops are available. return None如果还没有原生扩展保持ensure_native为空实现no-op即可——torch 基线会处理一切见 DeviceOps.ensure_native基类本身就是一个纯 torch 的空操作。等到后端开始携带原生代码时再参考 Part 2。核心接口速查表属性 / 方法是否必填用途DeviceSpec.device_type是设备类型字符串如cuda、musa、xpuDeviceSpec.backend_name是唯一标识一个具体后端实现的 LMCache 选择器DeviceSpec.torch_module_name是torch包上的属性名如cuda→torch.cudaDeviceSpec.ops_cls否返回该设备的DeviceOps子类基类返回DeviceOps本身纯 torch 基线DeviceSpec.is_available()是设备可用时返回TrueDeviceOps.ensure_native()否首次使用时调用一次覆盖它来绑定原生算子。基类为空操作如HpuDeviceOps直接继承不改写ops_cls的懒加载是刻意设计DeviceSpec被解析时不会把 torch 基线或原生.so拖进平台包的导入图见 device_spec.py 的注释说明。Step 4安装所选集成In-tree 后端从包含平台包的分支正常构建/安装 LMCache。合并发布后用户随 LMCache 一起获得该后端。外部后端用任意 PEP 517 前端构建厂商项目然后在与 LMCache 及推理引擎同一个 Python 环境中安装其 wheelpython -m pip install build python -m build --wheel python -m pip install dist/lmcache_foo_device-0.1.0-py3-none-any.whl安装任一集成后必须重启所有 LMCache 与推理引擎进程——平台注册表在进程生命周期内只构建并缓存一次_build_backend_registry使用lru_cache见 _device_detect.py。外部 wheel 的加载规则DeviceSpec.device_type是面向 torch 的设备类别如cuda多个后端可以有意识地共享同一个device_typeDeviceSpec.backend_name是 LMCache 专属的具体实现选择器必须是唯一、非空的小写字符串外部 wheel 的entry-point 名必须与DeviceSpec.backend_name匹配重复的backend_name在第一次确定性匹配后被忽略无法导入、解析到错误对象类型或构造时抛异常的插件会被记录日志并跳过其他设备不受影响entry-point 模块在平台注册表初始化时被导入。应直接导入lmcache.v1.platform.base.device_spec.DeviceSpec等基础接口把算子、原生库、IPC 包装和缓存上下文全部放到懒属性后面已安装的 entry-point 包是可执行代码只安装可信来源的 wheel。验证端到端跑通启动 LMCache server。下面示例中的 worker 走的是engine-driven 传输路径server 只有在--supported-transfer-mode为engine_driven或auto时才加载该路径默认值是lmcache_driven请注意显式指定lmcache server --l1-size-gb 10 --eviction-policy LRU --port 5555 \ --supported-transfer-mode engine_driven用 MP connector 运行 vLLM。如需指定具体的 torch 设备类别设置DEVICE_TYPE。共享同一设备类型的后端在恰好一个上报可用时会被自动选中若多个后端都上报可用则需要用LMCACHE_DEVICE_BACKEND精确选择实现export DEVICE_TYPEfoo # 可选选择 torch 面向的设备类型 export LMCACHE_DEVICE_BACKENDfoo # 可选多个可用后端并存时消歧 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-driven、其他设备 → engine-driven。上面示例显式设置engine_driven是为了让新接入的非 CUDA 设备无需额外能力检查即可工作。若要走lmcache_drivenIPC 零拷贝参见 Part 2。检查 LMCache 日志torch_dev..., torch_device_typefoo这确认你的DeviceSpec已被发现、torch 基线已生效。当设置了LMCACHE_DEVICE_BACKEND时LMCache 还会绑定backend_name匹配的后端这条日志来自 _device_detect.py 的get_torch_device()logger.info(torch_dev%s, torch_device_type%s, ...)。调试清单torch.foo.is_available()返回True外部 wheel 场景下importlib.metadata.entry_points(grouplmcache.device_plugins)包含foo且指向FooDeviceSpec若有其他可用后端共享device_typefoo设置LMCACHE_DEVICE_BACKENDfoo消歧若DEVICE_TYPE未被自动识别设置DEVICE_TYPEfoo强制指定面向 torch 的设备类别engine-driven 传输端到端可用检查 LMCache 日志确认走的是 SHM 还是 Pickle 子路径——两者都应成功存取store/retrieve正确性已验证TP1 / 多 worker 行为已验证。Part 2性能优化与高级传输模式基础功能验证通过后就可以针对设备做性能优化了。通过bind_native绑定设备专属算子DeviceOps基类把每个算子都委托给 torch_ops.py 的 torch 基线例如 multi_layer_block_kv_transfer 等所有方法都只是转发。厂商可以替换其中任意子集为设备专属实现。工作原理当ensure_native()调用self.bind_native(native_module)时该方法遍历模块的公开符号并作为实例属性重新绑定——覆盖掉那些委托给torch_ops的基类方法callers → lmcache.device_ops (DeviceOps 实例) │ bind_native() 覆盖层: ├── native.multi_layer_kv_transfer ← 厂商 CUDA/SYCL kernel ├── native.calculate_cdf ← 厂商 kernel └── (其他所有) ← torch_ops 基线从实现看device_ops.pybind_native会跳过下划线开头的私有符号把模块中每个type绑定为类型、每个callable绑定为算子并把EngineKVFormat同步到GPUKVFormat别名上。集成契约无论你如何构建 ops 模块都必须满足符号名一致被覆盖的每个函数必须使用与torch_ops.py完全相同的名字如multi_layer_block_kv_transfer调用签名一致位置/关键字参数、参数顺序与语义必须与基线一致——调用方只面向DeviceOps实例并不知道哪个后端应答可导入的 Python 模块原生模块必须能通过import导入实现方式不限纯 Python 文件、pybind11 扩展、ctypes 包装、RustPyO3模块等允许部分覆盖不必重实现每个函数未覆盖的部分继续用 torch 基线支持增量优化原生模块中的类型也会被绑定这正是StagingCopy、KernelGroupSpec等原生计划类型覆盖 ops_types.py 中桩类型的方式。实现要点multi_layer_block_kv_transfer与lmcache_memcpy_async是 engine-driven 与 LMCache-driven 两种传输的热点入口其他torch_ops.py函数按需覆盖即可如果在设备专属包装器内回退到通用路径如输入不受支持应直接调用lmcache.v1.platform.torch_ops里对应的函数以保持语义ensure_native()在DeviceSpec.get_ops()首次创建单例时被调用一次见 device_spec.py。允许软失败——记 warning 后不绑定返回即可。仓库内的三个参考实现展示了三种不同的绑定风格cuda/device_ops.pyensure_native()中import lmcache.cuda_ops并bind_native()扩展缺失时记 warning 并停留在 torch 基线软失败xpu/device_ops.py绑定 SYCLlmcache.xpu_ops扩展musa/device_ops.py不做bind_native而是直接以 OO 多态覆盖方法——例如multi_layer_block_kv_transfer会先用construct_musa_tensor_from_data_pointer把裸指针重建为 MUSA 张量视图再优先尝试native_kv_transfer失败则退回 torch 传输record_completion_on_stream/record_event_on_stream则先同步 MUSA 流再调用基类TorchMUSA 尚无 CUDA host-callback ABI。高级传输模式LMCache-drivenIPC 零拷贝默认传输模式是AUTO路由器严格按device_type分发——device_type cuda走LMCacheDrivenTransferContextIPC 零拷贝其余全部走EngineDrivenTransferContext。支持 IPC 句柄传输的非 CUDA 设备仍可显式 opt-in 到 LMCache-driven见下文。ROCm 特例在 PyTorch 下 ROCm GPU 上报device_type cuda。LMCache 注册了一个backend_name rocm的独立RocmDeviceSpec但复用 CUDA 平台的算子、缓存上下文和 IPC 包装。CUDA 与 ROCm 的可用性检查互斥因此 AUTO 模式无需额外配置即可选对后端。当调用方或LMCACHE_MP_TRANSFER_MODE显式请求lmcache_driven时_build_lmcache_driven_contextworker_transfer.py会执行两道硬性检查——两者都必须成功否则工厂抛出ValueError绝不静默降级你的DeviceSpec子类必须通过ipc_wrapper_cls绑定一个DeviceIPCWrapper子类暴露wrapclassmethod。平台级的 resolve_kv_wrapper_factory 直接从已注册的 spec 上读取该绑定——没有单独的注册表或自动扫描若device_type没有注册 spec/wrapper 会抛ValueErrorDeviceSpec.is_handle_transfer_available()必须返回True基类默认即True设备不支持 IPC 句柄传输时才覆盖为False例如 neuron/init.py。另外LMCache-driven 的 server 模块还要求后端旁必须提供BaseCacheContext子类in-tree 如lmcache/v1/platform/foo/cache_context.py外部如lmcache_foo/cache_context.py以及匹配的DeviceSpec.create_cache_context覆盖。平台无关的工厂 create_cache_context 通过DeviceSpec注册表按device_type分发并调用该钩子基类默认的create_cache_context抛出NotImplementedError见 device_spec.py缺失覆盖会响亮报错而非静默回退到 CPU 路径。缓存上下文本身管理用于 IPC 传输的 KV 缓存布局与指针。pin_memory_backend的 host 侧 pinning 是可选的只影响 staging 缓冲性能不是启用 LMCache-driven 模式的前提。create_transfer_contextworker_transfer.py把路由集中在一处先校验所有 KV 张量共享同一 device type再按mode参数或LMCACHE_MP_TRANSFER_MODE环境变量默认AUTO决定走_build_lmcache_driven_context、_build_engine_driven_context还是按device_type cuda分发。Event IPC 能力LMCache-driven 的多进程句柄路径还要求一个平台事件 IPC 后端。该能力由DeviceSpec.event_ipc_backend声明刻意与DeviceOps、DeviceIPCWrapper分离基类DeviceSpec返回None——具体设备必须显式 opt-inCUDA 风格事件 API 可直接使用DefaultEventIPCBackend(event_module..., device_type...)事件 ABI 不同的设备应在 in-tree 或外部后端旁实现自己的EventIPCBackend具体设备 spec 应缓存该后端因为请求 future 可能反复查询此属性。后端契约覆盖事件创建、句柄导出/导入、事件记录、流等待、查询与同步见 event_ipc.py 中EventIPCBackend协议的全部方法。check_event_support(device)在请求设备上这些操作不可用时必须抛RuntimeError。默认后端会检查是否存在接受interprocessTrue且提供from_ipc_handle的 CUDA 风格 Event 类型对应 CUDA 的Event(interprocessTrue)、ipc_handle()、Event.from_ipc_handle、record、wait、query、synchronize。自定义后端应针对自己的事件 ABI 校验等价前提。对于 CUDA 风格设备按如下方式绑定并缓存默认后端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_moduletorch.foo, device_typeself.device_type, ) self._event_backend_cache backend return backend参考实现见 cuda/init.py 的_select_event_ipc_backend()isolated-IPC 开启时选用TimelineSemaphoreEventIPCBackend否则用DefaultEventIPCBackend(event_moduletorch.cuda, ...)。Event IPC 操作必须在不引入设备级全局同步的前提下保持生产者/消费者流顺序query_event必须保持非阻塞以便请求 future 安全地轮询完成状态。如果设备不支持 Event IPC保持基类event_ipc_backend实现不变返回None。engine-driven 模式不需要该能力LMCache-driven 模式会显式报错而不是回退到 CUDA 或当前进程中的加速器。该能力在 worker/server 注册时、构造设备感知的 completion future 之前被检查运行时分发见 get_event_ipc_backend解析出device_type后查注册表无 spec 或后端为None都抛RuntimeError。这阻止了不支持的平台进入无法安全地给 KV 缓存内存排序的异步传输路径。STORE/RETRIEVE 消息线上格式不变现有的事件句柄字节仍随请求与响应 payload 传输。完整覆盖DeviceSpec以开启 LMCache-drivenclass FooDeviceSpec(DeviceSpec): property def ipc_wrapper_cls(self): 绑定该设备的 DeviceIPCWrapper 子类。 懒导入仅在真正走 LMCache-driven 路径时才拉入加速器专属模块。 from .ipc_wrapper import FooIPCWrapper return FooIPCWrapper def is_handle_transfer_available(self) - bool: 返回 True 表示设备支持 IPC 句柄传输。 return True # 基类默认不支持时覆盖为 False property def pin_memory_backend(self): 返回 PinMemoryBackend 子类或 None。 可选只影响 host staging 性能。 return None # 默认 def create_cache_context(self, *args, **kwargs): 懒导入并实例化该设备的 BaseCacheContext。 LMCache-driven 模式必需基类默认抛 NotImplementedError。 from .cache_context import FooCacheContext return FooCacheContext(*args, **kwargs)通过把 Part 1 中 vLLMkv_connector_extra_config的lmcache.mp.mp_transfer_mode设为lmcache_driven或导出LMCACHE_MP_TRANSFER_MODElmcache_driven来 opt-in。若两道硬性检查任一失败工厂抛ValueError拒绝构造该 context——请切回engine_driven或auto。仓库参考索引主题路径Device spec 基类lmcache/v1/platform/base/device_spec.pyDevice ops 基类lmcache/v1/platform/base/device_ops.pyTorch ops 基线lmcache/v1/platform/torch_ops.pyOps 类型与枚举lmcache/v1/platform/ops_types.pyEvent IPC 基类lmcache/v1/platform/base/event_ipc.py后端加载resolve_device_ops/_detect_devicelmcache/v1/platform/_device_detect.py 与 lmcache/v1/platform/init.py包级device_ops解析lmcache/init.pyCache context 基类lmcache/v1/platform/base/cache_context.pyCache context 工厂lmcache/v1/platform/cache_context.py参考DeviceSpecCUDAlmcache/v1/platform/cuda/init.py参考DeviceOpsCUDAbind_nativelmcache/v1/platform/cuda/device_ops.py参考DeviceOpsXPUbind_nativelmcache/v1/platform/xpu/device_ops.py参考DeviceOpsMUSA方法覆盖lmcache/v1/platform/musa/device_ops.py参考DeviceSpecNeuron/Trainium仅 engine-drivenlmcache/v1/platform/neuron/init.pyEngine-driven 调用点lmcache/v1/multiprocess/transfer_context/worker_transfer.pyEngineDrivenTransferContext、create_transfer_context关键设计原则回顾插件化而非分叉无论 in-tree 还是外部 wheel后端都只实现DeviceSpec/DeviceOps两个接口运行时统一走lmcache.device_plugins入口或子包发现核心代码零改动torch 基线兜底没有原生 kernel 也能跑通全部功能性能优化是增量式的——覆盖哪个算子就优化哪个算子失败要响亮LMCache-driven 的硬性检查ipc_wrapper_cls、is_handle_transfer_available、create_cache_context全部 fail-closed缺失能力就抛异常拒绝构造绝不静默降级到错误路径能力显式声明Event IPC、IPC 句柄传输、pin memory 等能力都由DeviceSpec显式 opt-in避免不支持的设备误入无法安全排序的异步传输路径。接入新设备时建议按本文顺序推进先按 Part 1 用 torch 基线跑通端到端功能并核对调试清单再按 Part 2 逐步绑定原生算子、实现 Event IPC 与 IPC 包装、最终启用lmcache_driven零拷贝传输。【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCache创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考