diff --git a/python/sglang/srt/models/deepseek_ocr.py b/python/sglang/srt/models/deepseek_ocr.py index 1b70614f6..859b4ebf2 100644 --- a/python/sglang/srt/models/deepseek_ocr.py +++ b/python/sglang/srt/models/deepseek_ocr.py @@ -1689,30 +1689,41 @@ class DeepseekOCRForCausalLM(nn.Module): else self.vision_model.dtype ) has_local_crops = self._collect_mm_flag(mm_items, "has_local_crops") - pixel_values = torch.stack([item.feature for item in mm_items], dim=0).type( - target_dtype - ) - images_crop = ( - torch.stack([item.images_crop for item in mm_items], dim=0) - .type(target_dtype) - .to(device=pixel_values.device) - ) - images_spatial_crop = ( - torch.cat([item.images_spatial_crop for item in mm_items], dim=0) - .type(torch.long) - .to(device=pixel_values.device) - ) + # Different images may have a different number of local crop patches + # (images_crop shape [1, num_patches, 3, H, W] varies per item), so we + # cannot stack them into a single tensor. Process each item separately + # and concatenate the resulting feature sequences. + vision_feature_lists: List[torch.Tensor] = [] + for idx, item in enumerate(mm_items): + pixel_values = item.feature.unsqueeze(0).type(target_dtype) + images_crop = ( + item.images_crop.unsqueeze(0) + .type(target_dtype) + .to(device=pixel_values.device) + ) + images_spatial_crop = item.images_spatial_crop.type(torch.long).to( + device=pixel_values.device + ) + if images_spatial_crop.dim() == 2: + images_spatial_crop = images_spatial_crop.unsqueeze(0) - assert images_crop.dim() == 6 - assert images_spatial_crop.dim() == 3 + assert images_crop.dim() == 6 + assert images_spatial_crop.dim() == 3 + + item_has_local_crops = ( + [has_local_crops[idx]] if has_local_crops is not None else None + ) + + vision_feature_lists.extend( + self._pixel_values_to_embedding( + pixel_values=pixel_values, + images_crop=images_crop, + images_spatial_crop=images_spatial_crop, + has_local_crops=item_has_local_crops, + ) + ) - vision_feature_lists = self._pixel_values_to_embedding( - pixel_values=pixel_values, - images_crop=images_crop, - images_spatial_crop=images_spatial_crop, - has_local_crops=has_local_crops, - ) vision_features = torch.cat(vision_feature_lists, dim=0).type(target_dtype) return vision_features diff --git a/test/registered/xpu/test_deepseek_ocr.py b/test/registered/xpu/test_deepseek_ocr.py index 33e38c789..5e8792829 100644 --- a/test/registered/xpu/test_deepseek_ocr.py +++ b/test/registered/xpu/test_deepseek_ocr.py @@ -6,9 +6,13 @@ import json import os import unittest from pathlib import Path +from unittest.mock import patch import requests +import torch +from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem +from sglang.srt.models.deepseek_ocr import DeepseekOCRForCausalLM from sglang.srt.utils import kill_process_tree from sglang.srt.utils.hf_transformers import get_tokenizer from sglang.test.ci.ci_register import register_xpu_ci @@ -109,5 +113,93 @@ class TestDeepSeekOCR(CustomTestCase): self.run_decode() +def _device_available(device: str) -> bool: + if device == "cpu": + return True + if device == "cuda": + return torch.cuda.is_available() + if device == "xpu": + return hasattr(torch, "xpu") and torch.xpu.is_available() + return False + + +class TestDeepSeekOCRProcessImageInputBatchedUnequalCrops(CustomTestCase): + """Regression: `_process_image_input` used to `torch.stack` `images_crop` + across `mm_items`, which crashed when two OCR requests in the same batch + had different `num_patches` (e.g. 6 vs 4). Guards against reintroducing + the stack. Runs on each device the host provides (CPU / XPU / CUDA).""" + + def _make_item(self, device, num_patches, tiles_w, tiles_h): + item = MultimodalDataItem(modality=Modality.IMAGE) + item.feature = torch.zeros(3, 640, 640, dtype=torch.float32, device=device) + item.images_crop = torch.zeros( + 1, num_patches, 3, 640, 640, dtype=torch.float32, device=device + ) + item.images_spatial_crop = torch.tensor( + [[[tiles_w, tiles_h]]], dtype=torch.long, device=device + ) + item.has_local_crops = True + return item + + def _run_batched_unequal_num_patches(self, device: str): + if not _device_available(device): + self.skipTest(f"device {device!r} not available on this host") + torch_device = torch.device(device) + + items = [ + self._make_item(torch_device, num_patches=6, tiles_w=3, tiles_h=2), + self._make_item(torch_device, num_patches=4, tiles_w=2, tiles_h=2), + ] + + def fake_pixel_values_to_embedding( + pixel_values, images_crop, images_spatial_crop, has_local_crops + ): + self.assertEqual(pixel_values.shape[0], 1) + self.assertEqual(images_crop.dim(), 6) + self.assertEqual(images_spatial_crop.dim(), 3) + self.assertEqual(pixel_values.device.type, torch_device.type) + self.assertEqual(images_crop.device.type, torch_device.type) + self.assertEqual(images_spatial_crop.device.type, torch_device.type) + n_patches = images_crop.shape[2] + return [torch.zeros(n_patches, 8, dtype=torch.float32, device=torch_device)] + + instance = DeepseekOCRForCausalLM.__new__(DeepseekOCRForCausalLM) + instance.is_ocr2 = False + + stub_param = torch.zeros(1, dtype=torch.float32, device=torch_device) + + class _Stub: + dtype = torch.float32 + + @staticmethod + def parameters(): + return iter([stub_param]) + + instance.sam_model = _Stub() + instance.vision_model = _Stub() + + with patch.object( + DeepseekOCRForCausalLM, + "_pixel_values_to_embedding", + side_effect=fake_pixel_values_to_embedding, + autospec=False, + ): + out = DeepseekOCRForCausalLM._process_image_input(instance, items) + + # Feature sequences from both items must be concatenated in order. + # 6 rows for item A + 4 rows for item B = 10 rows. + self.assertEqual(out.shape, (10, 8)) + self.assertEqual(out.device.type, torch_device.type) + + def test_batched_unequal_num_patches_no_crash_cpu(self): + self._run_batched_unequal_num_patches("cpu") + + def test_batched_unequal_num_patches_no_crash_xpu(self): + self._run_batched_unequal_num_patches("xpu") + + def test_batched_unequal_num_patches_no_crash_cuda(self): + self._run_batched_unequal_num_patches("cuda") + + if __name__ == "__main__": unittest.main()