[FEAT] Shared mem pool based cuda ipc for multi-modal data transport (#11917)

Co-authored-by: kousakawang <wanghanpei@bytedance.com>
Co-authored-by: Yuan Luo <4908075+yuan-luo@users.noreply.github.com>
This commit is contained in:
kousakawang
2025-11-02 16:46:37 +08:00
committed by GitHub
parent a920b9dace
commit 7efd8b3d1f
8 changed files with 424 additions and 20 deletions

View File

@@ -82,6 +82,7 @@ from sglang.srt.sampling.sampling_params import SamplingParams
from sglang.srt.server_args import ServerArgs, get_global_server_args
from sglang.srt.utils import flatten_nested_list
from sglang.srt.utils.common import is_npu
from sglang.srt.utils.cuda_ipc_transport_utils import CudaIpcTensorTransportProxy
_is_npu = is_npu()
@@ -1365,6 +1366,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
pixel_values = getattr(mm_item, "feature", None)
if isinstance(pixel_values, torch.Tensor):
mm_item.feature = pixel_values.to(self.device, non_blocking=True)
elif isinstance(pixel_values, CudaIpcTensorTransportProxy):
mm_item.feature = pixel_values.reconstruct_on_target_device(
torch.cuda.current_device()
)
self.multimodal_inputs = multimodal_inputs
self.token_type_ids = token_type_ids_tensor
self.seq_lens_sum = sum(seq_lens)