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, 14:45 UTC
[V1] Zero-copy tensor/ndarray serialization/transmission (#13790)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Apr 10, 2025 at 12:23 UTC
dd143ef54137807fdb8f91b836c5ec6617dfb507
4 files changed
+217
-58
tests/v1/test_serial_utils.py
new
+80
@@ -0,0 +1,80 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from collections import UserDict
3
+from dataclasses import dataclass
4
+
5
+import numpy as np
6
+import torch
7
+
8
+from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
9
+
10
+
11
+class UnrecognizedType(UserDict):
12
+
13
+ def __init__(self, an_int: int):
14
+ super().__init__()
15
+ self.an_int = an_int
16
+
17
+
18
+@dataclass
19
+class MyType:
20
+ tensor1: torch.Tensor
21
+ a_string: str
22
+ list_of_tensors: list[torch.Tensor]
23
+ numpy_array: np.ndarray
24
+ unrecognized: UnrecognizedType
25
+
26
+
27
+def test_encode_decode():
28
+ """Test encode/decode loop with zero-copy tensors."""
29
+
30
+ obj = MyType(
31
+ tensor1=torch.randint(low=0,
32
+ high=100,
33
+ size=(1024, ),
34
+ dtype=torch.int32),
35
+ a_string="hello",
36
+ list_of_tensors=[
37
+ torch.rand((1, 10), dtype=torch.float32),
38
+ torch.rand((3, 5, 4000), dtype=torch.float64),
39
+ torch.tensor(1984), # test scalar too
40
+ ],
41
+ numpy_array=np.arange(512),
42
+ unrecognized=UnrecognizedType(33),
43
+ )
44
+
45
+ encoder = MsgpackEncoder()
46
+ decoder = MsgpackDecoder(MyType)
47
+
48
+ encoded = encoder.encode(obj)
49
+
50
+ # There should be the main buffer + 2 large tensor buffers
51
+ # + 1 large numpy array. "large" is <= 256 bytes.
52
+ # The two small tensors are encoded inline.
53
+ assert len(encoded) == 4
54
+
55
+ decoded: MyType = decoder.decode(encoded)
56
+
57
+ assert_equal(decoded, obj)
58
+
59
+ # Test encode_into case
60
+
61
+ preallocated = bytearray()
62
+
63
+ encoded2 = encoder.encode_into(obj, preallocated)
64
+
65
+ assert len(encoded2) == 4
66
+ assert encoded2[0] is preallocated
67
+
68
+ decoded2: MyType = decoder.decode(encoded2)
69
+
70
+ assert_equal(decoded2, obj)
71
+
72
+
73
+def assert_equal(obj1: MyType, obj2: MyType):
74
+ assert torch.equal(obj1.tensor1, obj2.tensor1)
75
+ assert obj1.a_string == obj2.a_string
76
+ assert all(
77
+ torch.equal(a, b)
78
+ for a, b in zip(obj1.list_of_tensors, obj2.list_of_tensors))
79
+ assert np.array_equal(obj1.numpy_array, obj2.numpy_array)
80
+ assert obj1.unrecognized.an_int == obj2.unrecognized.an_int
vllm/v1/engine/core.py
+4
-4
@@ -490,14 +490,14 @@ class EngineCoreProc(EngineCore):
490
491
while True:
492
# (RequestType, RequestData)
493
- type_frame, data_frame = socket.recv_multipart(copy=False)
493
+ type_frame, *data_frames = socket.recv_multipart(copy=False)
494
request_type = EngineCoreRequestType(bytes(type_frame.buffer))
495
496
# Deserialize the request data.
497
decoder = add_request_decoder if (
498
request_type
499
== EngineCoreRequestType.ADD) else generic_decoder
500
- request = decoder.decode(data_frame.buffer)
500
+ request = decoder.decode(data_frames)
501
502
# Push to input queue for core busy loop.
503
self.input_queue.put_nowait((request_type, request))
@@ -514,8 +514,8 @@ class EngineCoreProc(EngineCore):
514
while True:
515
outputs = self.output_queue.get()
516
outputs.engine_index = engine_index
517
- encoder.encode_into(outputs, buffer)
518
- socket.send(buffer, copy=False)
517
+ buffers = encoder.encode_into(outputs, buffer)
518
+ socket.send_multipart(buffers, copy=False)
519
520
521
ENGINE_PAUSED_OUTPUTS = EngineCoreOutputs(engine_paused=True)
vllm/v1/engine/core_client.py
+13
-13
@@ -26,7 +26,7 @@ from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,
26
EngineCoreRequestType, UtilityOutput)
27
from vllm.v1.engine.core import EngineCore, EngineCoreProc
28
from vllm.v1.executor.abstract import Executor
29
-from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
29
+from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder, bytestr
30
from vllm.v1.utils import BackgroundProcHandle
31
32
logger = init_logger(__name__)
@@ -505,8 +505,8 @@ class SyncMPClient(MPClient):
505
# shutdown signal, exit thread.
506
break
507
508
- frame = out_socket.recv(copy=False)
509
- outputs = decoder.decode(frame.buffer)
508
+ frames = out_socket.recv_multipart(copy=False)
509
+ outputs = decoder.decode(frames)
510
if outputs.utility_output:
511
_process_utility_output(outputs.utility_output,
512
utility_results)
@@ -529,7 +529,7 @@ class SyncMPClient(MPClient):
529
def _send_input(self, request_type: EngineCoreRequestType, request: Any):
530
# (Identity, RequestType, SerializedRequest)
531
msg = (self.core_engine.identity, request_type.value,
532
- self.encoder.encode(request))
532
+ *self.encoder.encode(request))
533
self.input_socket.send_multipart(msg, copy=False)
534
535
def call_utility(self, method: str, *args) -> Any:
@@ -633,8 +633,8 @@ class AsyncMPClient(MPClient):
633
634
async def process_outputs_socket():
635
while True:
636
- (frame, ) = await output_socket.recv_multipart(copy=False)
637
- outputs: EngineCoreOutputs = decoder.decode(frame.buffer)
636
+ frames = await output_socket.recv_multipart(copy=False)
637
+ outputs: EngineCoreOutputs = decoder.decode(frames)
638
if outputs.utility_output:
639
_process_utility_output(outputs.utility_output,
640
utility_results)
@@ -666,12 +666,12 @@ class AsyncMPClient(MPClient):
666
if engine is None:
667
engine = self.core_engine
668
669
- message = (request_type.value, self.encoder.encode(request))
669
+ message = (request_type.value, *self.encoder.encode(request))
670
return self._send_input_message(message, engine)
671
672
- def _send_input_message(self, message: tuple[bytes, bytes],
672
+ def _send_input_message(self, message: tuple[bytestr, ...],
673
engine: CoreEngine) -> Awaitable[None]:
674
- message = (engine.identity, ) + message # type: ignore[assignment]
674
+ message = (engine.identity, ) + message
675
return self.input_socket.send_multipart(message, copy=False)
676
677
async def call_utility_async(self, method: str, *args) -> Any:
@@ -684,8 +684,8 @@ class AsyncMPClient(MPClient):
684
call_id = uuid.uuid1().int >> 64
685
future = asyncio.get_running_loop().create_future()
686
self.utility_results[call_id] = future
687
- message = (EngineCoreRequestType.UTILITY.value,
688
- self.encoder.encode((call_id, method, args)))
687
+ message = (EngineCoreRequestType.UTILITY.value, *self.encoder.encode(
688
+ (call_id, method, args)))
689
await self._send_input_message(message, engine)
690
self._ensure_output_queue_task()
691
return await future
@@ -760,7 +760,7 @@ class DPAsyncMPClient(AsyncMPClient):
760
761
# Control message used for triggering dp idle mode loop.
762
self.start_dp_msg = (EngineCoreRequestType.START_DP.value,
763
- self.encoder.encode(None))
763
+ *self.encoder.encode(None))
764
765
self.num_engines_running = 0
766
self.reqs_in_flight: dict[str, CoreEngine] = {}
@@ -794,7 +794,7 @@ class DPAsyncMPClient(AsyncMPClient):
794
# tokenized.
795
request.prompt = None
796
797
- msg = (EngineCoreRequestType.ADD.value, self.encoder.encode(request))
797
+ msg = (EngineCoreRequestType.ADD.value, *self.encoder.encode(request))
798
799
chosen_engine = self.get_core_engine_for_request()
800
self.reqs_in_flight[request.request_id] = chosen_engine
vllm/v1/serial_utils.py
+120
-41
@@ -1,61 +1,140 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
import pickle
4
+from collections.abc import Sequence
5
+from inspect import isclass
6
from types import FunctionType
5
-from typing import Any, Optional
7
+from typing import Any, Optional, Union
8
9
import cloudpickle
10
+import numpy as np
11
import torch
12
+import zmq
13
from msgspec import msgpack
14
11
-CUSTOM_TYPE_TENSOR = 1
12
-CUSTOM_TYPE_PICKLE = 2
13
-CUSTOM_TYPE_CLOUDPICKLE = 3
15
+CUSTOM_TYPE_PICKLE = 1
16
+CUSTOM_TYPE_CLOUDPICKLE = 2
17
18
+# TODO calibrate this size
19
+INLINE_BUF_SIZE_THRESHOLD = 256
20
16
-class MsgpackEncoder:
17
- """Encoder with custom torch tensor serialization."""
21
+bytestr = Union[bytes, bytearray, memoryview, zmq.Frame]
22
19
- def __init__(self):
20
- self.encoder = msgpack.Encoder(enc_hook=custom_enc_hook)
23
22
- def encode(self, obj: Any) -> bytes:
23
- return self.encoder.encode(obj)
24
+class MsgpackEncoder:
25
+ """Encoder with custom torch tensor and numpy array serialization.
26
25
- def encode_into(self, obj: Any, buf: bytearray) -> None:
26
- self.encoder.encode_into(obj, buf)
27
+ Note that unlike vanilla `msgspec` Encoders, this interface is generally
28
+ not thread-safe when encoding tensors / numpy arrays.
29
+ """
30
+
31
+ def __init__(self):
32
+ self.encoder = msgpack.Encoder(enc_hook=self.enc_hook)
33
+ # This is used as a local stash of buffers that we can then access from
34
+ # our custom `msgspec` hook, `enc_hook`. We don't have a way to
35
+ # pass custom data to the hook otherwise.
36
+ self.aux_buffers: Optional[list[bytestr]] = None
37
+
38
+ def encode(self, obj: Any) -> Sequence[bytestr]:
39
+ try:
40
+ self.aux_buffers = bufs = [b'']
41
+ bufs[0] = self.encoder.encode(obj)
42
+ # This `bufs` list allows us to collect direct pointers to backing
43
+ # buffers of tensors and np arrays, and return them along with the
44
+ # top-level encoded buffer instead of copying their data into the
45
+ # new buffer.
46
+ return bufs
47
+ finally:
48
+ self.aux_buffers = None
49
+
50
+ def encode_into(self, obj: Any, buf: bytearray) -> Sequence[bytestr]:
51
+ try:
52
+ self.aux_buffers = [buf]
53
+ bufs = self.aux_buffers
54
+ self.encoder.encode_into(obj, buf)
55
+ return bufs
56
+ finally:
57
+ self.aux_buffers = None
58
+
59
+ def enc_hook(self, obj: Any) -> Any:
60
+ if isinstance(obj, torch.Tensor):
61
+ return self._encode_ndarray(obj.numpy())
62
+
63
+ # Fall back to pickle for object or void kind ndarrays.
64
+ if isinstance(obj, np.ndarray) and obj.dtype.kind not in ('O', 'V'):
65
+ return self._encode_ndarray(obj)
66
+
67
+ if isinstance(obj, FunctionType):
68
+ # `pickle` is generally faster than cloudpickle, but can have
69
+ # problems serializing methods.
70
+ return msgpack.Ext(CUSTOM_TYPE_CLOUDPICKLE, cloudpickle.dumps(obj))
71
+
72
+ return msgpack.Ext(CUSTOM_TYPE_PICKLE,
73
+ pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL))
74
+
75
+ def _encode_ndarray(
76
+ self, obj: np.ndarray
77
+ ) -> tuple[str, tuple[int, ...], Union[int, memoryview]]:
78
+ assert self.aux_buffers is not None
79
+ if not obj.shape or obj.nbytes < INLINE_BUF_SIZE_THRESHOLD:
80
+ # Encode small arrays and scalars inline.
81
+ data = obj.data
82
+ else:
83
+ # Otherwise encode index of backing buffer.
84
+ obj = np.ascontiguousarray(obj)
85
+ data = len(self.aux_buffers)
86
+ self.aux_buffers.append(obj.data)
87
+ # We serialize the ndarray as a tuple of native types.
88
+ # The data is either inlined if small, or an index into a list of
89
+ # backing buffers that we've stashed in `aux_buffers`.
90
+ return obj.dtype.str, obj.shape, data
91
92
93
class MsgpackDecoder:
30
- """Decoder with custom torch tensor serialization."""
94
+ """Decoder with custom torch tensor and numpy array serialization.
95
+
96
+ Note that unlike vanilla `msgspec` Decoders, this interface is generally
97
+ not thread-safe when encoding tensors / numpy arrays.
98
+ """
99
100
def __init__(self, t: Optional[Any] = None):
101
args = () if t is None else (t, )
34
- self.decoder = msgpack.Decoder(*args, ext_hook=custom_ext_hook)
35
-
36
- def decode(self, obj: Any):
37
- return self.decoder.decode(obj)
38
-
39
-
40
-def custom_enc_hook(obj: Any) -> Any:
41
- if isinstance(obj, torch.Tensor):
42
- # NOTE(rob): it is fastest to use numpy + pickle
43
- # when serializing torch tensors.
44
- # https://gist.github.com/tlrmchlsmth/8067f1b24a82b6e2f90450e7764fa103 # noqa: E501
45
- return msgpack.Ext(CUSTOM_TYPE_TENSOR, pickle.dumps(obj.numpy()))
46
-
47
- if isinstance(obj, FunctionType):
48
- return msgpack.Ext(CUSTOM_TYPE_CLOUDPICKLE, cloudpickle.dumps(obj))
49
-
50
- return msgpack.Ext(CUSTOM_TYPE_PICKLE, pickle.dumps(obj))
51
-
52
-
53
-def custom_ext_hook(code: int, data: memoryview) -> Any:
54
- if code == CUSTOM_TYPE_TENSOR:
55
- return torch.from_numpy(pickle.loads(data))
56
- if code == CUSTOM_TYPE_PICKLE:
57
- return pickle.loads(data)
58
- if code == CUSTOM_TYPE_CLOUDPICKLE:
59
- return cloudpickle.loads(data)
60
-
61
- raise NotImplementedError(f"Extension type code {code} is not supported")
102
+ self.decoder = msgpack.Decoder(*args,
103
+ ext_hook=self.ext_hook,
104
+ dec_hook=self.dec_hook)
105
+ self.aux_buffers: Sequence[bytestr] = ()
106
+
107
+ def decode(self, bufs: Union[bytestr, Sequence[bytestr]]) -> Any:
108
+ if isinstance(bufs, (bytes, bytearray, memoryview, zmq.Frame)):
109
+ # TODO - This check can become `isinstance(bufs, bytestr)`
110
+ # as of Python 3.10.
111
+ return self.decoder.decode(bufs)
112
+
113
+ self.aux_buffers = bufs
114
+ try:
115
+ return self.decoder.decode(bufs[0])
116
+ finally:
117
+ self.aux_buffers = ()
118
+
119
+ def dec_hook(self, t: type, obj: Any) -> Any:
120
+ # Given native types in `obj`, convert to type `t`.
121
+ if isclass(t):
122
+ if issubclass(t, np.ndarray):
123
+ return self._decode_ndarray(obj)
124
+ if issubclass(t, torch.Tensor):
125
+ return torch.from_numpy(self._decode_ndarray(obj))
126
+ return obj
127
+
128
+ def _decode_ndarray(self, arr: Any) -> np.ndarray:
129
+ dtype, shape, data = arr
130
+ buffer = self.aux_buffers[data] if isinstance(data, int) else data
131
+ return np.ndarray(buffer=buffer, dtype=np.dtype(dtype), shape=shape)
132
+
133
+ def ext_hook(self, code: int, data: memoryview) -> Any:
134
+ if code == CUSTOM_TYPE_PICKLE:
135
+ return pickle.loads(data)
136
+ if code == CUSTOM_TYPE_CLOUDPICKLE:
137
+ return cloudpickle.loads(data)
138
+
139
+ raise NotImplementedError(
140
+ f"Extension type code {code} is not supported")