minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

feat: unify multimodal feature transport (#30904)

Browse Source
This commit is contained in:
Mick
2026-07-15 17:42:38 +08:00
committed by GitHub
parent f2c875d1c8
commit 947a14d617
7 changed files with 240 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/multimodal/processors/base_processor.py
+12 -4
View File
@@ -41,7 +41,6 @@ _is_cpu = is_cpu()
_is_npu = is_npu()
_is_xpu = is_xpu()
SGL_USE_CUDA_IPC = envs.SGLANG_USE_CUDA_IPC_TRANSPORT.get()
_IPC_POOL_HANDLE_CACHE = envs.SGLANG_USE_IPC_POOL_HANDLE_CACHE.get()
@@ -189,6 +188,15 @@ class BaseMultimodalProcessor(ABC):
self.server_args = server_args
self.transport_mode = transport_mode
self.keep_mm_feature_on_device = server_args.keep_mm_feature_on_device
configured_mm_feature_transport = getattr(
server_args, "mm_feature_transport", "cpu"
)
self.mm_feature_transport = (
configured_mm_feature_transport
if configured_mm_feature_transport in ("cpu", "cuda_ipc")
else "cpu"
)
self.use_cuda_ipc = self.mm_feature_transport == "cuda_ipc"
self.disable_fast_image_processor = server_args.disable_fast_image_processor
self.skip_tokenizer_init = server_args.skip_tokenizer_init
@@ -267,7 +275,7 @@ class BaseMultimodalProcessor(ABC):
skip_mm_pool = kwargs.get("skip_mm_pool", False)
if SGL_USE_CUDA_IPC and not skip_mm_pool:
if self.use_cuda_ipc and not skip_mm_pool:
# SGLANG_MM_FEATURE_CACHE_MB is the total pool budget across all
# tokenizer workers. Each worker gets an equal share so that adding
# workers doesn't multiply the GPU-side footprint.
@@ -488,7 +496,7 @@ class BaseMultimodalProcessor(ABC):
if not self.keep_mm_feature_on_device:
# move feature tensors to cpu
for feature_name in self.FEATURE_NAMES:
if SGL_USE_CUDA_IPC:
if self.use_cuda_ipc:
pass
else:
if feature_name in result and isinstance(
@@ -1473,7 +1481,7 @@ class BaseMultimodalProcessor(ABC):
4. copy
"""
if SGL_USE_CUDA_IPC:
if self.use_cuda_ipc:
# post-process, prepare for cuda-ipc transfer
for item in all_collected_items:
if isinstance(item.feature, torch.Tensor):
python/sglang/srt/multimodal/processors/ernie45_vl.py
+2 -5
View File
@@ -19,13 +19,10 @@ from sglang.srt.multimodal.processors.base_processor import (
from sglang.srt.multimodal.processors.base_processor import (
MultimodalSpecialTokens,
)
from sglang.srt.utils import get_bool_env_var, is_npu, logger
from sglang.srt.utils import is_npu, logger
_is_npu = is_npu()
SGL_USE_CUDA_IPC = get_bool_env_var("SGLANG_USE_CUDA_IPC_TRANSPORT")
IMAGE_FACTOR = 28
MIN_PIXELS = 4 * 28 * 28
# MAX_PIXELS = envs.SGLANG_IMAGE_MAX_PIXELS.get()
@@ -352,7 +349,7 @@ class Ernie4_5_VLImageProcessor(SGLangBaseProcessor):
if not self.keep_mm_feature_on_device:
# move feature tensors to cpu
for feature_name in self.FEATURE_NAMES:
if SGL_USE_CUDA_IPC:
if self.use_cuda_ipc:
pass
else:
if feature_name in result and isinstance(
python/sglang/srt/multimodal/processors/midashenglm.py
+1 -1
View File
@@ -70,7 +70,7 @@ class MiDashengLMMultimodalProcessor(BaseMultimodalProcessor):
**kwargs,
)
if not getattr(self.server_args, "keep_mm_feature_on_device", False):
if not self.keep_mm_feature_on_device and not self.use_cuda_ipc:
for feature_name in ["input_values"]:
if feature_name in result:
result[feature_name] = result[feature_name].cpu()
python/sglang/srt/multimodal/processors/moss_vl.py
+6 -40
View File
@@ -15,16 +15,12 @@ from sglang.srt.managers.schedule_batch import (
MultimodalProcessorOutput,
)
from sglang.srt.models.moss_vl import MossVLForConditionalGeneration
from sglang.srt.multimodal.processors.base_processor import (
SGL_USE_CUDA_IPC,
)
from sglang.srt.multimodal.processors.base_processor import (
BaseMultimodalProcessor as SGLangBaseProcessor,
)
from sglang.srt.multimodal.processors.base_processor import (
MultimodalSpecialTokens,
)
from sglang.srt.utils.cuda_ipc_transport_utils import CudaIpcTensorTransportProxy
class MossVLImageProcessor(SGLangBaseProcessor):
@@ -551,44 +547,14 @@ class MossVLImageProcessor(SGLangBaseProcessor):
if mm_items and vision_token_info:
mm_items[0].set("vision_token_info", vision_token_info[0])
if SGL_USE_CUDA_IPC:
if self.use_cuda_ipc:
for item in mm_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(item.feature, torch.Tensor):
item.feature = self._wrap_tensor_for_cuda_ipc(item.feature)
if isinstance(item.precomputed_embeddings, torch.Tensor):
item.precomputed_embeddings = self._wrap_tensor_for_cuda_ipc(
item.precomputed_embeddings
)
if isinstance(available_slice, torch.Tensor):
available_slice.copy_(
item.feature.reshape(-1).view(torch.int8),
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):
flattened = item.precomputed_embeddings.reshape(-1)
available_slice.copy_(
flattened.view(torch.int8),
non_blocking=True,
)
item.precomputed_embeddings = CudaIpcTensorTransportProxy(
data=available_slice,
info_data=item.precomputed_embeddings,
sync_buffer_meta=sync_flag,
)
return MultimodalProcessorOutput(
input_ids=input_ids.tolist(),
python/sglang/srt/server_args.py
+81 -1
View File
@@ -2223,9 +2223,15 @@ class ServerArgs:
bool,
"Adopt base image processor instead of fast image processor.",
] = False
mm_feature_transport: A[
Optional[Literal["cpu", "cuda_ipc"]],
"Transport multimodal features through CPU memory or a bounded CUDA IPC pool. "
"The default is CPU transport; CUDA IPC reserves GPU memory on the base GPU.",
] = None
keep_mm_feature_on_device: A[
bool,
"Keep multimodal feature tensors on device after processing to save D2H copy.",
"Deprecated. Use --mm-feature-transport=cuda_ipc for bounded GPU-resident "
"multimodal feature transport.",
] = False
# -------------------------------------------------------------------------
@@ -6109,7 +6115,81 @@ class ServerArgs:
"and min_new_tokens are unavailable."
)
def _handle_multimodal_feature_transport(self):
"""Resolve multimodal feature transport before tokenizer workers start.
CUDA IPC is deliberately opt-in: its fixed pool lives on ``base_gpu_id``
and reduces the memory left for model/KV-cache allocations. The legacy
flag and environment variable remain supported so existing deployments
continue to work, but both map to this single policy.
"""
requested_transport = self.mm_feature_transport
legacy_ipc_is_set = envs.SGLANG_USE_CUDA_IPC_TRANSPORT.is_set()
legacy_ipc_enabled = envs.SGLANG_USE_CUDA_IPC_TRANSPORT.get()
if self.keep_mm_feature_on_device:
if requested_transport == "cpu":
raise ValueError(
"--keep-mm-feature-on-device conflicts with "
"--mm-feature-transport=cpu. Use only "
"--mm-feature-transport=cuda_ipc."
)
requested_transport = "cuda_ipc"
logger.warning(
"--keep-mm-feature-on-device is deprecated; using "
"--mm-feature-transport=cuda_ipc instead."
)
if requested_transport is None:
if legacy_ipc_is_set:
requested_transport = "cuda_ipc" if legacy_ipc_enabled else "cpu"
logger.warning(
"SGLANG_USE_CUDA_IPC_TRANSPORT is deprecated; use "
"--mm-feature-transport=%s instead.",
requested_transport,
)
else:
requested_transport = "cpu"
elif legacy_ipc_is_set and legacy_ipc_enabled != (
requested_transport == "cuda_ipc"
):
logger.warning(
"--mm-feature-transport=%s overrides the conflicting legacy "
"SGLANG_USE_CUDA_IPC_TRANSPORT=%s setting.",
requested_transport,
int(legacy_ipc_enabled),
)
if requested_transport == "cuda_ipc":
if not is_cuda():
raise ValueError(
"--mm-feature-transport=cuda_ipc requires NVIDIA CUDA."
)
if self.nnodes != 1:
raise ValueError(
"--mm-feature-transport=cuda_ipc only supports a single node."
)
pool_budget_mb = envs.SGLANG_MM_FEATURE_CACHE_MB.get()
logger.info(
"Using CUDA IPC for multimodal features: reserving up to %d MiB "
"on base GPU %d across %d tokenizer worker(s). This reduces KV "
"cache headroom; a full pool falls back to CPU transport.",
pool_budget_mb,
self.base_gpu_id,
self.tokenizer_worker_num,
)
self.mm_feature_transport = requested_transport
# The bounded IPC pool owns device residency. Do not retain unpooled
# tensors after a pool miss, which would make HBM use request-dependent.
self.keep_mm_feature_on_device = False
envs.SGLANG_USE_CUDA_IPC_TRANSPORT.set(
"1" if requested_transport == "cuda_ipc" else "0"
)
def _handle_environment_variables(self):
self._handle_multimodal_feature_transport()
envs.SGLANG_ENABLE_TORCH_COMPILE.set("1" if self.enable_torch_compile else "0")
if self.mamba_ssm_dtype is not None:
envs.SGLANG_MAMBA_SSM_DTYPE.set(self.mamba_ssm_dtype)