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 4, 15:00 UTC
Auto-disable expandable_segments around cumem memory pool (#40812)
Signed-off-by: youkaichao <youkaichao@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
youkaichao committed
Apr 27, 2026 at 09:37 UTC
2ce95a761b9acd925d1bc69cfd1d4fc13de9e2b7
1 file changed
+40
-32
vllm/device_allocator/cumem.py
+40
-32
@@ -128,13 +128,6 @@ class CuMemAllocator:
128
return CuMemAllocator.instance
129
130
def __init__(self):
131
- conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
132
- assert "expandable_segments:True" not in conf, (
133
- "Expandable segments are not compatible with memory pool. "
134
- "Please track https://github.com/pytorch/pytorch/issues/147851 "
135
- "for the latest updates."
136
- )
137
-
131
self.pointer_to_data: dict[int, AllocationData] = {}
132
self.current_tag: str = CuMemAllocator.default_tag
133
self.allocator_and_pools: dict[str, Any] = {}
@@ -264,34 +257,49 @@ class CuMemAllocator:
257
258
assert isinstance(tag, str)
259
260
+ # Expandable segments are incompatible with the memory pool used for
261
+ # sleep mode (see https://github.com/pytorch/pytorch/issues/147851).
262
+ # If the user has enabled expandable segments via
263
+ # PYTORCH_CUDA_ALLOC_CONF, temporarily disable them for the duration
264
+ # of the memory pool context and restore on exit.
265
+ conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
266
+ expandable_was_enabled = "expandable_segments:True" in conf
267
+ if expandable_was_enabled:
268
+ torch.cuda.memory._set_allocator_settings("expandable_segments:False")
269
+
270
old_tag = self.current_tag
271
self.current_tag = tag
269
- with use_memory_pool_with_allocator(
270
- self.python_malloc_callback, self.python_free_callback
271
- ) as data:
272
- # start to hit another PyTorch bug in PyTorch 2.6,
273
- # possibly because of gc-related issue w.r.t. the allocator and
274
- # the memory pool.
275
- # to avoid the issue, we keep a reference of the data.
276
- # see https://github.com/pytorch/pytorch/issues/146431 .
277
- self.allocator_and_pools[tag] = data
278
- yield
279
- # PyTorch's bug, calling torch.cuda.empty_cache() will error
280
- # when using pluggable allocator, see
281
- # https://github.com/pytorch/pytorch/issues/145168 .
282
- # if we have some memory allocated and then freed,
283
- # the memory will not be released, e.g. in online quantization,
284
- # where the model is created in higher precision, and then
285
- # quantized in lower precision.
286
- # Find all unused allocations and manually release them.
287
- # TODO: we should expose `empty_cache` method in the memory pool.
288
- # TODO: ask for help from PyTorch team to expose this method.
289
- allocations = data[0].snapshot()
290
- for allocation in allocations:
291
- if allocation["allocated_size"] == 0:
292
- handle = self._python_free_callback(allocation["address"])
293
- unmap_and_release(handle)
272
+ try:
273
+ with use_memory_pool_with_allocator(
274
+ self.python_malloc_callback, self.python_free_callback
275
+ ) as data:
276
+ # start to hit another PyTorch bug in PyTorch 2.6,
277
+ # possibly because of gc-related issue w.r.t. the allocator
278
+ # and the memory pool.
279
+ # to avoid the issue, we keep a reference of the data.
280
+ # see https://github.com/pytorch/pytorch/issues/146431 .
281
+ self.allocator_and_pools[tag] = data
282
+ yield
283
+ # PyTorch's bug, calling torch.cuda.empty_cache() will error
284
+ # when using pluggable allocator, see
285
+ # https://github.com/pytorch/pytorch/issues/145168 .
286
+ # if we have some memory allocated and then freed,
287
+ # the memory will not be released, e.g. in online
288
+ # quantization, where the model is created in higher
289
+ # precision, and then quantized in lower precision.
290
+ # Find all unused allocations and manually release them.
291
+ # TODO: we should expose `empty_cache` method in the memory
292
+ # pool.
293
+ # TODO: ask for help from PyTorch team to expose this method.
294
+ allocations = data[0].snapshot()
295
+ for allocation in allocations:
296
+ if allocation["allocated_size"] == 0:
297
+ handle = self._python_free_callback(allocation["address"])
298
+ unmap_and_release(handle)
299
+ finally:
300
self.current_tag = old_tag
301
+ if expandable_was_enabled:
302
+ torch.cuda.memory._set_allocator_settings("expandable_segments:True")
303
304
def get_current_usage(self) -> int:
305
"""