A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 12:30 UTC
[CustomOp] Support object-level enable for CustomOp (#30547)
Signed-off-by: shen-shanshan <467638484@qq.com>
Shanshan Shen committed
Dec 15, 2025 at 11:02 UTC
738648fb81aa53639994bee81eb0daa19aeadf59
1 file changed
+7
-2
vllm/model_executor/custom_op.py
+7
-2
@@ -38,8 +38,9 @@ class CustomOp(nn.Module):
38
)
39
return super().__new__(op_cls_to_instantiate)
40
41
- def __init__(self):
41
+ def __init__(self, enforce_enable: bool = False):
42
super().__init__()
43
+ self._enforce_enable = enforce_enable
44
self._forward_method = self.dispatch_forward()
45
46
def forward(self, *args, **kwargs):
@@ -84,7 +85,11 @@ class CustomOp(nn.Module):
85
# NOTE(woosuk): Here we assume that vLLM was built for only one
86
# specific backend. Currently, we do not support dynamic dispatching.
87
compilation_config = get_cached_compilation_config()
87
- enabled = self.enabled()
88
+
89
+ # CustomOp object can be enforce enabled, e.g., enable device-specific
90
+ # kernels in ViT models when enabling graph mode. By default, it will
91
+ # follow the compilation_config to determine whether enable itself.
92
+ enabled = self._enforce_enable or self.enabled()
93
if enabled:
94
compilation_config.enabled_custom_ops.update([self.__class__.name])
95
else: