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 3, 01:45 UTC
[Bugfix] Make CuMemAllocator free callback stream-aware (#43020)
Signed-off-by: zixi-qi <zixi@inferact.ai> Co-authored-by: Claude <noreply@anthropic.com>
qizixi committed
May 21, 2026 at 20:36 UTC
18a27cc9a3641cc1dd3eae5113b75c7ccc029b5f
1 file changed
+6
vllm/device_allocator/cumem.py
+6
@@ -160,6 +160,12 @@ class CuMemAllocator:
160
data = self.pointer_to_data.pop(ptr)
161
if data.cpu_backup_tensor is not None:
162
data.cpu_backup_tensor = None
163
+ # Drain pending kernels before the C extension's cuMemUnmap.
164
+ # The pluggable allocator path doesn't defer reclaim like the
165
+ # regular caching allocator, so without this, in-flight work
166
+ # (e.g. quant helpers' transient tensors during weight loading)
167
+ # races the unmap and surfaces as CUDA_ERROR_ILLEGAL_ADDRESS.
168
+ torch.cuda.synchronize(data.handle[0])
169
logger.debug(
170
"Freed %s bytes for %s with address %s from cumem allocator",
171
data.handle[1],