diff --git a/python/sglang/srt/multimodal/processors/base_processor.py b/python/sglang/srt/multimodal/processors/base_processor.py index 90ffc00ac..2ba4e508c 100644 --- a/python/sglang/srt/multimodal/processors/base_processor.py +++ b/python/sglang/srt/multimodal/processors/base_processor.py @@ -6,6 +6,7 @@ import multiprocessing as mp import os import re from abc import ABC, abstractmethod +from contextlib import contextmanager from typing import ( Any, Dict, @@ -642,6 +643,24 @@ class BaseMultimodalProcessor(ABC): return "npu" return None + @contextmanager + def _temporary_fast_processor_cuda_pool(self, device: Optional[str]): + """Release fast-processor CUDA temporaries after CPU feature transport.""" + can_release = ( + device is not None + and torch.device(device).type == "cuda" + and not self.keep_mm_features_on_device + and not self.precompute_hash_before_cpu_transfer + ) + if not can_release: + yield + return + + with torch.cuda.device(device): + pool = torch.cuda.MemPool() + with torch.cuda.use_mem_pool(pool, device=device): + yield + def process_mm_data( self, input_text, @@ -690,14 +709,15 @@ class BaseMultimodalProcessor(ABC): if self.audio_config: kwargs.setdefault("audio_kwargs", {}).update(self.audio_config) + processor_device = None if ( hasattr(processor, "image_processor") and isinstance(processor.image_processor, BaseImageProcessor) and not self.disable_fast_image_processor ): - device = self._fast_image_processor_device(processor) - if device is not None: - kwargs["device"] = device + processor_device = self._fast_image_processor_device(processor) + if processor_device is not None: + kwargs["device"] = processor_device # Avoid double BOS when the chat template already wrote one. if self._tokenizer_auto_adds_specials and isinstance(input_text, str): @@ -705,24 +725,25 @@ class BaseMultimodalProcessor(ABC): if bos and input_text.startswith(bos): kwargs.setdefault("add_special_tokens", False) - result = processor.__call__( - text=[input_text], - padding=True, - return_tensors="pt", - **kwargs, - ) - # Deferred: the hash is computed on the GPU tensor first, and - # _precompute_hashes_before_cpu_transfer moves it down afterwards. - if ( - not self.keep_mm_features_on_device - and not self.precompute_hash_before_cpu_transfer - ): - # 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") + with self._temporary_fast_processor_cuda_pool(processor_device): + result = processor.__call__( + text=[input_text], + padding=True, + return_tensors="pt", + **kwargs, + ) + # Deferred: the hash is computed on the GPU tensor first, and + # _precompute_hashes_before_cpu_transfer moves it down afterwards. + if ( + not self.keep_mm_features_on_device + and not self.precompute_hash_before_cpu_transfer + ): + # 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") return result diff --git a/test/registered/unit/multimodal/test_processor_device_selection.py b/test/registered/unit/multimodal/test_processor_device_selection.py index 49458d071..2e08ca6ad 100644 --- a/test/registered/unit/multimodal/test_processor_device_selection.py +++ b/test/registered/unit/multimodal/test_processor_device_selection.py @@ -7,6 +7,8 @@ device has to come from what the worker was handed. """ import unittest +from contextlib import nullcontext +from types import SimpleNamespace from unittest.mock import patch from sglang.srt.multimodal.processors.base_processor import BaseMultimodalProcessor @@ -77,5 +79,83 @@ class TestFastImageProcessorDevice(CustomTestCase): self.assertIsNone(device) +class TestFastImageProcessorMemoryPool(CustomTestCase): + def _processor(self, *, transport="cpu", precompute_hash=False): + processor = _make(base_gpu_id=0) + processor.mm_feature_transport = transport + processor.precompute_hash_before_cpu_transfer = precompute_hash + return processor + + def test_pool_is_limited_to_immediate_cpu_transport(self): + cases = ( + (self._processor(), "cuda:0", True), + (self._processor(transport="cuda_ipc"), "cuda:0", False), + (self._processor(transport="cuda_vmm"), "cuda:0", False), + (self._processor(precompute_hash=True), "cuda:0", False), + (self._processor(), "cpu", False), + (self._processor(), None, False), + ) + for processor, device, expected in cases: + with ( + self.subTest(device=device, transport=processor.mm_feature_transport), + patch(f"{BASE}.torch.cuda.device", return_value=nullcontext()), + patch(f"{BASE}.torch.cuda.MemPool", return_value="pool") as mem_pool, + patch(f"{BASE}.torch.cuda.use_mem_pool", return_value=nullcontext()), + ): + with processor._temporary_fast_processor_cuda_pool(device): + pass + self.assertEqual(mem_pool.called, expected) + + def test_processor_call_uses_private_pool_until_cpu_copy_finishes(self): + class ImageProcessor: + pass + + class Feature: + def to(self, device): + events.append(("copy", device)) + + feature = Feature() + + class Processor: + image_processor = ImageProcessor() + tokenizer = SimpleNamespace(bos_token=None) + + def __call__(self, **kwargs): + events.append(("call", kwargs["device"])) + return {"pixel_values": feature} + + events = [] + processor = self._processor() + processor._processor = Processor() + processor._tokenizer = processor._processor.tokenizer + processor._tokenizer_auto_adds_specials = False + processor.disable_fast_image_processor = False + processor.image_config = {} + processor.video_config = {} + processor.audio_config = {} + processor.FEATURE_NAMES = ["pixel_values"] + + class PoolContext: + def __enter__(self): + events.append("enter") + + def __exit__(self, *args): + events.append("exit") + + with ( + patch(f"{BASE}.BaseImageProcessor", ImageProcessor), + patch(f"{BASE}.torch.cuda.device", return_value=nullcontext()), + patch(f"{BASE}.torch.cuda.MemPool", return_value="pool"), + patch(f"{BASE}.torch.cuda.use_mem_pool", return_value=PoolContext()), + patch(f"{BASE}.torch.Tensor", Feature), + ): + processor.process_mm_data("test", images=["image"]) + + self.assertEqual( + events, + ["enter", ("call", "cuda:0"), ("copy", "cpu"), "exit"], + ) + + if __name__ == "__main__": unittest.main()