[Bugfix] Reuse CUDA streams in packed weight transfer to cap reserved-memory waste (#52951)
Signed-off-by: Quentin Gallouédec <quentin.gallouedec@huggingface.co>
Quentin Gallouédec committed
Aug 24, 2026 at 02:54 UTC
26858770eccf9592073b502a57f4662aac358001
1 file changed
+11
-2
vllm/distributed/weight_transfer/packed_tensor.py
+11
-2
@@ -5,6 +5,7 @@
5
import math
6
from collections.abc import Callable, Iterator
7
from dataclasses import dataclass
8
+from functools import cache
9
from typing import Any
10
11
import torch
@@ -15,6 +16,14 @@ DEFAULT_PACKED_BUFFER_SIZE_BYTES = 1024 * 1024 * 1024 # 1GB
16
DEFAULT_PACKED_NUM_BUFFERS = 2
17
18
19
+# Streams are cached across calls: the caching allocator keys cached blocks by the
20
+# stream they were allocated on, so fresh streams at every call would strand the
21
+# packed buffers' reserved memory per call. See gh-52950.
22
+@cache
23
+def _get_streams(device_idx: int, num_buffers: int) -> tuple[torch.cuda.Stream, ...]:
24
+ return tuple(torch.cuda.Stream(device=device_idx) for _ in range(num_buffers))
25
+
26
+
27
def unpack_tensor(
28
packed_tensor: torch.Tensor,
29
names: list[str],
@@ -154,7 +163,7 @@ def packed_nccl_broadcast_producer(
163
Both producer and consumer must use the same value.
164
165
"""
157
- streams = [torch.cuda.Stream() for _ in range(num_buffers)]
166
+ streams = _get_streams(torch.accelerator.current_device_index(), num_buffers)
167
# Keep references to in-flight chunks so their packed_tensors
168
# aren't freed while an async broadcast is still reading them.
169
in_flight: list[PackedChunk | None] = [None] * num_buffers
@@ -209,7 +218,7 @@ def packed_nccl_broadcast_consumer(
218
"""
219
target_packed_tensor_size = buffer_size_bytes
220
212
- streams = [torch.cuda.Stream() for _ in range(num_buffers)]
221
+ streams = _get_streams(torch.accelerator.current_device_index(), num_buffers)
222
buffer_idx = 0
223
224
packing_tensor_meta_data: list[list[tuple[str, list[int], torch.dtype, int]]] = [