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, 21:00 UTC
[sleep mode] error out with expandable_segments (#14189)
Signed-off-by: youkaichao <youkaichao@gmail.com>
youkaichao committed
Mar 4, 2025 at 18:54 UTC
71c4b40562eb308d5fd93091373e4f913463a9eb
1 file changed
+7
vllm/device_allocator/cumem.py
+7
@@ -8,6 +8,7 @@
8
# not sure why, they are created from a different context.
9
# the only successful approach is to call cuda driver API in C.
10
import dataclasses
11
+import os
12
from contextlib import contextmanager
13
from typing import Any, Callable, Dict, Optional, Tuple, Union
14
@@ -140,6 +141,12 @@ class CuMemAllocator:
141
return CuMemAllocator.instance
142
143
def __init__(self):
144
+ conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
145
+ assert "expandable_segments:True" not in conf, \
146
+ ("Expandable segments are not compatible with memory pool. "
147
+ "Please track https://github.com/pytorch/pytorch/issues/147851 "
148
+ "for the latest updates.")
149
+
150
self.pointer_to_data: Dict[int, AllocationData] = {}
151
self.current_tag: str = CuMemAllocator.default_tag
152
self.allocator_and_pools: Dict[str, Any] = {}