xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-03-24 03:23:39 +08:00
Code Issues Packages Projects Releases Wiki Activity

[CustomOp] Support object-level enable for CustomOp (#30547)

Browse Source 
Signed-off-by: shen-shanshan <467638484@qq.com>
This commit is contained in:
Shanshan Shen 2025-12-15 11:02:09 +08:00 committed by GitHub
parent 917fdae5b2
commit 738648fb81
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
vllm/model_executor/custom_op.py
View File

@ -38,8 +38,9 @@ class CustomOp(nn.Module):
)
return super().__new__(op_cls_to_instantiate)
def __init__(self):
def __init__(self, enforce_enable: bool = False):
super().__init__()
self._enforce_enable = enforce_enable
self._forward_method = self.dispatch_forward()
def forward(self, *args, **kwargs):
@ -84,7 +85,11 @@ class CustomOp(nn.Module):
# NOTE(woosuk): Here we assume that vLLM was built for only one
# specific backend. Currently, we do not support dynamic dispatching.
compilation_config = get_cached_compilation_config()
enabled = self.enabled()
# CustomOp object can be enforce enabled, e.g., enable device-specific
# kernels in ViT models when enabling graph mode. By default, it will
# follow the compilation_config to determine whether enable itself.
enabled = self._enforce_enable or self.enabled()
if enabled:
compilation_config.enabled_custom_ops.update([self.__class__.name])
else: