diff --git a/components/src/dynamo/sglang/_compat.py b/components/src/dynamo/sglang/_compat.py index c0cd4199a89c..2bd48acec1d5 100644 --- a/components/src/dynamo/sglang/_compat.py +++ b/components/src/dynamo/sglang/_compat.py @@ -22,7 +22,7 @@ import inspect import logging -from functools import lru_cache +from functools import lru_cache, wraps from typing import Any logger = logging.getLogger(__name__) @@ -53,6 +53,49 @@ def ensure_sglang_top_level_exports() -> None: ensure_sglang_top_level_exports() +def ensure_sglang_tensor_image_size() -> None: + """Allow SGLang's image-token resolver to handle decoded image tensors. + + SGLang 0.5.13 and 0.5.14 assume every decoded image exposes the PIL + ``height``/``width`` attributes. Its CUDA JPEG decoder instead returns a + CHW tensor, causing multimodal requests to fall back to retokenization. + + Remove this compatibility override once the minimum supported SGLang + release handles tensor image dimensions itself. + """ + import torch + from sglang.srt.multimodal.processors.base_processor import BaseMultimodalProcessor + + original = getattr(BaseMultimodalProcessor, "resolve_image_token_counts", None) + if original is None or getattr( + original, "_dynamo_tensor_image_size_support", False + ): + return + + @wraps(original) + def resolve_image_token_counts(self: Any, images: list[Any]) -> list[int]: + if not any(isinstance(image, torch.Tensor) for image in images): + return original(self, images) + + image_sizes: list[tuple[int, int]] = [] + for image in images: + if isinstance(image, torch.Tensor): + if image.ndim < 2: + raise ValueError(f"Invalid image tensor shape: {image.shape}") + height, width = image.shape[-2:] + else: + height, width = image.height, image.width + image_sizes.append((int(height), int(width))) + + token_counts = self._processor._get_num_multimodal_tokens( + image_sizes=image_sizes + ).num_image_tokens + return [int(count) for count in token_counts] + + resolve_image_token_counts._dynamo_tensor_image_size_support = True # type: ignore[attr-defined] + BaseMultimodalProcessor.resolve_image_token_counts = resolve_image_token_counts + + @lru_cache(maxsize=32) def _get_async_generate_supported_kwarg_names( async_generate: Any, @@ -120,6 +163,7 @@ def enable_disjoint_streaming_output(server_args: Any) -> None: __all__ = [ "enable_disjoint_streaming_output", + "ensure_sglang_tensor_image_size", "ensure_sglang_top_level_exports", "filter_supported_async_generate_kwargs", ] diff --git a/components/src/dynamo/sglang/args.py b/components/src/dynamo/sglang/args.py index 8daca61cd9c2..3a99b748f161 100644 --- a/components/src/dynamo/sglang/args.py +++ b/components/src/dynamo/sglang/args.py @@ -28,7 +28,10 @@ ) from dynamo.common.utils.runtime import parse_endpoint from dynamo.runtime.logging import configure_dynamo_logging -from dynamo.sglang._compat import enable_disjoint_streaming_output +from dynamo.sglang._compat import ( + enable_disjoint_streaming_output, + ensure_sglang_tensor_image_size, +) from dynamo.sglang.backend_args import DynamoSGLangArgGroup, DynamoSGLangConfig configure_dynamo_logging() @@ -529,6 +532,8 @@ async def parse_args( ) else: server_args = ServerArgs.from_cli_args(parsed_args) + if server_args.get_model_config().is_multimodal: + ensure_sglang_tensor_image_size() if getattr(server_args, "schedule_low_priority_values_first", False): raise ValueError( diff --git a/components/src/dynamo/sglang/tests/test_sglang_unit.py b/components/src/dynamo/sglang/tests/test_sglang_unit.py index f271db417594..f226e7a37c8f 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_unit.py +++ b/components/src/dynamo/sglang/tests/test_sglang_unit.py @@ -11,6 +11,7 @@ from types import SimpleNamespace import pytest +import torch import yaml from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST @@ -19,6 +20,7 @@ from dynamo.common.constants import DisaggregationMode, EmbeddingTransferMode from dynamo.common.snapshot.constants import SNAPSHOT_CONTROL_DIR_ENV from dynamo.sglang._compat import ( + ensure_sglang_tensor_image_size, ensure_sglang_top_level_exports, filter_supported_async_generate_kwargs, ) @@ -119,6 +121,102 @@ def test_compat_restores_sglang_top_level_exports(): sgl.ServerArgs = original_server_args +def test_compat_supports_tensor_image_sizes_and_is_idempotent(caplog, monkeypatch): + from sglang.srt.multimodal.processors.base_processor import ( + BaseMultimodalProcessor, + BaseMultiModalProcessorOutput, + MultimodalSpecialTokens, + ) + + class Processor: + image_sizes = None + + def _get_num_multimodal_tokens(self, *, image_sizes): + self.image_sizes = image_sizes + return SimpleNamespace(num_image_tokens=[4]) + + class ConcreteMultimodalProcessor(BaseMultimodalProcessor): + async def process_mm_data_async(self, *args, **kwargs): + raise NotImplementedError + + original = BaseMultimodalProcessor.resolve_image_token_counts + try: + ensure_sglang_tensor_image_size() + installed = BaseMultimodalProcessor.resolve_image_token_counts + ensure_sglang_tensor_image_size() + + processor = object.__new__(ConcreteMultimodalProcessor) + processor._processor = Processor() + image_token_id = 99 + processor._process_and_collect_mm_items = lambda **kwargs: ( + [], + torch.tensor( + [20, image_token_id, image_token_id, image_token_id, image_token_id, 21] + ), + {}, + ) + base_output = BaseMultiModalProcessorOutput( + input_text="decoded prompt", + input_ids=[10, image_token_id, 11], + images=[torch.empty((3, 48, 80), dtype=torch.uint8)], + ) + mm_tokens = MultimodalSpecialTokens(image_token_id=image_token_id) + # SGLang defaults this on to preserve caller token IDs and expand only + # image placeholders instead of decoding and retokenizing the prompt. + monkeypatch.setenv("SGLANG_MM_AVOID_RETOKENIZE", "1") + + with caplog.at_level( + logging.WARNING, + logger="sglang.srt.multimodal.processors.base_processor", + ): + _, input_ids, _ = processor.process_and_combine_mm_data( + base_output, mm_tokens + ) + + assert installed is BaseMultimodalProcessor.resolve_image_token_counts + assert processor._processor.image_sizes == [(48, 80)] + assert input_ids.tolist() == [ + 10, + image_token_id, + image_token_id, + image_token_id, + image_token_id, + 11, + ] + assert not any( + "falling back to decode+retokenize" in record.message + for record in caplog.records + ) + finally: + BaseMultimodalProcessor.resolve_image_token_counts = original + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_multimodal", [False, True]) +async def test_tensor_image_size_compat_uses_resolved_model_capability( + monkeypatch, mock_sglang_cli, is_multimodal +): + server_args = SimpleNamespace( + disaggregation_mode="null", + dllm_algorithm=None, + kv_events_config=None, + get_model_config=lambda: SimpleNamespace(is_multimodal=is_multimodal), + ) + install_calls = [] + monkeypatch.setattr( + "dynamo.sglang.args.ServerArgs.from_cli_args", lambda _: server_args + ) + monkeypatch.setattr( + "dynamo.sglang.args.ensure_sglang_tensor_image_size", + lambda: install_calls.append(True), + ) + mock_sglang_cli(model="/tmp") + + await parse_args(sys.argv[1:]) + + assert install_calls == ([True] if is_multimodal else []) + + def test_compat_filters_async_generate_kwargs_for_older_engines(): class OldEngine: async def async_generate(self, input_ids=None, sampling_params=None):