[3/N][torch.compile] consolidate custom op logging (#10399)

Signed-off-by: youkaichao <youkaichao@gmail.com>
This commit is contained in:
youkaichao
2024-11-18 15:14:59 -08:00
committed by GitHub
parent 96d999fbe8
commit a03ea40792
3 changed files with 20 additions and 5 deletions

View File

@@ -61,10 +61,13 @@ class CustomOp(nn.Module):
def dispatch_forward(self):
# NOTE(woosuk): Here we assume that vLLM was built for only one
# specific backend. Currently, we do not support dynamic dispatching.
compilation_config = get_current_vllm_config().compilation_config
enabled = self.enabled()
logger.debug("custom op %s %s", self.__class__.name,
"enabled" if enabled else "disabled")
if enabled:
compilation_config.enabled_custom_ops.update([self.__class__.name])
else:
compilation_config.disabled_custom_ops.update(
[self.__class__.name])
if not enabled:
return self.forward_native