Fix DeepSeek-OCR batching crash on variable local-crop counts (#33214)

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 = (
+        # Different images may have a different number of local crop patches
-            torch.stack([item.images_crop for item in mm_items], dim=0)
+        # (images_crop shape [1, num_patches, 3, H, W] varies per item), so we
-            .type(target_dtype)
+        # cannot stack them into a single tensor. Process each item separately
-            .to(device=pixel_values.device)
+        # and concatenate the resulting feature sequences.
-        )
+        vision_feature_lists: List[torch.Tensor] = []
-        images_spatial_crop = (
+        for idx, item in enumerate(mm_items):
-            torch.cat([item.images_spatial_crop for item in mm_items], dim=0)
+            pixel_values = item.feature.unsqueeze(0).type(target_dtype)
-            .type(torch.long)
+            images_crop = (
-            .to(device=pixel_values.device)
+                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_crop.dim() == 6
-        assert images_spatial_crop.dim() == 3
+            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

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()