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 15, 18:30 UTC
[Frontend] Remove custom_cache_manager (#13791)
Signed-off-by: fulvius31 <asangior@redhat.com>
Alessandro Sangiorgi committed
Mar 19, 2025 at 11:13 UTC
374ee287d821306bc6977cdf5eaab98f9ca3f81e
3 files changed
+1
-71
vllm/executor/multiproc_worker_utils.py
-8
@@ -16,12 +16,8 @@ import torch
16
17
from vllm.config import VllmConfig
18
from vllm.logger import init_logger
19
-from vllm.triton_utils.importing import HAS_TRITON
19
from vllm.utils import _check_multiproc_method, get_mp_context, run_method
20
22
-if HAS_TRITON:
23
- from vllm.triton_utils import maybe_set_triton_cache_manager
24
-
21
logger = init_logger(__name__)
22
23
T = TypeVar('T')
@@ -314,7 +310,3 @@ def set_multiprocessing_worker_envs(parallel_config):
310
current_parallelism, default_omp_num_threads)
311
os.environ["OMP_NUM_THREADS"] = str(default_omp_num_threads)
312
torch.set_num_threads(default_omp_num_threads)
317
-
318
- # workaround for https://github.com/vllm-project/vllm/issues/6103
319
- if HAS_TRITON and parallel_config.world_size > 1:
320
- maybe_set_triton_cache_manager()
vllm/triton_utils/__init__.py
+1
-8
@@ -2,11 +2,4 @@
2
3
from vllm.triton_utils.importing import HAS_TRITON
4
5
-__all__ = ["HAS_TRITON"]
6
-
7
-if HAS_TRITON:
8
-
9
- from vllm.triton_utils.custom_cache_manager import (
10
- maybe_set_triton_cache_manager)
11
-
12
- __all__ += ["maybe_set_triton_cache_manager"]
5
+__all__ = ["HAS_TRITON"]
\ No newline at end of file
vllm/triton_utils/custom_cache_manager.py
deleted
-55
@@ -1,55 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-
3
-import os
4
-
5
-from triton.runtime.cache import (FileCacheManager, default_cache_dir,
6
- default_dump_dir, default_override_dir)
7
-
8
-from vllm.logger import init_logger
9
-
10
-logger = init_logger(__name__)
11
-
12
-
13
-def maybe_set_triton_cache_manager() -> None:
14
- """Set environment variable to tell Triton to use a
15
- custom cache manager"""
16
- cache_manger = os.environ.get("TRITON_CACHE_MANAGER", None)
17
- if cache_manger is None:
18
- manager = "vllm.triton_utils.custom_cache_manager:CustomCacheManager"
19
- logger.info("Setting Triton cache manager to: %s", manager)
20
- os.environ["TRITON_CACHE_MANAGER"] = manager
21
-
22
-
23
-class CustomCacheManager(FileCacheManager):
24
- """Re-implements Triton's cache manager, ensuring that a
25
- unique cache directory is created for each process. This is
26
- needed to avoid collisions when running with tp>1 and
27
- using multi-processing as the distributed backend.
28
-
29
- Note this issue was fixed by triton-lang/triton/pull/4295,
30
- but the fix is not yet included in triton==v3.0.0. However,
31
- it should be included in the subsequent version.
32
- """
33
-
34
- def __init__(self, key, override=False, dump=False):
35
- self.key = key
36
- self.lock_path = None
37
- if dump:
38
- self.cache_dir = default_dump_dir()
39
- self.cache_dir = os.path.join(self.cache_dir, self.key)
40
- self.lock_path = os.path.join(self.cache_dir, "lock")
41
- os.makedirs(self.cache_dir, exist_ok=True)
42
- elif override:
43
- self.cache_dir = default_override_dir()
44
- self.cache_dir = os.path.join(self.cache_dir, self.key)
45
- else:
46
- # create cache directory if it doesn't exist
47
- self.cache_dir = os.getenv("TRITON_CACHE_DIR",
48
- "").strip() or default_cache_dir()
49
- if self.cache_dir:
50
- self.cache_dir = f"{self.cache_dir}_{os.getpid()}"
51
- self.cache_dir = os.path.join(self.cache_dir, self.key)
52
- self.lock_path = os.path.join(self.cache_dir, "lock")
53
- os.makedirs(self.cache_dir, exist_ok=True)
54
- else:
55
- raise RuntimeError("Could not create or locate cache dir")