@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7b57a433da
Fork 0

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

[Model] Support Dots OCR (#24645)

Signed-off-by: Roger Wang <hey@rogerw.io> Co-authored-by: yinz-aizip <yinz@aizip.ai>

Roger Wang committed Sep 21, 2025 at 19:24 UTC 7b57a433daf94dc075ccf52267b82aae359d0d53
7 files changed +917
docs/models/supported_models.md
+1
@@ -352,6 +352,7 @@ th {
352 | `DeepseekV2ForCausalLM` | DeepSeek-V2 | `deepseek-ai/DeepSeek-V2`, `deepseek-ai/DeepSeek-V2-Chat`, etc. | ✅︎ | ✅︎ | ✅︎ |
353 | `DeepseekV3ForCausalLM` | DeepSeek-V3 | `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`, etc. | ✅︎ | ✅︎ | ✅︎ |
354 | `Dots1ForCausalLM` | dots.llm1 | `rednote-hilab/dots.llm1.base`, `rednote-hilab/dots.llm1.inst`, etc. | | ✅︎ | ✅︎ |
355 +| `DotsOCRForCausalLM` | dots_ocr | `rednote-hilab/dots.ocr` | | ✅︎ | ✅︎ |
356 | `Ernie4_5ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | ✅︎ | ✅︎ | ✅︎ |
357 | `Ernie4_5_MoeForCausalLM` | Ernie4.5MoE | `baidu/ERNIE-4.5-21B-A3B-PT`, `baidu/ERNIE-4.5-300B-A47B-PT`, etc. |✅︎| ✅︎ | ✅︎ |
358 | `ExaoneForCausalLM` | EXAONE-3 | `LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct`, etc. | ✅︎ | ✅︎ | ✅︎ |
examples/offline_inference/vision_language.py
+18
@@ -126,6 +126,23 @@ def run_chameleon(questions: list[str], modality: str) -> ModelRequestData:
126 )
127
128
129 +# Dots-OCR
130 +def run_dots_ocr(questions: list[str], modality: str) -> ModelRequestData:
131 + assert modality == "image"
132 +
133 + prompts = [f"<|img|><|imgpad|><|endofimg|>{question}" for question in questions]
134 + engine_args = EngineArgs(
135 + model="rednote-hilab/dots.ocr",
136 + limit_mm_per_prompt={modality: 1},
137 + trust_remote_code=True,
138 + )
139 +
140 + return ModelRequestData(
141 + engine_args=engine_args,
142 + prompts=prompts,
143 + )
144 +
145 +
146 def run_command_a_vision(questions: list[str], modality: str) -> ModelRequestData:
147 assert modality == "image"
148
@@ -1676,6 +1693,7 @@ model_example_map = {
1693 "aya_vision": run_aya_vision,
1694 "blip-2": run_blip2,
1695 "chameleon": run_chameleon,
1696 + "dots_ocr": run_dots_ocr,
1697 "command_a_vision": run_command_a_vision,
1698 "deepseek_vl_v2": run_deepseek_vl2,
1699 "ernie45_vl": run_ernie45_vl,
tests/models/registry.py
+2
@@ -448,6 +448,8 @@ _MULTIMODAL_EXAMPLE_MODELS = {
448 max_transformers_version="4.48", # noqa: E501
449 transformers_version_reason="HF model is not compatible.", # noqa: E501
450 hf_overrides={"architectures": ["DeepseekVLV2ForCausalLM"]}), # noqa: E501
451 + "DotsOCRForCausalLM": _HfExamplesInfo("rednote-hilab/dots.ocr",
452 + trust_remote_code=True),
453 "Emu3ForConditionalGeneration": _HfExamplesInfo("BAAI/Emu3-Chat-hf"),
454 "Ernie4_5_VLMoeForConditionalGeneration": _HfExamplesInfo("baidu/ERNIE-4.5-VL-28B-A3B-PT", # noqa: E501
455 trust_remote_code=True),
vllm/model_executor/models/dots_ocr.py new
+824
@@ -0,0 +1,824 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Iterable, Mapping
4 +from typing import Literal, Optional, TypedDict, Union
5 +
6 +import torch
7 +import torch.nn as nn
8 +import torch.nn.functional as F
9 +from torch.nn import LayerNorm
10 +from transformers.modeling_utils import PreTrainedModel
11 +from transformers.models.qwen2_vl import Qwen2VLProcessor
12 +
13 +from vllm.attention.layer import check_upstream_fa_availability
14 +from vllm.config import VllmConfig
15 +from vllm.model_executor.layers.activation import SiluAndMul
16 +from vllm.model_executor.layers.layernorm import RMSNorm
17 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,
18 + MergedColumnParallelLinear,
19 + QKVParallelLinear,
20 + RowParallelLinear)
21 +from vllm.model_executor.layers.quantization import QuantizationConfig
22 +from vllm.model_executor.models.interfaces import (MultiModalEmbeddings,
23 + SupportsMultiModal,
24 + SupportsPP)
25 +from vllm.model_executor.models.qwen2 import Qwen2ForCausalLM
26 +from vllm.model_executor.models.qwen2_vl import (Qwen2VLDummyInputsBuilder,
27 + Qwen2VLMultiModalProcessor,
28 + Qwen2VLProcessingInfo)
29 +from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
30 + init_vllm_registered_model,
31 + maybe_prefix,
32 + merge_multimodal_embeddings)
33 +from vllm.model_executor.models.vision import get_vit_attn_backend
34 +from vllm.multimodal import MULTIMODAL_REGISTRY
35 +from vllm.multimodal.inputs import MultiModalDataDict
36 +from vllm.platforms import _Backend
37 +from vllm.sequence import IntermediateTensors
38 +from vllm.transformers_utils.configs.dotsocr import (DotsOCRConfig,
39 + DotsVisionConfig)
40 +
41 +IMAGE_TOKEN = "<|imgpad|>"
42 +
43 +
44 +class DotsOCRImagePixelInputs(TypedDict):
45 + type: Literal["pixel_values", "image_grid_thw"]
46 +
47 + pixel_values: torch.Tensor
48 + image_grid_thw: torch.Tensor
49 +
50 +
51 +class DotsOCRImageEmbeddingInputs(TypedDict):
52 + type: Literal["image_embeds", "image_grid_thw"]
53 + image_embeds: torch.Tensor
54 + """Supported types:
55 + - List[`torch.Tensor`]: A list of tensors holding all images' features.
56 + Each tensor holds an image's features.
57 + - `torch.Tensor`: A tensor holding all images' features
58 + (concatenation of all images' feature tensors).
59 + Tensor shape: `(num_image_features, hidden_size)`
60 + - `num_image_features` varies based on
61 + the number and resolution of the images.
62 + - `hidden_size` must match the hidden size of language model backbone.
63 + """
64 +
65 + image_grid_thw: torch.Tensor
66 +
67 +
68 +DotsOCRImageInputs = Union[DotsOCRImagePixelInputs,
69 + DotsOCRImageEmbeddingInputs]
70 +
71 +
72 +class DotsOCRDummyInputsBuilder(Qwen2VLDummyInputsBuilder):
73 +
74 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
75 + num_images = mm_counts.get("image", 0)
76 + return IMAGE_TOKEN * num_images
77 +
78 + def get_dummy_mm_data(
79 + self,
80 + seq_len: int,
81 + mm_counts: Mapping[str, int],
82 + ) -> MultiModalDataDict:
83 + num_images = mm_counts.get("image", 0)
84 +
85 + target_width, target_height = self.info.get_image_size_with_most_features( # noqa: E501
86 + )
87 +
88 + return {
89 + "image":
90 + self._get_dummy_images(width=target_width,
91 + height=target_height,
92 + num_images=num_images),
93 + }
94 +
95 +
96 +class DotsOCRProcessingInfo(Qwen2VLProcessingInfo):
97 +
98 + def get_hf_config(self) -> DotsOCRConfig:
99 + config = self.ctx.get_hf_config()
100 + if not config.__class__.__name__ == 'DotsOCRConfig':
101 + raise TypeError(f"Expected DotsOCRConfig, got {type(config)}")
102 +
103 + if hasattr(config, "vision_config") and isinstance(
104 + config.vision_config, dict):
105 + config.vision_config = DotsVisionConfig(**config.vision_config)
106 +
107 + return config
108 +
109 + def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
110 + return {"image": None}
111 +
112 + def get_mm_max_tokens_per_item(
113 + self,
114 + seq_len: int,
115 + mm_counts: Mapping[str, int],
116 + ) -> Mapping[str, int]:
117 + max_image_tokens = self.get_max_image_tokens()
118 + return {"image": max_image_tokens}
119 +
120 + def get_hf_processor(
121 + self,
122 + **kwargs: object,
123 + ) -> Qwen2VLProcessor:
124 + self.get_tokenizer(
125 + ).image_token = IMAGE_TOKEN # Ensure image token is set
126 + processor = self.ctx.get_hf_processor(
127 + Qwen2VLProcessor,
128 + **kwargs,
129 + )
130 + processor.image_token = IMAGE_TOKEN
131 + processor.video_token = "<|video_pad|>"
132 + return processor
133 +
134 +
135 +def rotate_half(x):
136 + """Rotates half the hidden dims of the input."""
137 + x1 = x[..., :x.shape[-1] // 2]
138 + x2 = x[..., x.shape[-1] // 2:]
139 + return torch.cat((-x2, x1), dim=-1)
140 +
141 +
142 +def apply_rotary_pos_emb_vision(tensor: torch.Tensor,
143 + freqs: torch.Tensor) -> torch.Tensor:
144 + orig_dtype = tensor.dtype
145 + tensor = tensor.float()
146 +
147 + cos = freqs.cos()
148 + sin = freqs.sin()
149 +
150 + cos = cos.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
151 + sin = sin.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
152 +
153 + output = (tensor * cos) + (rotate_half(tensor) * sin)
154 +
155 + output = output.to(orig_dtype)
156 +
157 + return output
158 +
159 +
160 +class VisionRotaryEmbedding(nn.Module):
161 +
162 + def __init__(self, dim: int, theta: float = 10000.0) -> None:
163 + super().__init__()
164 + inv_freq = 1.0 / (theta
165 + **(torch.arange(0, dim, 2, dtype=torch.float) / dim))
166 + self.register_buffer("inv_freq", inv_freq, persistent=False)
167 +
168 + def forward(self, seqlen: int) -> torch.Tensor:
169 + seq = torch.arange(seqlen,
170 + device=self.inv_freq.device,
171 + dtype=self.inv_freq.dtype)
172 + freqs = torch.outer(seq, self.inv_freq)
173 + return freqs
174 +
175 +
176 +class PatchMerger(nn.Module):
177 +
178 + def __init__(
179 + self,
180 + dim: int,
181 + context_dim: int,
182 + spatial_merge_size: int = 2,
183 + pre_norm="layernorm",
184 + ) -> None:
185 + super().__init__()
186 + self.hidden_size = context_dim * (spatial_merge_size**2)
187 + self.pre_norm = pre_norm
188 + if self.pre_norm == "layernorm":
189 + self.ln_q = LayerNorm(context_dim, eps=1e-6)
190 + elif self.pre_norm == "rmsnorm":
191 + self.ln_q = RMSNorm(context_dim, eps=1e-6)
192 + else:
193 + print("no norm in patch merger")
194 +
195 + self.mlp = nn.Sequential(
196 + ColumnParallelLinear(self.hidden_size,
197 + self.hidden_size,
198 + bias=True,
199 + return_bias=False,
200 + disable_tp=True),
201 + nn.GELU(),
202 + RowParallelLinear(self.hidden_size,
203 + dim,
204 + bias=True,
205 + return_bias=False,
206 + disable_tp=True),
207 + )
208 +
209 + def forward(self, x: torch.Tensor) -> torch.Tensor:
210 + if self.pre_norm:
211 + x = self.mlp(self.ln_q(x).view(-1, self.hidden_size))
212 + else:
213 + x = self.mlp(x.view(-1, self.hidden_size))
214 + return x
215 +
216 +
217 +class DotsVisionAttention(nn.Module):
218 +
219 + def __init__(self,
220 + config,
221 + dim: int,
222 + num_heads: int = 16,
223 + bias: bool = True,
224 + *,
225 + quant_config: Optional[QuantizationConfig] = None,
226 + prefix: str = "") -> None:
227 + super().__init__()
228 + from vllm.distributed import (parallel_state,
229 + tensor_model_parallel_all_gather)
230 + from vllm.distributed import utils as dist_utils
231 +
232 + self.embed_dim = dim
233 + self.num_heads = num_heads
234 + self.head_dim = dim // num_heads
235 + self.tp_size = parallel_state.get_tensor_model_parallel_world_size()
236 + self.tp_rank = parallel_state.get_tensor_model_parallel_rank()
237 + self.num_heads_per_partition = dist_utils.divide(
238 + num_heads, self.tp_size)
239 +
240 + # qkv/proj follow Qwen2-VL style; bias controlled by arg
241 + self.qkv = QKVParallelLinear(hidden_size=dim,
242 + head_size=dim // num_heads,
243 + total_num_heads=num_heads,
244 + bias=bias,
245 + quant_config=quant_config,
246 + prefix=f"{prefix}.qkv")
247 + self.proj = RowParallelLinear(input_size=dim,
248 + output_size=dim,
249 + bias=bias,
250 + quant_config=quant_config,
251 + prefix=f"{prefix}.proj")
252 + self._all_gather = tensor_model_parallel_all_gather
253 + self._split_last = dist_utils.split_tensor_along_last_dim
254 +
255 + # Select attention backend
256 + self.attn_backend = get_vit_attn_backend(self.head_dim,
257 + torch.get_default_dtype())
258 + self.use_upstream_fa = False
259 + if self.attn_backend != _Backend.FLASH_ATTN and \
260 + check_upstream_fa_availability(torch.get_default_dtype()):
261 + self.attn_backend = _Backend.FLASH_ATTN
262 + self.use_upstream_fa = True
263 + if self.attn_backend not in {
264 + _Backend.FLASH_ATTN, _Backend.TORCH_SDPA, _Backend.XFORMERS,
265 + _Backend.ROCM_AITER_FA
266 + }:
267 + raise RuntimeError(
268 + f"Unsupported vision attention backend: {self.attn_backend}")
269 + self.is_flash_attn_backend = self.attn_backend in {
270 + _Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA
271 + }
272 +
273 + def _split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
274 + # qkv: [S, B, 3*dim]
275 + seq_len, bs, _ = qkv.shape
276 + if self.tp_size > 1:
277 + qkv = self._all_gather(qkv)
278 + q, k, v = qkv.chunk(3, dim=2)
279 + if self.tp_size > 1:
280 + q = self._split_last(q, num_partitions=self.tp_size)[self.tp_rank]
281 + k = self._split_last(k, num_partitions=self.tp_size)[self.tp_rank]
282 + v = self._split_last(v, num_partitions=self.tp_size)[self.tp_rank]
283 + new_shape = (seq_len, bs, self.num_heads_per_partition, self.head_dim)
284 + return (q.view(*new_shape), k.view(*new_shape), v.view(*new_shape))
285 +
286 + def forward(
287 + self,
288 + hidden_states: torch.Tensor,
289 + cu_seqlens: torch.Tensor,
290 + rotary_pos_emb: Optional[torch.Tensor] = None,
291 + *,
292 + max_seqlen: Optional[int] = None,
293 + seqlens: Optional[list[int]] = None,
294 + ) -> torch.Tensor:
295 + # [S, C] -> [S, B=1, C]
296 + x = hidden_states.unsqueeze(1)
297 + x, _ = self.qkv(x)
298 + q, k, v = self._split_qkv(x)
299 + bs = q.shape[1]
300 + # [S,B,H,D] -> [B,S,H,D]
301 + q = q.permute(1, 0, 2, 3).contiguous()
302 + k = k.permute(1, 0, 2, 3).contiguous()
303 + v = v.permute(1, 0, 2, 3).contiguous()
304 +
305 + if rotary_pos_emb is not None:
306 + qk_concat = torch.cat([q, k], dim=0)
307 + qk_rotated = apply_rotary_pos_emb_vision(qk_concat, rotary_pos_emb)
308 + q, k = torch.chunk(qk_rotated, 2, dim=0)
309 +
310 + if self.is_flash_attn_backend:
311 + if self.attn_backend == _Backend.ROCM_AITER_FA:
312 + from aiter import flash_attn_varlen_func
313 + else:
314 + if self.use_upstream_fa:
315 + from flash_attn import flash_attn_varlen_func
316 + else:
317 + from vllm.vllm_flash_attn import flash_attn_varlen_func
318 + q_ = q.reshape(bs * q.shape[1], q.shape[2], q.shape[3])
319 + k_ = k.reshape(bs * k.shape[1], k.shape[2], k.shape[3])
320 + v_ = v.reshape(bs * v.shape[1], v.shape[2], v.shape[3])
321 + output = flash_attn_varlen_func(q_,
322 + k_,
323 + v_,
324 + cu_seqlens_q=cu_seqlens,
325 + cu_seqlens_k=cu_seqlens,
326 + max_seqlen_q=max_seqlen,
327 + max_seqlen_k=max_seqlen,
328 + dropout_p=0.0,
329 + causal=False)
330 + context_layer = output.view(bs, -1, self.num_heads_per_partition,
331 + self.head_dim)
332 + elif self.attn_backend == _Backend.TORCH_SDPA:
333 + outputs = []
334 + for i in range(1, len(cu_seqlens)):
335 + s = int(cu_seqlens[i - 1])
336 + e = int(cu_seqlens[i])
337 + q_i = q[:, s:e].permute(0, 2, 1, 3)
338 + k_i = k[:, s:e].permute(0, 2, 1, 3)
339 + v_i = v[:, s:e].permute(0, 2, 1, 3)
340 + out_i = F.scaled_dot_product_attention(q_i,
341 + k_i,
342 + v_i,
343 + dropout_p=0.0)
344 + out_i = out_i.permute(0, 2, 1, 3)
345 + outputs.append(out_i)
346 + context_layer = torch.cat(outputs, dim=1) if outputs else q[:, :0]
347 + elif self.attn_backend == _Backend.XFORMERS:
348 + from xformers import ops as xops
349 + from xformers.ops.fmha.attn_bias import BlockDiagonalMask
350 + attn_bias = BlockDiagonalMask.from_seqlens(q_seqlen=seqlens,
351 + kv_seqlen=None,
352 + device=q.device)
353 + context_layer = xops.memory_efficient_attention_forward(
354 + q, k, v, attn_bias=attn_bias, p=0, scale=None)
355 + else:
356 + raise RuntimeError("Unsupported attention backend")
357 +
358 + # [B,S,H,D] -> [S,B,H*D] -> [S, C]
359 + context_layer = context_layer.permute(1, 0, 2, 3).contiguous()
360 + context_layer = context_layer.view(context_layer.shape[0], bs, -1)
361 + out, _ = self.proj(context_layer)
362 + return out.squeeze(1)
363 +
364 +
365 +class DotsSwiGLUFFN(nn.Module):
366 +
367 + def __init__(self,
368 + config,
369 + *,
370 + quant_config: Optional[QuantizationConfig] = None,
371 + prefix: str = ""):
372 + super().__init__()
373 + hidden_features = config.intermediate_size
374 + in_features = config.embed_dim
375 + bias = config.use_bias
376 +
377 + # Referenced aimv2.py AIMv2SwiGLUFFN
378 + self.fc13 = MergedColumnParallelLinear(in_features,
379 + [hidden_features] * 2,
380 + bias=bias,
381 + quant_config=quant_config,
382 + prefix=f"{prefix}.fc13",
383 + disable_tp=True)
384 + self.fc2 = RowParallelLinear(hidden_features,
385 + in_features,
386 + bias=bias,
387 + quant_config=quant_config,
388 + prefix=f"{prefix}.fc2",
389 + disable_tp=True)
390 + self.act_fn = SiluAndMul()
391 +
392 + def forward(self, x: torch.Tensor) -> torch.Tensor:
393 + x, _ = self.fc13(x)
394 + x = self.act_fn(x)
395 + x, _ = self.fc2(x)
396 + return x
397 +
398 + def load_weights(self, weights: Iterable[tuple[str,
399 + torch.Tensor]]) -> set[str]:
400 + params = dict(self.named_parameters())
401 + loaded: set[str] = set()
402 + for name, w in weights:
403 + # Map fc1 -> fc13 (shard 0)
404 + if name.startswith("fc1."):
405 + tgt = name.replace("fc1.", "fc13.")
406 + if tgt in params:
407 + params[tgt].weight_loader(params[tgt], w, 0)
408 + loaded.add(tgt)
409 + continue
410 + # Map fc3 -> fc13 (shard 1)
411 + if name.startswith("fc3."):
412 + tgt = name.replace("fc3.", "fc13.")
413 + if tgt in params:
414 + params[tgt].weight_loader(params[tgt], w, 1)
415 + loaded.add(tgt)
416 + continue
417 + # Pass-through for fc2 and others
418 + if name in params:
419 + params[name].weight_loader(params[name], w)
420 + loaded.add(name)
421 + return loaded
422 +
423 +
424 +class DotsPatchEmbed(nn.Module):
425 +
426 + def __init__(self, config):
427 + super().__init__()
428 + self.num_channels = config.num_channels
429 + self.patch_size = config.patch_size
430 + self.temporal_patch_size = config.temporal_patch_size
431 + self.embed_dim = config.embed_dim
432 + self.config = config
433 + self.proj = nn.Conv2d(
434 + config.num_channels,
435 + config.embed_dim,
436 + kernel_size=(config.patch_size, config.patch_size),
437 + stride=(config.patch_size, config.patch_size),
438 + )
439 + self.norm = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
440 +
441 + def forward(self, x: torch.Tensor, grid_thw=None) -> torch.Tensor:
442 + x = x.view(-1, self.num_channels, self.temporal_patch_size,
443 + self.patch_size, self.patch_size)[:, :, 0]
444 + x = self.proj(x).view(-1, self.embed_dim)
445 + x = self.norm(x)
446 + return x
447 +
448 +
449 +class DotsViTPreprocessor(nn.Module):
450 +
451 + def __init__(self, config):
452 + super().__init__()
453 + self.patch_h = config.patch_size
454 + self.patch_w = config.patch_size
455 + self.embed_dim = config.embed_dim
456 + self.config = config
457 + self.patchifier = DotsPatchEmbed(config)
458 +
459 + def forward(self, x: torch.Tensor, grid_thw=None) -> torch.Tensor:
460 + tokens = self.patchifier(x, grid_thw)
461 + return tokens
462 +
463 +
464 +class DotsVisionBlock(nn.Module):
465 +
466 + def __init__(self,
467 + config,
468 + *,
469 + quant_config: Optional[QuantizationConfig] = None,
470 + prefix: str = ""):
471 + super().__init__()
472 +
473 + self.attn = DotsVisionAttention(
474 + config,
475 + config.embed_dim,
476 + num_heads=config.num_attention_heads,
477 + bias=config.use_bias,
478 + quant_config=quant_config,
479 + prefix=f"{prefix}.attn",
480 + )
481 + self.norm1 = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
482 + self.mlp = DotsSwiGLUFFN(config,
483 + quant_config=quant_config,
484 + prefix=f"{prefix}.mlp")
485 + self.norm2 = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
486 +
487 + def forward(self,
488 + hidden_states: torch.Tensor,
489 + *,
490 + cu_seqlens: torch.Tensor,
491 + rotary_pos_emb: torch.Tensor,
492 + max_seqlen: Optional[int] = None,
493 + seqlens: Optional[list[int]] = None) -> torch.Tensor:
494 + hidden_states = hidden_states + self.attn(
495 + self.norm1(hidden_states),
496 + cu_seqlens=cu_seqlens,
497 + rotary_pos_emb=rotary_pos_emb,
498 + max_seqlen=max_seqlen,
499 + seqlens=seqlens,
500 + )
501 + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
502 + return hidden_states
503 +
504 +
505 +class DotsVisionTransformer(PreTrainedModel):
506 +
507 + def __init__(
508 + self,
509 + config: DotsVisionConfig,
510 + quant_config: Optional[QuantizationConfig] = None,
511 + *,
512 + num_hidden_layers_override: Optional[int] = None,
513 + require_post_norm: Optional[bool] = None,
514 + prefix: str = "",
515 + ) -> None:
516 + super().__init__(config)
517 + self.config = config
518 + self.spatial_merge_size = config.spatial_merge_size
519 +
520 + self.patch_embed = DotsViTPreprocessor(config)
521 +
522 + head_dim = config.embed_dim // config.num_attention_heads
523 + self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2)
524 + self.attn_backend = get_vit_attn_backend(
525 + head_size=head_dim, dtype=torch.get_default_dtype())
526 + if self.attn_backend != _Backend.FLASH_ATTN and \
527 + check_upstream_fa_availability(torch.get_default_dtype()):
528 + self.attn_backend = _Backend.FLASH_ATTN
529 +
530 + # Keep blocks for compatibility with other vision towers
531 + num_layers = (config.num_hidden_layers if num_hidden_layers_override
532 + is None else num_hidden_layers_override)
533 + self.blocks = nn.ModuleList([
534 + DotsVisionBlock(config,
535 + quant_config=quant_config,
536 + prefix=f"{prefix}.blocks.{i}")
537 + for i in range(num_layers)
538 + ])
539 + if require_post_norm is None:
540 + require_post_norm = (len(self.blocks) == config.num_hidden_layers)
541 + if require_post_norm and self.config.post_norm:
542 + self.post_trunk_norm = RMSNorm(config.embed_dim,
543 + eps=config.rms_norm_eps)
544 + else:
545 + self.post_trunk_norm = None
546 +
547 + self.merger = PatchMerger(
548 + dim=config.hidden_size,
549 + context_dim=config.embed_dim,
550 + spatial_merge_size=config.spatial_merge_size,
551 + )
552 +
553 + @property
554 + def dtype(self) -> torch.dtype:
555 + return self.patch_embed.patchifier.proj.weight.dtype
556 +
557 + @property
558 + def device(self) -> torch.device:
559 + return self.patch_embed.patchifier.proj.weight.device
560 +
561 + def get_pos_ids_by_grid(self, grid_thw):
562 + pos_ids = []
563 + for t, h, w in grid_thw:
564 + hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
565 + hpos_ids = hpos_ids.reshape(
566 + h // self.spatial_merge_size,
567 + self.spatial_merge_size,
568 + w // self.spatial_merge_size,
569 + self.spatial_merge_size,
570 + )
571 + hpos_ids = hpos_ids.permute(0, 2, 1, 3)
572 + hpos_ids = hpos_ids.flatten()
573 +
574 + wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
575 + wpos_ids = wpos_ids.reshape(
576 + h // self.spatial_merge_size,
577 + self.spatial_merge_size,
578 + w // self.spatial_merge_size,
579 + self.spatial_merge_size,
580 + )
581 + wpos_ids = wpos_ids.permute(0, 2, 1, 3)
582 + wpos_ids = wpos_ids.flatten()
583 + pos_ids.append(
584 + torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
585 +
586 + return pos_ids
587 +
588 + def rot_pos_emb(self, grid_thw):
589 + pos_ids = self.get_pos_ids_by_grid(grid_thw)
590 + pos_ids = torch.cat(pos_ids, dim=0)
591 + max_grid_size = grid_thw[:, 1:].max()
592 + rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
593 + rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
594 + return rotary_pos_emb
595 +
596 + def compute_attn_mask_seqlen(
597 + self, cu_seqlens: torch.Tensor
598 + ) -> tuple[Optional[int], Optional[list[int]]]:
599 + max_seqlen, seqlens = None, None
600 + if self.attn_backend == _Backend.FLASH_ATTN:
601 + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
602 + elif self.attn_backend == _Backend.XFORMERS:
603 + seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
604 + return max_seqlen, seqlens
605 +
606 + def forward(self, hidden_states: torch.Tensor,
607 + grid_thw: torch.Tensor) -> torch.Tensor:
608 + hidden_states = hidden_states.to(self.dtype)
609 + hidden_states = self.patch_embed(hidden_states, grid_thw)
610 +
611 + rotary_pos_emb = self.rot_pos_emb(grid_thw)
612 +
613 + cu_seqlens = torch.repeat_interleave(
614 + grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
615 + dim=0,
616 + dtype=grid_thw.dtype
617 + if torch.jit.is_tracing() else torch.int32,
618 + )
619 + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
620 +
621 + max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)
622 + for blk in self.blocks:
623 + hidden_states = blk(hidden_states,
624 + cu_seqlens=cu_seqlens,
625 + rotary_pos_emb=rotary_pos_emb,
626 + max_seqlen=max_seqlen,
627 + seqlens=seqlens)
628 +
629 + if self.post_trunk_norm is not None:
630 + hidden_states = self.post_trunk_norm(hidden_states)
631 +
632 + hidden_states = self.merger(hidden_states)
633 + return hidden_states
634 +
635 +
636 +@MULTIMODAL_REGISTRY.register_processor(
637 + Qwen2VLMultiModalProcessor,
638 + info=DotsOCRProcessingInfo,
639 + dummy_inputs=DotsOCRDummyInputsBuilder,
640 +)
641 +class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
642 + hf_to_vllm_mapper = WeightsMapper(
643 + orig_to_new_substr={
644 + ".attn.qkv_proj.": ".attn.qkv.",
645 + ".attn.out_proj.": ".attn.proj.",
646 + },
647 + orig_to_new_prefix={
648 + "lm_head.": "language_model.lm_head.",
649 + "model.": "language_model.model.",
650 + },
651 + )
652 +
653 + @classmethod
654 + def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:
655 + if modality.startswith("image"):
656 + return "<|img|><|imgpad|><|endofimg|>"
657 +
658 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
659 + super().__init__()
660 +
661 + self.config: DotsOCRConfig = vllm_config.model_config.hf_config
662 + self.quant_config = vllm_config.quant_config
663 + self.multimodal_config = vllm_config.model_config.multimodal_config
664 +
665 + if isinstance(self.config.vision_config, dict):
666 + vision_config = DotsVisionConfig(**self.config.vision_config)
667 + self.config.vision_config = vision_config
668 + else:
669 + vision_config = self.config.vision_config
670 +
671 + self.vision_tower = DotsVisionTransformer(
672 + vision_config,
673 + quant_config=self.quant_config,
674 + prefix=maybe_prefix(prefix, "vision_tower"),
675 + )
676 + self.language_model: Qwen2ForCausalLM = init_vllm_registered_model(
677 + vllm_config=vllm_config,
678 + hf_config=self.config,
679 + prefix=maybe_prefix(prefix, "language_model"),
680 + architectures=["Qwen2ForCausalLM"],
681 + )
682 +
683 + def _validate_and_reshape_mm_tensor(self, mm_input: object,
684 + name: str) -> torch.Tensor:
685 + if not isinstance(mm_input, (torch.Tensor, list)):
686 + raise ValueError(f"Incorrect type of {name}. "
687 + f"Got type: {type(mm_input)}")
688 + if isinstance(mm_input, torch.Tensor):
689 + if mm_input.ndim == 2:
690 + return mm_input
691 + if mm_input.ndim != 3:
692 + raise ValueError(f"{name} should be 2D or batched 3D tensor. "
693 + f"Got ndim: {mm_input.ndim} "
694 + f"(shape={mm_input.shape})")
695 + return torch.concat(list(mm_input))
696 + else:
697 + return torch.concat(mm_input)
698 +
699 + def _parse_and_validate_image_input(
700 + self, **kwargs: object) -> Optional[DotsOCRImageInputs]:
701 + pixel_values = kwargs.pop("pixel_values", None)
702 + image_embeds = kwargs.pop("image_embeds", None)
703 + image_grid_thw = kwargs.pop("image_grid_thw", None)
704 +
705 + if pixel_values is None and image_embeds is None:
706 + return None
707 +
708 + if pixel_values is not None:
709 + pixel_values = self._validate_and_reshape_mm_tensor(
710 + pixel_values, "image pixel values")
711 + image_grid_thw = self._validate_and_reshape_mm_tensor(
712 + image_grid_thw, "image grid_thw")
713 +
714 + if not isinstance(pixel_values, (torch.Tensor, list)):
715 + raise ValueError("Incorrect type of image pixel values. "
716 + f"Got type: {type(pixel_values)}")
717 +
718 + return DotsOCRImagePixelInputs(type="pixel_values",
719 + pixel_values=pixel_values,
720 + image_grid_thw=image_grid_thw)
721 +
722 + if image_embeds is not None:
723 + image_embeds = self._validate_and_reshape_mm_tensor(
724 + image_embeds, "image embeds")
725 + image_grid_thw = self._validate_and_reshape_mm_tensor(
726 + image_grid_thw, "image grid_thw")
727 +
728 + if not isinstance(image_embeds, torch.Tensor):
729 + raise ValueError("Incorrect type of image embeddings. "
730 + f"Got type: {type(image_embeds)}")
731 + return DotsOCRImageEmbeddingInputs(type="image_embeds",
732 + image_embeds=image_embeds,
733 + image_grid_thw=image_grid_thw)
734 +
735 + def _process_image_input(
736 + self, image_input: DotsOCRImageInputs) -> tuple[torch.Tensor, ...]:
737 + grid_thw = image_input["image_grid_thw"]
738 + assert grid_thw.ndim == 2
739 + grid_thw_list = grid_thw.tolist()
740 +
741 + if image_input["type"] == "image_embeds":
742 + image_embeds = image_input["image_embeds"].type(
743 + self.vision_tower.dtype)
744 + else:
745 + pixel_values = image_input["pixel_values"].type(
746 + self.vision_tower.dtype)
747 + image_embeds = self.vision_tower(
748 + pixel_values, grid_thw)[:, :self.config.hidden_size]
749 +
750 + # Split concatenated embeddings for each image item.
751 + merge_size = self.vision_tower.spatial_merge_size
752 + sizes = (torch.tensor(grid_thw_list, dtype=torch.long).prod(-1) //
753 + (merge_size * merge_size)).tolist()
754 +
755 + return image_embeds.split(sizes)
756 +
757 + def get_language_model(self) -> torch.nn.Module:
758 + return self.language_model
759 +
760 + def get_multimodal_embeddings(
761 + self, **kwargs: object) -> Optional[MultiModalEmbeddings]:
762 + image_input = self._parse_and_validate_image_input(**kwargs)
763 + if image_input is None:
764 + return []
765 + vision_embeddings = self._process_image_input(image_input)
766 + return vision_embeddings
767 +
768 + def get_input_embeddings(
769 + self,
770 + input_ids: torch.Tensor,
771 + multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
772 + ) -> torch.Tensor:
773 + inputs_embeds = self.language_model.get_input_embeddings(input_ids)
774 + if multimodal_embeddings is not None:
775 + inputs_embeds = merge_multimodal_embeddings(
776 + input_ids,
777 + inputs_embeds,
778 + multimodal_embeddings,
779 + self.config.image_token_id,
780 + )
781 +
782 + return inputs_embeds
783 +
784 + def forward(
785 + self,
786 + input_ids: Optional[torch.Tensor],
787 + positions: torch.Tensor,
788 + intermediate_tensors: Optional[IntermediateTensors] = None,
789 + inputs_embeds: Optional[torch.Tensor] = None,
790 + **kwargs,
791 + ) -> Union[torch.Tensor, IntermediateTensors]:
792 + if intermediate_tensors is not None:
793 + inputs_embeds = None
794 + elif inputs_embeds is None and kwargs.get("pixel_values") is not None:
795 + image_input = self._parse_and_validate_image_input(**kwargs)
796 + if image_input is None:
797 + inputs_embeds = None
798 + else:
799 + assert input_ids is not None
800 + inputs_embeds = self.get_multimodal_embeddings(
801 + input_ids,
802 + image_input=image_input,
803 + )
804 + input_ids = None
805 +
806 + hidden_states = self.language_model(
807 + input_ids=input_ids,
808 + positions=positions,
809 + intermediate_tensors=intermediate_tensors,
810 + inputs_embeds=inputs_embeds,
811 + )
812 +
813 + return hidden_states
814 +
815 + def compute_logits(
816 + self,
817 + hidden_states: torch.Tensor,
818 + ) -> Optional[torch.Tensor]:
819 + return self.language_model.compute_logits(hidden_states)
820 +
821 + def load_weights(self, weights: Iterable[tuple[str,
822 + torch.Tensor]]) -> set[str]:
823 + loader = AutoWeightsLoader(self)
824 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/registry.py
+1
@@ -219,6 +219,7 @@ _MULTIMODAL_MODELS = {
219 "ChameleonForConditionalGeneration": ("chameleon", "ChameleonForConditionalGeneration"), # noqa: E501
220 "Cohere2VisionForConditionalGeneration": ("cohere2_vision", "Cohere2VisionForConditionalGeneration"), # noqa: E501
221 "DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"),
222 + "DotsOCRForCausalLM": ("dots_ocr", "DotsOCRForCausalLM"),
223 "Ernie4_5_VLMoeForConditionalGeneration": ("ernie45_vl", "Ernie4_5_VLMoeForConditionalGeneration"), # noqa: E501
224 "FuyuForCausalLM": ("fuyu", "FuyuForCausalLM"),
225 "Gemma3ForConditionalGeneration": ("gemma3_mm", "Gemma3ForConditionalGeneration"), # noqa: E501
vllm/transformers_utils/configs/__init__.py
+2
@@ -9,6 +9,7 @@ Model configs may be defined in this directory for the following reasons:
9
10 from vllm.transformers_utils.configs.chatglm import ChatGLMConfig
11 from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekVLV2Config
12 +from vllm.transformers_utils.configs.dotsocr import DotsOCRConfig
13 from vllm.transformers_utils.configs.eagle import EAGLEConfig
14 # RWConfig is for the original tiiuae/falcon-40b(-instruct) and
15 # tiiuae/falcon-7b(-instruct) models. Newer Falcon models will use the
@@ -36,6 +37,7 @@ from vllm.transformers_utils.configs.ultravox import UltravoxConfig
37 __all__ = [
38 "ChatGLMConfig",
39 "DeepseekVLV2Config",
40 + "DotsOCRConfig",
41 "EAGLEConfig",
42 "RWConfig",
43 "JAISConfig",
vllm/transformers_utils/configs/dotsocr.py new
+69
@@ -0,0 +1,69 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Any, Optional
4 +
5 +from transformers.configuration_utils import PretrainedConfig
6 +from transformers.models.qwen2 import Qwen2Config
7 +
8 +
9 +class DotsVisionConfig(PretrainedConfig):
10 + model_type: str = "dots_vit"
11 +
12 + def __init__(
13 + self,
14 + embed_dim: int = 1536, # vision encoder embed size
15 + hidden_size: int = 1536, # after merger hidden size
16 + intermediate_size: int = 4224,
17 + num_hidden_layers: int = 42,
18 + num_attention_heads: int = 12,
19 + num_channels: int = 3,
20 + patch_size: int = 14,
21 + spatial_merge_size: int = 2,
22 + temporal_patch_size: int = 1,
23 + rms_norm_eps: float = 1e-5,
24 + use_bias: bool = False,
25 + attn_implementation="flash_attention_2",
26 + initializer_range=0.02,
27 + init_merger_std=0.02,
28 + is_causal=False, # ve causal forward
29 + post_norm=True,
30 + gradient_checkpointing=False,
31 + **kwargs: Any,
32 + ):
33 + super().__init__(**kwargs)
34 + self.embed_dim = embed_dim
35 + self.hidden_size = hidden_size
36 + self.intermediate_size = intermediate_size
37 + self.num_hidden_layers = num_hidden_layers
38 + self.num_attention_heads = num_attention_heads
39 + self.num_channels = num_channels
40 + self.patch_size = patch_size
41 + self.spatial_merge_size = spatial_merge_size
42 + self.temporal_patch_size = temporal_patch_size
43 + self.rms_norm_eps = rms_norm_eps
44 + self.use_bias = use_bias
45 + self.attn_implementation = attn_implementation
46 + self.initializer_range = initializer_range
47 + self.init_merger_std = init_merger_std
48 + self.is_causal = is_causal
49 + self.post_norm = post_norm
50 + self.gradient_checkpointing = gradient_checkpointing
51 +
52 +
53 +class DotsOCRConfig(Qwen2Config):
54 + model_type = "dots_ocr"
55 +
56 + def __init__(self,
57 + image_token_id=151665,
58 + video_token_id=151656,
59 + vision_config: Optional[dict] = None,
60 + *args,
61 + **kwargs):
62 + super().__init__(*args, **kwargs)
63 + self.image_token_id = image_token_id
64 + self.video_token_id = video_token_id
65 + self.vision_config = DotsVisionConfig(**(vision_config or {}))
66 +
67 + def save_pretrained(self, save_directory, **kwargs):
68 + self._auto_class = None
69 + super().save_pretrained(save_directory, **kwargs)