biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[PluggableLayer][1/N] Define PluggableLayer (#32331)

Browse Source 
Signed-off-by: whx-sjtu <2952154980@qq.com>
This commit is contained in:
whx
2026-01-21 00:19:21 +08:00
committed by GitHub
parent 7901109ea5
commit 4ca62a0dbd
4 changed files with 99 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
docs/design/custom_op.md
View File

@@ -8,15 +8,6 @@ This document will introduce how CustomOp works in vLLM and how to implement a n
`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.
??? code
```python
class CustomOp(nn.Module):
op_registry: dict[str, type["CustomOp"]] = {}
op_registry_oot: dict[str, type["CustomOp"]] = {}
```
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.
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.