插件系统¶

社区频繁要求能够通过自定义功能来扩展 vLLM。为了实现这一点，vLLM 提供了一个插件系统，允许用户在不修改 vLLM 代码库的情况下添加自定义功能。本文档解释了插件在 vLLM 中是如何工作的，以及如何为 vLLM 创建插件。

vLLM 中的插件系统如何工作¶

插件是 vLLM 执行的用户注册代码。鉴于 vLLM 的架构（请参阅架构概述），可能涉及多个进程，尤其是在使用各种并行化技术的分布式推理时。为了成功启用插件，vLLM 创建的每个进程都需要加载插件。这是通过 vllm.plugins 模块中的 load_plugins_by_group 函数完成的。

vLLM 如何发现插件¶

vLLM 的插件系统使用标准的 Python entry_points 机制。该机制允许开发人员在他们的 Python 包中注册函数供其他包使用。插件示例

代码

# inside `setup.py` file
from setuptools import setup

setup(name='vllm_add_dummy_model',
    version='0.1',
    packages=['vllm_add_dummy_model'],
    entry_points={
        'vllm.general_plugins':
        ["register_dummy_model = vllm_add_dummy_model:register"]
    })

# inside `vllm_add_dummy_model.py` file
def register():
    from vllm import ModelRegistry

    if "MyLlava" not in ModelRegistry.get_supported_archs():
        ModelRegistry.register_model(
            "MyLlava",
            "vllm_add_dummy_model.my_llava:MyLlava",
        )

有关将入口点添加到您的包的更多信息，请查看官方文档。

每个插件有三个部分

插件组：入口点组的名称。vLLM 使用入口点组 vllm.general_plugins 来注册通用插件。这是 setup.py 文件中 entry_points 的键。对于 vLLM 的通用插件，请始终使用 vllm.general_plugins。
插件名称：插件的名称。这是 entry_points 字典中值。在上面的示例中，插件名称是 register_dummy_model。可以通过 VLLM_PLUGINS 环境变量按名称过滤插件。如果要仅加载特定插件，请将 VLLM_PLUGINS 设置为插件名称。
插件值：要在插件系统中注册的函数或模块的完全限定名。在上面的示例中，插件值是 vllm_add_dummy_model:register，它引用了 vllm_add_dummy_model 模块中名为 register 的函数。

支持的插件类型¶

通用插件（组名为 vllm.general_plugins）：这些插件的主要用途是将自定义的、非核心的（out-of-the-tree）模型注册到 vLLM 中。这是通过调用 ModelRegistry.register_model 在插件函数内部注册模型来实现的。
平台插件（组名为 vllm.platform_plugins）：这些插件的主要用途是将自定义的、非核心的（out-of-the-tree）平台注册到 vLLM 中。当当前环境中不支持该平台时，插件函数应返回 None；当支持该平台时，则返回平台类的完全限定名。
IO 处理器插件（组名为 vllm.io_processor_plugins）：这些插件的主要用途是为池化模型注册自定义的模型提示和模型输出的预/后处理。插件函数返回 IOProcessor 类的完全限定名。
状态记录器插件（组名为 vllm.stat_logger_plugins）：这些插件的主要用途是将自定义的、非核心的（out-of-the-tree）日志记录器注册到 vLLM 中。入口点应为继承自 StatLoggerBase 的类。

编写插件的指南¶

可重入性：入口点中指定的函数应该是可重入的，这意味着它可以被多次调用而不会引起问题。这是必要的，因为该函数在某些进程中可能会被多次调用。

平台插件指南¶

创建一个平台插件项目，例如 vllm_add_dummy_platform。项目结构应如下所示

vllm_add_dummy_platform/
├── vllm_add_dummy_platform/
│   ├── __init__.py
│   ├── my_dummy_platform.py
│   ├── my_dummy_worker.py
│   ├── my_dummy_attention.py
│   ├── my_dummy_device_communicator.py
│   ├── my_dummy_custom_ops.py
├── setup.py

在 setup.py 文件中，添加以下入口点

setup(
    name="vllm_add_dummy_platform",
    ...
    entry_points={
        "vllm.platform_plugins": [
            "my_dummy_platform = vllm_add_dummy_platform:register"
        ]
    },
    ...
)

请确保 vllm_add_dummy_platform:register 是一个可调用的函数，并返回平台类的完全限定名。例如

def register():
    return "vllm_add_dummy_platform.my_dummy_platform.MyDummyPlatform"

在 my_dummy_platform.py 中实现平台类 MyDummyPlatform。平台类应继承自 vllm.platforms.interface.Platform。请遵循接口逐个实现函数。以下是一些重要的函数和属性，至少应实现这些：
- _enum：此属性是来自 PlatformEnum 的设备枚举。通常，它应该是 PlatformEnum.OOT，表示该平台是外部的。
- device_type：此属性应返回 PyTorch 使用的设备类型。例如，"cpu"、"cuda" 等。
- device_name：此属性通常设置为与 device_type 相同。它主要用于日志记录。
- check_and_update_config：此函数在 vLLM 初始化过程的早期被调用。它用于插件更新 vLLM 配置。例如，块大小、图模式配置等可以在此函数中更新。最重要的事情是，必须在此函数中设置 worker_cls，以便 vLLM 知道要在工作进程中使用哪个工作类。
- get_attn_backend_cls：此函数应返回注意力后端类的完全限定名。
- get_device_communicator_cls：此函数应返回设备通信器的完全限定名。
在 my_dummy_worker.py 中实现工作类 MyDummyWorker。工作类应继承自 WorkerBase。请遵循接口逐个实现函数。基本上，基类中的所有接口都应该实现，因为它们在 vLLM 的不同地方被调用。为了确保模型可以执行，需要实现的基本函数是：
- init_device：调用此函数来设置工作进程的设备。
- initialize_cache：调用此函数来设置工作进程的缓存配置。
- load_model：调用此函数将模型权重加载到设备。
- get_kv_cache_spec：调用此函数来生成模型的 KV 缓存规范。
- determine_available_memory：调用此函数来分析模型的峰值内存使用情况，以确定可以使用多少内存用于 KV 缓存而不会导致 OOM（内存不足）。
- initialize_from_config：调用此函数，使用指定的 kv_cache_config 分配设备 KV 缓存。
- execute_model：每一步都调用此函数来推断模型。
可以实现的附加函数有：
- 如果插件想支持睡眠模式功能，请实现 sleep 和 wakeup 函数。
- 如果插件想支持图模式功能，请实现 compile_or_warm_up_model 函数。
- 如果插件想支持投机解码功能，请实现 take_draft_token_ids 函数。
- 如果插件想支持 LORA 功能，请实现 add_lora、remove_lora、list_loras 和 pin_lora 函数。
- 如果插件想支持数据并行功能，请实现 execute_dummy_batch 函数。
请参考工作类基类 WorkerBase 以了解更多可实现的函数。
在 my_dummy_attention.py 中实现注意力后端类 MyDummyAttention。注意力后端类应继承自 AttentionBackend。它用于在您的设备上计算注意力。以 vllm.v1.attention.backends 为例，它包含许多注意力后端实现。
为高性能实现自定义算子。大多数算子都可以通过 PyTorch 的原生实现运行，但性能可能不佳。在这种情况下，您可以为您的插件实现特定的自定义算子。目前，vLLM 支持以下类型的自定义算子：
- PyTorch 算子有三种类型的 PyTorch 算子：
  - communicator ops：设备通信器算子。例如 all-reduce、all-gather 等。请在 my_dummy_device_communicator.py 中实现设备通信器类 MyDummyDeviceCommunicator。设备通信器类应继承自 DeviceCommunicatorBase。
  - common ops：通用算子。例如 matmul、softmax 等。请通过注册 OOT（out-of-tree）的方式实现通用算子。更多详细信息请参阅 CustomOp 类。
  - csrc ops：C++ 算子。这类算子是用 C++ 实现的，并注册为 torch 自定义算子。请按照 csrc 模块和 vllm._custom_ops 来实现您的算子。
- Triton 算子：自定义方式目前不适用于 Triton 算子。
（可选）实现其他可插拔模块，如 LORA、图后端、量化、Mamba 注意力后端等。

兼容性保证¶

vLLM 保证已记录的插件接口，例如 ModelRegistry.register_model，将始终可用于插件注册模型。但是，插件开发者有责任确保他们的插件与他们目标 vLLM 版本兼容。例如，"vllm_add_dummy_model.my_llava:MyLlava" 应该与插件目标 vLLM 版本兼容。

模型/模块的接口在 vLLM 的开发过程中可能会发生变化。如果您看到任何弃用日志信息，请将您的插件升级到最新版本。

弃用公告¶

弃用

Platform.get_attn_backend_cls 中的 use_v1 参数已弃用。它已在 v0.13.0 中移除。
vllm.attention 中的 _Backend 已弃用。它已在 v0.13.0 中移除。请使用 vllm.attention.backends.registry.register_backend 将新的注意力后端添加到 AttentionBackendEnum 中。