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
[Security] Serialize using safetensors instead of pickle in Mooncake Pipe (#14228)
Signed-off-by: KuntaiDu <kuntai@uchicago.edu>
Kuntai Du committed
Mar 4, 2025 at 15:10 UTC
288ca110f68d23909728627d3100e5a8db820aa2
1 file changed
+6
-6
vllm/distributed/kv_transfer/kv_pipe/mooncake_pipe.py
+6
-6
@@ -2,13 +2,14 @@
2
3
import json
4
import os
5
-import pickle
5
from concurrent.futures import ThreadPoolExecutor
6
from dataclasses import dataclass
7
from typing import Optional, Union
8
9
import torch
10
import zmq
11
+from safetensors.torch import load as safetensors_load
12
+from safetensors.torch import save as safetensors_save
13
14
from vllm.config import KVTransferConfig
15
from vllm.distributed.kv_transfer.kv_pipe.base import KVPipeBase
@@ -237,14 +238,13 @@ class MooncakePipe(KVPipeBase):
238
return hash(tensor.data_ptr())
239
240
def _send_impl(self, tensor: torch.Tensor) -> None:
240
- """Implement the tensor sending logic."""
241
- value_bytes = pickle.dumps(tensor)
242
- self.transfer_engine.send_bytes(value_bytes)
241
+ """Implement the tensor sending logic using safetensors."""
242
+ self.transfer_engine.send_bytes(safetensors_save({"tensor": tensor}))
243
244
def _recv_impl(self) -> torch.Tensor:
245
- """Implement the tensor receiving logic."""
245
+ """Implement the tensor receiving logic using safetensors."""
246
data = self.transfer_engine.recv_bytes()
247
- return pickle.loads(data)
247
+ return safetensors_load(data)["tensor"].to(self.device)
248
249
def send_tensor(self, tensor: Optional[torch.Tensor]) -> None:
250
"""Send tensor to the target process."""