biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 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
1 changed files with 7 additions and 2 deletions
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) return super().__new__(op_cls_to_instantiate)
def __init__(self): def __init__(self, enforce_enable: bool = False):
super().__init__() super().__init__()
self._enforce_enable = enforce_enable
self._forward_method = self.dispatch_forward() self._forward_method = self.dispatch_forward()
def forward(self, *args, **kwargs): 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 # NOTE(woosuk): Here we assume that vLLM was built for only one
# specific backend. Currently, we do not support dynamic dispatching. # specific backend. Currently, we do not support dynamic dispatching.
compilation_config = get_cached_compilation_config() 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: if enabled:
compilation_config.enabled_custom_ops.update([self.__class__.name]) compilation_config.enabled_custom_ops.update([self.__class__.name])
else: else: