[VLM] avoid CUDA placement on non-CUDA platforms (#38750)

This commit is contained in:
Vincent Liu
2026-09-19 11:13:22 +08:00
committed by GitHub
parent c3aa09b0db
commit 090263eff6
3 changed files with 85 additions and 13 deletions
@@ -23,6 +23,7 @@ import torch
from PIL import Image
from transformers import BaseImageProcessor
from sglang.srt import platforms
from sglang.srt.managers.schedule_batch import (
Modality,
MultimodalDataItem,
@@ -721,19 +722,23 @@ class BaseMultimodalProcessor(ABC):
return processor, _tokenizer_of(processor)
def _preprocessing_competes_with_the_scheduler(self) -> bool:
"""Whether image preprocessing submits its work to the serving GPU.
"""Whether image preprocessing contends with the serving accelerator.
The fast image processor runs inside the tokenizer process but on
``cuda:{base_gpu_id}`` -- the device the scheduler serves from. A second
preprocessing worker there is one more competitor for that device rather
than added parallelism.
The fast image processor runs inside the tokenizer process but may run on
the same accelerator as the scheduler. A second preprocessing worker there
adds device contention rather than CPU preprocessing parallelism.
"""
if _is_cpu or get_exec().deterministic.rl_on_policy_target is not None:
return False
if self.disable_fast_image_processor:
return False
image_processor = getattr(self._processor, "image_processor", None)
return isinstance(image_processor, BaseImageProcessor)
if not isinstance(image_processor, BaseImageProcessor):
return False
if _is_xpu or _is_npu:
return True
platform = platforms.current_platform
return platform.is_cuda_alike()
def _resolve_auto_mm_processor_worker_num(self) -> int:
"""The worker count to use when the user did not ask for one.
@@ -769,9 +774,12 @@ class BaseMultimodalProcessor(ABC):
if _is_xpu:
return "xpu"
if not _is_npu:
platform = platforms.current_platform
if not platform.is_cuda_alike():
return None
# Per-worker placement travels as a constructor argument, and
# this record is that argument.
return f"cuda:{server_args.base_gpu_id}"
return f"{platform.device_type}:{server_args.base_gpu_id}"
if processor.__class__.__name__ == "MiniMaxVLProcessor":
# MiniMax's image/video processors create 10-dim tensors during
# patch extraction, exceeding the Ascend 8-dim limit; patch them
@@ -75,6 +75,21 @@ class TestMmProcessConfigValidation(CustomTestCase):
class TestBaseProcessorConfigExtraction(CustomTestCase):
"""Verify BaseMultimodalProcessor.__init__ extracts configs from server_args."""
def _patch_platform(self, cuda_alike, device_type):
platforms = SimpleNamespace(
current_platform=SimpleNamespace(
is_cuda_alike=lambda: cuda_alike,
device_type=device_type,
)
)
return patch.multiple(
"sglang.srt.multimodal.processors.base_processor",
_is_cpu=False,
_is_xpu=False,
_is_npu=False,
platforms=platforms,
)
def _make_processor(
self,
mm_process_config,
@@ -210,12 +225,23 @@ class TestBaseProcessorConfigExtraction(CustomTestCase):
9.30 -> 4.02 req/s on GB300 for full-page images."""
from transformers import BaseImageProcessor
with self._patch_platform(cuda_alike=True, device_type="cuda"):
proc = self._make_processor(
{}, image_processor=MagicMock(spec=BaseImageProcessor)
)
self.assertEqual(proc.mm_processor_worker_num, 1)
self.assertIsNone(proc.mm_processor_executor)
def test_non_accelerator_fast_processor_gets_two_workers(self):
from transformers import BaseImageProcessor
with self._patch_platform(cuda_alike=False, device_type="custom"):
proc = self._make_processor(
{}, image_processor=MagicMock(spec=BaseImageProcessor)
)
self.assertEqual(proc.mm_processor_worker_num, 2)
self.assertIsNotNone(proc.mm_processor_executor)
def test_explicit_request_overrides_the_path_decision(self):
"""The server argument wins: an operator who measured their own workload
can still ask for concurrency on the GPU path."""
@@ -240,7 +266,10 @@ class TestBaseProcessorConfigExtraction(CustomTestCase):
BaseMultimodalProcessor,
)
with patch.object(BaseMultimodalProcessor, "auto_mm_processor_worker_num", 3):
with (
patch.object(BaseMultimodalProcessor, "auto_mm_processor_worker_num", 3),
self._patch_platform(cuda_alike=True, device_type="cuda"),
):
proc = self._make_processor(
{}, image_processor=MagicMock(spec=BaseImageProcessor)
)
@@ -51,10 +51,16 @@ class TestFastImageProcessorDevice(CustomTestCase):
reset_context()
self.addCleanup(reset_context)
def _device(self, processor, **platform):
def _device(self, processor, cuda_alike=True, device_type="cuda", **platform):
flags = {"_is_cpu": False, "_is_xpu": False, "_is_npu": False}
flags.update(platform)
with patch.multiple(BASE, **flags):
platforms = SimpleNamespace(
current_platform=SimpleNamespace(
is_cuda_alike=lambda: cuda_alike,
device_type=device_type,
)
)
with patch.multiple(BASE, platforms=platforms, **flags):
return processor._fast_image_processor_device(_Processor())
def test_device_follows_the_instance_base_gpu_id(self):
@@ -92,6 +98,23 @@ class TestFastImageProcessorDevice(CustomTestCase):
device = processor._fast_image_processor_device(Glm4vProcessor())
self.assertIsNone(device)
def test_platform_respects_cuda_compatibility(self):
processor = _make(base_gpu_id=3)
for cuda_alike, device_type, expected in (
(False, "other", None),
(True, "cuda", "cuda:3"),
(True, "musa", "musa:3"),
):
with self.subTest(cuda_alike=cuda_alike, device_type=device_type):
self.assertEqual(
self._device(
processor,
cuda_alike=cuda_alike,
device_type=device_type,
),
expected,
)
class TestFastImageProcessorMemoryPool(CustomTestCase):
def setUp(self):
@@ -161,6 +184,18 @@ class TestFastImageProcessorMemoryPool(CustomTestCase):
events.append("exit")
with (
patch.multiple(
BASE,
_is_cpu=False,
_is_xpu=False,
_is_npu=False,
platforms=SimpleNamespace(
current_platform=SimpleNamespace(
is_cuda_alike=lambda: True,
device_type="cuda",
)
),
),
patch(f"{BASE}.BaseImageProcessor", ImageProcessor),
patch(f"{BASE}.torch.cuda.device", return_value=nullcontext()),
patch(f"{BASE}.torch.cuda.MemPool", return_value="pool"),