diff --git a/python/sglang/srt/multimodal/processors/base_processor.py b/python/sglang/srt/multimodal/processors/base_processor.py index eef87d25c..fdf9cb55c 100644 --- a/python/sglang/srt/multimodal/processors/base_processor.py +++ b/python/sglang/srt/multimodal/processors/base_processor.py @@ -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): diff --git a/python/sglang/srt/multimodal/processors/ernie45_vl.py b/python/sglang/srt/multimodal/processors/ernie45_vl.py index f7d0a9e1a..7177bf932 100644 --- a/python/sglang/srt/multimodal/processors/ernie45_vl.py +++ b/python/sglang/srt/multimodal/processors/ernie45_vl.py @@ -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( diff --git a/python/sglang/srt/multimodal/processors/midashenglm.py b/python/sglang/srt/multimodal/processors/midashenglm.py index 985be22fd..aa58a5cf8 100644 --- a/python/sglang/srt/multimodal/processors/midashenglm.py +++ b/python/sglang/srt/multimodal/processors/midashenglm.py @@ -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() diff --git a/python/sglang/srt/multimodal/processors/moss_vl.py b/python/sglang/srt/multimodal/processors/moss_vl.py index 409d9eb23..565b0ec09 100644 --- a/python/sglang/srt/multimodal/processors/moss_vl.py +++ b/python/sglang/srt/multimodal/processors/moss_vl.py @@ -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(), diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index c000a88eb..6d3cec64c 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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) diff --git a/test/registered/unit/managers/test_mm_process_config.py b/test/registered/unit/managers/test_mm_process_config.py index 71628f55a..1c854c09e 100644 --- a/test/registered/unit/managers/test_mm_process_config.py +++ b/test/registered/unit/managers/test_mm_process_config.py @@ -1,4 +1,5 @@ import unittest +from types import SimpleNamespace from unittest.mock import MagicMock, patch from sglang.srt.server_args import ServerArgs @@ -103,6 +104,62 @@ class TestBaseProcessorConfigExtraction(unittest.TestCase): self.assertEqual(proc.audio_config, {}) +class TestMultimodalFeatureTransportRuntime(unittest.TestCase): + @staticmethod + def _server_args(mm_feature_transport): + return SimpleNamespace( + mm_feature_transport=mm_feature_transport, + keep_mm_feature_on_device=False, + disable_fast_image_processor=False, + skip_tokenizer_init=False, + mm_process_config={}, + tokenizer_worker_num=1, + base_gpu_id=2, + ) + + @staticmethod + def _processor(): + processor = MagicMock() + processor.tokenizer.encode.return_value = [] + return processor + + def test_cuda_ipc_pool_uses_resolved_server_arg(self): + # The processor module can be imported before this instance is built; + # transport policy must still resolve from the instance's ServerArgs. + from sglang.srt.multimodal.processors import base_processor + + with patch.object( + base_processor.BaseMultimodalProcessor, "__abstractmethods__", set() + ), patch.object(base_processor, "MmItemMemoryPool") as memory_pool: + processor = base_processor.BaseMultimodalProcessor( + hf_config=MagicMock(), + server_args=self._server_args("cuda_ipc"), + _processor=self._processor(), + transport_mode=None, + ) + + self.assertEqual(processor.mm_feature_transport, "cuda_ipc") + self.assertTrue(processor.use_cuda_ipc) + memory_pool.assert_called_once() + + def test_cpu_transport_does_not_allocate_ipc_pool(self): + from sglang.srt.multimodal.processors import base_processor + + with patch.object( + base_processor.BaseMultimodalProcessor, "__abstractmethods__", set() + ), patch.object(base_processor, "MmItemMemoryPool") as memory_pool: + processor = base_processor.BaseMultimodalProcessor( + hf_config=MagicMock(), + server_args=self._server_args("cpu"), + _processor=self._processor(), + transport_mode=None, + ) + + self.assertEqual(processor.mm_feature_transport, "cpu") + self.assertFalse(processor.use_cuda_ipc) + memory_pool.assert_not_called() + + class TestProcessMmDataKwargs(unittest.TestCase): """Verify process_mm_data injects per-modality kwargs correctly.""" @@ -114,6 +171,7 @@ class TestProcessMmDataKwargs(unittest.TestCase): server_args = MagicMock() server_args.mm_process_config = mm_process_config + server_args.mm_feature_transport = "cpu" server_args.disable_fast_image_processor = True server_args.keep_mm_feature_on_device = True server_args.skip_tokenizer_init = False @@ -135,6 +193,8 @@ class TestProcessMmDataKwargs(unittest.TestCase): proc.server_args = server_args proc.keep_mm_feature_on_device = server_args.keep_mm_feature_on_device + proc.mm_feature_transport = server_args.mm_feature_transport + proc.use_cuda_ipc = False proc.disable_fast_image_processor = server_args.disable_fast_image_processor proc.skip_tokenizer_init = server_args.skip_tokenizer_init proc._processor = mock_processor @@ -217,6 +277,7 @@ class TestOverrideProcessorsConfigInjection(unittest.TestCase): """Create an override processor with mocked dependencies.""" server_args = MagicMock() server_args.mm_process_config = mm_process_config + server_args.mm_feature_transport = "cpu" server_args.disable_fast_image_processor = True server_args.keep_mm_feature_on_device = False server_args.skip_tokenizer_init = False @@ -232,6 +293,8 @@ class TestOverrideProcessorsConfigInjection(unittest.TestCase): proc.server_args = server_args proc.keep_mm_feature_on_device = server_args.keep_mm_feature_on_device + proc.mm_feature_transport = server_args.mm_feature_transport + proc.use_cuda_ipc = False proc.disable_fast_image_processor = server_args.disable_fast_image_processor proc.skip_tokenizer_init = server_args.skip_tokenizer_init proc._processor = mock_hf_processor @@ -317,6 +380,7 @@ class TestDoubleBosGuard(unittest.TestCase): server_args = MagicMock() server_args.mm_process_config = {} + server_args.mm_feature_transport = "cpu" server_args.disable_fast_image_processor = True server_args.keep_mm_feature_on_device = True diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index f7029e221..353efb421 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -59,6 +59,80 @@ class TestPrepareServerArgs(CustomTestCase): os.unlink(config_file) +class TestMultimodalFeatureTransport(CustomTestCase): + @patch("sglang.srt.server_args.is_cuda", return_value=True) + def test_cuda_ipc_is_explicit_and_bounded(self, _mock_is_cuda): + server_args = ServerArgs( + model_path="dummy", + mm_feature_transport="cuda_ipc", + tokenizer_worker_num=4, + base_gpu_id=2, + ) + + with patch.dict(os.environ, {"SGLANG_USE_CUDA_IPC_TRANSPORT": "0"}): + with self.assertLogs(server_args_module.logger, level="INFO") as logs: + server_args._handle_multimodal_feature_transport() + + self.assertEqual(server_args.mm_feature_transport, "cuda_ipc") + self.assertTrue(envs.SGLANG_USE_CUDA_IPC_TRANSPORT.get()) + + output = "\n".join(logs.output) + self.assertIn("base GPU 2", output) + self.assertIn("4 tokenizer worker", output) + + @patch("sglang.srt.server_args.is_cuda", return_value=True) + def test_legacy_keep_flag_maps_to_cuda_ipc(self, _mock_is_cuda): + server_args = ServerArgs(model_path="dummy", keep_mm_feature_on_device=True) + + with patch.dict(os.environ, {"SGLANG_USE_CUDA_IPC_TRANSPORT": "0"}): + with self.assertLogs(server_args_module.logger, level="WARNING") as logs: + server_args._handle_multimodal_feature_transport() + + self.assertEqual(server_args.mm_feature_transport, "cuda_ipc") + self.assertFalse(server_args.keep_mm_feature_on_device) + self.assertTrue(envs.SGLANG_USE_CUDA_IPC_TRANSPORT.get()) + + self.assertIn("deprecated", logs.output[0]) + + @patch("sglang.srt.server_args.is_cuda", return_value=True) + def test_explicit_cpu_overrides_legacy_environment(self, _mock_is_cuda): + server_args = ServerArgs(model_path="dummy", mm_feature_transport="cpu") + + with patch.dict(os.environ, {"SGLANG_USE_CUDA_IPC_TRANSPORT": "1"}): + with self.assertLogs(server_args_module.logger, level="WARNING") as logs: + server_args._handle_multimodal_feature_transport() + + self.assertEqual(server_args.mm_feature_transport, "cpu") + self.assertFalse(envs.SGLANG_USE_CUDA_IPC_TRANSPORT.get()) + + self.assertIn("overrides", logs.output[0]) + + def test_default_transport_is_cpu(self): + server_args = ServerArgs(model_path="dummy") + + with patch.dict(os.environ, {"SGLANG_USE_CUDA_IPC_TRANSPORT": "0"}): + server_args._handle_multimodal_feature_transport() + + self.assertEqual(server_args.mm_feature_transport, "cpu") + self.assertFalse(envs.SGLANG_USE_CUDA_IPC_TRANSPORT.get()) + + @patch("sglang.srt.server_args.is_cuda", return_value=False) + def test_cuda_ipc_rejects_non_nvidia_platforms(self, _mock_is_cuda): + server_args = ServerArgs(model_path="dummy", mm_feature_transport="cuda_ipc") + + with self.assertRaisesRegex(ValueError, "requires NVIDIA CUDA"): + server_args._handle_multimodal_feature_transport() + + @patch("sglang.srt.server_args.is_cuda", return_value=True) + def test_cuda_ipc_rejects_multi_node(self, _mock_is_cuda): + server_args = ServerArgs( + model_path="dummy", mm_feature_transport="cuda_ipc", nnodes=2 + ) + + with self.assertRaisesRegex(ValueError, "single node"): + server_args._handle_multimodal_feature_transport() + + class TestMambaCacheStochasticRounding(unittest.TestCase): def test_rejects_fp32_ssm_cache(self): server_args = ServerArgs(