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
[New Model]Donut model (#23229)
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
汪志鹏 committed
Aug 24, 2025 at 20:52 UTC
416f05929ac66f5ae364936b70087fc60cacee4b
11 files changed
+1240
-3
docs/models/supported_models.md
+1
@@ -615,6 +615,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
615
| `ChameleonForConditionalGeneration` | Chameleon | T + I | `facebook/chameleon-7b`, etc. | | ✅︎ | ✅︎ |
616
| `Cohere2VisionForConditionalGeneration` | Command A Vision | T + I<sup>+</sup> | `CohereLabs/command-a-vision-07-2025`, etc. | | ✅︎ | ✅︎ |
617
| `DeepseekVLV2ForCausalLM`<sup>^</sup> | DeepSeek-VL2 | T + I<sup>+</sup> | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ | ✅︎ |
618
+| `DonutForConditionalGeneration`<sup>^</sup> | Donut | T + I | `ByteDance/Dolphin`, `naver-clova-ix/donut-base-finetuned-docvqa`, etc. | | | |
619
| `Florence2ForConditionalGeneration` | Florence-2 | T + I | `microsoft/Florence-2-base`, `microsoft/Florence-2-large`, etc. | | | |
620
| `FuyuForCausalLM` | Fuyu | T + I | `adept/fuyu-8b`, etc. | | ✅︎ | ✅︎ |
621
| `Gemma3ForConditionalGeneration` | Gemma 3 | T + I<sup>+</sup> | `google/gemma-3-4b-it`, `google/gemma-3-27b-it`, etc. | ✅︎ | ✅︎ | ⚠️ |
examples/offline_inference/dolphin.py
new
+311
@@ -0,0 +1,311 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import argparse
5
+import copy
6
+import os
7
+from dataclasses import dataclass
8
+
9
+import cv2
10
+import numpy as np
11
+import regex as re
12
+from PIL import Image
13
+from transformers import DonutProcessor
14
+
15
+from vllm import LLM, SamplingParams
16
+from vllm.inputs import ExplicitEncoderDecoderPrompt, TextPrompt, TokensPrompt
17
+from vllm.multimodal.utils import fetch_image
18
+
19
+
20
+# Copied from https://github.com/bytedance/Dolphin/utils/utils.py
21
+@dataclass
22
+class ImageDimensions:
23
+ original_w: int
24
+ original_h: int
25
+ padded_w: int
26
+ padded_h: int
27
+
28
+
29
+# Copied from https://github.com/bytedance/Dolphin/utils/utils.py
30
+def map_to_original_coordinates(
31
+ x1, y1, x2, y2, dims: ImageDimensions
32
+) -> tuple[int, int, int, int]:
33
+ try:
34
+ top = (dims.padded_h - dims.original_h) // 2
35
+ left = (dims.padded_w - dims.original_w) // 2
36
+ orig_x1 = max(0, x1 - left)
37
+ orig_y1 = max(0, y1 - top)
38
+ orig_x2 = min(dims.original_w, x2 - left)
39
+ orig_y2 = min(dims.original_h, y2 - top)
40
+ if orig_x2 <= orig_x1:
41
+ orig_x2 = min(orig_x1 + 1, dims.original_w)
42
+ if orig_y2 <= orig_y1:
43
+ orig_y2 = min(orig_y1 + 1, dims.original_h)
44
+ return int(orig_x1), int(orig_y1), int(orig_x2), int(orig_y2)
45
+ except Exception as e:
46
+ print(f"map_to_original_coordinates error: {str(e)}")
47
+ return 0, 0, min(100, dims.original_w), min(100, dims.original_h)
48
+
49
+
50
+# Copied from https://github.com/bytedance/Dolphin/utils/utils.py
51
+def adjust_box_edges(image, boxes: list[list[float]], max_pixels=15, threshold=0.2):
52
+ if isinstance(image, str):
53
+ image = cv2.imread(image)
54
+ img_h, img_w = image.shape[:2]
55
+ new_boxes = []
56
+ for box in boxes:
57
+ best_box = copy.deepcopy(box)
58
+
59
+ def check_edge(img, current_box, i, is_vertical):
60
+ edge = current_box[i]
61
+ gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
62
+ _, binary = cv2.threshold(
63
+ gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU
64
+ )
65
+ if is_vertical:
66
+ line = binary[current_box[1] : current_box[3] + 1, edge]
67
+ else:
68
+ line = binary[edge, current_box[0] : current_box[2] + 1]
69
+ transitions = np.abs(np.diff(line))
70
+ return np.sum(transitions) / len(transitions)
71
+
72
+ edges = [(0, -1, True), (2, 1, True), (1, -1, False), (3, 1, False)]
73
+ current_box = copy.deepcopy(box)
74
+ current_box[0] = min(max(current_box[0], 0), img_w - 1)
75
+ current_box[1] = min(max(current_box[1], 0), img_h - 1)
76
+ current_box[2] = min(max(current_box[2], 0), img_w - 1)
77
+ current_box[3] = min(max(current_box[3], 0), img_h - 1)
78
+
79
+ for i, direction, is_vertical in edges:
80
+ best_score = check_edge(image, current_box, i, is_vertical)
81
+ if best_score <= threshold:
82
+ continue
83
+ for step in range(max_pixels):
84
+ current_box[i] += direction
85
+ if i == 0 or i == 2:
86
+ current_box[i] = min(max(current_box[i], 0), img_w - 1)
87
+ else:
88
+ current_box[i] = min(max(current_box[i], 0), img_h - 1)
89
+ score = check_edge(image, current_box, i, is_vertical)
90
+ if score < best_score:
91
+ best_score = score
92
+ best_box = copy.deepcopy(current_box)
93
+ if score <= threshold:
94
+ break
95
+ new_boxes.append(best_box)
96
+ return new_boxes
97
+
98
+
99
+# Copied from https://github.com/bytedance/Dolphin/utils/utils.py
100
+def process_coordinates(coords, padded_image, dims: ImageDimensions, previous_box=None):
101
+ try:
102
+ x1, y1 = int(coords[0] * dims.padded_w), int(coords[1] * dims.padded_h)
103
+ x2, y2 = int(coords[2] * dims.padded_w), int(coords[3] * dims.padded_h)
104
+ x1, y1, x2, y2 = (
105
+ max(0, min(x1, dims.padded_w - 1)),
106
+ max(0, min(y1, dims.padded_h - 1)),
107
+ max(0, min(x2, dims.padded_w)),
108
+ max(0, min(y2, dims.padded_h)),
109
+ )
110
+ if x2 <= x1:
111
+ x2 = min(x1 + 1, dims.padded_w)
112
+ if y2 <= y1:
113
+ y2 = min(y1 + 1, dims.padded_h)
114
+ new_boxes = adjust_box_edges(padded_image, [[x1, y1, x2, y2]])
115
+ x1, y1, x2, y2 = new_boxes[0]
116
+ x1, y1, x2, y2 = (
117
+ max(0, min(x1, dims.padded_w - 1)),
118
+ max(0, min(y1, dims.padded_h - 1)),
119
+ max(0, min(x2, dims.padded_w)),
120
+ max(0, min(y2, dims.padded_h)),
121
+ )
122
+ if x2 <= x1:
123
+ x2 = min(x1 + 1, dims.padded_w)
124
+ if y2 <= y1:
125
+ y2 = min(y1 + 1, dims.padded_h)
126
+ if previous_box is not None:
127
+ prev_x1, prev_y1, prev_x2, prev_y2 = previous_box
128
+ if (x1 < prev_x2 and x2 > prev_x1) and (y1 < prev_y2 and y2 > prev_y1):
129
+ y1 = prev_y2
130
+ y1 = min(y1, dims.padded_h - 1)
131
+ if y2 <= y1:
132
+ y2 = min(y1 + 1, dims.padded_h)
133
+ new_previous_box = [x1, y1, x2, y2]
134
+ orig_x1, orig_y1, orig_x2, orig_y2 = map_to_original_coordinates(
135
+ x1, y1, x2, y2, dims
136
+ )
137
+ return x1, y1, x2, y2, orig_x1, orig_y1, orig_x2, orig_y2, new_previous_box
138
+ except Exception as e:
139
+ print(f"process_coordinates error: {str(e)}")
140
+ orig_x1, orig_y1, orig_x2, orig_y2 = (
141
+ 0,
142
+ 0,
143
+ min(100, dims.original_w),
144
+ min(100, dims.original_h),
145
+ )
146
+ return 0, 0, 100, 100, orig_x1, orig_y1, orig_x2, orig_y2, [0, 0, 100, 100]
147
+
148
+
149
+# Copied from https://github.com/bytedance/Dolphin/utils/utils.py
150
+def prepare_image(image) -> tuple[np.ndarray, ImageDimensions]:
151
+ try:
152
+ image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
153
+ original_h, original_w = image_cv.shape[:2]
154
+ max_size = max(original_h, original_w)
155
+ top = (max_size - original_h) // 2
156
+ bottom = max_size - original_h - top
157
+ left = (max_size - original_w) // 2
158
+ right = max_size - original_w - left
159
+ padded_image = cv2.copyMakeBorder(
160
+ image_cv, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(0, 0, 0)
161
+ )
162
+ padded_h, padded_w = padded_image.shape[:2]
163
+ dimensions = ImageDimensions(
164
+ original_w=original_w,
165
+ original_h=original_h,
166
+ padded_w=padded_w,
167
+ padded_h=padded_h,
168
+ )
169
+ return padded_image, dimensions
170
+ except Exception as e:
171
+ print(f"prepare_image error: {str(e)}")
172
+ h, w = image.height, image.width
173
+ dimensions = ImageDimensions(original_w=w, original_h=h, padded_w=w, padded_h=h)
174
+ return np.zeros((h, w, 3), dtype=np.uint8), dimensions
175
+
176
+
177
+# Copied from https://github.com/bytedance/Dolphin/utils/utils.py
178
+def parse_layout_string(bbox_str):
179
+ """Parse layout string using regular expressions"""
180
+ pattern = r"\[(\d*\.?\d+),\s*(\d*\.?\d+),\s*(\d*\.?\d+),\s*(\d*\.?\d+)\]\s*(\w+)"
181
+ matches = re.finditer(pattern, bbox_str)
182
+
183
+ parsed_results = []
184
+ for match in matches:
185
+ coords = [float(match.group(i)) for i in range(1, 5)]
186
+ label = match.group(5).strip()
187
+ parsed_results.append((coords, label))
188
+
189
+ return parsed_results
190
+
191
+
192
+model_id = "ByteDance/Dolphin"
193
+
194
+# The input image size for Dolphin is 896 x 896,
195
+# and the patch_size is 4 x 4.
196
+# Therefore, the initial number of patches is:
197
+# Height: 896 / 4 = 224 patches
198
+# Width: 896 / 4 = 224 patches
199
+
200
+# The Dolphin model uses a staged downsampling approach,
201
+# defined by the "depths": [2, 2, 14, 2] configuration.
202
+# Before entering stages 2, 3, and 4, a "Patch Merging" operation is performed,
203
+# which halves the feature map's dimensions (dividing both height and width by 2).
204
+# Before Stage 2: The size changes from 224 x 224 to (224/2) x (224/2) = 112 x 112.
205
+# Before Stage 3: The size changes from 112 x 112 to (112/2) x (112/2) = 56 x 56.
206
+# Before Stage 4: The size changes from 56 x 56 to (56/2) x (56/2) = 28 x 28.
207
+
208
+# Because vLLM needs to fill the image features with an encoder_prompt,
209
+# and the encoder_prompt will have `<pad>` tokens added when tokenized,
210
+# we need to construct an encoder_prompt with a length of 28 x 28 - 1 = 783.
211
+encoder_prompt = "".join(["0"] * 783)
212
+sampling_params = SamplingParams(
213
+ temperature=0.0,
214
+ max_tokens=2048,
215
+)
216
+
217
+processor = DonutProcessor.from_pretrained(model_id)
218
+llm = LLM(
219
+ model=model_id,
220
+ dtype="float16",
221
+ max_num_seqs=8,
222
+ hf_overrides={"architectures": ["DonutForConditionalGeneration"]},
223
+)
224
+
225
+parser = argparse.ArgumentParser()
226
+parser.add_argument(
227
+ "--image_path", type=str, default=None, help="Path to a local image file."
228
+)
229
+args = parser.parse_args()
230
+
231
+if args.image_path:
232
+ if not os.path.exists(args.image_path):
233
+ raise FileNotFoundError(f"Error: File not found at {args.image_path}")
234
+ image = Image.open(args.image_path).convert("RGB")
235
+else:
236
+ image = fetch_image(
237
+ "https://huggingface.co/datasets/hf-internal-testing/example-documents/resolve/main/jpeg_images/0.jpg"
238
+ )
239
+
240
+
241
+prompt = "Parse the reading order of this document. "
242
+decoder_prompt = f"<s>{prompt}<Answer/>"
243
+decoder_prompt_tokens = TokensPrompt(
244
+ prompt_token_ids=processor.tokenizer(decoder_prompt, add_special_tokens=False)[
245
+ "input_ids"
246
+ ]
247
+)
248
+enc_dec_prompt = ExplicitEncoderDecoderPrompt(
249
+ encoder_prompt=TextPrompt(prompt=encoder_prompt, multi_modal_data={"image": image}),
250
+ decoder_prompt=decoder_prompt_tokens,
251
+)
252
+layout_outputs = llm.generate(prompts=enc_dec_prompt, sampling_params=sampling_params)
253
+layout_result_str = layout_outputs[0].outputs[0].text
254
+print(f"Layout analysis output:\n{layout_result_str}")
255
+
256
+padded_image, dims = prepare_image(image)
257
+layout_results = parse_layout_string(layout_result_str)
258
+text_table_elements = []
259
+previous_box = None
260
+reading_order = 0
261
+for bbox_coords, label in layout_results:
262
+ if label == "fig":
263
+ continue
264
+ try:
265
+ x1, y1, x2, y2, orig_x1, orig_y1, orig_x2, orig_y2, previous_box = (
266
+ process_coordinates(bbox_coords, padded_image, dims, previous_box)
267
+ )
268
+ cropped = padded_image[y1:y2, x1:x2]
269
+ if cropped.size > 0 and cropped.shape[0] > 3 and cropped.shape[1] > 3:
270
+ pil_crop = Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))
271
+ prompt_ocr = (
272
+ "Parse the table in the image. "
273
+ if label == "tab"
274
+ else "Read text in the image. "
275
+ )
276
+ text_table_elements.append(
277
+ {
278
+ "crop": pil_crop,
279
+ "prompt": prompt_ocr,
280
+ "reading_order": reading_order,
281
+ }
282
+ )
283
+ reading_order += 1
284
+ except Exception as e:
285
+ print(f"Error processing bbox (label: {label}): {str(e)}")
286
+ continue
287
+
288
+if text_table_elements:
289
+ batch_prompts = []
290
+ for elem in text_table_elements:
291
+ decoder_prompt_str = f"<s>{elem['prompt']}<Answer/>"
292
+ decoder_prompt_tokens = TokensPrompt(
293
+ prompt_token_ids=processor.tokenizer(
294
+ decoder_prompt_str, add_special_tokens=False
295
+ )["input_ids"]
296
+ )
297
+ enc_dec_prompt = ExplicitEncoderDecoderPrompt(
298
+ encoder_prompt=TextPrompt(
299
+ prompt=encoder_prompt, multi_modal_data={"image": elem["crop"]}
300
+ ),
301
+ decoder_prompt=decoder_prompt_tokens,
302
+ )
303
+ batch_prompts.append(enc_dec_prompt)
304
+ batch_outputs = llm.generate(prompts=batch_prompts, sampling_params=sampling_params)
305
+ for i, output in enumerate(batch_outputs):
306
+ text_table_elements[i]["text"] = output.outputs[0].text.strip()
307
+
308
+print("------" * 8)
309
+text_table_elements.sort(key=lambda x: x["reading_order"])
310
+for elem in text_table_elements:
311
+ print(elem.get("text", ""))
examples/offline_inference/encoder_decoder_multimodal.py
+46
@@ -13,6 +13,7 @@ from typing import NamedTuple
13
from vllm import LLM, EngineArgs, PromptType, SamplingParams
14
from vllm.assets.audio import AudioAsset
15
from vllm.assets.image import ImageAsset
16
+from vllm.multimodal.utils import fetch_image
17
from vllm.utils import FlexibleArgumentParser
18
19
@@ -21,6 +22,50 @@ class ModelRequestData(NamedTuple):
22
prompts: Sequence[PromptType]
23
24
25
+def run_donut():
26
+ engine_args = EngineArgs(
27
+ model="naver-clova-ix/donut-base-finetuned-docvqa",
28
+ max_num_seqs=2,
29
+ limit_mm_per_prompt={"image": 1},
30
+ dtype="float16",
31
+ hf_overrides={"architectures": ["DonutForConditionalGeneration"]},
32
+ )
33
+
34
+ # The input image size for donut-base-finetuned-docvqa is 2560 x 1920,
35
+ # and the patch_size is 4 x 4.
36
+ # Therefore, the initial number of patches is:
37
+ # Height: 1920 / 4 = 480 patches
38
+ # Width: 2560 / 4 = 640 patches
39
+ # The Swin model uses a staged downsampling approach,
40
+ # defined by the "depths": [2, 2, 14, 2] configuration.
41
+ # Before entering stages 2, 3, and 4, a "Patch Merging" operation is performed,
42
+ # which halves the feature map's dimensions (dividing both height and width by 2).
43
+ # Before Stage 2: The size changes from 480 x 640 to (480/2) x (640/2) = 240 x 320.
44
+ # Before Stage 3: The size changes from 240 x 320 to (240/2) x (320/2) = 120 x 160.
45
+ # Before Stage 4: The size changes from 120 x 160 to (120/2) x (160/2) = 60 x 80.
46
+ # Because vLLM needs to fill the image features with an encoder_prompt,
47
+ # and the encoder_prompt will have `<pad>` tokens added when tokenized,
48
+ # we need to construct an encoder_prompt with a length of 60 x 80 - 1 = 4799.
49
+ prompts = [
50
+ {
51
+ "encoder_prompt": {
52
+ "prompt": "".join(["$"] * 4799),
53
+ "multi_modal_data": {
54
+ "image": fetch_image(
55
+ "https://huggingface.co/datasets/hf-internal-testing/example-documents/resolve/main/jpeg_images/0.jpg"
56
+ ) # noqa: E501
57
+ },
58
+ },
59
+ "decoder_prompt": "<s_docvqa><s_question>What time is the coffee break?</s_question><s_answer>", # noqa: E501
60
+ },
61
+ ]
62
+
63
+ return ModelRequestData(
64
+ engine_args=engine_args,
65
+ prompts=prompts,
66
+ )
67
+
68
+
69
def run_florence2():
70
engine_args = EngineArgs(
71
model="microsoft/Florence-2-large",
@@ -118,6 +163,7 @@ def run_whisper():
163
164
165
model_example_map = {
166
+ "donut": run_donut,
167
"florence2": run_florence2,
168
"mllama": run_mllama,
169
"whisper": run_whisper,
tests/models/multimodal/processing/test_common.py
+2
@@ -160,6 +160,7 @@ def _test_processing_correctness(
160
# incorrect token ids. So we need use `add_special_tokens=False` here
161
# to leave bos_token to be added by the processor.
162
_ADD_SPECIAL_TOKENS_OVERRIDES = {
163
+ "donut": False,
164
"mllama": False,
165
"ovis": False,
166
"ovis2_5": False,
@@ -270,6 +271,7 @@ def _test_processing_correctness_one(
271
"facebook/chameleon-7b",
272
"CohereLabs/command-a-vision-07-2025",
273
"deepseek-ai/deepseek-vl2-tiny",
274
+ "naver-clova-ix/donut-base-finetuned-docvqa",
275
"microsoft/Florence-2-base",
276
"adept/fuyu-8b",
277
"google/gemma-3-4b-it",
tests/models/registry.py
+3
@@ -513,6 +513,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {
513
is_available_online=False,
514
),
515
# [Encoder-decoder]
516
+ "DonutForConditionalGeneration": _HfExamplesInfo("naver-clova-ix/donut-base-finetuned-docvqa", # noqa: E501
517
+ hf_overrides={"architectures": ["DonutForConditionalGeneration"], "model_type": "donut"}, # noqa: E501
518
+ extras={"dolphin": "ByteDance/Dolphin"}), # noqa: E501
519
# Florence-2 uses BartFastTokenizer which can't be loaded from AutoTokenizer
520
# Therefore, we borrow the BartTokenizer from the original Bart model
521
"Florence2ForConditionalGeneration": _HfExamplesInfo("microsoft/Florence-2-base", # noqa: E501
vllm/engine/llm_engine.py
+1
-1
@@ -1822,7 +1822,7 @@ class LLMEngine:
1822
assert isinstance(mm_processor, EncDecMultiModalProcessor)
1823
1824
if mm_processor.pad_dummy_encoder_prompt:
1825
- return # Skip encoder length check for Whisper
1825
+ return # Skip encoder length check for Whisper and Donut
1826
1827
if model_config.is_multimodal_model:
1828
suggestion = (
vllm/model_executor/models/donut.py
new
+398
@@ -0,0 +1,398 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import math
5
+from collections.abc import Iterable, Mapping, Sequence
6
+from typing import Literal, Optional, TypedDict, Union
7
+
8
+import torch
9
+import torch.nn as nn
10
+from transformers import BatchFeature, NougatProcessor
11
+
12
+from vllm.config import VllmConfig
13
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
14
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
15
+from vllm.model_executor.models.bart import BartParallelLMHead, MBartDecoder
16
+from vllm.model_executor.models.interfaces import (MultiModalEmbeddings,
17
+ SupportsMultiModal,
18
+ SupportsV0Only)
19
+from vllm.model_executor.models.swin import SwinModel
20
+from vllm.model_executor.models.utils import (AutoWeightsLoader,
21
+ _flatten_embeddings, flatten_bn)
22
+from vllm.model_executor.sampling_metadata import SamplingMetadata
23
+from vllm.multimodal import MULTIMODAL_REGISTRY
24
+from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
25
+ MultiModalKwargsItems)
26
+from vllm.multimodal.parse import MultiModalDataItems
27
+from vllm.multimodal.processing import (BaseProcessingInfo,
28
+ EncDecMultiModalProcessor,
29
+ PromptIndexTargets, PromptInsertion,
30
+ PromptUpdate)
31
+from vllm.multimodal.profiling import BaseDummyInputsBuilder
32
+
33
+
34
+class MBartDecoderWrapper(nn.Module):
35
+
36
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
37
+ super().__init__()
38
+ config = vllm_config.model_config.hf_config
39
+ cache_config = vllm_config.cache_config
40
+ quant_config = vllm_config.quant_config
41
+
42
+ self.decoder = MBartDecoder(config,
43
+ cache_config,
44
+ quant_config=quant_config,
45
+ prefix=f"{prefix}.decoder")
46
+
47
+ def forward(self, *args, **kwargs):
48
+ return self.decoder(*args, **kwargs)
49
+
50
+
51
+class DonutLanguageForConditionalGeneration(nn.Module, SupportsV0Only):
52
+
53
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
54
+ super().__init__()
55
+
56
+ config = vllm_config.model_config.hf_config
57
+
58
+ self.config = config
59
+ self.model = MBartDecoderWrapper(vllm_config=vllm_config,
60
+ prefix=f"{prefix}.model")
61
+ embed_scale = math.sqrt(
62
+ config.d_model) if config.scale_embedding else 1.0
63
+
64
+ self.vocab_size = config.vocab_size
65
+ self.lm_head = BartParallelLMHead(self.vocab_size,
66
+ config.d_model,
67
+ embed_scale=embed_scale)
68
+
69
+ self.logits_processor = LogitsProcessor(self.vocab_size,
70
+ config.vocab_size)
71
+
72
+ def forward(
73
+ self,
74
+ input_ids: torch.Tensor,
75
+ positions: torch.Tensor,
76
+ inputs_embeds: torch.Tensor,
77
+ **kwargs,
78
+ ) -> torch.Tensor:
79
+ r"""
80
+ Args:
81
+ input_ids
82
+ torch.Tensor of *decoder* input token ids.
83
+ positions
84
+ torch.Tensor of *decoder* position indices.
85
+ Returns:
86
+ Output torch.Tensor
87
+ """
88
+
89
+ return self.model(decoder_input_ids=input_ids,
90
+ decoder_positions=positions,
91
+ encoder_hidden_states=inputs_embeds)
92
+
93
+ def compute_logits(
94
+ self,
95
+ hidden_states: torch.Tensor,
96
+ sampling_metadata: SamplingMetadata,
97
+ ) -> Optional[torch.Tensor]:
98
+ logits = self.logits_processor(self.lm_head, hidden_states,
99
+ sampling_metadata)
100
+ return logits
101
+
102
+ def load_weights(self, weights: Iterable[tuple[str,
103
+ torch.Tensor]]) -> set[str]:
104
+ stacked_params_mapping = [
105
+ # (param_name, shard_name, shard_id)
106
+ ("qkv_proj", "q_proj", "q"),
107
+ ("qkv_proj", "k_proj", "k"),
108
+ ("qkv_proj", "v_proj", "v"),
109
+ ]
110
+
111
+ params_dict = dict(self.named_parameters())
112
+ loaded_params: set[str] = set()
113
+ for name, loaded_weight in weights:
114
+ for (param_name, weight_name, shard_id) in stacked_params_mapping:
115
+ if weight_name not in name:
116
+ continue
117
+ name = name.replace(weight_name, param_name)
118
+ param = params_dict[name]
119
+ weight_loader = param.weight_loader
120
+ weight_loader(param, loaded_weight, shard_id)
121
+ break
122
+ else:
123
+ if "final_logits_bias" in name:
124
+ continue
125
+ # if self.config.tie_word_embeddings and "embed_tokens" in name:
126
+ # continue
127
+ param = params_dict[name]
128
+ weight_loader = getattr(param, "weight_loader",
129
+ default_weight_loader)
130
+ weight_loader(param, loaded_weight)
131
+ loaded_params.add(name)
132
+ return loaded_params
133
+
134
+
135
+class DonutImagePixelInputs(TypedDict):
136
+ type: Literal["pixel_values"]
137
+ data: torch.Tensor
138
+ """Shape: (batch_size, num_channel, height, width)"""
139
+
140
+
141
+class DonutProcessingInfo(BaseProcessingInfo):
142
+
143
+ def get_hf_config(self):
144
+ return self.ctx.get_hf_config()
145
+
146
+ def get_hf_processor(self):
147
+ return self.ctx.get_hf_processor()
148
+
149
+ def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
150
+ return {"image": 1}
151
+
152
+ def get_num_image_tokens(self) -> int:
153
+ return 1
154
+
155
+
156
+class DonutDummyInputsBuilder(BaseDummyInputsBuilder[DonutProcessingInfo]):
157
+
158
+ def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
159
+ return ""
160
+
161
+ def get_dummy_mm_data(
162
+ self,
163
+ seq_len: int,
164
+ mm_counts: Mapping[str, int],
165
+ ) -> MultiModalDataDict:
166
+ num_images = mm_counts.get("image", 0)
167
+
168
+ target_width, target_height = self.info.get_hf_config(
169
+ ).encoder.image_size
170
+
171
+ return {
172
+ "image":
173
+ self._get_dummy_images(width=target_width,
174
+ height=target_height,
175
+ num_images=num_images)
176
+ }
177
+
178
+
179
+class DonutMultiModalProcessor(EncDecMultiModalProcessor[DonutProcessingInfo]):
180
+
181
+ def _hf_processor_applies_updates(
182
+ self,
183
+ prompt_text: str,
184
+ mm_items: MultiModalDataItems,
185
+ hf_processor_mm_kwargs: Mapping[str, object],
186
+ tokenization_kwargs: Mapping[str, object],
187
+ ) -> bool:
188
+ return False
189
+
190
+ def create_encoder_prompt(
191
+ self,
192
+ prompt: Union[str, list[int]],
193
+ mm_data: MultiModalDataDict,
194
+ ) -> Union[str, list[int]]:
195
+ return prompt
196
+
197
+ def create_decoder_prompt(
198
+ self,
199
+ prompt: Union[str, list[int]],
200
+ mm_data: MultiModalDataDict,
201
+ ) -> Union[str, list[int]]:
202
+ return prompt
203
+
204
+ @property
205
+ def pad_dummy_encoder_prompt(self) -> bool:
206
+ return True
207
+
208
+ def _call_hf_processor(
209
+ self,
210
+ prompt: str,
211
+ mm_data: Mapping[str, object],
212
+ mm_kwargs: Mapping[str, object],
213
+ tok_kwargs: Mapping[str, object],
214
+ ) -> BatchFeature:
215
+ hf_processor = self.info.get_hf_processor()
216
+ if mm_data:
217
+ processed_outputs = super()._call_hf_processor(
218
+ prompt, mm_data, mm_kwargs, tok_kwargs)
219
+ if isinstance(hf_processor, NougatProcessor):
220
+ processed_outputs["input_ids"] = processed_outputs["labels"]
221
+ else:
222
+ tokenizer = hf_processor.tokenizer
223
+ processed_outputs = tokenizer(prompt,
224
+ add_special_tokens=False,
225
+ return_tensors="pt")
226
+ return processed_outputs
227
+
228
+ def _get_mm_fields_config(
229
+ self,
230
+ hf_inputs: BatchFeature,
231
+ hf_processor_mm_kwargs: Mapping[str, object],
232
+ ) -> Mapping[str, MultiModalFieldConfig]:
233
+ return dict(pixel_values=MultiModalFieldConfig.batched("image"))
234
+
235
+ def _get_prompt_updates(
236
+ self,
237
+ mm_items: MultiModalDataItems,
238
+ hf_processor_mm_kwargs: Mapping[str, object],
239
+ out_mm_kwargs: MultiModalKwargsItems,
240
+ ) -> Sequence[PromptUpdate]:
241
+ hf_processor = self.info.get_hf_processor()
242
+ tokenizer = hf_processor.tokenizer
243
+ pad_token_id = tokenizer.pad_token_id
244
+ num_image_tokens = self.info.get_num_image_tokens()
245
+ image_tokens = [pad_token_id] * num_image_tokens
246
+
247
+ return [
248
+ PromptInsertion(
249
+ modality="image",
250
+ target=PromptIndexTargets.start(),
251
+ insertion=image_tokens,
252
+ )
253
+ ]
254
+
255
+
256
+@MULTIMODAL_REGISTRY.register_processor(DonutMultiModalProcessor,
257
+ info=DonutProcessingInfo,
258
+ dummy_inputs=DonutDummyInputsBuilder)
259
+class DonutForConditionalGeneration(nn.Module, SupportsMultiModal,
260
+ SupportsV0Only):
261
+
262
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
263
+ super().__init__()
264
+ config = vllm_config.model_config.hf_config
265
+ processor_config = vllm_config.model_config.hf_image_processor_config
266
+
267
+ self.config = config
268
+ self.vision_config = config.encoder
269
+ self.processor_config = processor_config
270
+ self.encoder = SwinModel(config=config.encoder)
271
+
272
+ self.decoder = DonutLanguageForConditionalGeneration(
273
+ vllm_config=vllm_config.with_hf_config(config.decoder),
274
+ prefix=f"{prefix}.decoder",
275
+ )
276
+ self.pad_token_id = config.pad_token_id
277
+
278
+ def _validate_pixel_values(
279
+ self, data: Union[torch.Tensor, list[torch.Tensor]]
280
+ ) -> Union[torch.Tensor, list[torch.Tensor]]:
281
+
282
+ # size = self.processor_config["size"]
283
+ h, w = self.config.encoder.image_size
284
+ expected_dims = (3, h, w)
285
+
286
+ def _validate_shape(d: torch.Tensor):
287
+ actual_dims = tuple(d.shape)
288
+
289
+ if actual_dims != expected_dims:
290
+ raise ValueError(
291
+ "The expected shape of pixel values per batch "
292
+ f"is {expected_dims}. You supplied {actual_dims}.")
293
+
294
+ for d in data:
295
+ _validate_shape(d)
296
+
297
+ return data
298
+
299
+ def _parse_and_validate_image_input(self, **kwargs: object):
300
+ pixel_values: Optional[Union[list[list[torch.Tensor]],
301
+ list[torch.Tensor],
302
+ torch.Tensor]] = kwargs.pop(
303
+ "pixel_values", None)
304
+ image_embeds: Optional[Union[list[list[torch.Tensor]],
305
+ list[torch.Tensor],
306
+ torch.Tensor]] = kwargs.pop(
307
+ "image_embeds", None)
308
+
309
+ if pixel_values is None and image_embeds is None:
310
+ return None
311
+
312
+ if pixel_values is not None and image_embeds is not None:
313
+ raise ValueError(
314
+ "Both pixel values and image embeds are provided.")
315
+
316
+ if pixel_values is not None:
317
+ return DonutImagePixelInputs(
318
+ type="pixel_values",
319
+ data=self._validate_pixel_values(
320
+ flatten_bn(pixel_values, concat=True)),
321
+ )
322
+
323
+ if image_embeds is not None:
324
+ raise NotImplementedError
325
+
326
+ raise AssertionError("This line should be unreachable.")
327
+
328
+ def _process_image_input(
329
+ self, image_input: DonutImagePixelInputs) -> torch.Tensor:
330
+ assert image_input["type"] == "pixel_values"
331
+ pixel_values = image_input["data"]
332
+ dtype = next(self.encoder.parameters()).dtype
333
+ pixel_values = pixel_values.to(dtype)
334
+ return self.encoder(pixel_values)
335
+
336
+ def get_language_model(self) -> torch.nn.Module:
337
+ return self.decoder
338
+
339
+ def get_multimodal_embeddings(
340
+ self, **kwargs: object) -> Optional[MultiModalEmbeddings]:
341
+ image_input = self._parse_and_validate_image_input(**kwargs)
342
+ if image_input is None:
343
+ return None
344
+ vision_embeddings = self._process_image_input(image_input)
345
+ return vision_embeddings
346
+
347
+ def get_input_embeddings(
348
+ self,
349
+ input_ids: torch.Tensor,
350
+ multimodal_embeddings: MultiModalEmbeddings,
351
+ ) -> torch.Tensor:
352
+ return _flatten_embeddings(multimodal_embeddings)
353
+
354
+ def forward(
355
+ self,
356
+ input_ids: torch.Tensor,
357
+ positions: torch.Tensor,
358
+ *,
359
+ encoder_input_ids: torch.Tensor,
360
+ encoder_positions: torch.Tensor,
361
+ **kwargs,
362
+ ) -> torch.Tensor:
363
+ r"""
364
+ Args:
365
+ input_ids
366
+ torch.Tensor of *decoder* input token ids.
367
+ positions
368
+ torch.Tensor of *decoder* position indices.
369
+ encoder_input_ids
370
+ torch.Tensor of *encoder* input token ids.
371
+ encoder_positions
372
+ torch.Tensor of *encoder* position indices
373
+ Returns:
374
+ Output torch.Tensor
375
+ """
376
+
377
+ inputs_embeds = None
378
+ if encoder_input_ids.numel() > 0:
379
+ vision_embeddings = self.get_multimodal_embeddings(**kwargs)
380
+ inputs_embeds = self.get_input_embeddings(encoder_input_ids,
381
+ vision_embeddings)
382
+
383
+ hidden_states = self.decoder(input_ids,
384
+ positions,
385
+ inputs_embeds=inputs_embeds)
386
+ return hidden_states
387
+
388
+ def compute_logits(
389
+ self,
390
+ hidden_states: torch.Tensor,
391
+ sampling_metadata: SamplingMetadata,
392
+ ) -> Optional[torch.Tensor]:
393
+ return self.decoder.compute_logits(hidden_states, sampling_metadata)
394
+
395
+ def load_weights(self, weights: Iterable[tuple[str,
396
+ torch.Tensor]]) -> set[str]:
397
+ loader = AutoWeightsLoader(self)
398
+ return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+1
@@ -252,6 +252,7 @@ _MULTIMODAL_MODELS = {
252
"Tarsier2ForConditionalGeneration": ("qwen2_vl", "Tarsier2ForConditionalGeneration"), # noqa: E501
253
"VoxtralForConditionalGeneration": ("voxtral", "VoxtralForConditionalGeneration"), # noqa: E501
254
# [Encoder-decoder]
255
+ "DonutForConditionalGeneration": ("donut", "DonutForConditionalGeneration"),
256
"Florence2ForConditionalGeneration": ("florence2", "Florence2ForConditionalGeneration"), # noqa: E501
257
"MllamaForConditionalGeneration": ("mllama", "MllamaForConditionalGeneration"), # noqa: E501
258
"Llama4ForConditionalGeneration": ("mllama4", "Llama4ForConditionalGeneration"), # noqa: E501
vllm/model_executor/models/swin.py
new
+475
@@ -0,0 +1,475 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from collections.abc import Iterable
5
+from typing import Optional
6
+
7
+import torch
8
+import torch.nn as nn
9
+from transformers import SwinConfig
10
+from transformers.models.swin.modeling_swin import SwinEmbeddings
11
+from transformers.models.swin.modeling_swin import SwinLayer as HFSwinLayer
12
+from transformers.models.swin.modeling_swin import SwinPatchMerging
13
+from transformers.pytorch_utils import meshgrid
14
+
15
+from vllm.model_executor.layers.activation import get_act_fn
16
+from vllm.model_executor.layers.linear import (ColumnParallelLinear,
17
+ QKVParallelLinear,
18
+ RowParallelLinear)
19
+from vllm.model_executor.layers.quantization import QuantizationConfig
20
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
21
+
22
+
23
+class SwinSelfAttention(nn.Module):
24
+
25
+ def __init__(
26
+ self,
27
+ config: SwinConfig,
28
+ dim: int,
29
+ num_heads: int,
30
+ window_size: int,
31
+ quant_config: Optional[QuantizationConfig] = None,
32
+ prefix: str = "",
33
+ ) -> None:
34
+ super().__init__()
35
+ if dim % num_heads != 0:
36
+ raise ValueError(
37
+ f"The hidden size ({dim}) is not a multiple of the number of "
38
+ f"attention heads ({num_heads})")
39
+
40
+ self.num_attention_heads = num_heads
41
+ self.attention_head_size = int(dim / num_heads)
42
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
43
+ self.window_size = (window_size if isinstance(window_size, Iterable)
44
+ else (window_size, window_size))
45
+ self.scale = self.attention_head_size**-0.5
46
+
47
+ self.relative_position_bias_table = nn.Parameter(
48
+ torch.zeros(
49
+ (2 * self.window_size[0] - 1) * (2 * self.window_size[1] - 1),
50
+ num_heads))
51
+
52
+ # get pair-wise relative position index for each token inside the window
53
+ coords_h = torch.arange(self.window_size[0])
54
+ coords_w = torch.arange(self.window_size[1])
55
+ coords = torch.stack(meshgrid([coords_h, coords_w], indexing="ij"))
56
+ coords_flatten = torch.flatten(coords, 1)
57
+ relative_coords = coords_flatten[:, :, None] - coords_flatten[:,
58
+ None, :]
59
+ relative_coords = relative_coords.permute(1, 2, 0).contiguous()
60
+ relative_coords[:, :, 0] += self.window_size[0] - 1
61
+ relative_coords[:, :, 1] += self.window_size[1] - 1
62
+ relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
63
+ relative_position_index = relative_coords.sum(-1)
64
+
65
+ self.relative_position_index = nn.Parameter(relative_position_index,
66
+ requires_grad=False)
67
+
68
+ self.qkv = QKVParallelLinear(
69
+ hidden_size=dim,
70
+ head_size=self.attention_head_size,
71
+ total_num_heads=self.num_attention_heads,
72
+ bias=config.qkv_bias,
73
+ quant_config=quant_config,
74
+ prefix=f"{prefix}.qkv",
75
+ )
76
+
77
+ def transpose_for_scores(self, x):
78
+ new_x_shape = x.size()[:-1] + (self.num_attention_heads,
79
+ self.attention_head_size)
80
+ x = x.view(new_x_shape)
81
+ return x.permute(0, 2, 1, 3)
82
+
83
+ def _get_rel_pos_bias(self) -> torch.Tensor:
84
+ relative_position_bias = self.relative_position_bias_table[
85
+ self.relative_position_index.view(-1)]
86
+ relative_position_bias = relative_position_bias.view(
87
+ self.window_size[0] * self.window_size[1],
88
+ self.window_size[0] * self.window_size[1], -1)
89
+ relative_position_bias = relative_position_bias.permute(
90
+ 2, 0, 1).contiguous()
91
+ return relative_position_bias.unsqueeze(0)
92
+
93
+ def forward(
94
+ self,
95
+ hidden_states: torch.Tensor,
96
+ attention_mask: Optional[torch.FloatTensor] = None,
97
+ head_mask: Optional[torch.FloatTensor] = None,
98
+ output_attentions: Optional[bool] = False,
99
+ ) -> tuple[torch.Tensor, ...]:
100
+ batch_size, dim, num_channels = hidden_states.shape
101
+
102
+ qkv_output, _ = self.qkv(hidden_states)
103
+ query_layer, key_layer, value_layer = qkv_output.chunk(3, dim=-1)
104
+
105
+ key_layer = self.transpose_for_scores(key_layer)
106
+ value_layer = self.transpose_for_scores(value_layer)
107
+ query_layer = self.transpose_for_scores(query_layer)
108
+
109
+ attention_scores = self._get_rel_pos_bias()
110
+ if attention_mask is not None:
111
+ mask_shape = attention_mask.shape[0]
112
+ attention_mask_expanded = attention_mask.view(
113
+ 1, mask_shape, 1, dim,
114
+ dim).expand(batch_size // mask_shape, mask_shape,
115
+ self.num_attention_heads, dim, dim)
116
+ attention_scores = attention_scores + \
117
+ attention_mask_expanded.unsqueeze(
118
+ 1).unsqueeze(0)
119
+ attention_scores = attention_scores.view(-1,
120
+ self.num_attention_heads,
121
+ dim, dim)
122
+
123
+ context_layer = torch.nn.functional.scaled_dot_product_attention(
124
+ query_layer,
125
+ key_layer,
126
+ value_layer,
127
+ attn_mask=attention_scores,
128
+ dropout_p=0.,
129
+ )
130
+ attention_probs = None
131
+
132
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
133
+ new_context_layer_shape = context_layer.size()[:-2] + (
134
+ self.all_head_size, )
135
+ context_layer = context_layer.view(new_context_layer_shape)
136
+
137
+ outputs = (context_layer,
138
+ attention_probs) if output_attentions else (context_layer, )
139
+
140
+ return outputs
141
+
142
+
143
+class SwinSelfOutput(nn.Module):
144
+
145
+ def __init__(
146
+ self,
147
+ config: SwinConfig,
148
+ dim: int,
149
+ quant_config: Optional[QuantizationConfig] = None,
150
+ prefix: str = "",
151
+ ) -> None:
152
+ super().__init__()
153
+ self.dense = RowParallelLinear(
154
+ input_size=dim,
155
+ output_size=dim,
156
+ quant_config=quant_config,
157
+ prefix=f"{prefix}.dense",
158
+ )
159
+
160
+ def forward(self, hidden_states: torch.Tensor,
161
+ input_tensor: torch.Tensor) -> torch.Tensor:
162
+ hidden_states, _ = self.dense(hidden_states)
163
+
164
+ return hidden_states
165
+
166
+
167
+class SwinAttention(nn.Module):
168
+
169
+ def __init__(self,
170
+ config: SwinConfig,
171
+ dim: int,
172
+ num_heads: int,
173
+ window_size: int,
174
+ quant_config: Optional[QuantizationConfig] = None,
175
+ prefix: str = "") -> None:
176
+ super().__init__()
177
+ self.self = SwinSelfAttention(config,
178
+ dim,
179
+ num_heads,
180
+ window_size,
181
+ quant_config=quant_config,
182
+ prefix=f"{prefix}.self")
183
+ self.output = SwinSelfOutput(config,
184
+ dim,
185
+ quant_config=quant_config,
186
+ prefix=f"{prefix}.output")
187
+ self.pruned_heads = set()
188
+
189
+ def forward(
190
+ self,
191
+ hidden_states: torch.Tensor,
192
+ attention_mask: Optional[torch.FloatTensor] = None,
193
+ head_mask: Optional[torch.FloatTensor] = None,
194
+ output_attentions: Optional[bool] = False,
195
+ ) -> tuple[torch.Tensor]:
196
+ self_outputs = self.self(hidden_states, attention_mask, head_mask,
197
+ output_attentions)
198
+ attention_output = self.output(self_outputs[0], hidden_states)
199
+ outputs = (attention_output, ) + self_outputs[1:]
200
+ return outputs
201
+
202
+
203
+class SwinIntermediate(nn.Module):
204
+
205
+ def __init__(self,
206
+ config: SwinConfig,
207
+ dim: int,
208
+ quant_config: Optional[QuantizationConfig] = None,
209
+ prefix: str = "") -> None:
210
+ super().__init__()
211
+ self.dense = ColumnParallelLinear(dim,
212
+ int(config.mlp_ratio * dim),
213
+ quant_config=quant_config,
214
+ prefix=f"{prefix}.dense")
215
+ self.intermediate_act_fn = get_act_fn(config.hidden_act)
216
+
217
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
218
+ hidden_states, _ = self.dense(hidden_states)
219
+ hidden_states = self.intermediate_act_fn(hidden_states)
220
+ return hidden_states
221
+
222
+
223
+class SwinOutput(nn.Module):
224
+
225
+ def __init__(self,
226
+ config: SwinConfig,
227
+ dim: int,
228
+ quant_config: Optional[QuantizationConfig] = None,
229
+ prefix: str = "") -> None:
230
+ super().__init__()
231
+ self.dense = RowParallelLinear(int(config.mlp_ratio * dim),
232
+ dim,
233
+ quant_config=quant_config,
234
+ prefix=f"{prefix}.dense")
235
+
236
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
237
+ hidden_states, _ = self.dense(hidden_states)
238
+ return hidden_states
239
+
240
+
241
+class SwinLayer(HFSwinLayer):
242
+
243
+ def __init__(
244
+ self,
245
+ config: SwinConfig,
246
+ dim: int,
247
+ input_resolution: int,
248
+ num_heads: int,
249
+ drop_path_rate: float = 0.0,
250
+ shift_size: int = 0,
251
+ quant_config: Optional[QuantizationConfig] = None,
252
+ prefix: str = "",
253
+ ) -> None:
254
+ super().__init__(
255
+ config=config,
256
+ dim=dim,
257
+ input_resolution=input_resolution,
258
+ num_heads=num_heads,
259
+ drop_path_rate=drop_path_rate,
260
+ shift_size=shift_size,
261
+ )
262
+
263
+ self.attention = SwinAttention(config,
264
+ dim,
265
+ num_heads,
266
+ window_size=self.window_size,
267
+ quant_config=quant_config,
268
+ prefix=f"{prefix}.attention")
269
+ self.intermediate = SwinIntermediate(config,
270
+ dim,
271
+ quant_config=quant_config,
272
+ prefix=f"{prefix}.intermediate")
273
+ self.output = SwinOutput(config,
274
+ dim,
275
+ quant_config=quant_config,
276
+ prefix=f"{prefix}.output")
277
+
278
+
279
+class SwinStage(nn.Module):
280
+
281
+ def __init__(
282
+ self,
283
+ config: SwinConfig,
284
+ dim: int,
285
+ input_resolution: int,
286
+ depth: int,
287
+ num_heads: int,
288
+ drop_path: list[float],
289
+ downsample: Optional[SwinPatchMerging] = None,
290
+ quant_config: Optional[QuantizationConfig] = None,
291
+ prefix: str = "",
292
+ ) -> None:
293
+ super().__init__()
294
+ self.config = config
295
+ self.dim = dim
296
+ self.blocks = nn.ModuleList([
297
+ SwinLayer(config=config,
298
+ dim=dim,
299
+ input_resolution=input_resolution,
300
+ num_heads=num_heads,
301
+ drop_path_rate=drop_path[layer_idx],
302
+ shift_size=0 if
303
+ (layer_idx % 2 == 0) else config.window_size // 2,
304
+ quant_config=quant_config,
305
+ prefix=f"{prefix}.blocks.{layer_idx}")
306
+ for layer_idx in range(depth)
307
+ ])
308
+
309
+ # patch merging layer
310
+ if downsample is not None:
311
+ self.downsample = downsample(input_resolution,
312
+ dim=dim,
313
+ norm_layer=nn.LayerNorm)
314
+ else:
315
+ self.downsample = None
316
+
317
+ self.pointing = False
318
+
319
+ def forward(
320
+ self,
321
+ hidden_states: torch.Tensor,
322
+ input_dimensions: tuple[int, int],
323
+ head_mask: Optional[torch.FloatTensor] = None,
324
+ output_attentions: Optional[bool] = False,
325
+ always_partition: Optional[bool] = False,
326
+ ) -> tuple[torch.Tensor]:
327
+ height, width = input_dimensions
328
+ for i, layer_module in enumerate(self.blocks):
329
+ layer_head_mask = head_mask[i] if head_mask is not None else None
330
+
331
+ layer_outputs = layer_module(hidden_states, input_dimensions,
332
+ layer_head_mask, output_attentions,
333
+ always_partition)
334
+
335
+ hidden_states = layer_outputs[0]
336
+
337
+ hidden_states_before_downsampling = hidden_states
338
+ if self.downsample is not None:
339
+ height_downsampled, width_downsampled = (height + 1) // 2, (width +
340
+ 1) // 2
341
+ output_dimensions = (height, width, height_downsampled,
342
+ width_downsampled)
343
+ hidden_states = self.downsample(hidden_states_before_downsampling,
344
+ input_dimensions)
345
+ else:
346
+ output_dimensions = (height, width, height, width)
347
+
348
+ stage_outputs = (hidden_states, hidden_states_before_downsampling,
349
+ output_dimensions)
350
+
351
+ if output_attentions:
352
+ stage_outputs += layer_outputs[1:]
353
+ return stage_outputs
354
+
355
+
356
+class SwinEncoder(nn.Module):
357
+
358
+ def __init__(
359
+ self,
360
+ config: SwinConfig,
361
+ grid_size: int,
362
+ quant_config: Optional[QuantizationConfig] = None,
363
+ prefix: str = "",
364
+ ) -> None:
365
+ super().__init__()
366
+ self.num_layers = len(config.depths)
367
+ self.config = config
368
+ dpr = [
369
+ x.item() for x in torch.linspace(
370
+ 0, config.drop_path_rate, sum(config.depths), device="cpu")
371
+ ]
372
+ self.layers = nn.ModuleList([
373
+ SwinStage(config=config,
374
+ dim=int(config.embed_dim * 2**layer_idx),
375
+ input_resolution=(grid_size[0] // (2**layer_idx),
376
+ grid_size[1] // (2**layer_idx)),
377
+ depth=config.depths[layer_idx],
378
+ num_heads=config.num_heads[layer_idx],
379
+ drop_path=dpr[sum(config.depths[:layer_idx]
380
+ ):sum(config.depths[:layer_idx + 1])],
381
+ downsample=SwinPatchMerging if
382
+ (layer_idx < self.num_layers - 1) else None,
383
+ quant_config=quant_config,
384
+ prefix=f"{prefix}.layers.{layer_idx}")
385
+ for layer_idx in range(self.num_layers)
386
+ ])
387
+
388
+ def forward(
389
+ self,
390
+ hidden_states: torch.Tensor,
391
+ input_dimensions: tuple[int, int],
392
+ head_mask: Optional[torch.FloatTensor] = None,
393
+ output_attentions: Optional[bool] = False,
394
+ always_partition: Optional[bool] = False,
395
+ ) -> tuple[torch.Tensor]:
396
+ for i, layer_module in enumerate(self.layers):
397
+ layer_head_mask = head_mask[i] if head_mask is not None else None
398
+
399
+ layer_outputs = layer_module(hidden_states, input_dimensions,
400
+ layer_head_mask, output_attentions,
401
+ always_partition)
402
+
403
+ hidden_states = layer_outputs[0]
404
+ output_dimensions = layer_outputs[2]
405
+
406
+ input_dimensions = (output_dimensions[-2], output_dimensions[-1])
407
+
408
+ return hidden_states
409
+
410
+
411
+class SwinModel(nn.Module):
412
+ config_class: SwinConfig
413
+
414
+ def __init__(
415
+ self,
416
+ config: SwinConfig,
417
+ quant_config: Optional[QuantizationConfig] = None,
418
+ prefix: str = "",
419
+ ) -> None:
420
+ super().__init__()
421
+ self.config = config
422
+ self.num_layers = len(config.depths)
423
+ self.num_features = int(config.embed_dim * 2**(self.num_layers - 1))
424
+
425
+ self.embeddings = SwinEmbeddings(config)
426
+ self.encoder = SwinEncoder(config,
427
+ self.embeddings.patch_grid,
428
+ quant_config=quant_config,
429
+ prefix=f"{prefix}.encoder")
430
+
431
+ def forward(
432
+ self,
433
+ pixel_values: Optional[torch.FloatTensor] = None,
434
+ head_mask: Optional[torch.FloatTensor] = None,
435
+ output_attentions: Optional[bool] = None,
436
+ ) -> tuple[torch.Tensor]:
437
+ embedding_output, input_dimensions = self.embeddings(pixel_values)
438
+
439
+ encoder_outputs = self.encoder(
440
+ embedding_output,
441
+ input_dimensions,
442
+ head_mask=head_mask,
443
+ output_attentions=output_attentions,
444
+ )
445
+
446
+ return encoder_outputs
447
+
448
+ def load_weights(self, weights: Iterable[tuple[str,
449
+ torch.Tensor]]) -> set[str]:
450
+ stacked_params_mapping = [
451
+ ("qkv", "query", "q"),
452
+ ("qkv", "key", "k"),
453
+ ("qkv", "value", "v"),
454
+ ]
455
+ params_dict = dict(self.named_parameters())
456
+ loaded_params: set[str] = set()
457
+
458
+ for name, loaded_weight in weights:
459
+
460
+ for (param_name, weight_name, shard_id) in stacked_params_mapping:
461
+ if weight_name not in name:
462
+ continue
463
+ name = name.replace(weight_name, param_name)
464
+
465
+ param = params_dict[name]
466
+ weight_loader = param.weight_loader
467
+ weight_loader(param, loaded_weight, shard_id)
468
+ break
469
+ else:
470
+ param = params_dict[name]
471
+ weight_loader = getattr(param, "weight_loader",
472
+ default_weight_loader)
473
+ weight_loader(param, loaded_weight)
474
+ loaded_params.add(name)
475
+ return loaded_params
vllm/multimodal/profiling.py
+1
-1
@@ -209,7 +209,7 @@ class MultiModalProfiler(Generic[_I]):
209
if processor.pad_dummy_encoder_prompt:
210
num_tokens_to_pad = max(total_len, seq_len) - total_len
211
encoder_prompt_token_ids.extend([0] * num_tokens_to_pad)
212
- # NOTE: Whisper allows total_len > seq_len.
212
+ # NOTE: Whisper and Donut allows total_len > seq_len.
213
elif total_len > seq_len and not envs.VLLM_USE_V1:
214
# `max_num_batched_tokens` is defined by `SchedulerConfig`
215
logger.warning_once(
vllm/v1/engine/processor.py
+1
-1
@@ -389,7 +389,7 @@ class Processor:
389
assert isinstance(mm_processor, EncDecMultiModalProcessor)
390
391
if mm_processor.pad_dummy_encoder_prompt:
392
- return # Skip encoder length check for Whisper
392
+ return # Skip encoder length check for Whisper and Donut
393
394
if model_config.is_multimodal_model:
395
suggestion = (