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 13, 05:45 UTC
[BugFix] Fix multi-modal async scheduling race condition (#28706)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Nov 14, 2025 at 01:11 UTC
bc3e43069aadb1fa301a9f60a22872b6ec4453b9
3 files changed
+43
-31
vllm/distributed/device_communicators/shm_object_storage.py
+3
-3
@@ -342,8 +342,8 @@ class MsgpackSerde(ObjectSerde):
342
from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
343
344
self.encoder = MsgpackEncoder()
345
- self.tensor_decoder = MsgpackDecoder(torch.Tensor)
346
- self.mm_decoder = MsgpackDecoder(MultiModalKwargsItem)
345
+ self.tensor_decoder = MsgpackDecoder(torch.Tensor, share_mem=False)
346
+ self.mm_decoder = MsgpackDecoder(MultiModalKwargsItem, share_mem=False)
347
self._mm_kwargs_item_cls = MultiModalKwargsItem
348
349
def serialize(self, value: Any) -> tuple[bytes | list[bytes], int, bytes, int]:
@@ -368,7 +368,7 @@ class MsgpackSerde(ObjectSerde):
368
# pickle.loads do not read past the end of a pickled object
369
# within a large buffer, so we can skip storing the metadata size
370
type_name, nbytes, len_arr = pickle.loads(data_view)
371
- serialized_data = bytearray(data_view[-nbytes:])
371
+ serialized_data = data_view[-nbytes:]
372
373
if type_name == torch.Tensor.__name__:
374
obj = []
vllm/v1/serial_utils.py
+19
-7
@@ -31,6 +31,7 @@ from vllm.multimodal.inputs import (
31
MultiModalSharedField,
32
NestedTensors,
33
)
34
+from vllm.utils.platform_utils import is_pin_memory_available
35
from vllm.v1.engine import UtilityResult
36
from vllm.v1.utils import tensor_data
37
@@ -282,7 +283,9 @@ class MsgpackDecoder:
283
not thread-safe when encoding tensors / numpy arrays.
284
"""
285
285
- def __init__(self, t: Any | None = None):
286
+ def __init__(self, t: Any | None = None, share_mem: bool = True):
287
+ self.share_mem = share_mem
288
+ self.pin_tensors = is_pin_memory_available()
289
args = () if t is None else (t,)
290
self.decoder = msgpack.Decoder(
291
*args, ext_hook=self.ext_hook, dec_hook=self.dec_hook
@@ -347,21 +350,30 @@ class MsgpackDecoder:
350
# zero-copy decode. We assume the ndarray will not be kept around,
351
# as it now locks the whole received message buffer in memory.
352
buffer = self.aux_buffers[data] if isinstance(data, int) else data
350
- return np.frombuffer(buffer, dtype=dtype).reshape(shape)
353
+ arr = np.frombuffer(buffer, dtype=dtype)
354
+ if not self.share_mem:
355
+ arr = arr.copy()
356
+ return arr.reshape(shape)
357
358
def _decode_tensor(self, arr: Any) -> torch.Tensor:
359
dtype, shape, data = arr
354
- # Copy from inline representation, to decouple the memory storage
355
- # of the message from the original buffer. And also make Torch
356
- # not complain about a readonly memoryview.
357
- buffer = self.aux_buffers[data] if isinstance(data, int) else bytearray(data)
360
+ is_aux = isinstance(data, int)
361
+ buffer = self.aux_buffers[data] if is_aux else data
362
+ buffer = buffer if isinstance(buffer, memoryview) else memoryview(buffer)
363
torch_dtype = getattr(torch, dtype)
364
assert isinstance(torch_dtype, torch.dtype)
360
- if not buffer: # torch.frombuffer doesn't like empty buffers
365
+ if not buffer.nbytes: # torch.frombuffer doesn't like empty buffers
366
assert 0 in shape
367
return torch.empty(shape, dtype=torch_dtype)
368
# Create uint8 array
369
arr = torch.frombuffer(buffer, dtype=torch.uint8)
370
+ # Clone ensures tensor is backed by pytorch-owned memory for safe
371
+ # future async CPU->GPU transfer.
372
+ # Pin larger tensors for more efficient CPU->GPU transfer.
373
+ if not is_aux:
374
+ arr = arr.clone()
375
+ elif not self.share_mem:
376
+ arr = arr.pin_memory() if self.pin_tensors else arr.clone()
377
# Convert back to proper shape & type
378
return arr.view(torch_dtype).view(shape)
379
vllm/v1/worker/gpu_model_runner.py
+21
-21
@@ -2590,28 +2590,28 @@ class GPUModelRunner(
2590
)
2591
)
2592
2593
- dp_rank = self.parallel_config.data_parallel_rank
2594
- if ubatch_slices:
2595
- assert num_tokens_across_dp is not None
2596
- num_input_tokens = int(num_tokens_across_dp[dp_rank].item())
2597
- self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens)
2598
- elif num_tokens_across_dp is not None:
2599
- num_input_tokens = int(num_tokens_across_dp[dp_rank].item())
2600
- else:
2601
- num_input_tokens = self._get_num_input_tokens(
2602
- scheduler_output.total_num_scheduled_tokens
2603
- )
2593
+ dp_rank = self.parallel_config.data_parallel_rank
2594
+ if ubatch_slices:
2595
+ assert num_tokens_across_dp is not None
2596
+ num_input_tokens = int(num_tokens_across_dp[dp_rank].item())
2597
+ self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens)
2598
+ elif num_tokens_across_dp is not None:
2599
+ num_input_tokens = int(num_tokens_across_dp[dp_rank].item())
2600
+ else:
2601
+ num_input_tokens = self._get_num_input_tokens(
2602
+ scheduler_output.total_num_scheduled_tokens
2603
+ )
2604
2605
- (
2606
- input_ids,
2607
- inputs_embeds,
2608
- positions,
2609
- intermediate_tensors,
2610
- model_kwargs,
2611
- ec_connector_output,
2612
- ) = self._preprocess(
2613
- scheduler_output, num_input_tokens, intermediate_tensors
2614
- )
2605
+ (
2606
+ input_ids,
2607
+ inputs_embeds,
2608
+ positions,
2609
+ intermediate_tensors,
2610
+ model_kwargs,
2611
+ ec_connector_output,
2612
+ ) = self._preprocess(
2613
+ scheduler_output, num_input_tokens, intermediate_tensors
2614
+ )
2615
2616
uniform_decode = (
2617
max_num_scheduled_tokens == self.uniform_decode_query_len