[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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user