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 2, 03:30 UTC
[Platform] Add is_cumem_allocator_available (#43838)
Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
wangxiyuan committed
Jun 3, 2026 at 10:54 UTC
02a01496fc65b495b0c76699387f3509b07cbec5
2 files changed
+12
-11
vllm/config/model.py
+4
-11
@@ -80,16 +80,6 @@ else:
80
81
logger = init_logger(__name__)
82
83
-
84
-def is_cumem_allocator_available() -> bool:
85
- try:
86
- from vllm.device_allocator.cumem import cumem_available
87
- except ImportError:
88
- return False
89
-
90
- return cumem_available
91
-
92
-
83
RunnerOption = Literal["auto", RunnerType]
84
ConvertType = Literal["none", "embed", "classify"]
85
ConvertOption = Literal["auto", ConvertType]
@@ -542,7 +532,10 @@ class ModelConfig:
532
"Enabling cumem allocator because sleep mode requires it."
533
)
534
self.enable_cumem_allocator = True
545
- if self.enable_cumem_allocator and not is_cumem_allocator_available():
535
+ if (
536
+ self.enable_cumem_allocator
537
+ and not current_platform.is_cumem_allocator_available()
538
+ ):
539
raise ValueError("cumem allocator is not supported on current platform.")
540
541
hf_config = get_config(
vllm/platforms/interface.py
+8
@@ -199,6 +199,14 @@ class Platform:
199
# all ROCm platforms for now.
200
return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM)
201
202
+ def is_cumem_allocator_available(self) -> bool:
203
+ try:
204
+ from vllm.device_allocator.cumem import cumem_available
205
+ except ImportError:
206
+ return False
207
+
208
+ return cumem_available
209
+
210
@classmethod
211
def get_pass_manager_cls(cls) -> str:
212
"""