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
[BugFix] Handle non-contiguous tensors properly when serializing (#16492)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Apr 11, 2025 at 17:54 UTC
41cc883c293eebc0aecd5b078f19517f44a6d871
2 files changed
+30
-11
tests/v1/test_serial_utils.py
+18
-4
@@ -22,6 +22,10 @@ class MyType:
22
list_of_tensors: list[torch.Tensor]
23
numpy_array: np.ndarray
24
unrecognized: UnrecognizedType
25
+ small_f_contig_tensor: torch.Tensor
26
+ large_f_contig_tensor: torch.Tensor
27
+ small_non_contig_tensor: torch.Tensor
28
+ large_non_contig_tensor: torch.Tensor
29
30
31
def test_encode_decode():
@@ -40,6 +44,10 @@ def test_encode_decode():
44
],
45
numpy_array=np.arange(512),
46
unrecognized=UnrecognizedType(33),
47
+ small_f_contig_tensor=torch.rand(5, 4).t(),
48
+ large_f_contig_tensor=torch.rand(1024, 4).t(),
49
+ small_non_contig_tensor=torch.rand(2, 4)[:, 1:3],
50
+ large_non_contig_tensor=torch.rand(1024, 512)[:, 10:20],
51
)
52
53
encoder = MsgpackEncoder()
@@ -47,10 +55,10 @@ def test_encode_decode():
55
56
encoded = encoder.encode(obj)
57
50
- # There should be the main buffer + 2 large tensor buffers
51
- # + 1 large numpy array. "large" is <= 256 bytes.
58
+ # There should be the main buffer + 4 large tensor buffers
59
+ # + 1 large numpy array. "large" is <= 512 bytes.
60
# The two small tensors are encoded inline.
53
- assert len(encoded) == 4
61
+ assert len(encoded) == 6
62
63
decoded: MyType = decoder.decode(encoded)
64
@@ -62,7 +70,7 @@ def test_encode_decode():
70
71
encoded2 = encoder.encode_into(obj, preallocated)
72
65
- assert len(encoded2) == 4
73
+ assert len(encoded2) == 6
74
assert encoded2[0] is preallocated
75
76
decoded2: MyType = decoder.decode(encoded2)
@@ -78,3 +86,9 @@ def assert_equal(obj1: MyType, obj2: MyType):
86
for a, b in zip(obj1.list_of_tensors, obj2.list_of_tensors))
87
assert np.array_equal(obj1.numpy_array, obj2.numpy_array)
88
assert obj1.unrecognized.an_int == obj2.unrecognized.an_int
89
+ assert torch.equal(obj1.small_f_contig_tensor, obj2.small_f_contig_tensor)
90
+ assert torch.equal(obj1.large_f_contig_tensor, obj2.large_f_contig_tensor)
91
+ assert torch.equal(obj1.small_non_contig_tensor,
92
+ obj2.small_non_contig_tensor)
93
+ assert torch.equal(obj1.large_non_contig_tensor,
94
+ obj2.large_non_contig_tensor)
vllm/v1/serial_utils.py
+12
-7
@@ -14,9 +14,10 @@ from msgspec import msgpack
14
15
CUSTOM_TYPE_PICKLE = 1
16
CUSTOM_TYPE_CLOUDPICKLE = 2
17
+CUSTOM_TYPE_RAW_VIEW = 3
18
19
# TODO calibrate this size
19
-INLINE_BUF_SIZE_THRESHOLD = 256
20
+MIN_NOCOPY_BUF_SIZE = 512
21
22
bytestr = Union[bytes, bytearray, memoryview, zmq.Frame]
23
@@ -76,14 +77,16 @@ class MsgpackEncoder:
77
self, obj: np.ndarray
78
) -> tuple[str, tuple[int, ...], Union[int, memoryview]]:
79
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
80
+ arr_data = obj.data if obj.data.c_contiguous else obj.tobytes()
81
+ if not obj.shape or obj.nbytes < MIN_NOCOPY_BUF_SIZE:
82
+ # Encode small arrays and scalars inline. Using this extension type
83
+ # ensures we can avoid copying when decoding.
84
+ data = msgpack.Ext(CUSTOM_TYPE_RAW_VIEW, arr_data)
85
else:
83
- # Otherwise encode index of backing buffer.
84
- obj = np.ascontiguousarray(obj)
86
+ # Otherwise encode index of backing buffer to avoid copy.
87
data = len(self.aux_buffers)
86
- self.aux_buffers.append(obj.data)
88
+ self.aux_buffers.append(arr_data)
89
+
90
# We serialize the ndarray as a tuple of native types.
91
# The data is either inlined if small, or an index into a list of
92
# backing buffers that we've stashed in `aux_buffers`.
@@ -131,6 +134,8 @@ class MsgpackDecoder:
134
return np.ndarray(buffer=buffer, dtype=np.dtype(dtype), shape=shape)
135
136
def ext_hook(self, code: int, data: memoryview) -> Any:
137
+ if code == CUSTOM_TYPE_RAW_VIEW:
138
+ return data
139
if code == CUSTOM_TYPE_PICKLE:
140
return pickle.loads(data)
141
if code == CUSTOM_TYPE_CLOUDPICKLE: