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

feat(vlm): auto-select cuda vmm on multi-node mnnvl (#33936)

Browse Source
This commit is contained in:
Mick
2026-08-08 16:00:58 +08:00
committed by GitHub
parent db3898fec1
commit d747bd052e
11 changed files with 245 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/server_args/test_server_args.py
+120
View File
@@ -181,6 +181,16 @@ class TestMultimodalFeatureTransport(CustomTestCase):
self.assertIn("deprecated", logs.output[0])
def test_legacy_keep_flag_rejects_explicit_cuda_vmm(self):
server_args = ServerArgs(
model_path="dummy",
keep_mm_feature_on_device=True,
mm_feature_transport="cuda_vmm",
)
with self.assertRaisesRegex(ValueError, "conflicts.*cuda_vmm"):
server_args._handle_multimodal_feature_transport()
@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")
@@ -231,6 +241,91 @@ class TestMultimodalFeatureTransport(CustomTestCase):
self.assertIn("auto-resolved to cuda_ipc", "\n".join(logs.output))
@patch("sglang.srt.server_args.os.path.exists", return_value=True)
@patch("sglang.srt.server_args.is_mnnvl_fabric_device", return_value=True)
@patch("sglang.srt.server_args.is_cuda", return_value=True)
@patch(
"sglang.srt.model_loader.utils.supports_cuda_vmm_feature_transport",
return_value=True,
)
def test_default_transport_is_cuda_vmm_for_supported_multinode_mnnvl(
self,
_mock_supports_cuda_vmm,
_mock_is_cuda,
_mock_is_mnnvl,
_mock_path_exists,
):
server_args = ServerArgs(model_path="dummy", nnodes=2)
self._set_model_type(server_args, is_multimodal=True)
with patch.dict(os.environ, {}, clear=False):
envs.SGLANG_USE_CUDA_IPC_TRANSPORT.clear()
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_vmm")
self.assertFalse(envs.SGLANG_USE_CUDA_IPC_TRANSPORT.get())
output = "\n".join(logs.output)
self.assertIn("auto-resolved to cuda_vmm", output)
self.assertIn("CUDA FABRIC", output)
@patch("sglang.srt.server_args.os.path.exists", return_value=True)
@patch("sglang.srt.server_args.is_mnnvl_fabric_device", return_value=True)
@patch("sglang.srt.server_args.is_cuda", return_value=True)
@patch(
"sglang.srt.model_loader.utils.supports_cuda_vmm_feature_transport",
return_value=False,
)
def test_default_transport_is_cpu_for_unsupported_multinode_model(
self,
_mock_supports_cuda_vmm,
_mock_is_cuda,
_mock_is_mnnvl,
_mock_path_exists,
):
server_args = ServerArgs(model_path="dummy", nnodes=2)
self._set_model_type(server_args, is_multimodal=True)
with self.assertLogs(server_args_module.logger, level="INFO") as logs:
server_args._handle_multimodal_feature_transport()
self.assertEqual(server_args.mm_feature_transport, "cpu")
self.assertIn("has not opted into CUDA VMM", "\n".join(logs.output))
@patch("sglang.srt.server_args.os.path.exists", return_value=False)
@patch("sglang.srt.server_args.is_mnnvl_fabric_device", return_value=True)
@patch("sglang.srt.server_args.is_cuda", return_value=True)
def test_default_transport_is_cpu_without_imex_channel(
self, _mock_is_cuda, _mock_is_mnnvl, _mock_path_exists
):
server_args = ServerArgs(model_path="dummy", nnodes=2)
self._set_model_type(server_args, is_multimodal=True)
with patch.dict(os.environ, {}, clear=False):
envs.SGLANG_USE_CUDA_IPC_TRANSPORT.clear()
with self.assertLogs(server_args_module.logger, level="INFO") as logs:
server_args._handle_multimodal_feature_transport()
self.assertEqual(server_args.mm_feature_transport, "cpu")
self.assertIn("no IMEX channel", "\n".join(logs.output))
@patch("sglang.srt.server_args.is_mnnvl_fabric_device", return_value=False)
@patch("sglang.srt.server_args.is_cuda", return_value=True)
def test_default_transport_is_cpu_for_multinode_non_mnnvl(
self, _mock_is_cuda, _mock_is_mnnvl
):
server_args = ServerArgs(model_path="dummy", nnodes=2)
self._set_model_type(server_args, is_multimodal=True)
with patch.dict(os.environ, {}, clear=False):
envs.SGLANG_USE_CUDA_IPC_TRANSPORT.clear()
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=True)
def test_default_transport_is_cuda_ipc_for_language_only_model(self, _mock_is_cuda):
server_args = ServerArgs(model_path="dummy", language_only=True)
@@ -259,6 +354,31 @@ class TestMultimodalFeatureTransport(CustomTestCase):
with self.assertRaisesRegex(ValueError, "single node"):
server_args._handle_multimodal_feature_transport()
@patch("sglang.srt.server_args.is_cuda", return_value=True)
def test_cuda_vmm_is_explicit_and_uses_shared_budget(self, _mock_is_cuda):
server_args = ServerArgs(
model_path="dummy",
mm_feature_transport="cuda_vmm",
nnodes=2,
tokenizer_worker_num=2,
)
with (
patch.dict(os.environ, {"SGLANG_USE_CUDA_IPC_TRANSPORT": "1"}),
envs.SGLANG_MM_FEATURE_CACHE_MB.override(256),
):
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_vmm")
self.assertFalse(envs.SGLANG_USE_CUDA_IPC_TRANSPORT.get())
output = "\n".join(logs.output)
self.assertIn("CUDA FABRIC", output)
self.assertIn("256 MiB", output)
self.assertIn("2 tokenizer worker", output)
self.assertIn("falls back to inline CPU", output)
@patch("sglang.srt.server_args.is_cuda", return_value=False)
def test_cuda_vmm_rejects_non_nvidia_platforms(self, _mock_is_cuda):
server_args = ServerArgs(model_path="dummy", mm_feature_transport="cuda_vmm")