@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 160 lines 11.8 KB
Rendered Raw
1 # Plugin System
2
3 The community frequently requests the ability to extend vLLM with custom features. To facilitate this, vLLM includes a plugin system that allows users to add custom features without modifying the vLLM codebase. This document explains how plugins work in vLLM and how to create a plugin for vLLM.
4
5 ## How Plugins Work in vLLM
6
7 Plugins are user-registered code that vLLM executes. Given vLLM's architecture (see [Arch Overview](arch_overview.md)), multiple processes may be involved, especially when using distributed inference with various parallelism techniques. To enable plugins successfully, every process created by vLLM needs to load the plugin. This is done by the [load_plugins_by_group][vllm.plugins.load_plugins_by_group] function in the `vllm.plugins` module.
8
9 ## How vLLM Discovers Plugins
10
11 vLLM's plugin system uses the standard Python `entry_points` mechanism. This mechanism allows developers to register functions in their Python packages for use by other packages. An example of a plugin:
12
13 ??? code
14
15 ```python
16 # inside `setup.py` file
17 from setuptools import setup
18
19 setup(name='vllm_add_dummy_model',
20 version='0.1',
21 packages=['vllm_add_dummy_model'],
22 entry_points={
23 'vllm.general_plugins':
24 ["register_dummy_model = vllm_add_dummy_model:register"]
25 })
26
27 # inside `vllm_add_dummy_model/__init__.py` file
28 def register():
29 from vllm import ModelRegistry
30
31 if "MyLlava" not in ModelRegistry.get_supported_archs():
32 ModelRegistry.register_model(
33 "MyLlava",
34 "vllm_add_dummy_model.my_llava:MyLlava",
35 )
36 ```
37
38 For more information on adding entry points to your package, please check the [official documentation](https://setuptools.pypa.io/en/latest/userguide/entry_point.html).
39
40 Every plugin has three parts:
41
42 1. **Plugin group**: The name of the entry point group. vLLM uses the entry point group `vllm.general_plugins` to register general plugins. This is the key of `entry_points` in the `setup.py` file. Always use `vllm.general_plugins` for vLLM's general plugins.
43 2. **Plugin name**: The name of the plugin. This is the value in the dictionary of the `entry_points` dictionary. In the example above, the plugin name is `register_dummy_model`. Plugins can be filtered by their names using the `VLLM_PLUGINS` environment variable. To load only a specific plugin, set `VLLM_PLUGINS` to the plugin name.
44 3. **Plugin value**: The fully qualified name of the function or module to register in the plugin system. In the example above, the plugin value is `vllm_add_dummy_model:register`, which refers to a function named `register` in the `vllm_add_dummy_model` module.
45
46 ## Types of supported plugins
47
48 - **General plugins** (with group name `vllm.general_plugins`): The primary use case for these plugins is to register custom, out-of-the-tree models into vLLM. This is done by calling `ModelRegistry.register_model` to register the model inside the plugin function. For an example of an official model plugin, see the [bart-plugin](https://github.com/vllm-project/bart-plugin) which adds support for `BartForConditionalGeneration`.
49
50 - **Platform plugins** (with group name `vllm.platform_plugins`): The primary use case for these plugins is to register custom, out-of-the-tree platforms into vLLM. The plugin function should return `None` when the platform is not supported in the current environment, or the platform class's fully qualified name when the platform is supported.
51
52 - **IO Processor plugins** (with group name `vllm.io_processor_plugins`): The primary use case for these plugins is to register custom pre-/post-processing of the model prompt and model output for pooling models. The plugin function returns the IOProcessor's class fully qualified name.
53
54 - **Stat logger plugins** (with group name `vllm.stat_logger_plugins`): The primary use case for these plugins is to register custom, out-of-the-tree loggers into vLLM. The entry point should be a class that subclasses StatLoggerBase.
55
56 - **Endpoint plugins** (with group name `vllm.endpoint_plugins`): The primary use case for these plugins is to register custom, out-of-the-tree HTTP routes on the OpenAI compatible API server. Unlike the other plugin groups above, endpoint plugins are loaded only in the API server front end process and are **not loaded by default**. See [Endpoint Plugins](endpoint_plugins.md) for the interface and [Security](../usage/security.md#endpoint-plugins) for the opt-in and trust model.
57
58 ## Guidelines for Writing Plugins
59
60 - **Being re-entrant**: The function specified in the entry point should be re-entrant, meaning it can be called multiple times without causing issues. This is necessary because the function might be called multiple times in some processes.
61
62 ### Platform plugins guidelines
63
64 1. Create a platform plugin project, for example, `vllm_add_dummy_platform`. The project structure should look like this:
65
66 ```shell
67 vllm_add_dummy_platform/
68 ├── vllm_add_dummy_platform/
69 │ ├── __init__.py
70 │ ├── my_dummy_platform.py
71 │ ├── my_dummy_worker.py
72 │ ├── my_dummy_attention.py
73 │ ├── my_dummy_device_communicator.py
74 │ ├── my_dummy_custom_ops.py
75 ├── setup.py
76 ```
77
78 2. In the `setup.py` file, add the following entry point:
79
80 ```python
81 setup(
82 name="vllm_add_dummy_platform",
83 ...
84 entry_points={
85 "vllm.platform_plugins": [
86 "my_dummy_platform = vllm_add_dummy_platform:register"
87 ]
88 },
89 ...
90 )
91 ```
92
93 Please make sure `vllm_add_dummy_platform:register` is a callable function and returns the platform class's fully qualified name. for example:
94
95 ```python
96 def register():
97 return "vllm_add_dummy_platform.my_dummy_platform.MyDummyPlatform"
98 ```
99
100 3. Implement the platform class `MyDummyPlatform` in `my_dummy_platform.py`. The platform class should inherit from `vllm.platforms.interface.Platform`. Please follow the interface to implement the functions one by one. There are some important functions and properties that should be implemented at least:
101
102 - `_enum`: This property is the device enumeration from [PlatformEnum][vllm.platforms.interface.PlatformEnum]. Usually, it should be `PlatformEnum.OOT`, which means the platform is out-of-tree.
103 - `device_type`: This property should return the type of the device which pytorch uses. For example, `"cpu"`, `"cuda"`, etc.
104 - `device_name`: This property is set the same as `device_type` usually. It's mainly used for logging purposes.
105 - `check_and_update_config`: This function is called very early in the vLLM's initialization process. It's used for plugins to update the vllm configuration. For example, the block size, graph mode config, etc., can be updated in this function. The most important thing is that the **worker_cls** should be set in this function to let vLLM know which worker class to use for the worker process.
106 - `get_attn_backend_cls`: This function should return the attention backend class's fully qualified name.
107 - `get_device_communicator_cls`: This function should return the device communicator class's fully qualified name.
108
109 4. Implement the worker class `MyDummyWorker` in `my_dummy_worker.py`. The worker class should inherit from [WorkerBase][vllm.v1.worker.worker_base.WorkerBase]. Please follow the interface to implement the functions one by one. Basically, all interfaces in the base class should be implemented, since they are called here and there in vLLM. To make sure a model can be executed, the basic functions should be implemented are:
110
111 - `init_device`: This function is called to set up the device for the worker.
112 - `initialize_cache`: This function is called to set cache config for the worker.
113 - `load_model`: This function is called to load the model weights to device.
114 - `get_kv_cache_spec`: This function is called to generate the kv cache spec for the model.
115 - `determine_available_memory`: This function is called to profiles the peak memory usage of the model to determine how much memory can be used for KV cache without OOMs.
116 - `initialize_from_config`: This function is called to allocate device KV cache with the specified kv_cache_config
117 - `execute_model`: This function is called every step to inference the model.
118
119 Additional functions that can be implemented are:
120
121 - If the plugin wants to support sleep mode feature, please implement the `sleep` and `wakeup` functions.
122 - If the plugin wants to support graph mode feature, please implement the `compile_or_warm_up_model` function.
123 - If the plugin wants to support speculative decoding feature, please implement the `take_draft_token_ids` function.
124 - If the plugin wants to support lora feature, please implement the `add_lora`,`remove_lora`,`list_loras` and `pin_lora` functions.
125 - If the plugin wants to support data parallelism feature, please implement the `execute_dummy_batch` functions.
126
127 Please look at the worker base class [WorkerBase][vllm.v1.worker.worker_base.WorkerBase] for more functions that can be implemented.
128
129 5. Implement the attention backend class `MyDummyAttention` in `my_dummy_attention.py`. The attention backend class should inherit from [AttentionBackend][vllm.v1.attention.backend.AttentionBackend]. It's used to calculate attentions with your device. Take `vllm.v1.attention.backends` as examples, it contains many attention backend implementations.
130
131 6. Implement custom ops for high performance. Most ops can be run by pytorch native implementation, while the performance may not be good. In this case, you can implement specific custom ops for your plugins. Currently, there are kinds of custom ops vLLM supports:
132
133 - pytorch ops
134 there are 3 kinds of pytorch ops:
135
136 - `communicator ops`: Device communicator op. Such as all-reduce, all-gather, etc.
137 Please implement the device communicator class `MyDummyDeviceCommunicator` in `my_dummy_device_communicator.py`. The device communicator class should inherit from [DeviceCommunicatorBase][vllm.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase].
138 - `common ops`: Common ops. Such as matmul, softmax, etc.
139 Please implement the common ops by register oot way. See more detail in [CustomOp][vllm.model_executor.custom_op.CustomOp] class.
140 - `csrc ops`: C++ ops. This kind of ops are implemented in C++ and are registered as torch custom ops.
141 Following csrc module and `vllm._custom_ops` to implement your ops.
142
143 - triton ops
144 Custom way doesn't work for triton ops now.
145
146 7. (optional) Implement other pluggable modules, such as lora, graph backend, quantization, mamba attention backend, etc.
147
148 ## Compatibility Guarantee
149
150 vLLM guarantees the interface of documented plugins, such as `ModelRegistry.register_model`, will always be available for plugins to register models. However, it is the responsibility of plugin developers to ensure their plugins are compatible with the version of vLLM they are targeting. For example, `"vllm_add_dummy_model.my_llava:MyLlava"` should be compatible with the version of vLLM that the plugin targets.
151
152 The interface for the model/module may change during vLLM's development. If you see any deprecation log info, please upgrade your plugin to the latest version.
153
154 ## Deprecation announcement
155
156 !!! warning "Deprecations"
157 - `use_v1` parameter in `Platform.get_attn_backend_cls` is deprecated. It has been removed in v0.13.0.
158 - `_Backend` in `vllm.attention` is deprecated. It has been removed in v0.13.0. Please use `vllm.v1.attention.backends.registry.register_backend` to add new attention backend to `AttentionBackendEnum` instead.
159 - `seed_everything` platform interface is deprecated. It has been removed in v0.16.0. Please use `vllm.utils.torch_utils.set_random_seed` instead.
160 - `prompt` in `Platform.validate_request` is deprecated. It has been removed in v0.18.0.