From b6f71d5850eadd97de8a7d569e7271a58a9ffa2e Mon Sep 17 00:00:00 2001 From: Mick Date: Sun, 24 May 2026 19:48:06 +0800 Subject: [PATCH] [VLM] avoid extra cuda-ipc staging for preprocessed input (#26096) --- python/sglang/srt/managers/schedule_batch.py | 62 +++++++------ .../multimodal/processors/base_processor.py | 92 ++++++++----------- test/manual/vlm/test_mm_utils.py | 52 +++++++++++ 3 files changed, 127 insertions(+), 79 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index e64d0e945..ce81a8cad 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -347,18 +347,23 @@ class MultimodalDataItem: ret.validate() return ret - def reconstruct(self): - if not isinstance(self.feature, CudaIpcTensorTransportProxy): - return + def has_cuda_ipc_proxy(self): + return ( + isinstance(self.feature, CudaIpcTensorTransportProxy) + or isinstance(self.precomputed_embeddings, CudaIpcTensorTransportProxy) + or any( + isinstance(value, CudaIpcTensorTransportProxy) + for value in self.model_specific_data.values() + ) + ) - reconstruct_device = torch.cuda.current_device() + def reconstruct(self, target_device: int): + """materialize cuda ipc proxy tensors in-place on target_device""" if isinstance(self.feature, CudaIpcTensorTransportProxy): - self.feature = self.feature.reconstruct_on_target_device(reconstruct_device) + self.feature = self.feature.reconstruct_on_target_device(target_device) if isinstance(self.precomputed_embeddings, CudaIpcTensorTransportProxy): self.precomputed_embeddings = ( - self.precomputed_embeddings.reconstruct_on_target_device( - reconstruct_device - ) + self.precomputed_embeddings.reconstruct_on_target_device(target_device) ) for extra_key in self.model_specific_data: if isinstance( @@ -366,17 +371,18 @@ class MultimodalDataItem: ): extra_data = self.model_specific_data[ extra_key - ].reconstruct_on_target_device(reconstruct_device) + ].reconstruct_on_target_device(target_device) self.model_specific_data[extra_key] = extra_data @dataclasses.dataclass class MultimodalProcessorOutput: - """Raw output from multimodal processors, before pad/hash computation. + """Raw output from multimodal processors before scheduler-side preparation (pad, hash). This is the typed replacement for the dict previously returned by - ``BaseMultimodalProcessor.process_mm_data_async``. Unlike - ``MultimodalInputs``, items here do NOT carry pad_value or hash yet. + ``BaseMultimodalProcessor.process_mm_data_async``. Preprocessed inputs may + already carry ``pad_value`` and ``hash`` to avoid hashing the same tensor once + per scheduler TP rank. """ mm_items: List[MultimodalDataItem] @@ -496,16 +502,16 @@ class MultimodalInputs: @staticmethod def from_processor_output(obj: "MultimodalProcessorOutput"): mm_items = obj.mm_items + assert isinstance(mm_items, list) + mm_items = [item for item in mm_items if item.is_valid()] + + # try reconstructing from cuda-ipc + reconstruct_device = None for mm_item in mm_items: - mm_item.reconstruct() - - ret = MultimodalInputs( - mm_items=mm_items, - padded_input_ids=obj.padded_input_ids, - ) - - assert isinstance(ret.mm_items, list) - ret.mm_items = [item for item in ret.mm_items if item.is_valid()] + if mm_item.has_cuda_ipc_proxy(): + if reconstruct_device is None: + reconstruct_device = torch.cuda.current_device() + mm_item.reconstruct(reconstruct_device) if envs.SGLANG_MM_BUFFER_SIZE_MB.get() > 0: # Multi-modal feature hashing optimization: @@ -522,19 +528,23 @@ class MultimodalInputs: if not is_feature_buffer_initialized(): init_feature_buffer(device) reset_buffer_offset() - for item in ret.mm_items: + for item in mm_items: if item.feature is not None: if isinstance(item.feature, torch.Tensor): item.feature = try_add_to_buffer(item.feature) - for item in ret.mm_items: + for item in mm_items: item.set_pad_value() if envs.SGLANG_MM_BUFFER_SIZE_MB.get() > 0: - for item in ret.mm_items: + for item in mm_items: if item.feature is not None: item.feature = item.feature.to("cpu", non_blocking=True) + mm_inputs = MultimodalInputs( + mm_items=mm_items, + padded_input_ids=obj.padded_input_ids, + ) optional_args = [ "mrope_positions", "mrope_position_delta", @@ -554,9 +564,9 @@ class MultimodalInputs: for arg in optional_args: val = getattr(obj, arg, None) if val is not None: - setattr(ret, arg, val) + setattr(mm_inputs, arg, val) - return ret + return mm_inputs def contains_image_inputs(self) -> bool: return any(item.is_image() for item in self.mm_items) diff --git a/python/sglang/srt/multimodal/processors/base_processor.py b/python/sglang/srt/multimodal/processors/base_processor.py index bd92d1e42..d3bb63f5f 100644 --- a/python/sglang/srt/multimodal/processors/base_processor.py +++ b/python/sglang/srt/multimodal/processors/base_processor.py @@ -1193,6 +1193,32 @@ class BaseMultimodalProcessor(ABC): return input_ids.flatten().to(dtype=torch.long) return torch.tensor(input_ids, dtype=torch.long).flatten() + def _wrap_tensor_for_cuda_ipc(self, tensor: torch.Tensor): + """helper function to turn a tensor into a cuda-ipc tensor""" + if not tensor.is_cuda: + return tensor + + sync_flag, available_slice, byte_offset = ( + self.cudaipc_mmfeature_pool.return_a_slice_tensor_with_flag(tensor) + ) + if isinstance(available_slice, torch.Tensor): + available_slice.copy_(tensor.view(torch.int8).view(-1), non_blocking=True) + return CudaIpcTensorTransportProxy( + data=available_slice, + info_data=tensor, + sync_buffer_meta=sync_flag, + pool_ipc_handle=( + self.cudaipc_mmfeature_pool._pool_ipc_handle + if _IPC_POOL_HANDLE_CACHE + else None + ), + pool_byte_offset=byte_offset, + pool_device_index=self.cudaipc_mmfeature_pool._pool_device_index, + ) + if self.server_args.keep_mm_feature_on_device: + return tensor + return tensor.cpu() + def process_and_combine_mm_data( self, base_output: BaseMultiModalProcessorOutput, @@ -1304,6 +1330,13 @@ class BaseMultimodalProcessor(ABC): all_collected_items = get_new_expanded_mm_items(all_collected_items) + for item in all_collected_items: + if item.format in ( + MultimodalInputFormat.PROCESSOR_OUTPUT, + MultimodalInputFormat.PRECOMPUTED_EMBEDDING, + ): + item.set_pad_value() + """ solution for cuda-ipc memory-leak: 1. memory-pool: each time get a slice from memory-pool and use it as transport-data (with async lock guard) @@ -1313,60 +1346,13 @@ class BaseMultimodalProcessor(ABC): """ if SGL_USE_CUDA_IPC: - # post-process + # post-process, prepare for cuda-ipc transfer for item in all_collected_items: - if isinstance(item.feature, torch.Tensor) and item.feature.is_cuda: - sync_flag, available_slice, byte_offset = ( - self.cudaipc_mmfeature_pool.return_a_slice_tensor_with_flag( - item.feature - ) + if isinstance(item.feature, torch.Tensor): + item.feature = self._wrap_tensor_for_cuda_ipc(item.feature) + if isinstance(item.precomputed_embeddings, torch.Tensor): + item.precomputed_embeddings = self._wrap_tensor_for_cuda_ipc( + item.precomputed_embeddings ) - if isinstance(available_slice, torch.Tensor): - available_slice.copy_( - item.feature.view(torch.int8).view(-1), non_blocking=True - ) - item.feature = CudaIpcTensorTransportProxy( - data=available_slice, - info_data=item.feature, - sync_buffer_meta=sync_flag, - pool_ipc_handle=( - self.cudaipc_mmfeature_pool._pool_ipc_handle - if _IPC_POOL_HANDLE_CACHE - else None - ), - pool_byte_offset=byte_offset, - pool_device_index=self.cudaipc_mmfeature_pool._pool_device_index, - ) - elif not self.server_args.keep_mm_feature_on_device: - item.feature = item.feature.cpu() - elif ( - isinstance(item.precomputed_embeddings, torch.Tensor) - and item.precomputed_embeddings.is_cuda - ): - - sync_flag, available_slice, byte_offset = ( - self.cudaipc_mmfeature_pool.return_a_slice_tensor_with_flag( - item.precomputed_embeddings - ) - ) - if isinstance(available_slice, torch.Tensor): - available_slice.copy_( - item.precomputed_embeddings.view(torch.int8).view(-1), - non_blocking=True, - ) - item.precomputed_embeddings = CudaIpcTensorTransportProxy( - data=available_slice, - info_data=item.precomputed_embeddings, - sync_buffer_meta=sync_flag, - pool_ipc_handle=( - self.cudaipc_mmfeature_pool._pool_ipc_handle - if _IPC_POOL_HANDLE_CACHE - else None - ), - pool_byte_offset=byte_offset, - pool_device_index=self.cudaipc_mmfeature_pool._pool_device_index, - ) - elif not self.server_args.keep_mm_feature_on_device: - item.precomputed_embeddings = item.precomputed_embeddings.cpu() return all_collected_items, input_ids, ret diff --git a/test/manual/vlm/test_mm_utils.py b/test/manual/vlm/test_mm_utils.py index c526c4324..9302c40a6 100644 --- a/test/manual/vlm/test_mm_utils.py +++ b/test/manual/vlm/test_mm_utils.py @@ -45,6 +45,58 @@ class TestMultimodalInputsFromDict(unittest.TestCase): self.assertTrue(torch.equal(mm_inputs.mm_items[0].feature, feature_tensor)) proxy_feature.reconstruct_on_target_device.assert_called_once_with(0) + def test_materialize_precomputed_embedding_proxy_without_feature(self): + embedding_tensor = torch.tensor([[1.0, 2.0]], dtype=torch.float32) + proxy_embedding = _make_proxy_with_reconstruct_result(embedding_tensor) + mm_item = MultimodalDataItem( + modality=Modality.IMAGE, + offsets=[(0, 1)], + precomputed_embeddings=proxy_embedding, + ) + + with ( + patch.object(schedule_batch.torch.cuda, "is_available", return_value=True), + patch.object(schedule_batch.torch.cuda, "current_device", return_value=0), + patch.object( + schedule_batch.envs.SGLANG_MM_BUFFER_SIZE_MB, "get", return_value=0 + ), + ): + mm_inputs = MultimodalInputs.from_dict({"mm_items": [mm_item]}) + + self.assertTrue( + torch.equal( + mm_inputs.mm_items[0].precomputed_embeddings, + embedding_tensor, + ) + ) + proxy_embedding.reconstruct_on_target_device.assert_called_once_with(0) + + def test_materialize_model_specific_proxy_without_feature(self): + grid_tensor = torch.tensor([[1, 2, 3]], dtype=torch.int64) + proxy_grid = _make_proxy_with_reconstruct_result(grid_tensor) + mm_item = MultimodalDataItem( + modality=Modality.IMAGE, + offsets=[(0, 1)], + model_specific_data={"image_grid_thw": proxy_grid}, + ) + + with ( + patch.object(schedule_batch.torch.cuda, "is_available", return_value=True), + patch.object(schedule_batch.torch.cuda, "current_device", return_value=0), + patch.object( + schedule_batch.envs.SGLANG_MM_BUFFER_SIZE_MB, "get", return_value=0 + ), + ): + mm_inputs = MultimodalInputs.from_dict({"mm_items": [mm_item]}) + + self.assertTrue( + torch.equal( + mm_inputs.mm_items[0].model_specific_data["image_grid_thw"], + grid_tensor, + ) + ) + proxy_grid.reconstruct_on_target_device.assert_called_once_with(0) + if __name__ == "__main__": unittest.main(verbosity=2)