外掛系統¶

社群頻繁要求能夠透過自定義功能來擴充套件 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 中。