[VLM] avoid extra cuda-ipc staging for preprocessed input (#26096)

python/sglang/srt/managers/schedule_batch.py

@@ -347,18 +347,23 @@ class MultimodalDataItem:
         ret.validate()
         return ret
 
-    def reconstruct(self):
+    def has_cuda_ipc_proxy(self):
-        if not isinstance(self.feature, CudaIpcTensorTransportProxy):
+        return (
-            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(
+                self.precomputed_embeddings.reconstruct_on_target_device(target_device)
-                    reconstruct_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
+    ``BaseMultimodalProcessor.process_mm_data_async``.  Preprocessed inputs may
-    ``MultimodalInputs``, items here do NOT carry pad_value or hash yet.
+    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()
+            if mm_item.has_cuda_ipc_proxy():
-
+                if reconstruct_device is None:
-        ret = MultimodalInputs(
+                    reconstruct_device = torch.cuda.current_device()
-            mm_items=mm_items,
+                mm_item.reconstruct(reconstruct_device)
-            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 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)

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:
+                if isinstance(item.feature, torch.Tensor):
-                    sync_flag, available_slice, byte_offset = (
+                    item.feature = self._wrap_tensor_for_cuda_ipc(item.feature)
-                        self.cudaipc_mmfeature_pool.return_a_slice_tensor_with_flag(
+                if isinstance(item.precomputed_embeddings, torch.Tensor):
-                            item.feature
+                    item.precomputed_embeddings = self._wrap_tensor_for_cuda_ipc(
-                        )
-                    )
-                    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

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)