A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Bugfix] Fix deepseek-ocr multi-image inference and add `merge_by_field_config=True` with tensor schema support (#27361)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Oct 23, 2025 at 08:15 UTC
2566dca2a9e4e24c941845905e0ebad62441a1fa
4 files changed
+112
-66
examples/offline_inference/vision_language_multi_image.py
+48
-2
@@ -44,6 +44,7 @@ class ModelRequestData(NamedTuple):
44
stop_token_ids: list[int] | None = None
45
chat_template: str | None = None
46
lora_requests: list[LoRARequest] | None = None
47
+ sampling_params: SamplingParams | None = None
48
49
50
# NOTE: The default `max_num_seqs` and `max_model_len` may result in OOM on
@@ -201,6 +202,46 @@ def load_deepseek_vl2(question: str, image_urls: list[str]) -> ModelRequestData:
202
)
203
204
205
+def load_deepseek_ocr(question: str, image_urls: list[str]) -> ModelRequestData:
206
+ from vllm.model_executor.models.deepseek_ocr import NGramPerReqLogitsProcessor
207
+
208
+ model_name = "deepseek-ai/DeepSeek-OCR"
209
+
210
+ engine_args = EngineArgs(
211
+ model=model_name,
212
+ max_num_seqs=2,
213
+ limit_mm_per_prompt={"image": len(image_urls)},
214
+ logits_processors=[NGramPerReqLogitsProcessor],
215
+ )
216
+
217
+ placeholder = "<image>\n" * len(image_urls)
218
+ prompt = placeholder + question
219
+
220
+ # The following sampling params config is taken from
221
+ # the official Deepseek-OCR inference example.
222
+ # (IMPORTANT) Use the custom logits processor and avoid skipping
223
+ # special tokens for this model for the optimal OCR performance.
224
+ sampling_params = SamplingParams(
225
+ temperature=0.0,
226
+ max_tokens=8192,
227
+ # ngram logit processor args
228
+ extra_args=dict(
229
+ ngram_size=30,
230
+ window_size=90,
231
+ # whitelist: <td>, </td>
232
+ whitelist_token_ids={128821, 128822},
233
+ ),
234
+ skip_special_tokens=False,
235
+ )
236
+
237
+ return ModelRequestData(
238
+ engine_args=engine_args,
239
+ prompt=prompt,
240
+ image_data=[fetch_image(url) for url in image_urls],
241
+ sampling_params=sampling_params,
242
+ )
243
+
244
+
245
def load_gemma3(question: str, image_urls: list[str]) -> ModelRequestData:
246
model_name = "google/gemma-3-4b-it"
247
@@ -1253,6 +1294,7 @@ model_example_map = {
1294
"bee": load_bee,
1295
"command_a_vision": load_command_a_vision,
1296
"deepseek_vl_v2": load_deepseek_vl2,
1297
+ "deepseek_ocr": load_deepseek_ocr,
1298
"gemma3": load_gemma3,
1299
"h2ovl_chat": load_h2ovl,
1300
"hyperclovax_seed_vision": load_hyperclovax_seed_vision,
@@ -1325,8 +1367,12 @@ def run_chat(model: str, question: str, image_urls: list[str], seed: int | None)
1367
engine_args = asdict(req_data.engine_args) | {"seed": seed}
1368
llm = LLM(**engine_args)
1369
1328
- sampling_params = SamplingParams(
1329
- temperature=0.0, max_tokens=256, stop_token_ids=req_data.stop_token_ids
1370
+ sampling_params = (
1371
+ SamplingParams(
1372
+ temperature=0.0, max_tokens=256, stop_token_ids=req_data.stop_token_ids
1373
+ )
1374
+ if req_data.sampling_params is None
1375
+ else req_data.sampling_params
1376
)
1377
outputs = llm.chat(
1378
[
tests/models/multimodal/processing/test_common.py
+1
@@ -332,6 +332,7 @@ def _test_processing_correctness_one(
332
"facebook/chameleon-7b",
333
"CohereLabs/command-a-vision-07-2025",
334
"deepseek-ai/deepseek-vl2-tiny",
335
+ "deepseek-ai/DeepSeek-OCR",
336
"baidu/ERNIE-4.5-VL-28B-A3B-PT",
337
"adept/fuyu-8b",
338
"google/gemma-3-4b-it",
vllm/model_executor/models/deepseek_ocr.py
+58
-55
@@ -4,6 +4,7 @@
4
5
import math
6
from collections.abc import Iterable, Mapping, Sequence
7
+from typing import Annotated, Literal
8
9
import torch
10
import torch.nn as nn
@@ -53,6 +54,7 @@ from vllm.transformers_utils.processors.deepseek_ocr import (
54
count_tiles,
55
)
56
from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config
57
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
58
from vllm.v1.sample.logits_processor import (
59
AdapterLogitsProcessor,
60
RequestLogitsProcessor,
@@ -65,6 +67,28 @@ from .deepseek_vl2 import MlpProjector
67
_IMAGE_TOKEN = "<image>"
68
69
70
+class DeepseekOCRImagePixelInputs(TensorSchema):
71
+ """
72
+ Dimensions:
73
+ - b: Batch size
74
+ - n: Number of images
75
+ - p: Number of patches
76
+ - base_size: Base size of the processor
77
+ - image_size: Image size of the processor
78
+ """
79
+
80
+ type: Literal["pixel_values"]
81
+ data: Annotated[
82
+ torch.Tensor,
83
+ TensorShape("bn", 3, "base_size", "base_size", dynamic_dims={"bnp"}),
84
+ ]
85
+ images_crop: Annotated[
86
+ torch.Tensor,
87
+ TensorShape("bnp", 3, "image_size", "image_size", dynamic_dims={"bnp"}),
88
+ ]
89
+ images_spatial_crop: Annotated[torch.Tensor, TensorShape("bn", 2)]
90
+
91
+
92
class NoRepeatNGramLogitsProcessor:
93
def __init__(
94
self,
@@ -260,10 +284,15 @@ class DeepseekOCRMultiModalProcessor(
284
hf_inputs: BatchFeature,
285
hf_processor_mm_kwargs: Mapping[str, object],
286
) -> Mapping[str, MultiModalFieldConfig]:
287
+ images_spatial_crop = hf_inputs.get("images_spatial_crop", torch.empty((0, 2)))
288
+ is_tiled = (images_spatial_crop[:, 0] > 1) | (images_spatial_crop[:, 1] > 1)
289
+ patches_per_image = torch.where(is_tiled, images_spatial_crop.prod(dim=-1), 0)
290
return dict(
291
pixel_values=MultiModalFieldConfig.batched("image"),
292
images_spatial_crop=MultiModalFieldConfig.batched("image"),
266
- images_crop=MultiModalFieldConfig.batched("image"),
293
+ images_crop=MultiModalFieldConfig.flat_from_sizes(
294
+ "image", patches_per_image
295
+ ),
296
)
297
298
def _get_prompt_updates(
@@ -302,35 +331,6 @@ class DeepseekOCRMultiModalProcessor(
331
)
332
]
333
305
- # TODO(Isotr0py): Check if we still need this workaround for
306
- # deepseek-ocr processor.
307
- # def _cached_apply_hf_processor(
308
- # self,
309
- # prompt: str | list[int],
310
- # mm_data_items: MultiModalDataItems,
311
- # hf_processor_mm_kwargs: Mapping[str, object],
312
- # tokenization_kwargs: Mapping[str, object],
313
- # mm_uuids: MultiModalUUIDDict | None = None,
314
- # ) -> tuple[list[int], MultiModalKwargs, bool]:
315
- # # The processor logic is different for len(images) <= 2 vs > 2
316
- # # Since the processing cache assumes that the processor output is
317
- # # invariant of how many images are passed per prompt, we only
318
- # # perform caching for the most common case
319
- # if mm_data_items.get_count("image", strict=False) > 2:
320
- # # This code path corresponds to the cache being disabled
321
- # return self._apply_hf_processor_main(
322
- # prompt=prompt,
323
- # mm_items=mm_data_items,
324
- # hf_processor_mm_kwargs=hf_processor_mm_kwargs,
325
- # enable_hf_prompt_update=True,
326
- # )
327
-
328
- # return super()._cached_apply_hf_processor(
329
- # prompt=prompt,
330
- # mm_data_items=mm_data_items,
331
- # hf_processor_mm_kwargs=hf_processor_mm_kwargs,
332
- # )
333
-
334
335
@MULTIMODAL_REGISTRY.register_processor(
336
DeepseekOCRMultiModalProcessor,
@@ -338,6 +338,8 @@ class DeepseekOCRMultiModalProcessor(
338
dummy_inputs=DeepseekOCRDummyInputsBuilder,
339
)
340
class DeepseekOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
341
+ merge_by_field_config = True
342
+
343
hf_to_vllm_mapper = WeightsMapper(
344
orig_to_new_prefix={
345
# map prefix for language backbone
@@ -389,6 +391,7 @@ class DeepseekOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
391
self.vision_model = DeepCLIPVisionTransformer(
392
config=clip_vision_config,
393
quant_config=quant_config,
394
+ prefix=maybe_prefix(prefix, "vision_model"),
395
)
396
397
self.projector = MlpProjector(self.projector_config)
@@ -426,7 +429,9 @@ class DeepseekOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
429
self.language_model.make_empty_intermediate_tensors
430
)
431
429
- def _parse_and_validate_image_input(self, **kwargs: object):
432
+ def _parse_and_validate_image_input(
433
+ self, **kwargs: object
434
+ ) -> DeepseekOCRImagePixelInputs | None:
435
pixel_values = kwargs.pop("pixel_values", None)
436
images_spatial_crop = kwargs.pop("images_spatial_crop", None)
437
images_crop = kwargs.pop("images_crop", None)
@@ -435,23 +440,16 @@ class DeepseekOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
440
return None
441
442
if pixel_values is not None:
438
- if not isinstance(pixel_values, (torch.Tensor, list)):
439
- raise ValueError(
440
- f"Incorrect type of pixel values. Got type: {type(pixel_values)}"
441
- )
442
-
443
- if not isinstance(images_spatial_crop, (torch.Tensor, list)):
444
- raise ValueError(
445
- "Incorrect type of image sizes. "
446
- f"Got type: {type(images_spatial_crop)}"
447
- )
448
-
449
- if not isinstance(images_crop, (torch.Tensor, list)):
450
- raise ValueError(
451
- f"Incorrect type of image crop. Got type: {type(images_crop)}"
452
- )
453
-
454
- return [pixel_values, images_crop, images_spatial_crop]
443
+ base_size = self.vision_config.image_size
444
+ return DeepseekOCRImagePixelInputs(
445
+ type="pixel_values",
446
+ data=pixel_values,
447
+ images_crop=images_crop,
448
+ images_spatial_crop=images_spatial_crop,
449
+ resolve_bindings={
450
+ "base_size": base_size,
451
+ },
452
+ )
453
454
raise AssertionError("This line should be unreachable.")
455
@@ -518,10 +516,13 @@ class DeepseekOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
516
) -> NestedTensors:
517
images_in_this_batch = []
518
519
+ is_tiled = (images_spatial_crop[:, 0] > 1) | (images_spatial_crop[:, 1] > 1)
520
+ patches_per_image = torch.where(is_tiled, images_spatial_crop.prod(dim=-1), 0)
521
+ images_crop = images_crop.split(patches_per_image.tolist())
522
for jdx in range(images_spatial_crop.size(0)):
522
- patches = images_crop[jdx][0].to(torch.bfloat16)
523
- image_ori = pixel_values[jdx]
524
- crop_shape = images_spatial_crop[jdx][0]
523
+ patches = images_crop[jdx]
524
+ image_ori = pixel_values[[jdx]]
525
+ crop_shape = images_spatial_crop[jdx]
526
527
global_features = self._encode_global_features(image_ori)
528
local_features = self._encode_local_features(patches, crop_shape)
@@ -540,10 +541,12 @@ class DeepseekOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
541
542
return images_in_this_batch
543
543
- def _process_image_input(self, image_input) -> torch.Tensor:
544
- pixel_values = image_input[0].to(torch.bfloat16)
545
- images_crop = image_input[1]
546
- images_spatial_crop = image_input[2].to(dtype=torch.long)
544
+ def _process_image_input(
545
+ self, image_input: DeepseekOCRImagePixelInputs
546
+ ) -> torch.Tensor:
547
+ pixel_values = image_input.data
548
+ images_crop = image_input.images_crop
549
+ images_spatial_crop = image_input.images_spatial_crop.to(dtype=torch.long)
550
551
vision_features = self._pixel_values_to_embedding(
552
pixel_values=pixel_values,
vllm/transformers_utils/processors/deepseek_ocr.py
+5
-9
@@ -411,20 +411,16 @@ class DeepseekOCRProcessor(ProcessorMixin):
411
images_seq_mask = images_seq_mask[:-1]
412
413
if len(images_list) == 0:
414
- pixel_values = torch.zeros((1, 3, self.base_size, self.base_size))
415
- images_spatial_crop = torch.zeros((1, 1), dtype=torch.long)
416
- images_crop = torch.zeros(
417
- (1, 3, self.image_size, self.image_size)
418
- ).unsqueeze(0)
414
+ pixel_values = torch.zeros((0, 3, self.base_size, self.base_size))
415
+ images_spatial_crop = torch.zeros((0, 2), dtype=torch.long)
416
+ images_crop = torch.zeros((0, 3, self.image_size, self.image_size))
417
else:
418
pixel_values = torch.stack(images_list, dim=0)
419
images_spatial_crop = torch.tensor(images_spatial_crop, dtype=torch.long)
420
if images_crop_list:
423
- images_crop = torch.stack(images_crop_list, dim=0).unsqueeze(0)
421
+ images_crop = torch.stack(images_crop_list, dim=0)
422
else:
425
- images_crop = torch.zeros(
426
- (1, 3, self.image_size, self.image_size)
427
- ).unsqueeze(0)
423
+ images_crop = torch.zeros((0, 3, self.image_size, self.image_size))
424
425
input_ids = input_ids.unsqueeze(0)
426