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 14, 05:30 UTC
Optimize configuration access with LRU cache in custom ops (#22204)
Signed-off-by: zitian zhao <zitian.zhao@tencentmusic.com>
ZiTian.Zhao committed
Aug 5, 2025 at 12:43 UTC
4b3e4474d73ae9cf0d6c8315570fdffd71037d08
2 files changed
+13
-5
vllm/config.py
+9
-1
@@ -15,7 +15,7 @@ from collections.abc import Mapping
15
from contextlib import contextmanager
16
from dataclasses import (MISSING, Field, asdict, field, fields, is_dataclass,
17
replace)
18
-from functools import cached_property
18
+from functools import cached_property, lru_cache
19
from importlib.util import find_spec
20
from typing import (TYPE_CHECKING, Any, Callable, ClassVar, Literal, Optional,
21
Protocol, TypeVar, Union, cast, get_args)
@@ -5123,6 +5123,14 @@ def set_current_vllm_config(vllm_config: VllmConfig,
5123
finally:
5124
_current_vllm_config = old_vllm_config
5125
_current_prefix = old_prefix
5126
+ # Clear the compilation config cache when context changes
5127
+ get_cached_compilation_config.cache_clear()
5128
+
5129
+
5130
+@lru_cache(maxsize=1)
5131
+def get_cached_compilation_config():
5132
+ """Cache config to avoid repeated calls to get_current_vllm_config()"""
5133
+ return get_current_vllm_config().compilation_config
5134
5135
5136
def get_current_vllm_config() -> VllmConfig:
vllm/model_executor/custom_op.py
+4
-4
@@ -5,7 +5,7 @@ from typing import Optional
5
6
import torch.nn as nn
7
8
-from vllm.config import get_current_vllm_config
8
+from vllm.config import get_cached_compilation_config
9
from vllm.logger import init_logger
10
from vllm.platforms import current_platform
11
@@ -86,7 +86,7 @@ class CustomOp(nn.Module):
86
def dispatch_forward(self):
87
# NOTE(woosuk): Here we assume that vLLM was built for only one
88
# specific backend. Currently, we do not support dynamic dispatching.
89
- compilation_config = get_current_vllm_config().compilation_config
89
+ compilation_config = get_cached_compilation_config()
90
enabled = self.enabled()
91
if enabled:
92
compilation_config.enabled_custom_ops.update([self.__class__.name])
@@ -115,7 +115,7 @@ class CustomOp(nn.Module):
115
@classmethod
116
def enabled(cls) -> bool:
117
# if no name, then it was not registered
118
- compilation_config = get_current_vllm_config().compilation_config
118
+ compilation_config = get_cached_compilation_config()
119
custom_ops = compilation_config.custom_ops
120
if not hasattr(cls, "name"):
121
logger.warning_once(
@@ -138,7 +138,7 @@ class CustomOp(nn.Module):
138
Specifying 'all' or 'none' in custom_op takes precedence.
139
"""
140
from vllm.config import CompilationLevel
141
- compilation_config = get_current_vllm_config().compilation_config
141
+ compilation_config = get_cached_compilation_config()
142
default_on = (compilation_config.level < CompilationLevel.PIECEWISE
143
or not compilation_config.use_inductor)
144
count_none = compilation_config.custom_ops.count("none")