chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

11
python/sglang/srt/managers/mm_utils.py
View File

@@ -13,6 +13,7 @@ from torch import nn
from sglang.srt.layers.multimodal import gpu_tensor_hash
from sglang.srt.managers.schedule_batch import (
CudaIpcTensorTransportProxy,
Modality,
MultimodalDataItem,
MultimodalInputs,
@@ -77,7 +78,6 @@ class TransportProxyTensor(torch.Tensor):
"tensor_data": None,
"ipc_extra": None,
}
transport_mode = self._metadata.get("transport_mode", "default")
if transport_mode == "cuda_ipc" and self.is_cuda:
@@ -91,6 +91,7 @@ class TransportProxyTensor(torch.Tensor):
"dtype": self.dtype,
"stride": self.stride(),
"device_index": self.device.index,
"storage_offset": self.storage_offset(),
}
state["tensor_data"] = None
except Exception as e:
@@ -113,12 +114,13 @@ class TransportProxyTensor(torch.Tensor):
if transport_mode == "cuda_ipc" and state["ipc_extra"] is not None:
ipc_extra = state["ipc_extra"]
handle, shape, dtype, stride, source_device_index = (
handle, shape, dtype, stride, source_device_index, s_offset = (
ipc_extra["handle"],
ipc_extra["shape"],
ipc_extra["dtype"],
ipc_extra["stride"],
ipc_extra["device_index"],
ipc_extra["storage_offset"],
)
try:
@@ -127,7 +129,7 @@ class TransportProxyTensor(torch.Tensor):
storage = torch.UntypedStorage._new_shared_cuda(*handle)
reconstructed_tensor = torch.empty(
0, dtype=dtype, device=target_device
).set_(storage, storage_offset=0, size=shape, stride=stride)
).set_(storage, storage_offset=s_offset, size=shape, stride=stride)
self.set_(reconstructed_tensor)
except Exception as e:
print(f"Error: Failed to deserialize from CUDA IPC handle ({e}).")
@@ -811,4 +813,7 @@ def hash_feature(f):
return data_hash(arr_bytes)
elif isinstance(f, torch.Tensor):
return tensor_hash([f])
elif isinstance(f, CudaIpcTensorTransportProxy):
reconstruct_t = f.reconstruct_on_target_device(torch.cuda.current_device())
return tensor_hash([reconstruct_t])
return data_hash(f)