@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d73a9457a5

Fork 0

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, 08:30 UTC

[Core] Improve Tensor serialisation (#18774)

Signed-off-by: Lukas Geiger <lukas.geiger94@gmail.com>

Lukas Geiger committed May 28, 2025 at 02:46 UTC d73a9457a56e99d33676f668611bfb232357f1b9

1 file changed +9 -8

vllm/v1/serial_utils.py

+9 -8

@@ -158,10 +158,8 @@ class MsgpackEncoder:

158 self, obj: torch.Tensor

159 ) -> tuple[str, tuple[int, ...], Union[int, memoryview]]:

160 assert self.aux_buffers is not None

161 - # this creates a copy of the tensor if it's not already contiguous

162 - obj = obj.contiguous()

161 # view the tensor as a 1D array of bytes

164 - arr = obj.view((obj.numel(), )).view(torch.uint8).numpy()

162 + arr = obj.flatten().view(torch.uint8).numpy()

163 if obj.nbytes < self.size_threshold:

164 # Smaller tensors are encoded inline, just like ndarrays.

165 data = msgpack.Ext(CUSTOM_TYPE_RAW_VIEW, arr.data)

@@ -169,7 +167,7 @@ class MsgpackEncoder:

167 # Otherwise encode index of backing buffer to avoid copy.

168 data = len(self.aux_buffers)

169 self.aux_buffers.append(arr.data)

172 - dtype = str(obj.dtype)[6:] # remove 'torch.' prefix

170 + dtype = str(obj.dtype).removeprefix("torch.")

171 return dtype, obj.shape, data

172

173 def _encode_nested_tensors(self, nt: NestedTensors) -> Any:

@@ -245,7 +243,7 @@ class MsgpackDecoder:

243 # zero-copy decode. We assume the ndarray will not be kept around,

244 # as it now locks the whole received message buffer in memory.

245 buffer = self.aux_buffers[data] if isinstance(data, int) else data

248 - return np.ndarray(buffer=buffer, dtype=np.dtype(dtype), shape=shape)

246 + return np.frombuffer(buffer, dtype=dtype).reshape(shape)

247

248 def _decode_tensor(self, arr: Any) -> torch.Tensor:

249 dtype, shape, data = arr

@@ -254,12 +252,15 @@ class MsgpackDecoder:

252 # not complain about a readonly memoryview.

253 buffer = self.aux_buffers[data] if isinstance(data, int) \

254 else bytearray(data)

257 - # Create numpy wrapper around the bytes

258 - arr = np.ndarray(buffer=buffer, dtype=np.uint8, shape=(len(buffer), ))

255 torch_dtype = getattr(torch, dtype)

256 assert isinstance(torch_dtype, torch.dtype)

257 + if not buffer: # torch.frombuffer doesn't like empty buffers

258 + assert 0 in shape

259 + return torch.empty(shape, dtype=torch_dtype)

260 + # Create uint8 array

261 + arr = torch.frombuffer(buffer, dtype=torch.uint8)

262 # Convert back to proper shape & type

262 - return torch.from_numpy(arr).view(torch_dtype).view(shape)

263 + return arr.view(torch_dtype).view(shape)

264

265 def _decode_mm_items(self, obj: list) -> list[MultiModalKwargsItem]:

266 decoded_items = []