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, 17:01 UTC
Fix RunAI streamer tensor buffer reuse during weight loading (#43464)
Signed-off-by: bbartels <benjamin@bartels.dev>
Benjamin Bartels committed
May 28, 2026 at 03:16 UTC
05eec7120ea76383f099b8ad0f731e7da2abc75d
2 files changed
+24
-1
tests/model_executor/model_loader/runai_streamer_loader/test_weight_utils.py
+22
@@ -6,6 +6,7 @@ import tempfile
6
7
import huggingface_hub.constants
8
import torch
9
+from safetensors.torch import save_file
10
11
from vllm.model_executor.model_loader.weight_utils import (
12
download_weights_from_hf,
@@ -14,6 +15,27 @@ from vllm.model_executor.model_loader.weight_utils import (
15
)
16
17
18
+def test_runai_safetensors_weights_iterator_clones_reused_buffers(
19
+ tmp_path, monkeypatch
20
+):
21
+ monkeypatch.setenv("RUNAI_STREAMER_MEMORY_LIMIT", "0")
22
+ weights_file = tmp_path / "model.safetensors"
23
+ expected_tensors = {
24
+ "first": torch.tensor([1.0, 2.0]),
25
+ "second": torch.tensor([3.0, 4.0]),
26
+ }
27
+ save_file(expected_tensors, weights_file)
28
+
29
+ actual_tensors = dict(
30
+ runai_safetensors_weights_iterator([str(weights_file)], False)
31
+ )
32
+
33
+ assert actual_tensors.keys() == expected_tensors.keys()
34
+ assert actual_tensors["first"].data_ptr() != actual_tensors["second"].data_ptr()
35
+ for name, expected_tensor in expected_tensors.items():
36
+ assert torch.equal(actual_tensors[name], expected_tensor)
37
+
38
+
39
def test_runai_model_loader():
40
with tempfile.TemporaryDirectory() as tmpdir:
41
huggingface_hub.constants.HF_HUB_OFFLINE = False
vllm/model_executor/model_loader/weight_utils.py
+2
-1
@@ -1090,7 +1090,8 @@ def runai_safetensors_weights_iterator(
1090
mininterval=2,
1091
)
1092
1093
- yield from tensor_iter
1093
+ for name, tensor in tensor_iter:
1094
+ yield name, tensor.clone()
1095
1096
1097
def _init_fastsafetensors_loader(