chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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

77
python/sglang/srt/multimodal/processors/base_processor.py
View File

@@ -13,10 +13,24 @@ from PIL import Image
from transformers import BaseImageProcessorFast
from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
from sglang.srt.utils import is_npu, load_audio, load_image, load_video, logger
from sglang.srt.utils import (
get_bool_env_var,
is_npu,
load_audio,
load_image,
load_video,
logger,
)
from sglang.srt.utils.cuda_ipc_transport_utils import (
MM_FEATURE_CACHE_SIZE,
CudaIpcTensorTransportProxy,
MmItemMemoryPool,
)
_is_npu = is_npu()
SGL_USE_CUDA_IPC = get_bool_env_var("SGLANG_USE_CUDA_IPC_TRANSPORT")
@dataclasses.dataclass
class BaseMultiModalProcessorOutput:
@@ -210,6 +224,9 @@ class BaseMultimodalProcessor(ABC):
"input_features",
]
if SGL_USE_CUDA_IPC:
self.cudaipc_mmfeature_pool = MmItemMemoryPool(MM_FEATURE_CACHE_SIZE)
def process_mm_data(
self, input_text, images=None, videos=None, audios=None, **kwargs
) -> dict:
@@ -254,10 +271,13 @@ class BaseMultimodalProcessor(ABC):
if not self.server_args.keep_mm_feature_on_device:
# move feature tensors to cpu
for feature_name in self.FEATURE_NAMES:
if feature_name in result and isinstance(
result[feature_name], torch.Tensor
):
result[feature_name] = result[feature_name].to("cpu")
if SGL_USE_CUDA_IPC:
pass
else:
if feature_name in result and isinstance(
result[feature_name], torch.Tensor
):
result[feature_name] = result[feature_name].to("cpu")
return result
@@ -663,4 +683,51 @@ class BaseMultimodalProcessor(ABC):
mm_token_id=mm_token_id,
)
"""
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)
2. if can not get a slice , transport normal tensor
3. copy tensor in scheduler and release it (use position mark)
4. copy
"""
if SGL_USE_CUDA_IPC:
# post-process
for item in all_collected_items:
if isinstance(item.feature, torch.Tensor) and item.feature.is_cuda:
sync_flag, available_slice = (
self.cudaipc_mmfeature_pool.return_a_slice_tensor_with_flag(
item.feature
)
)
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,
)
elif (
isinstance(item.precomputed_embeddings, torch.Tensor)
and item.precomputed_embeddings.is_cuda
):
sync_flag, available_slice = (
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,
)
return all_collected_items, input_ids, ret