From 7efd8b3d1f2c598ab077790db43e2911408fdbd2 Mon Sep 17 00:00:00 2001 From: kousakawang Date: Sun, 2 Nov 2025 16:46:37 +0800 Subject: [PATCH] [FEAT] Shared mem pool based cuda ipc for multi-modal data transport (#11917) Co-authored-by: kousakawang Co-authored-by: Yuan Luo <4908075+yuan-luo@users.noreply.github.com> --- python/sglang/srt/managers/mm_utils.py | 11 +- python/sglang/srt/managers/schedule_batch.py | 5 + .../sglang/srt/managers/tokenizer_manager.py | 29 +- python/sglang/srt/models/qwen2.py | 1 + .../multimodal/processors/base_processor.py | 77 ++++- python/sglang/srt/server_args.py | 5 + .../srt/utils/cuda_ipc_transport_utils.py | 314 ++++++++++++++++++ test/srt/models/test_vlm_models.py | 2 + 8 files changed, 424 insertions(+), 20 deletions(-) create mode 100644 python/sglang/srt/utils/cuda_ipc_transport_utils.py diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index 60283080b..90ba08eb2 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -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) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 11f133839..2a26797e9 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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) diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 070b8b63d..8d3dd9af8 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -75,7 +75,11 @@ from sglang.srt.managers.scheduler_input_blocker import input_blocker_guard_regi from sglang.srt.managers.tokenizer_communicator_mixin import TokenizerCommunicatorMixin from sglang.srt.metrics.collector import TokenizerMetricsCollector from sglang.srt.sampling.sampling_params import SamplingParams -from sglang.srt.server_args import PortArgs, ServerArgs +from sglang.srt.server_args import ( + PortArgs, + ServerArgs, + set_global_server_args_for_tokenizer, +) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.tracing.trace import ( trace_get_proc_propagate_context, @@ -106,6 +110,16 @@ asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) logger = logging.getLogger(__name__) +def _determine_tensor_transport_mode(server_args: ServerArgs) -> TensorTransportMode: + is_cross_node = server_args.dist_init_addr + + if is_cross_node: + # Fallback to default CPU transport for multi-node + return "default" + else: + return "cuda_ipc" + + @dataclasses.dataclass class ReqState: """Store the state a request.""" @@ -183,6 +197,8 @@ class TokenizerManager(TokenizerCommunicatorMixin): ) # Initialize tokenizer and processor + set_global_server_args_for_tokenizer(server_args) + if self.model_config.is_multimodal: import_processors("sglang.srt.multimodal.processors") try: @@ -920,7 +936,6 @@ class TokenizerManager(TokenizerCommunicatorMixin): batch_req = BatchTokenizedEmbeddingReqInput(batch=tokenized_objs) self.send_to_scheduler.send_pyobj(batch_req) - # Create states for each individual request in the batch for i, tokenized_obj in enumerate(tokenized_objs): tmp_obj = obj[i] @@ -2204,16 +2219,6 @@ class ServerStatus(Enum): UnHealthy = "UnHealthy" -def _determine_tensor_transport_mode(server_args: ServerArgs) -> TensorTransportMode: - is_cross_node = server_args.dist_init_addr - - if is_cross_node: - # Fallback to default CPU transport for multi-node - return "default" - else: - return "cuda_ipc" - - async def print_exception_wrapper(func): """ Sometimes an asyncio function does not print exception. diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py index 98a3b20d9..75ceff821 100644 --- a/python/sglang/srt/models/qwen2.py +++ b/python/sglang/srt/models/qwen2.py @@ -340,6 +340,7 @@ class Qwen2Model(nn.Module): input_embeds: torch.Tensor = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> Union[torch.Tensor, PPProxyTensors]: + if self.pp_group.is_first_rank: if input_embeds is None: hidden_states = self.embed_tokens(input_ids) diff --git a/python/sglang/srt/multimodal/processors/base_processor.py b/python/sglang/srt/multimodal/processors/base_processor.py index f7bebf860..c06cf345a 100644 --- a/python/sglang/srt/multimodal/processors/base_processor.py +++ b/python/sglang/srt/multimodal/processors/base_processor.py @@ -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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 8d106e483..4faed41d3 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -3889,6 +3889,11 @@ def set_global_server_args_for_scheduler(server_args: ServerArgs): _global_server_args = server_args +def set_global_server_args_for_tokenizer(server_args: ServerArgs): + global _global_server_args + _global_server_args = server_args + + def get_global_server_args() -> ServerArgs: if _global_server_args is None: raise ValueError("Global server args is not set yet!") diff --git a/python/sglang/srt/utils/cuda_ipc_transport_utils.py b/python/sglang/srt/utils/cuda_ipc_transport_utils.py new file mode 100644 index 000000000..e29265b2a --- /dev/null +++ b/python/sglang/srt/utils/cuda_ipc_transport_utils.py @@ -0,0 +1,314 @@ +import fcntl +import logging +from multiprocessing import shared_memory +from typing import Tuple + +import numpy as np +import torch + +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import get_int_env_var + +logger = logging.getLogger(__name__) + +MM_FEATURE_CACHE_SIZE = ( + 2 * 1024 * 1024 * 1024 + if not get_int_env_var("SGLANG_MM_FEATURE_CACHE_MB") + else get_int_env_var("SGLANG_MM_FEATURE_CACHE_MB") * 1024 * 1024 +) + +SHM_LOCK_FILE = "/tmp/shm_wr_lock.lock" + + +class ShmSyncBuffer: + def __init__(self, byte_size: int = 4): + self.buffer = shared_memory.SharedMemory(create=True, size=byte_size) + self.buffer_wrapper = np.ndarray(1, dtype=np.float32, buffer=self.buffer.buf) + self.buffer_wrapper *= 0 + self.meta_data = { + "handle": self.buffer.name, + "shape": self.buffer_wrapper.shape, + "dtype": str(self.buffer_wrapper.dtype), + } + + def __del__(self): + if isinstance(self.buffer, shared_memory.SharedMemory): + self.buffer.close() + self.buffer.unlink() + + +class MmItemMemoryChunk: + def __init__(self, area: Tuple, sync_buffer: ShmSyncBuffer): + self.area = area + self.sync_flag = sync_buffer + + @property + def mem_size(self): + return self.area[1] - self.area[0] + + @property + def start(self): + return self.area[0] + + @property + def end(self): + return self.area[1] + + def try_to_recycle(self) -> bool: + try: + tp_num = get_global_server_args().tp_size + except: + logger.info( + "get_global_server_args has not been inited , skip this turn 's recycle" + ) + tp_num = -1 + if self.sync_flag.buffer_wrapper.item() == float(tp_num): + self.sync_flag.buffer_wrapper *= 0 + return True + + return False + + +class MmItemMemoryPool: + def __init__(self, memory_size): + self.memory_pool = torch.empty( + memory_size, dtype=torch.int8, device="cuda" + ).contiguous() + + self.sync_flag_list = [] + + init_chunk = MmItemMemoryChunk((0, memory_size), self.pop_sync_buffer()) + self.available_chunks = [init_chunk] + self.occupied_chunks = [] + + def clear_sync_flag_list(self): + # call each chunk's __del__ + self.sync_flag_list.clear() + + def pop_sync_buffer(self): + if len(self.sync_flag_list) == 0: + try: + new_sync_buffer = ShmSyncBuffer() + return new_sync_buffer + except: + logger.info("allocate shm buffer failed") + raise RuntimeError + else: + return self.sync_flag_list.pop() + + def push_sync_buffer(self, sync_buffer): + self.sync_flag_list.append(sync_buffer) + + def get_available_chunk(self, src_tensor: torch.Tensor) -> MmItemMemoryChunk: + # find currently available_chunks contain a available chunk or not + # if not, return None + src_tensor_size = src_tensor.numel() * src_tensor.element_size() + min_size = self.memory_pool.numel() * self.memory_pool.element_size() + 1 + selected_chunk = None + for chunk in self.available_chunks: + if chunk.mem_size >= src_tensor_size: + if chunk.mem_size < min_size: + min_size = chunk.mem_size + selected_chunk = chunk + + if selected_chunk: + occupied_chunk_area = ( + selected_chunk.start, + selected_chunk.start + src_tensor_size, + ) + occupied_chunk_sync_flag = selected_chunk.sync_flag + new_occupied_chunk = MmItemMemoryChunk( + occupied_chunk_area, occupied_chunk_sync_flag + ) + + self.occupied_chunks.append(new_occupied_chunk) + self.available_chunks.remove(selected_chunk) + + available_split_chunk_area = (new_occupied_chunk.end, selected_chunk.end) + # add a new chunk + if available_split_chunk_area[0] != available_split_chunk_area[1]: + split_available_chunk = MmItemMemoryChunk( + available_split_chunk_area, self.pop_sync_buffer() + ) + self.available_chunks.append(split_available_chunk) + + return new_occupied_chunk + + return None + + def return_a_slice_tensor_with_flag(self, src_tensor: torch.Tensor): + self.recycle_chunks() + self.merge_chunks() + + available_chunk = self.get_available_chunk(src_tensor) + if available_chunk is not None: + return ( + available_chunk.sync_flag.meta_data, + self.memory_pool[available_chunk.start : available_chunk.end], + ) + return None, None + + def recycle_chunks(self): + + new_occupied_chunks = [] + for chunk in self.occupied_chunks: + if chunk.try_to_recycle(): + self.available_chunks.append(chunk) + else: + new_occupied_chunks.append(chunk) + self.occupied_chunks = new_occupied_chunks + + def merge_chunks(self): + # merge_all_available_chunks + merged_chunks = [] + for chunk in sorted(self.available_chunks, key=lambda x: x.start): + if len(merged_chunks) == 0: + merged_chunks.append(chunk) + else: + if chunk.start == merged_chunks[-1].end: + to_merge_chunk = merged_chunks.pop() + to_merge_chunk_sync = to_merge_chunk.sync_flag + merged_chunk_area = (to_merge_chunk.start, chunk.end) + merged_chunks.append( + MmItemMemoryChunk(merged_chunk_area, to_merge_chunk_sync) + ) + self.push_sync_buffer(chunk.sync_flag) + else: + merged_chunks.append(chunk) + + self.available_chunks = merged_chunks + + +class CudaIpcTensorTransportProxy: + """ + A torch.tensor's proxy used to do inter-process data-sharing + including: + + torch.tensor(on gpu)'s cuda-ipc-hande infos + a shm sync buffer's meta data which is used to sync between different process + """ + + def __init__( + self, + data: torch.Tensor, + info_data: torch.Tensor, + sync_buffer_meta, + ): + + if (not isinstance(data, torch.Tensor)) or ( + not isinstance(info_data, torch.Tensor) + ): + raise TypeError( + f"Input 'data' must be a torch.Tensor, but got {type(data)}" + ) + + self.proxy_state = self.get_proxy_state(data, info_data) + self.reconstruct_tensor = None + self.sync_data_meta = sync_buffer_meta + self.sync_buffer = None + + @property + def get_sync_flag(self): + if not self.sync_buffer: + shm_name = self.sync_data_meta["handle"] + self.sync_buffer = shared_memory.SharedMemory(name=shm_name) + + shape = self.sync_data_meta["shape"] + dtype = self.sync_data_meta["dtype"] + return np.ndarray(shape, dtype=dtype, buffer=self.sync_buffer.buf) + + def close_shm(self): + self.sync_buffer.close() + self.sync_buffer = None + + def get_proxy_state(self, data, info_data): + # acquire all serialize metadata from _metadata + state = {} + + try: + storage = data.untyped_storage() + handle = storage._share_cuda_() + + state["ipc_extra"] = { + "handle": handle, + "shape": data.shape, + "dtype": data.dtype, + "stride": data.stride(), + "device_index": data.device.index, + "storage_offset": data.storage_offset(), + "recons_shape": info_data.shape, + "recons_dtype": info_data.dtype, + } + state["tensor_data"] = None + except Exception as e: + # Failed to get CUDA IPC handle (possibly tp). Falling back to default transport. + state["ipc_extra"] = None + state["tensor_data"] = data + + return state + + def reconstruct_on_target_device(self, rebuild_device_idx): + rebuild_device = torch.device(f"cuda:{rebuild_device_idx}") + if ( + isinstance(self.reconstruct_tensor, torch.Tensor) + and self.reconstruct_tensor.device == rebuild_device + ): + return self.reconstruct_tensor + + if self.proxy_state["ipc_extra"]: + ipc_extra = self.proxy_state["ipc_extra"] + ( + handle, + shape, + dtype, + stride, + source_device_index, + s_offset, + recons_shape, + recons_dtype, + ) = ( + ipc_extra["handle"], + ipc_extra["shape"], + ipc_extra["dtype"], + ipc_extra["stride"], + ipc_extra["device_index"], + ipc_extra["storage_offset"], + ipc_extra["recons_shape"], + ipc_extra["recons_dtype"], + ) + + try: + target_device = torch.device(f"cuda:{source_device_index}") + with torch.cuda.device(target_device): + storage = torch.UntypedStorage._new_shared_cuda(*handle) + slice_tensor = torch.empty( + 0, dtype=dtype, device=target_device + ).set_(storage, storage_offset=s_offset, size=shape, stride=stride) + + reconstructed_tensor = torch.empty( + recons_shape, dtype=recons_dtype, device=rebuild_device + ).contiguous() + reconstructed_tensor.view(torch.int8).view(-1).copy_(slice_tensor) + + open(SHM_LOCK_FILE, "a").close() + # write the shm_sync_buffer with a file lock + with open(SHM_LOCK_FILE, "w+") as f: + fcntl.flock(f, fcntl.LOCK_EX) + sync_flag = self.get_sync_flag + sync_flag += 1 + fcntl.flock(f, fcntl.LOCK_UN) + + self.close_shm() + + except Exception as e: + logger.info(f"Error: Failed to deserialize from CUDA IPC handle ({e}).") + raise e + elif isinstance(self.proxy_state["tensor_data"], torch.Tensor): + reconstructed_tensor = self.proxy_state["tensor_data"].to( + rebuild_device, non_blocking=True + ) + else: + raise TypeError("invalid proxy_state") + + self.reconstruct_tensor = reconstructed_tensor + return self.reconstruct_tensor diff --git a/test/srt/models/test_vlm_models.py b/test/srt/models/test_vlm_models.py index 597b0e2ae..383b2d57f 100644 --- a/test/srt/models/test_vlm_models.py +++ b/test/srt/models/test_vlm_models.py @@ -145,6 +145,8 @@ class TestVLMModels(CustomTestCase): process_env = os.environ.copy() if custom_env: process_env.update(custom_env) + # if test vlm with cuda_ipc feature, open this env_var + process_env["SGLANG_USE_CUDA_IPC_TRANSPORT"] = "1" # Prepare stdout/stderr redirection if needed stdout_file = None