From 4b668bc6cce3a65d797f38a19a7161843b299b96 Mon Sep 17 00:00:00 2001 From: youkaichao Date: Sat, 16 Nov 2024 18:10:41 -0800 Subject: [PATCH 1/2] consolidate custom op logging Signed-off-by: youkaichao --- vllm/config.py | 7 +++++++ vllm/model_executor/custom_op.py | 8 +++++--- vllm/plugins/__init__.py | 4 ++++ 3 files changed, 16 insertions(+), 3 deletions(-) diff --git a/vllm/config.py b/vllm/config.py index 7e37edbe594b1..9e704fc63fb31 100644 --- a/vllm/config.py +++ b/vllm/config.py @@ -2161,6 +2161,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: Set[str] = PrivateAttr + disabled_custom_ops: Set[str] = PrivateAttr + def model_post_init(self, __context: Any) -> None: self.level = envs.VLLM_TORCH_COMPILE_LEVEL @@ -2182,6 +2186,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 = set() + self.disabled_custom_ops = set() + def init_during_runtime(self): """To complete the initialization of config, we need to know the compile context, which is only available diff --git a/vllm/model_executor/custom_op.py b/vllm/model_executor/custom_op.py index 6ae7d7cf6964f..fcbd47ec2ade8 100644 --- a/vllm/model_executor/custom_op.py +++ b/vllm/model_executor/custom_op.py @@ -61,10 +61,12 @@ 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.add(self.__class__.name) + else: + compilation_config.disabled_custom_ops.add(self.__class__.name) if not enabled: return self.forward_native diff --git a/vllm/plugins/__init__.py b/vllm/plugins/__init__.py index c20b9ec891d5d..5dcc13feeca4c 100644 --- a/vllm/plugins/__init__.py +++ b/vllm/plugins/__init__.py @@ -92,6 +92,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 From 4bbd10c723b22fdbc1a58831f07d4c3c081ae732 Mon Sep 17 00:00:00 2001 From: youkaichao Date: Sat, 16 Nov 2024 18:16:30 -0800 Subject: [PATCH 2/2] use Counter Signed-off-by: youkaichao --- vllm/config.py | 13 +++++++------ vllm/model_executor/custom_op.py | 5 +++-- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/vllm/config.py b/vllm/config.py index 9e704fc63fb31..faba148bb0718 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 @@ -2162,8 +2163,8 @@ class CompilationConfig(BaseModel): capture_sizes: List[int] = PrivateAttr # keep track of enabled and disabled custom ops - enabled_custom_ops: Set[str] = PrivateAttr - disabled_custom_ops: Set[str] = PrivateAttr + 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 @@ -2186,8 +2187,8 @@ def model_post_init(self, __context: Any) -> None: func = __import__(module).__dict__[func_name] self.inductor_compile_config[k] = func - self.enabled_custom_ops = set() - self.disabled_custom_ops = set() + self.enabled_custom_ops = Counter() + self.disabled_custom_ops = Counter() def init_during_runtime(self): """To complete the initialization of config, diff --git a/vllm/model_executor/custom_op.py b/vllm/model_executor/custom_op.py index fcbd47ec2ade8..b07966f2ab7d0 100644 --- a/vllm/model_executor/custom_op.py +++ b/vllm/model_executor/custom_op.py @@ -64,9 +64,10 @@ def dispatch_forward(self): compilation_config = get_current_vllm_config().compilation_config enabled = self.enabled() if enabled: - compilation_config.enabled_custom_ops.add(self.__class__.name) + compilation_config.enabled_custom_ops.update([self.__class__.name]) else: - compilation_config.disabled_custom_ops.add(self.__class__.name) + compilation_config.disabled_custom_ops.update( + [self.__class__.name]) if not enabled: return self.forward_native