@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 316 lines 12.7 KB
Rendered Raw
1 # CustomOp
2
3 `CustomOp` is an abstract class used for dispatching the forward method of various operations to the appropriate backend. It also offers a mechanism for both vLLM and OOT (Out-Of-Tree) plugins to register their custom operations.
4
5 This document will introduce how CustomOp works in vLLM and how to implement a new `CustomOp`.
6
7 ## How CustomOp Works in vLLM
8
9 `CustomOp` manages two dictionaries of all custom ops (i.e., op classes, indexed by registered name) in its class, for vLLM and OOT plugins respectively.
10
11 We can use `@CustomOp.register("op_name")` to register an op class to the `CustomOp` system. After this, the `op_name` and its class will be added into the `op_registry` dictionary. In addition, We can also register an OOT op by `@CustomOp.register_oot("op_name")`. We will introduce this mechanism in detail later.
12
13 When a `CustomOp` is called (i.e., call its `forward()` method), if it is enabled (i.e., with `--compilation_config.custom_ops '["+op_name"]'`), it will automatically dispatch the forward method to the appropriate backend according to `current_platform`. Otherwise (i.e., it is disabled), it will only call the `forward_native()` method to use PyTorch-native implementation of this forward method.
14
15 - **CPU platform:** dispatch to `forward_cpu()`.
16 - **CUDA platform:** dispatch to `forward_cuda()`.
17 - **ROCm platform:** dispatch to `forward_hip()`. If `forward_hip()` is not implemented, it will use `forward_cuda()` as a fallback.
18 - **XPU platform:** dispatch to `forward_xpu()`.
19 - **TPU platform:** dispatch to `forward_tpu()`.
20 - **OOT platform:** dispatch to `forward_oot()`. This will only be called on OOT platforms.
21 - **Default:** dispatch to `forward_native()` as a final fallback for all platforms.
22
23 !!! note
24 Note that the dispatching logic might not be absolute because of class inheritance. Derived class might override the behavior.
25
26 Furthermore, vLLM decides whether to enable or disable a `CustomOp` based on `compilation_config.custom_ops`. To be specific, if a `CustomOp` is not registered in `compilation_config.custom_ops` (i.e., uses the default config), it will be enabled if `compilation_config.custom_ops` contains `all`, or will be disabled if it contains `none`.
27
28 !!! note
29 Note that `all` and `none` cannot coexist in `compilation_config.custom_ops`.
30
31 By default, if `compilation_config.backend == "inductor"` and `compilation_config.mode != CompilationMode.NONE`, a `none` will be appended into `compilation_config.custom_ops`, otherwise a `all` will be appended. In other words, this means `CustomOp` will be disabled in some platforms (i.e., those use `inductor` as default backend for `torch.compile`) when running with torch compile mode. In this case, Inductor generates (fused) Triton kernels for those disabled custom ops.
32
33 !!! note
34 For multi-modal models, vLLM has enforced the enabling of some custom ops to use device-specific deep-optimized kernels for better performance in ViT part, such as `MMEncoderAttention` and `ApplyRotaryEmb`. We can also pass a `enforce_enable=True` param to the `__init__()` method of the `CustomOp` to enforce enable itself at object-level.
35
36 Note that this `enforce_enable` mechanism will be removed after we add a separate `compilation_config` for multi-modal part.
37
38 ## How to Customise Your Configuration for CustomOp
39
40 vLLM also offers fine-grained control over which custom ops to enable or disable for users, by manually passing a `--compilation_config.custom_ops '["..."]'` when launching a server.
41
42 For example:
43
44 - Use `--compilation_config.custom_ops '["all"]'` to enable all custom ops.
45 - Use `--compilation_config.custom_ops '["none"]'` to disable all custom ops.
46 - Use `--compilation_config.custom_ops '["all,-op1"]'` to enable all custom ops except op1 (i.e., prefixed with a `-` means "disable").
47 - Use `--compilation_config.custom_ops '["none,+op1,+op2"]'` to only enable op1 and op2 (i.e., prefixed with a `+` means "enable").
48
49 ## Types of Supported CustomOp in vLLM
50
51 **1. Attention:**
52
53 ```python
54 --8<-- "vllm/model_executor/layers/mla.py:multi_head_latent_attention"
55
56 ```
57
58 **2. Activation:**
59
60 ```python
61 --8<-- "vllm/model_executor/layers/activation.py:silu_and_mul"
62
63 --8<-- "vllm/model_executor/layers/activation.py:mul_and_silu"
64
65 --8<-- "vllm/model_executor/layers/activation.py:gelu_new"
66
67 --8<-- "vllm/model_executor/layers/activation.py:gelu_fast"
68
69 --8<-- "vllm/model_executor/layers/activation.py:quick_gelu"
70
71 --8<-- "vllm/model_executor/layers/activation.py:gelu_and_mul"
72
73 --8<-- "vllm/model_executor/layers/activation.py:gelu_and_mul_sparse"
74
75 --8<-- "vllm/model_executor/layers/activation.py:relu2"
76
77 --8<-- "vllm/model_executor/layers/activation.py:xielu"
78
79 --8<-- "vllm/model_executor/layers/activation.py:swigluoai_and_mul"
80
81 --8<-- "vllm/model_executor/layers/activation.py:fatrelu_and_mul"
82 ```
83
84 **3. MM-Conv:**
85
86 ```python
87 --8<-- "vllm/model_executor/layers/conv.py:conv2d"
88
89 --8<-- "vllm/model_executor/layers/conv.py:conv3d"
90 ```
91
92 **4. Embedding:**
93
94 ```python
95 --8<-- "vllm/model_executor/layers/vocab_parallel_embedding.py:vocab_parallel_embedding"
96
97 --8<-- "vllm/model_executor/layers/vocab_parallel_embedding.py:parallel_lm_head"
98 ```
99
100 **5. Linear:**
101
102 ```python
103 --8<-- "vllm/model_executor/layers/linear.py:row_parallel_linear"
104
105 --8<-- "vllm/model_executor/layers/linear.py:column_parallel_linear"
106
107 --8<-- "vllm/model_executor/layers/linear.py:replicated_linear"
108 ```
109
110 **6. Logits Processor:**
111
112 ```python
113 --8<-- "vllm/model_executor/layers/logits_processor.py:logits_processor"
114 ```
115
116 **7. Mamba:**
117
118 ```python
119 --8<-- "vllm/model_executor/layers/mamba/mamba_mixer.py:mamba_mixer"
120
121 --8<-- "vllm/model_executor/layers/mamba/mamba_mixer2.py:mamba_mixer2"
122
123 --8<-- "vllm/model_executor/layers/mamba/mamba_mixer2.py:mixer2_gated_rms_norm"
124
125 --8<-- "vllm/model_executor/layers/mamba/short_conv.py:short_conv"
126 ```
127
128 **8. MoE:**
129
130 ```python
131 --8<-- "vllm/model_executor/layers/fused_moe/layer.py:fused_moe"
132
133 --8<-- "vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py:modular_fused_moe"
134
135 --8<-- "vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py:unquantized_fused_moe"
136
137 --8<-- "vllm/model_executor/models/transformers/moe.py:transformers_fused_moe"
138
139 --8<-- "vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py:grouped_topk"
140 ```
141
142 **9. Norm:**
143
144 ```python
145 --8<-- "vllm/model_executor/layers/layernorm.py:rms_norm"
146
147 --8<-- "vllm/model_executor/layers/layernorm.py:rms_norm_gated"
148
149 --8<-- "vllm/model_executor/layers/layernorm.py:gemma_rms_norm"
150 ```
151
152 **10. Quantization:**
153
154 ```python
155 --8<-- "vllm/model_executor/layers/quantization/input_quant_fp8.py:quant_fp8"
156 ```
157
158 **11. Rope:**
159
160 ```python
161 --8<-- "vllm/model_executor/layers/rotary_embedding/base.py:rotary_embedding"
162
163 --8<-- "vllm/model_executor/layers/rotary_embedding/dual_chunk_rope.py:dual_chunk_rotary_embedding"
164
165 --8<-- "vllm/model_executor/layers/rotary_embedding/common.py:apply_rotary_emb"
166 ```
167
168 **12. Encoder:**
169
170 ```python
171 --8<-- "vllm/model_executor/models/deepencoder2.py:qwen2_decoder"
172
173 --8<-- "vllm/model_executor/layers/attention/mm_encoder_attention.py:mm_encoder_attn"
174
175 --8<-- "vllm/model_executor/models/deepencoder.py:rel_pos_attention"
176 ```
177
178 ## Guidelines for Implementing a New CustomOp
179
180 ### Implement a New CustomOp in vLLM
181
182 This part is a tutorial of how to implement a New `CustomOp` in vLLM.
183
184 Steps:
185
186 1. Implement a new op class, which extends from `CustomOp` base class.
187 2. Add the `@CustomOp.register("op_name")` decorator on this op class to register it into `CustomOp` system.
188 3. Implement different `forward_xxx()` method according to your needs.
189
190 Taking `MMEncoderAttention` as an example:
191
192 ??? code
193
194 ```python
195 @CustomOp.register("mm_encoder_attn")
196 class MMEncoderAttention(CustomOp):
197
198 def __init__(
199 self,
200 num_heads: int,
201 head_size: int,
202 scale: float | None = None,
203 num_kv_heads: int | None = None,
204 prefix: str = "",
205 multimodal_config: MultiModalConfig | None = None,
206 ) -> None:
207 super().__init__()
208 # Init...
209
210 def forward_native(
211 self,
212 query: torch.Tensor,
213 key: torch.Tensor,
214 value: torch.Tensor,
215 cu_seqlens: torch.Tensor | None = None,
216 max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
217 ) -> torch.Tensor:
218 # Call TORCH_SDPA implementation...
219
220 def forward_cuda(
221 self,
222 query: torch.Tensor,
223 key: torch.Tensor,
224 value: torch.Tensor,
225 cu_seqlens: torch.Tensor | None = None,
226 max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
227 ) -> torch.Tensor:
228 # Call FA or TORCH_SDPA implementation...
229
230 def forward_cpu(
231 self,
232 query: torch.Tensor,
233 key: torch.Tensor,
234 value: torch.Tensor,
235 cu_seqlens: torch.Tensor | None = None,
236 max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
237 ) -> torch.Tensor:
238 # Call TORCH_SDPA implementation...
239
240 def forward_xpu(
241 self,
242 query: torch.Tensor,
243 key: torch.Tensor,
244 value: torch.Tensor,
245 cu_seqlens: torch.Tensor | None = None,
246 max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
247 ) -> torch.Tensor:
248 # Call FA implementation...
249
250 def forward_tpu(
251 self,
252 query: torch.Tensor,
253 key: torch.Tensor,
254 value: torch.Tensor,
255 cu_seqlens: torch.Tensor | None = None,
256 max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
257 ) -> torch.Tensor:
258 # Call PALLAS implementation...
259 ```
260
261 ### Register a New CustomOp in OOT Device Plugins
262
263 Currently, thanks to [vLLM's hardware-plugin mechanism](./plugin_system.md), there are various OOT device plugins emerging out to enable vLLM seamlessly runs on different hardwares. You can also find more details about this mechanism at [Introducing vLLM Hardware Plugin, Best Practice from Ascend NPU](https://blog.vllm.ai/2025/05/12/hardware-plugin.html).
264
265 - **Official device plugins:** [vllm-ascend](https://github.com/vllm-project/vllm-ascend) (for Huawei Ascend NPU), [vllm-spyre](https://github.com/vllm-project/vllm-spyre)
266 (for Spyre), [vllm-gaudi](https://github.com/vllm-project/vllm-gaudi) (for Intel Gaudi), [vllm-neuron](https://github.com/vllm-project/vllm-neuron) (for AWS Neuron), [vllm-meta](https://github.com/vllm-project/vllm-metal) (for Apple Silicon), etc.
267 - **Non-official device plugins:** [vllm-metax](https://github.com/MetaX-MACA/vLLM-metax) (for MetaX GPU), [vllm-kunlun](https://github.com/baidu/vLLM-Kunlun) (for Baidu Kunlun XPU), [vllm-musa](https://github.com/MooreThreads/vllm-musa) (for Moore Threads GPU), etc.
268
269 In this case, `CustomOp` can enable these hardware manufacturers to seamlessly replace vLLM's operations with their deep-optimized kernels for specific devices at runtime, by just registering an OOT `CustomOp` and implementing the `forward_oot()` method.
270
271 Now, this part will show you how to register an OOT `CustomOp` for a device plugin.
272
273 Taking `MMEncoderAttention` as an example:
274
275 1. Implement a `CustomMMEncoderAttention` class which extends from `MMEncoderAttention` and implement its `forward_oot()` method.
276 2. Register your `CustomMMEncoderAttention` into vLLM to replace `MMEncoderAttention`.
277
278 ??? code
279
280 ```python
281 from vllm.model_executor.layers.attention import MMEncoderAttention
282 from vllm.model_executor.custom_op import CustomOp
283
284
285 @CustomOp.register_oot("MMEncoderAttention")
286 class CustomMMEncoderAttention(MMEncoderAttention):
287
288 def __init__(...):
289 super().__init__(...)
290
291 def forward_oot(...):
292 # Call optimized device-specific kernels.
293 ...
294 ```
295
296 In this case, a new item `{"MMEncoderAttention": CustomMMEncoderAttention}` will be added into `op_registry_oot`. When initializing a `MMEncoderAttention` op object, if the class name (i.e., `MMEncoderAttention`) is contained in the keys of `op_registry_oot`, vLLM will replace it with our registered class (i.e., `CustomMMEncoderAttention`) and instantiate it.
297
298 After that, when this `MMEncoderAttention` op is called, your `forward_oot()` will be called if it is enabled. Thus, you will get expected performance on your hardwares without directly modify vLLM.
299
300 In addition, you can also register all your `CustomOp` at one place for better management.
301
302 ??? code
303
304 ```python
305 from vllm.model_executor.custom_op import CustomOp
306
307
308 REGISTERED_CUSTOM_OPS = {
309 "CustomOP1": YourCustomOp1,
310 "CustomOP2": YourCustomOp2,
311 "CustomOP3": YourCustomOp3,
312 }
313
314 for op_name, op_cls in REGISTERED_CUSTOM_OPS.items():
315 CustomOp.register_oot(_decorated_op_cls=op_cls, name=op_name)
316 ```