跳到内容

vLLM 的插件系统

社区经常请求扩展 vLLM 的自定义功能。为了方便这一点,vLLM 包含一个插件系统,允许用户在不修改 vLLM 代码库的情况下添加自定义功能。本文档解释了插件在 vLLM 中如何工作以及如何为 vLLM 创建插件。

vLLM 中插件的工作原理

插件是用户注册的由 vLLM 执行的代码。考虑到 vLLM 的架构(参见 架构概览),可能涉及多个进程,尤其是在使用分布式推理和各种并行技术时。为了成功启用插件,vLLM 创建的每个进程都需要加载插件。这是通过 vllm.plugins 模块中的 load_general_plugins 函数完成的。vLLM 创建的每个进程在开始任何工作之前都会调用此函数。

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",
        )

有关向您的包添加入口点的更多信息,请查阅官方文档

每个插件包含三个部分

  1. 插件组:入口点组的名称。vLLM 使用入口点组 vllm.general_plugins 来注册通用插件。这是 setup.py 文件中 entry_points 的键。vLLM 的通用插件总是使用 vllm.general_plugins
  2. 插件名称:插件的名称。这是 entry_points 字典中的值。在上面的示例中,插件名称是 register_dummy_model。可以使用 VLLM_PLUGINS 环境变量按名称过滤插件。要只加载特定的插件,请将 VLLM_PLUGINS 设置为插件名称。
  3. 插件值:要在插件系统中注册的函数的完全限定名称。在上面的示例中,插件值是 vllm_add_dummy_model:register,它指向 vllm_add_dummy_model 模块中名为 register 的函数。

支持的插件类型

  • 通用插件(组名 vllm.general_plugins):这些插件的主要用例是将自定义的、非内置的模型注册到 vLLM 中。这是通过在插件函数内部调用 ModelRegistry.register_model 来注册模型实现的。

  • 平台插件(组名 vllm.platform_plugins):这些插件的主要用例是将自定义的、非内置的平台注册到 vLLM 中。当当前环境中不支持该平台时,插件函数应返回 None;当支持该平台时,应返回平台类的完全限定名称。

编写插件的指南

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

兼容性保证

vLLM 保证已文档化的插件接口,例如 ModelRegistry.register_model,将始终可用于插件注册模型。然而,插件开发者有责任确保他们的插件与他们目标版本的 vLLM 兼容。例如,"vllm_add_dummy_model.my_llava:MyLlava" 应该与插件目标版本的 vLLM 兼容。模型接口在 vLLM 开发过程中可能会发生变化。