Skip to content

Commit 29daef4

Browse files
p88hnjhill
andcommitted
Apply suggestions from code review
Co-authored-by: Nick Hill <[email protected]> Signed-off-by: Staszek Pasko <[email protected]>
1 parent 49941ea commit 29daef4

File tree

1 file changed

+7
-5
lines changed

1 file changed

+7
-5
lines changed

vllm/v1/serial_utils.py

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -140,15 +140,15 @@ def _encode_tensor(
140140
# this creates a copy of the tensor
141141
obj = obj.contiguous() if not obj.is_contiguous() else obj
142142
# view the tensor as a 1D array of bytes
143-
arr = obj.view([obj.numel()]).view(torch.uint8).numpy()
143+
arr = obj.view((obj.numel(),)).view(torch.uint8).numpy()
144144
if obj.nbytes < self.size_threshold:
145145
data = msgpack.Ext(CUSTOM_TYPE_RAW_VIEW, arr.data)
146146
else:
147147
# Otherwise encode index of backing buffer to avoid copy.
148148
data = len(self.aux_buffers)
149149
self.aux_buffers.append(arr.data)
150-
dt = str(obj.dtype)[6:] # remove 'torch.' prefix
151-
return dt, obj.shape, data
150+
dtype = str(obj.dtype)[6:] # remove 'torch.' prefix
151+
return dtype, obj.shape, data
152152

153153
def _encode_nested_tensors(self, nt: NestedTensors) -> Any:
154154
if isinstance(nt, torch.Tensor):
@@ -228,8 +228,10 @@ def _decode_tensor(self, arr: Any) -> torch.Tensor:
228228
# the returned memory is non-writeable.
229229
buffer = self.aux_buffers[data] if isinstance(data, int) \
230230
else bytearray(data)
231-
arr = np.ndarray(buffer=buffer, dtype=np.uint8, shape=[len(buffer)])
232-
return torch.from_numpy(arr).view(getattr(torch, dtype)).view(shape)
231+
arr = np.ndarray(buffer=buffer, dtype=np.uint8, shape=(len(buffer),))
232+
torch_dtype = getattr(torch, dtype)
233+
assert isinstance(torch_dtype, torch.dtype)
234+
return torch.from_numpy(arr).view(torch_dtype).view(shape)
233235

234236
def _decode_mm_items(self, obj: list) -> list[MultiModalKwargsItem]:
235237
decoded_items = []

0 commit comments

Comments
 (0)