From 9eea0b89efa584adefcfd80649c54794f2153e76 Mon Sep 17 00:00:00 2001
From: youkaichao <youkaichao@gmail.com>
Date: Mon, 18 Nov 2024 15:14:59 -0800
Subject: [PATCH] [3/N][torch.compile] consolidate custom op logging (#10399)

Signed-off-by: youkaichao <youkaichao@gmail.com>
---
 vllm/config.py                   | 12 ++++++++++--
 vllm/model_executor/custom_op.py |  9 ++++++---
 vllm/plugins/__init__.py         |  4 ++++
 3 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/vllm/config.py b/vllm/config.py
index 14017bbdb3cf2..ea9ec43cc5a15 100644
--- a/vllm/config.py
+++ b/vllm/config.py
@@ -4,8 +4,9 @@
 import warnings
 from dataclasses import dataclass, field, replace
 from pathlib import Path
-from typing import (TYPE_CHECKING, Any, Callable, ClassVar, Dict, Final, List,
-                    Literal, Mapping, Optional, Set, Tuple, Type, Union)
+from typing import (TYPE_CHECKING, Any, Callable, ClassVar, Counter, Dict,
+                    Final, List, Literal, Mapping, Optional, Set, Tuple, Type,
+                    Union)
 
 import torch
 from pydantic import BaseModel, Field, PrivateAttr
@@ -2169,6 +2170,10 @@ class CompilationConfig(BaseModel):
     compile_sizes: List[int] = PrivateAttr
     capture_sizes: List[int] = PrivateAttr
 
+    # keep track of enabled and disabled custom ops
+    enabled_custom_ops: Counter[str] = PrivateAttr
+    disabled_custom_ops: Counter[str] = PrivateAttr
+
     def model_post_init(self, __context: Any) -> None:
         self.level = envs.VLLM_TORCH_COMPILE_LEVEL
 
@@ -2190,6 +2195,9 @@ def model_post_init(self, __context: Any) -> None:
             func = __import__(module).__dict__[func_name]
             self.inductor_compile_config[k] = func
 
+        self.enabled_custom_ops = Counter()
+        self.disabled_custom_ops = Counter()
+
     def init_backend(self) -> Union[str, Callable]:
         if self.level == CompilationLevel.NO_COMPILATION:
             raise ValueError("No compilation level is set.")
diff --git a/vllm/model_executor/custom_op.py b/vllm/model_executor/custom_op.py
index 6ae7d7cf6964f..b07966f2ab7d0 100644
--- a/vllm/model_executor/custom_op.py
+++ b/vllm/model_executor/custom_op.py
@@ -61,10 +61,13 @@ def forward_hpu(self, *args, **kwargs):
     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
diff --git a/vllm/plugins/__init__.py b/vllm/plugins/__init__.py
index a0c73a752b5e8..c5182139db50b 100644
--- a/vllm/plugins/__init__.py
+++ b/vllm/plugins/__init__.py
@@ -80,6 +80,10 @@ def set_current_vllm_config(vllm_config: VllmConfig):
         _current_vllm_config = vllm_config
         yield
     finally:
+        logger.debug("enabled custom ops: %s",
+                     vllm_config.compilation_config.enabled_custom_ops)
+        logger.debug("disabled custom ops: %s",
+                     vllm_config.compilation_config.disabled_custom_ops)
         _current_vllm_config = old_vllm_config