插件系统¶

社区经常提出通过自定义功能扩展 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/__init__.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",
        )

有关向您的包添加入口点（entry points）的更多信息，请查看官方文档。

每个插件由三部分组成

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

支持的插件类型¶

通用插件（General plugins）（组名为 vllm.general_plugins）：这些插件的主要用例是将自定义的、非原生模型注册到 vLLM 中。这是通过在插件函数中调用 ModelRegistry.register_model 来注册模型的。关于官方模型插件的示例，请参阅 bart-plugin，它添加了对 BartForConditionalGeneration 的支持。
平台插件（Platform plugins）（组名为 vllm.platform_plugins）：这些插件的主要用例是将自定义的、非原生平台注册到 vLLM 中。当当前环境不支持该平台时，插件函数应返回 None；当支持该平台时，则返回平台类的完全限定名。
IO 处理器插件（IO Processor plugins）（组名为 vllm.io_processor_plugins）：这些插件的主要用例是为池化（pooling）模型注册自定义的模型 Prompt 和输出的预处理/后处理。插件函数返回 IOProcessor 类的完全限定名。
状态记录器插件（Stat logger plugins）（组名为 vllm.stat_logger_plugins）：这些插件的主要用例是将自定义的、非原生记录器注册到 vLLM 中。入口点应为 StatLoggerBase 的子类。

插件编写指南¶

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

平台插件指南¶

创建一个平台插件项目，例如 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，表示该平台为非原生（out-of-tree）平台。
- device_type：此属性应返回 PyTorch 使用的设备类型。例如 "cpu", "cuda" 等。
- device_name：此属性通常与 device_type 设置相同。它主要用于日志记录。
- check_and_update_config：此函数在 vLLM 初始化过程的早期被调用。它用于插件更新 vLLM 配置。例如，可以在此函数中更新块大小、图模式配置等。最重要的是，必须在此函数中设置 worker_cls，以便让 vLLM 知道工作进程应使用哪个工作类（worker class）。
- 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：调用此函数以分析模型的峰值内存占用，从而确定在不发生 OOM 的情况下可用于 KV 缓存的内存量。
- initialize_from_config：调用此函数以使用指定的 kv_cache_config 分配设备 KV 缓存。
- execute_model：在每一步推理时调用此函数以运行模型。
可以实现的额外函数包括：
- 如果插件想要支持睡眠模式（sleep mode）功能，请实现 sleep 和 wakeup 函数。
- 如果插件想要支持图模式（graph mode）功能，请实现 compile_or_warm_up_model 函数。
- 如果插件想要支持投机解码（speculative decoding）功能，请实现 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 算子：有 3 种 PyTorch 算子：
  - communicator ops（通信器算子）：设备通信器算子。例如 all-reduce, all-gather 等。请在 my_dummy_device_communicator.py 中实现设备通信器类 MyDummyDeviceCommunicator。该类应继承自 DeviceCommunicatorBase。
  - common ops（通用算子）：通用算子，例如 matmul, softmax 等。请通过 OOT 注册方式实现通用算子。详情请参阅 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.v1.attention.backends.registry.register_backend 向 AttentionBackendEnum 添加新的注意力后端。
seed_everything 平台接口已被弃用，已在 v0.16.0 版本中移除。请改用 vllm.utils.torch_utils.set_random_seed。
Platform.validate_request 中的 prompt 已被弃用，已在 v0.18.0 版本中移除。