Skip to content

Commit 15da497

Browse files
committed
[Platform] add pre_register_and_update function
Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>
1 parent 29f1d47 commit 15da497

File tree

3 files changed

+41
-1
lines changed

3 files changed

+41
-1
lines changed

vllm/config.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3057,7 +3057,8 @@ class VllmConfig:
30573057
kv_transfer_config: KVTransferConfig = field(default=None,
30583058
init=True) # type: ignore
30593059
# some opaque config, only used to provide additional information
3060-
# for the hash computation, mainly used for testing and debugging.
3060+
# for the hash computation, mainly used for testing, debugging or out of
3061+
# tree config registration.
30613062
additional_config: SupportsHash = field(default=None,
30623063
init=True) # type: ignore
30633064
instance_id: str = ""

vllm/engine/arg_utils.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
from vllm.executor.executor_base import ExecutorBase
2121
from vllm.logger import init_logger
2222
from vllm.model_executor.layers.quantization import QUANTIZATION_METHODS
23+
from vllm.plugins import load_general_plugins
2324
from vllm.transformers_utils.utils import check_gguf_file
2425
from vllm.usage.usage_lib import UsageContext
2526
from vllm.utils import FlexibleArgumentParser, StoreBoolean
@@ -203,6 +204,8 @@ class EngineArgs:
203204

204205
calculate_kv_scales: Optional[bool] = None
205206

207+
additional_config: Optional[Dict[str, Any]] = None
208+
206209
def __post_init__(self):
207210
if not self.tokenizer:
208211
self.tokenizer = self.model
@@ -984,6 +987,14 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
984987
'be loaded from the model checkpoint if available. '
985988
'Otherwise, the scales will default to 1.0.')
986989

990+
parser.add_argument(
991+
"--additional-config",
992+
type=json.loads,
993+
default=None,
994+
help="Additional config for specified platform in JSON format. "
995+
"Different platforms may support different configs. Make sure the "
996+
"configs are valid for the platform you are using. The input format"
997+
" is like '{\"config_key\":\"config_value\"}'")
987998
return parser
988999

9891000
@classmethod
@@ -1044,6 +1055,9 @@ def create_load_config(self) -> LoadConfig:
10441055
def create_engine_config(self,
10451056
usage_context: Optional[UsageContext] = None
10461057
) -> VllmConfig:
1058+
from vllm.platforms import current_platform
1059+
current_platform.pre_register_and_update()
1060+
10471061
if envs.VLLM_USE_V1:
10481062
self._override_v1_engine_args(usage_context)
10491063

@@ -1287,6 +1301,7 @@ def create_engine_config(self,
12871301
prompt_adapter_config=prompt_adapter_config,
12881302
compilation_config=self.compilation_config,
12891303
kv_transfer_config=self.kv_transfer_config,
1304+
additional_config=self.additional_config,
12901305
)
12911306

12921307
if envs.VLLM_USE_V1:
@@ -1347,6 +1362,12 @@ def add_cli_args(parser: FlexibleArgumentParser,
13471362
parser.add_argument('--disable-log-requests',
13481363
action='store_true',
13491364
help='Disable logging requests.')
1365+
# Initialize plugin to update the parser, for example, The plugin may
1366+
# adding a new kind of quantization method to --quantization argument or
1367+
# a new device to --device argument.
1368+
load_general_plugins()
1369+
from vllm.platforms import current_platform
1370+
current_platform.pre_register_and_update(parser)
13501371
return parser
13511372

13521373

vllm/platforms/interface.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,10 @@
1313

1414
if TYPE_CHECKING:
1515
from vllm.config import VllmConfig
16+
from vllm.utils import FlexibleArgumentParser
1617
else:
1718
VllmConfig = None
19+
FlexibleArgumentParser = None
1820

1921
logger = init_logger(__name__)
2022

@@ -222,6 +224,22 @@ def seed_everything(cls, seed: int) -> None:
222224
np.random.seed(seed)
223225
torch.manual_seed(seed)
224226

227+
@classmethod
228+
def pre_register_and_update(cls,
229+
parser: Optional[FlexibleArgumentParser] = None
230+
) -> None:
231+
"""
232+
Do some pre-registeration or update action for the current platform.
233+
234+
This function is called before global VllmConfig is initialized or cli
235+
arguments are parsed. It's used for out-of-tree platforms to register or
236+
update the configuration.
237+
238+
For example, the out-of-tree quantization config can be imported and
239+
registered here dynamically.
240+
"""
241+
pass
242+
225243
@classmethod
226244
def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
227245
"""

0 commit comments

Comments
 (0)