@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 03fd76c570
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 10, 12:45 UTC

[Model] Add LFM2-VL model support (#31758)

Signed-off-by: Tianshu Yu <tianshuyu.formal@gmail.com> Signed-off-by: Cyrus Leung <cyrus.tl.leung@gmail.com> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>

tianshu-Michael-yu committed Jan 8, 2026 at 05:00 UTC 03fd76c57002747ded9bee5e7a5a699475cb4e9f
6 files changed +1266 -1
docs/models/supported_models.md
+1
@@ -703,6 +703,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
703 | `KeyeVL1_5ForConditionalGeneration` | Keye-VL-1_5-8B | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-1_5-8B` | ✅︎ | ✅︎ |
704 | `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I<sup>+</sup> | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | ✅︎ |
705 | `LightOnOCRForConditionalGeneration` | LightOnOCR-1B | T + I<sup>+</sup> | `lightonai/LightOnOCR-1B`, etc | ✅︎ | ✅︎ |
706 +| `Lfm2VlForConditionalGeneration` | LFM2-VL | T + I<sup>+</sup> | `LiquidAI/LFM2-VL-450M`, `LiquidAI/LFM2-VL-3B`, `LiquidAI/LFM2-VL-8B-A1B`, etc. | ✅︎ | ✅︎ |
707 | `Llama4ForConditionalGeneration` | Llama 4 | T + I<sup>+</sup> | `meta-llama/Llama-4-Scout-17B-16E-Instruct`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8`, `meta-llama/Llama-4-Maverick-17B-128E-Instruct`, etc. | ✅︎ | ✅︎ |
708 | `Llama_Nemotron_Nano_VL` | Llama Nemotron Nano VL | T + I<sup>E+</sup> | `nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1` | ✅︎ | ✅︎ |
709 | `LlavaForConditionalGeneration` | LLaVA-1.5, Pixtral (HF Transformers) | T + I<sup>E+</sup> | `llava-hf/llava-1.5-7b-hf`, `TIGER-Lab/Mantis-8B-siglip-llama3` (see note), `mistral-community/pixtral-12b`, etc. | ✅︎ | ✅︎ |
examples/offline_inference/vision_language.py
+33 -1
@@ -15,7 +15,7 @@ from dataclasses import asdict
15 from typing import NamedTuple
16
17 from huggingface_hub import snapshot_download
18 -from transformers import AutoTokenizer
18 +from transformers import AutoProcessor, AutoTokenizer
19
20 from vllm import LLM, EngineArgs, SamplingParams
21 from vllm.assets.image import ImageAsset
@@ -875,6 +875,37 @@ def run_lightonocr(questions: list[str], modality: str) -> ModelRequestData:
875 )
876
877
878 +def run_lfm2_vl(questions: list[str], modality: str) -> ModelRequestData:
879 + assert modality == "image"
880 +
881 + model_name = "LiquidAI/LFM2-VL-450M"
882 +
883 + engine_args = EngineArgs(
884 + model=model_name,
885 + max_model_len=4096,
886 + limit_mm_per_prompt={modality: 1},
887 + )
888 +
889 + processor = AutoProcessor.from_pretrained(model_name)
890 + messages = [
891 + [
892 + {
893 + "role": "user",
894 + "content": [{"type": "image"}, {"type": "text", "text": question}],
895 + }
896 + ]
897 + for question in questions
898 + ]
899 + prompts = processor.apply_chat_template(
900 + messages, tokenize=False, add_generation_prompt=True
901 + )
902 +
903 + return ModelRequestData(
904 + engine_args=engine_args,
905 + prompts=prompts,
906 + )
907 +
908 +
909 def run_llama4(questions: list[str], modality: str) -> ModelRequestData:
910 assert modality == "image"
911
@@ -1849,6 +1880,7 @@ model_example_map = {
1880 "keye_vl1_5": run_keye_vl1_5,
1881 "kimi_vl": run_kimi_vl,
1882 "lightonocr": run_lightonocr,
1883 + "lfm2_vl": run_lfm2_vl,
1884 "llama4": run_llama4,
1885 "llava": run_llava,
1886 "llava-next": run_llava_next,
tests/models/registry.py
+4
@@ -742,6 +742,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {
742 "LightOnOCRForConditionalGeneration": _HfExamplesInfo(
743 "lightonai/LightOnOCR-1B-1025"
744 ),
745 + "Lfm2VlForConditionalGeneration": _HfExamplesInfo(
746 + "LiquidAI/LFM2-VL-450M",
747 + min_transformers_version="5.0.0",
748 + ),
749 "Llama4ForConditionalGeneration": _HfExamplesInfo(
750 "meta-llama/Llama-4-Scout-17B-16E-Instruct",
751 max_model_len=10240,
vllm/model_executor/models/lfm2_vl.py new
+732
@@ -0,0 +1,732 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import itertools
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
11 +from transformers import BatchFeature
12 +from transformers.activations import ACT2FN
13 +from transformers.models.lfm2_vl import Lfm2VlProcessor
14 +from transformers.models.lfm2_vl.configuration_lfm2_vl import Lfm2VlConfig
15 +from transformers.models.lfm2_vl.image_processing_lfm2_vl_fast import (
16 + Lfm2VlImageProcessorFast,
17 + find_closest_aspect_ratio,
18 + round_by_factor,
19 +)
20 +
21 +from vllm.config import VllmConfig
22 +from vllm.config.multimodal import BaseDummyOptions
23 +from vllm.forward_context import set_forward_context
24 +from vllm.model_executor.layers.mamba.mamba_utils import (
25 + MambaStateDtypeCalculator,
26 + MambaStateShapeCalculator,
27 +)
28 +from vllm.model_executor.models.module_mapping import MultiModelKeys
29 +from vllm.multimodal import MULTIMODAL_REGISTRY
30 +from vllm.multimodal.inputs import (
31 + MultiModalDataDict,
32 + MultiModalFieldConfig,
33 + MultiModalKwargsItems,
34 +)
35 +from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
36 +from vllm.multimodal.processing import (
37 + BaseMultiModalProcessor,
38 + BaseProcessingInfo,
39 + PromptReplacement,
40 + PromptUpdateDetails,
41 +)
42 +from vllm.multimodal.profiling import BaseDummyInputsBuilder
43 +from vllm.sequence import IntermediateTensors
44 +from vllm.utils.tensor_schema import TensorSchema, TensorShape
45 +
46 +from .interfaces import (
47 + IsHybrid,
48 + MultiModalEmbeddings,
49 + SupportsLoRA,
50 + SupportsMultiModal,
51 + SupportsPP,
52 +)
53 +from .siglip2 import Siglip2Model
54 +from .utils import (
55 + AutoWeightsLoader,
56 + WeightsMapper,
57 + init_vllm_registered_model,
58 + maybe_prefix,
59 +)
60 +
61 +
62 +class Lfm2VLImagePixelInputs(TensorSchema):
63 + """
64 + Dimensions:
65 + - b: Number of images in the prompt
66 + - bn: Batch size * number of images
67 + - d: Number of dimensions
68 + - fd: Number of features per dimension
69 + """
70 +
71 + type: Literal["pixel_values"] = "pixel_values"
72 + pixel_values: Annotated[torch.Tensor, TensorShape("bn", "d", "fd")]
73 + spatial_shapes: Annotated[torch.Tensor, TensorShape("bn", 2)]
74 + num_patches: Annotated[torch.Tensor, TensorShape("b")]
75 +
76 +
77 +LFM2VLImageInputs = Lfm2VLImagePixelInputs
78 +
79 +
80 +class Lfm2VLProcessingInfo(BaseProcessingInfo):
81 + def get_hf_config(self):
82 + return self.ctx.get_hf_config(Lfm2VlConfig)
83 +
84 + def get_hf_processor(self, **kwargs):
85 + return self.ctx.get_hf_processor(Lfm2VlProcessor, **kwargs)
86 +
87 + def get_image_processor(self, **kwargs: object) -> Lfm2VlImageProcessorFast:
88 + return self.get_hf_processor(**kwargs).image_processor
89 +
90 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:
91 + return {"image": None}
92 +
93 + def get_image_size_with_most_features(self) -> ImageSize:
94 + processor = self.get_image_processor()
95 + max_image_tokens = processor.max_image_tokens
96 + encoder_patch_size = processor.encoder_patch_size
97 + downsample_factor = processor.downsample_factor
98 + max_pixels = max_image_tokens * (encoder_patch_size**2) * (downsample_factor**2)
99 + side = int(math.sqrt(max_pixels))
100 + return ImageSize(width=side, height=side)
101 +
102 + def _is_image_too_large(
103 + self,
104 + height: int,
105 + width: int,
106 + max_image_tokens: int,
107 + encoder_patch_size: int,
108 + downsample_factor: int,
109 + max_pixels_tolerance: float,
110 + ) -> bool:
111 + """Check if the image is too large to be processed as one tile."""
112 + total_factor = encoder_patch_size * downsample_factor
113 +
114 + h_bar = max(encoder_patch_size, round_by_factor(height, total_factor))
115 + w_bar = max(encoder_patch_size, round_by_factor(width, total_factor))
116 + return (
117 + h_bar * w_bar
118 + > max_image_tokens
119 + * encoder_patch_size**2
120 + * downsample_factor**2
121 + * max_pixels_tolerance
122 + )
123 +
124 + def smart_resize(
125 + self,
126 + height: int,
127 + width: int,
128 + downsample_factor: int,
129 + min_image_tokens: int,
130 + max_image_tokens: int,
131 + encoder_patch_size: int,
132 + ) -> tuple[int, int]:
133 + total_factor = encoder_patch_size * downsample_factor
134 + smart_resize_min_pixels = (
135 + min_image_tokens * encoder_patch_size**2 * downsample_factor**2
136 + )
137 + smart_resize_max_pixels = (
138 + max_image_tokens * encoder_patch_size**2 * downsample_factor**2
139 + )
140 +
141 + h_bar = max(total_factor, round_by_factor(height, total_factor))
142 + w_bar = max(total_factor, round_by_factor(width, total_factor))
143 +
144 + if h_bar * w_bar > smart_resize_max_pixels:
145 + beta = math.sqrt((height * width) / smart_resize_max_pixels)
146 + h_bar = max(
147 + total_factor, math.floor(height / beta / total_factor) * total_factor
148 + )
149 + w_bar = max(
150 + total_factor, math.floor(width / beta / total_factor) * total_factor
151 + )
152 + elif h_bar * w_bar < smart_resize_min_pixels:
153 + beta = math.sqrt(smart_resize_min_pixels / (height * width))
154 + h_bar = math.ceil(height * beta / total_factor) * total_factor
155 + w_bar = math.ceil(width * beta / total_factor) * total_factor
156 +
157 + return w_bar, h_bar
158 +
159 + def _target_ratios(self, min_tiles: int, max_tiles: int) -> list[tuple[int, int]]:
160 + ratios = [
161 + (w, h)
162 + for n in range(min_tiles, max_tiles + 1)
163 + for w in range(1, n + 1)
164 + for h in range(1, n + 1)
165 + if min_tiles <= w * h <= max_tiles
166 + ]
167 + return sorted(set(ratios), key=lambda x: x[0] * x[1])
168 +
169 + def _get_grid_layout(
170 + self,
171 + height: int,
172 + width: int,
173 + min_tiles: int,
174 + max_tiles: int,
175 + tile_size: int,
176 + ) -> tuple[int, int]:
177 + aspect_ratio = width / height
178 + target_ratios = self._target_ratios(min_tiles, max_tiles)
179 + # find best matching grid configuration
180 + grid_width, grid_height = find_closest_aspect_ratio(
181 + aspect_ratio, target_ratios, width, height, tile_size
182 + )
183 + total_patches = grid_width * grid_height
184 + return grid_width, grid_height, total_patches
185 +
186 + def _get_image_feature_grid_size(
187 + self,
188 + image_width: int,
189 + image_height: int,
190 + processor: Lfm2VlProcessor | None,
191 + ) -> tuple[int, int]:
192 + if processor is None:
193 + processor = self.get_image_processor()
194 +
195 + downsample_factor = processor.image_processor.downsample_factor
196 + encoder_patch_size = processor.image_processor.encoder_patch_size
197 + max_pixels_tolerance = processor.image_processor.max_pixels_tolerance
198 + min_tiles = processor.image_processor.min_tiles
199 + max_tiles = processor.image_processor.max_tiles
200 + max_image_tokens = processor.image_processor.max_image_tokens
201 + tile_size = processor.image_processor.tile_size
202 +
203 + do_image_splitting = not min_tiles == max_tiles == 1
204 + is_image_large = self._is_image_too_large(
205 + height=image_height,
206 + width=image_width,
207 + max_image_tokens=max_image_tokens,
208 + encoder_patch_size=encoder_patch_size,
209 + downsample_factor=downsample_factor,
210 + max_pixels_tolerance=max_pixels_tolerance,
211 + )
212 +
213 + # Big image will be cropped into patches and small images are just resized
214 + if is_image_large and do_image_splitting:
215 + grid_width, grid_height, total_patches = self._get_grid_layout(
216 + image_height,
217 + image_width,
218 + min_tiles=min_tiles,
219 + max_tiles=max_tiles,
220 + tile_size=tile_size,
221 + )
222 + else:
223 + grid_width = grid_height = total_patches = 1
224 +
225 + if grid_width * grid_height != 1: # Thumbnail
226 + total_patches += 1
227 +
228 + return grid_width, grid_height, total_patches
229 +
230 + def get_num_patches(
231 + self,
232 + *,
233 + image_width: int,
234 + image_height: int,
235 + processor: Lfm2VlProcessor | None,
236 + ) -> int:
237 + _, _, total_patches = self._get_image_feature_grid_size(
238 + image_width=image_width,
239 + image_height=image_height,
240 + processor=processor,
241 + )
242 + return total_patches
243 +
244 + def get_image_repl(
245 + self,
246 + image_width: int,
247 + image_height: int,
248 + spatial_shapes: torch.Tensor,
249 + processor: Lfm2VlProcessor | None,
250 + ) -> str:
251 + if processor is None:
252 + processor = self.get_hf_processor()
253 +
254 + grid_placeholder = "<|img_row_{n_h}_col_{n_w}|>"
255 + image_token = processor.image_token
256 + image_start_token = processor.image_start_token
257 + image_end_token = processor.image_end_token
258 + image_thumbnail_token = processor.image_thumbnail_token
259 +
260 + num_thumbnail_tokens, num_tokens_per_tile = self.get_num_image_tokens(
261 + spatial_shapes=spatial_shapes,
262 + processor=processor,
263 + )
264 + tile_img_placeholder = grid_placeholder + (image_token * num_tokens_per_tile)
265 +
266 + grid_w, grid_h, _ = self._get_image_feature_grid_size(
267 + image_width=image_width,
268 + image_height=image_height,
269 + processor=processor,
270 + )
271 +
272 + if grid_w > 1 or grid_h > 1:
273 + tiles_placeholder: list[str] = [
274 + tile_img_placeholder.format(n_h=i + 1, n_w=j + 1)
275 + for i in range(grid_h)
276 + for j in range(grid_w)
277 + ]
278 +
279 + if num_thumbnail_tokens > 0:
280 + tiles_placeholder.append(
281 + image_thumbnail_token + (image_token * num_thumbnail_tokens)
282 + )
283 + else:
284 + tiles_placeholder = [image_token * num_thumbnail_tokens]
285 +
286 + placeholder = "".join(
287 + itertools.chain([image_start_token], tiles_placeholder, [image_end_token])
288 + )
289 + return placeholder
290 +
291 + def get_num_image_tokens(
292 + self,
293 + *,
294 + spatial_shapes: torch.Tensor,
295 + processor: Lfm2VlProcessor | None,
296 + ) -> tuple[int, int]:
297 + tile_size = processor.image_processor.tile_size
298 + downsample_factor = processor.image_processor.downsample_factor
299 + encoder_patch_size = processor.image_processor.encoder_patch_size
300 + num_thumbnail_tokens = spatial_shapes[-1].prod() // (downsample_factor**2)
301 + num_patches_tile = tile_size // encoder_patch_size
302 + dwn_num_patches_tile = math.ceil(num_patches_tile / downsample_factor)
303 + num_tiles_tokens = dwn_num_patches_tile * dwn_num_patches_tile
304 + return num_thumbnail_tokens, num_tiles_tokens
305 +
306 +
307 +class Lfm2VLDummyInputsBuilder(BaseDummyInputsBuilder[Lfm2VLProcessingInfo]):
308 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
309 + num_images = mm_counts.get("image", 0)
310 + processor = self.info.get_hf_processor()
311 + image_token = processor.image_token
312 + return image_token * num_images
313 +
314 + def get_dummy_mm_data(
315 + self,
316 + seq_len: int,
317 + mm_counts: Mapping[str, int],
318 + mm_options: Mapping[str, BaseDummyOptions] | None = None,
319 + ) -> MultiModalDataDict:
320 + num_images = mm_counts.get("image", 0)
321 +
322 + target_width, target_height = self.info.get_image_size_with_most_features()
323 +
324 + image_overrides = mm_options.get("image") if mm_options else None
325 +
326 + return {
327 + "image": self._get_dummy_images(
328 + width=target_width,
329 + height=target_height,
330 + num_images=num_images,
331 + overrides=image_overrides,
332 + ),
333 + }
334 +
335 +
336 +class Lfm2VLMultiModalProcessor(BaseMultiModalProcessor[Lfm2VLProcessingInfo]):
337 + def _call_hf_processor(
338 + self,
339 + prompt: str,
340 + mm_data: Mapping[str, object],
341 + mm_kwargs: Mapping[str, object],
342 + tok_kwargs: Mapping[str, object],
343 + ) -> BatchFeature:
344 + # Text-only input not supported in composite processor
345 + if not (images := mm_data.get("images", [])):
346 + prompt_ids = self.info.get_tokenizer().encode(prompt)
347 + prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids)
348 + return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
349 +
350 + processed_outputs = super()._call_hf_processor(
351 + prompt,
352 + mm_data,
353 + mm_kwargs,
354 + tok_kwargs,
355 + )
356 +
357 + parsed_images = (
358 + self._get_data_parser()
359 + .parse_mm_data({"image": images})
360 + .get_items("image", ImageProcessorItems)
361 + )
362 + image_sizes = [
363 + parsed_images.get_image_size(i) for i in range(len(parsed_images))
364 + ]
365 + hf_processor = self.info.get_hf_processor(**mm_kwargs)
366 +
367 + num_patches = [
368 + self.info.get_num_patches(
369 + image_width=size.width,
370 + image_height=size.height,
371 + processor=hf_processor,
372 + )
373 + for size in image_sizes
374 + ]
375 + processed_outputs["num_patches"] = torch.tensor(num_patches)
376 +
377 + return processed_outputs
378 +
379 + def _get_mm_fields_config(
380 + self,
381 + hf_inputs: BatchFeature,
382 + hf_processor_mm_kwargs: Mapping[str, object],
383 + ) -> Mapping[str, MultiModalFieldConfig]:
384 + num_patches = hf_inputs.get("num_patches", torch.empty(0))
385 +
386 + return dict[str, MultiModalFieldConfig](
387 + pixel_values=MultiModalFieldConfig.flat_from_sizes("image", num_patches),
388 + spatial_shapes=MultiModalFieldConfig.flat_from_sizes(
389 + "image", num_patches, keep_on_cpu=True
390 + ),
391 + num_patches=MultiModalFieldConfig.batched("image", keep_on_cpu=True),
392 + )
393 +
394 + def _get_prompt_updates(
395 + self,
396 + mm_items: MultiModalDataItems,
397 + hf_processor_mm_kwargs: Mapping[str, object],
398 + out_mm_kwargs: MultiModalKwargsItems,
399 + ) -> Sequence[PromptReplacement]:
400 + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
401 + image_token = hf_processor.image_token
402 +
403 + def get_image_replacement_lfm2vl(item_idx: int):
404 + images = mm_items.get_items("image", ImageProcessorItems)
405 + image_size = images.get_image_size(item_idx)
406 + out_item = out_mm_kwargs["image"][item_idx]
407 + spatial_shapes = out_item["spatial_shapes"].data
408 + assert isinstance(spatial_shapes, torch.Tensor)
409 + image_repl = self.info.get_image_repl(
410 + image_width=image_size.width,
411 + image_height=image_size.height,
412 + spatial_shapes=spatial_shapes,
413 + processor=hf_processor,
414 + )
415 + return PromptUpdateDetails.select_text(
416 + image_repl,
417 + embed_text=image_token,
418 + )
419 +
420 + return [
421 + PromptReplacement(
422 + modality="image",
423 + target=image_token,
424 + replacement=get_image_replacement_lfm2vl,
425 + )
426 + ]
427 +
428 +
429 +class Lfm2VLMultiModalProjector(nn.Module):
430 + def __init__(
431 + self, config: Lfm2VlConfig, use_data_parallel: bool = False, prefix: str = ""
432 + ):
433 + super().__init__()
434 + self.use_data_parallel = use_data_parallel
435 +
436 + in_channels = config.vision_config.hidden_size * (config.downsample_factor**2)
437 + self.factor = config.downsample_factor
438 + self.projector_use_layernorm = config.projector_use_layernorm
439 + if self.projector_use_layernorm:
440 + self.layer_norm = nn.LayerNorm(in_channels)
441 + self.linear_1 = nn.Linear(
442 + in_channels,
443 + config.projector_hidden_size,
444 + bias=config.projector_bias,
445 + )
446 + self.act = ACT2FN[config.projector_hidden_act]
447 + self.linear_2 = nn.Linear(
448 + config.projector_hidden_size,
449 + config.text_config.hidden_size,
450 + bias=config.projector_bias,
451 + )
452 +
453 + def forward(self, image_features: torch.Tensor):
454 + image_features = self.pixel_unshuffle(image_features)
455 + if self.projector_use_layernorm:
456 + image_features = self.layer_norm(image_features)
457 + hidden_states = self.linear_1(image_features)
458 + hidden_states = self.act(hidden_states)
459 + hidden_states = self.linear_2(hidden_states)
460 + return hidden_states
461 +
462 + def pixel_unshuffle(self, hidden_states: torch.Tensor):
463 + batch_size, width, height, channels = hidden_states.size()
464 + hidden_states = hidden_states.reshape(
465 + batch_size, width, height // self.factor, channels * self.factor
466 + )
467 + hidden_states = hidden_states.permute(0, 2, 1, 3)
468 + hidden_states = hidden_states.reshape(
469 + batch_size,
470 + height // self.factor,
471 + width // self.factor,
472 + channels * self.factor**2,
473 + )
474 + hidden_states = hidden_states.permute(0, 2, 1, 3)
475 + return hidden_states
476 +
477 +
478 +@MULTIMODAL_REGISTRY.register_processor(
479 + Lfm2VLMultiModalProcessor,
480 + info=Lfm2VLProcessingInfo,
481 + dummy_inputs=Lfm2VLDummyInputsBuilder,
482 +)
483 +class Lfm2VLForConditionalGeneration(
484 + nn.Module, SupportsMultiModal, SupportsLoRA, SupportsPP, IsHybrid
485 +):
486 + merge_by_field_config = True
487 +
488 + hf_to_vllm_mapper = WeightsMapper(
489 + orig_to_new_prefix={
490 + "lm_head.": "language_model.lm_head.",
491 + "model.language_model.": "language_model.model.",
492 + "model.vision_tower.": "vision_tower.",
493 + "model.multi_modal_projector.": "multi_modal_projector.",
494 + }
495 + )
496 +
497 + @classmethod
498 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:
499 + if modality.startswith("image"):
500 + return "<image>"
501 +
502 + raise ValueError("Only image modality is supported")
503 +
504 + @classmethod
505 + def get_mamba_state_dtype_from_config(
506 + cls,
507 + vllm_config: "VllmConfig",
508 + ) -> tuple[torch.dtype, ...]:
509 + return MambaStateDtypeCalculator.short_conv_state_dtype(
510 + vllm_config.model_config.dtype,
511 + vllm_config.cache_config.mamba_cache_dtype,
512 + )
513 +
514 + @classmethod
515 + def get_mamba_state_shape_from_config(
516 + cls,
517 + vllm_config: "VllmConfig",
518 + ) -> tuple[tuple[int, int]]:
519 + """Calculate shapes for LFM2's convolutional cache.
520 +
521 + Args:
522 + vllm_config: vLLM config
523 +
524 + Returns:
525 + Tuple containing:
526 + - conv_state_shape: Shape for convolutional state cache
527 + """
528 + parallel_config = vllm_config.parallel_config
529 + hf_language_config = vllm_config.model_config.hf_config.text_config
530 +
531 + return MambaStateShapeCalculator.short_conv_state_shape(
532 + tp_world_size=parallel_config.tensor_parallel_size,
533 + intermediate_size=hf_language_config.hidden_size,
534 + conv_kernel=hf_language_config.conv_L_cache,
535 + )
536 +
537 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
538 + super().__init__()
539 + config: Lfm2VlConfig = vllm_config.model_config.hf_config
540 + multimodal_config = vllm_config.model_config.multimodal_config
541 + vision_config = config.vision_config
542 + quant_config = vllm_config.quant_config
543 +
544 + self.config = config
545 + self.vllm_config = vllm_config
546 + self.multimodal_config = multimodal_config
547 + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
548 +
549 + if vision_config.model_type == "siglip2_vision_model":
550 + self.vision_tower = Siglip2Model(
551 + config=vision_config,
552 + quant_config=quant_config,
553 + multimodal_config=multimodal_config,
554 + prefix=maybe_prefix(prefix, "vision_tower"),
555 + )
556 + else:
557 + raise ValueError(
558 + f"Unsupported visual tokenizer model_type: {vision_config.model_type}"
559 + )
560 +
561 + self.multi_modal_projector = Lfm2VLMultiModalProjector(
562 + config=config,
563 + use_data_parallel=self.use_data_parallel,
564 + prefix=f"{prefix}.multi_modal_projector",
565 + )
566 +
567 + self.language_model = init_vllm_registered_model(
568 + vllm_config=vllm_config,
569 + hf_config=config.text_config,
570 + prefix=maybe_prefix(prefix, "language"),
571 + architectures=config.text_config.architectures,
572 + )
573 +
574 + self.make_empty_intermediate_tensors = (
575 + self.language_model.make_empty_intermediate_tensors
576 + )
577 +
578 + def get_language_model(self) -> torch.nn.Module:
579 + return self.language_model
580 +
581 + def _parse_and_validate_image_input(
582 + self, **kwargs: object
583 + ) -> LFM2VLImageInputs | None:
584 + pixel_values = kwargs.pop("pixel_values", None)
585 + spatial_shapes = kwargs.pop("spatial_shapes", None)
586 + num_patches = kwargs.pop("num_patches", None)
587 + if pixel_values is None:
588 + return None
589 +
590 + return LFM2VLImageInputs(
591 + type="pixel_values",
592 + pixel_values=pixel_values,
593 + spatial_shapes=spatial_shapes,
594 + num_patches=num_patches,
595 + )
596 +
597 + def image_pixels_to_features(
598 + self,
599 + pixel_values: torch.FloatTensor,
600 + spatial_shapes: torch.Tensor,
601 + ) -> torch.Tensor:
602 + pixel_values = pixel_values.to(
603 + dtype=self.vision_tower.vision_model.embeddings.patch_embedding.weight.dtype
604 + ) # fp16 compatibility
605 +
606 + # LFM2-VL's HF processor pads patch sequences with trailing zeros.
607 + # Derive the valid-patch mask from spatial_shapes instead of carrying
608 + # pixel_attention_mask through the vLLM multimodal pipeline.
609 + max_seq_len = pixel_values.shape[1]
610 + lengths_cpu = (spatial_shapes[:, 0] * spatial_shapes[:, 1]).to(
611 + dtype=torch.int32
612 + )
613 + max_seqlen = (
614 + lengths_cpu.max().reshape(1).to(device=pixel_values.device)
615 + if lengths_cpu.numel()
616 + else torch.tensor([0], dtype=torch.int32, device=pixel_values.device)
617 + )
618 + lengths = lengths_cpu.to(device=pixel_values.device)
619 + packed_mask = (
620 + torch.arange(max_seq_len, device=pixel_values.device)[None, :]
621 + < lengths[:, None]
622 + )
623 + cu_seqlens = torch.zeros(
624 + lengths.shape[0] + 1,
625 + dtype=torch.int32,
626 + device=lengths.device,
627 + )
628 + cu_seqlens[1:] = torch.cumsum(lengths, dim=0)
629 +
630 + with set_forward_context(None, self.vllm_config):
631 + vision_outputs = self.vision_tower(
632 + pixel_values=pixel_values,
633 + spatial_shapes=spatial_shapes,
634 + packed_mask=packed_mask,
635 + cu_seqlens=cu_seqlens,
636 + max_seqlen=max_seqlen,
637 + )
638 + image_outputs = getattr(vision_outputs, "last_hidden_state", vision_outputs)
639 +
640 + image_features = []
641 +
642 + # spatial_shapes is on CPU (keep_on_cpu=True), so .tolist() is instant
643 + spatial_shapes_list = spatial_shapes.tolist()
644 + for img_idx, (feature_org_h, feature_org_w) in enumerate(spatial_shapes_list):
645 + feature_len = feature_org_h * feature_org_w
646 + feature = image_outputs[img_idx, :feature_len]
647 +
648 + # reshape to original height and width
649 + feature = feature.reshape(1, feature_org_h, feature_org_w, -1)
650 +
651 + # project the image representation
652 + img_embedding = self.multi_modal_projector(feature)
653 +
654 + # flatten here to handle variable length in naflex
655 + img_embedding = img_embedding.reshape(-1, img_embedding.size(-1))
656 + image_features.append(img_embedding)
657 +
658 + return image_features
659 +
660 + def _process_image_input(
661 + self,
662 + image_input: LFM2VLImageInputs,
663 + ) -> torch.Tensor | list[torch.Tensor]:
664 + pixel_values = image_input["pixel_values"]
665 + spatial_shapes = image_input["spatial_shapes"]
666 + num_patches = image_input["num_patches"]
667 +
668 + image_features = self.image_pixels_to_features(
669 + pixel_values,
670 + spatial_shapes=spatial_shapes,
671 + )
672 +
673 + # Group patches by image - num_patches is on CPU (keep_on_cpu=True)
674 + # so .tolist() is instant with no DtoH sync
675 + num_patches_list = num_patches.tolist()
676 + batched_features: list[torch.Tensor] = []
677 + patch_idx = 0
678 + for count in num_patches_list:
679 + # Slice the list of patch tensors for this image
680 + image_patches = image_features[patch_idx : patch_idx + count]
681 + # Concatenate patches for this image
682 + batched_features.append(torch.cat(image_patches, dim=0))
683 + patch_idx += count
684 +
685 + return batched_features
686 +
687 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
688 + image_input = self._parse_and_validate_image_input(**kwargs)
689 + if image_input is None:
690 + return []
691 +
692 + return self._process_image_input(image_input)
693 +
694 + def forward(
695 + self,
696 + input_ids: torch.Tensor,
697 + positions: torch.Tensor,
698 + intermediate_tensors: IntermediateTensors | None = None,
699 + inputs_embeds: torch.Tensor | None = None,
700 + **kwargs: object,
701 + ) -> torch.Tensor | IntermediateTensors:
702 + if intermediate_tensors is not None:
703 + inputs_embeds = None
704 +
705 + hidden_states = self.language_model(
706 + input_ids=input_ids,
707 + positions=positions,
708 + intermediate_tensors=intermediate_tensors,
709 + inputs_embeds=inputs_embeds,
710 + )
711 + return hidden_states
712 +
713 + def compute_logits(
714 + self,
715 + hidden_states: torch.Tensor,
716 + ) -> torch.Tensor | None:
717 + logits = self.language_model.compute_logits(hidden_states)
718 + return logits
719 +
720 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
721 + loader = AutoWeightsLoader(self)
722 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
723 +
724 + def get_mm_mapping(self) -> MultiModelKeys:
725 + """
726 + Get the module prefix in multimodal models
727 + """
728 + return MultiModelKeys.from_string_field(
729 + language_model="language_model",
730 + connector="multi_modal_projector",
731 + tower_model="vision_tower",
732 + )
vllm/model_executor/models/registry.py
+1
@@ -349,6 +349,7 @@ _MULTIMODAL_MODELS = {
349 "lightonocr",
350 "LightOnOCRForConditionalGeneration",
351 ),
352 + "Lfm2VlForConditionalGeneration": ("lfm2_vl", "Lfm2VLForConditionalGeneration"),
353 "Llama_Nemotron_Nano_VL": ("nemotron_vl", "LlamaNemotronVLChatModel"),
354 "Llama4ForConditionalGeneration": ("mllama4", "Llama4ForConditionalGeneration"), # noqa: E501
355 "LlavaForConditionalGeneration": ("llava", "LlavaForConditionalGeneration"),
vllm/model_executor/models/siglip2.py new
+495
@@ -0,0 +1,495 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Implementation of Siglip2VisionModel intended to be only used
4 +within a vision language model."""
5 +
6 +from collections.abc import Iterable
7 +
8 +import torch
9 +from torch import nn
10 +from torch.nn import functional as F
11 +from transformers import Siglip2VisionConfig
12 +
13 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
14 +from vllm.compilation.decorators import support_torch_compile
15 +from vllm.config import MultiModalConfig
16 +from vllm.distributed import get_tensor_model_parallel_world_size
17 +from vllm.model_executor.layers.activation import get_act_fn
18 +from vllm.model_executor.layers.linear import (
19 + ColumnParallelLinear,
20 + QKVParallelLinear,
21 + RowParallelLinear,
22 +)
23 +from vllm.model_executor.layers.quantization import QuantizationConfig
24 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader
25 +
26 +from .vision import should_torch_compile_mm_vit
27 +
28 +
29 +class Siglip2VisionEmbeddings(nn.Module):
30 + def __init__(self, config: Siglip2VisionConfig):
31 + super().__init__()
32 + self.config = config
33 + self.embed_dim = config.hidden_size
34 + self.patch_size = config.patch_size
35 + self.patch_embedding = nn.Linear(
36 + in_features=config.num_channels * self.patch_size * self.patch_size,
37 + out_features=self.embed_dim,
38 + )
39 + self.num_patches = config.num_patches
40 + self.position_embedding_size = int(self.num_patches**0.5)
41 + self.position_embedding = nn.Embedding(self.num_patches, self.embed_dim)
42 +
43 + @staticmethod
44 + def resize_positional_embeddings(
45 + positional_embeddings: torch.Tensor,
46 + spatial_shapes: torch.LongTensor,
47 + max_length: int,
48 + ) -> torch.Tensor:
49 + """
50 + Resize positional embeddings to image-specific size and pad to a fixed size.
51 +
52 + Args:
53 + positional_embeddings (`torch.Tensor`):
54 + Position embeddings of shape (height, width, embed_dim)
55 + spatial_shapes (`torch.LongTensor`):
56 + Spatial shapes of shape (batch_size, 2) to resize the positional
57 + embeddings to
58 + max_length (`int`):
59 + Maximum length of the positional embeddings to pad resized
60 + positional embeddings to
61 +
62 + Returns:
63 + `torch.Tensor`: Embeddings of shape (batch_size, max_length, embed_dim)
64 + """
65 + batch_size = spatial_shapes.shape[0]
66 + embed_dim = positional_embeddings.shape[-1]
67 + source_dtype = positional_embeddings.dtype
68 +
69 + resulted_positional_embeddings = torch.empty(
70 + (batch_size, max_length, embed_dim),
71 + device=positional_embeddings.device,
72 + dtype=source_dtype,
73 + )
74 +
75 + # (height, width, embed_dim) -> (1, embed_dim, height, width) for interpolation
76 + positional_embeddings = positional_embeddings.permute(2, 0, 1).unsqueeze(0)
77 +
78 + # Upcast to float32 on CPU because antialias is not supported for
79 + # bfloat16/float16 on CPU
80 + if positional_embeddings.device.type == "cpu":
81 + positional_embeddings = positional_embeddings.to(torch.float32)
82 +
83 + for i in range(batch_size):
84 + # (1, dim, height, width) -> (1, dim, target_height, target_width)
85 + height, width = spatial_shapes[i]
86 + resized_embeddings = F.interpolate(
87 + positional_embeddings,
88 + size=(height, width),
89 + mode="bilinear",
90 + align_corners=False,
91 + antialias=True,
92 + )
93 +
94 + # (1, dim, target_height, target_width) ->
95 + # (target_height * target_width, dim)
96 + resized_embeddings = resized_embeddings.reshape(
97 + embed_dim, height * width
98 + ).transpose(0, 1)
99 +
100 + # Cast to original dtype
101 + resized_embeddings = resized_embeddings.to(source_dtype)
102 +
103 + resulted_positional_embeddings[i, : height * width] = resized_embeddings
104 + resulted_positional_embeddings[i, height * width :] = resized_embeddings[0]
105 +
106 + return resulted_positional_embeddings
107 +
108 + def forward(
109 + self, pixel_values: torch.FloatTensor, spatial_shapes: torch.LongTensor
110 + ) -> torch.Tensor:
111 + """
112 + Args:
113 + pixel_values (`torch.FloatTensor`):
114 + Pixel values of shape (batch_size, max_num_patches,
115 + num_channels * patch_size * patch_size)
116 + spatial_shapes (`list[tuple[int, int]]`):
117 + Spatial shapes of shape (batch_size, 2) to resize the positional
118 + embeddings to
119 + """
120 +
121 + # Apply patch embeddings to already patchified pixel values
122 + target_dtype = self.patch_embedding.weight.dtype
123 + patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))
124 +
125 + # Get positional resized and padded positional embeddings
126 + positional_embeddings = self.position_embedding.weight.reshape(
127 + self.position_embedding_size, self.position_embedding_size, -1
128 + )
129 + resized_positional_embeddings = self.resize_positional_embeddings(
130 + positional_embeddings, spatial_shapes, max_length=pixel_values.shape[1]
131 + )
132 +
133 + # Add positional embeddings to patch embeddings
134 + embeddings = patch_embeds + resized_positional_embeddings
135 + return embeddings
136 +
137 +
138 +class Siglip2Attention(nn.Module):
139 + """Multi-headed attention from 'Attention Is All You Need' paper"""
140 +
141 + def __init__(
142 + self,
143 + config: Siglip2VisionConfig,
144 + quant_config: QuantizationConfig | None = None,
145 + multimodal_config: MultiModalConfig | None = None,
146 + prefix: str = "",
147 + ):
148 + super().__init__()
149 + self.config = config
150 + self.embed_dim = config.hidden_size
151 + self.num_heads = config.num_attention_heads
152 + self.head_dim = self.embed_dim // self.num_heads
153 + if self.head_dim * self.num_heads != self.embed_dim:
154 + raise ValueError(
155 + f"embed_dim must be divisible by num_heads "
156 + f"(got `embed_dim`: {self.embed_dim} and `num_heads`:"
157 + f" {self.num_heads})."
158 + )
159 + self.scale = self.head_dim**-0.5
160 + self.dropout = config.attention_dropout
161 +
162 + use_data_parallel = (
163 + multimodal_config is not None
164 + and multimodal_config.mm_encoder_tp_mode == "data"
165 + )
166 + tp_size = 1 if use_data_parallel else get_tensor_model_parallel_world_size()
167 + assert self.num_heads % tp_size == 0
168 + self.num_heads_per_partition = self.num_heads // tp_size
169 +
170 + self.qkv_proj = QKVParallelLinear(
171 + hidden_size=self.embed_dim,
172 + head_size=self.head_dim,
173 + total_num_heads=self.num_heads,
174 + quant_config=quant_config,
175 + prefix=f"{prefix}.qkv_proj",
176 + disable_tp=use_data_parallel,
177 + )
178 + self.out_proj = RowParallelLinear(
179 + input_size=self.embed_dim,
180 + output_size=self.embed_dim,
181 + quant_config=quant_config,
182 + prefix=f"{prefix}.out_proj",
183 + disable_tp=use_data_parallel,
184 + )
185 + self.attn = MMEncoderAttention(
186 + num_heads=self.num_heads_per_partition,
187 + head_size=self.head_dim,
188 + scale=self.scale,
189 + prefix=f"{prefix}.attn",
190 + multimodal_config=multimodal_config,
191 + )
192 +
193 + def forward(
194 + self,
195 + hidden_states: torch.Tensor,
196 + cu_seqlens: torch.Tensor,
197 + max_seqlen: int | torch.Tensor,
198 + ) -> torch.Tensor:
199 + qkv, _ = self.qkv_proj(
200 + hidden_states
201 + ) # batch_size, q_len, 3 * num_heads_per_partition * head_dim
202 + bsz, q_len, _ = qkv.shape
203 + query_states, key_states, value_states = qkv.chunk(3, dim=-1)
204 + query_states = query_states.view(
205 + bsz, q_len, self.num_heads_per_partition, self.head_dim
206 + )
207 + key_states = key_states.view(
208 + bsz, q_len, self.num_heads_per_partition, self.head_dim
209 + )
210 + value_states = value_states.view(
211 + bsz, q_len, self.num_heads_per_partition, self.head_dim
212 + )
213 +
214 + # Use unified MultiHeadAttention implementation
215 + out = self.attn(
216 + query=query_states,
217 + key=key_states,
218 + value=value_states,
219 + cu_seqlens=cu_seqlens,
220 + max_seqlen=max_seqlen,
221 + )
222 + out = out.reshape(bsz, q_len, -1)
223 + attn_output, _ = self.out_proj(out)
224 + return attn_output
225 +
226 +
227 +class Siglip2MLP(nn.Module):
228 + def __init__(
229 + self,
230 + config: Siglip2VisionConfig,
231 + quant_config: QuantizationConfig | None = None,
232 + multimodal_config: MultiModalConfig | None = None,
233 + prefix: str = "",
234 + ):
235 + super().__init__()
236 + self.config = config
237 + self.activation_fn = get_act_fn(config.hidden_act)
238 + use_data_parallel = (
239 + multimodal_config is not None
240 + and multimodal_config.mm_encoder_tp_mode == "data"
241 + )
242 + self.fc1 = ColumnParallelLinear(
243 + config.hidden_size,
244 + config.intermediate_size,
245 + quant_config=quant_config,
246 + prefix=f"{prefix}.fc1",
247 + disable_tp=use_data_parallel,
248 + )
249 + self.fc2 = RowParallelLinear(
250 + config.intermediate_size,
251 + config.hidden_size,
252 + quant_config=quant_config,
253 + prefix=f"{prefix}.fc2",
254 + disable_tp=use_data_parallel,
255 + )
256 +
257 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
258 + hidden_states, _ = self.fc1(hidden_states)
259 + hidden_states = self.activation_fn(hidden_states)
260 + hidden_states, _ = self.fc2(hidden_states)
261 + return hidden_states
262 +
263 +
264 +@support_torch_compile(
265 + dynamic_arg_dims={"hidden_states": [0, 1], "cu_seqlens": 0},
266 + enable_if=should_torch_compile_mm_vit,
267 +)
268 +class Siglip2EncoderLayer(nn.Module):
269 + def __init__(
270 + self,
271 + config: Siglip2VisionConfig,
272 + quant_config: QuantizationConfig | None = None,
273 + multimodal_config: MultiModalConfig | None = None,
274 + prefix: str = "",
275 + ):
276 + super().__init__()
277 + self.embed_dim = config.hidden_size
278 + self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
279 + self.self_attn = Siglip2Attention(
280 + config,
281 + quant_config=quant_config,
282 + multimodal_config=multimodal_config,
283 + prefix=f"{prefix}.self_attn",
284 + )
285 + self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
286 + self.mlp = Siglip2MLP(
287 + config,
288 + quant_config=quant_config,
289 + multimodal_config=multimodal_config,
290 + prefix=f"{prefix}.mlp",
291 + )
292 +
293 + def forward(
294 + self,
295 + hidden_states: torch.Tensor,
296 + cu_seqlens: torch.Tensor,
297 + max_seqlen: int | torch.Tensor,
298 + ) -> torch.Tensor:
299 + """
300 + Args:
301 + hidden_states: Input tensor of shape (batch, seq_len, embed_dim).
302 + cu_seqlens: Cumulative sequence lengths tensor.
303 + max_seqlen: Maximum sequence length.
304 + """
305 + residual = hidden_states
306 +
307 + hidden_states = self.layer_norm1(hidden_states)
308 + hidden_states = self.self_attn(
309 + hidden_states=hidden_states,
310 + cu_seqlens=cu_seqlens,
311 + max_seqlen=max_seqlen,
312 + )
313 + hidden_states = residual + hidden_states
314 +
315 + residual = hidden_states
316 + hidden_states = self.layer_norm2(hidden_states)
317 + hidden_states = self.mlp(hidden_states)
318 + hidden_states = residual + hidden_states
319 + return hidden_states
320 +
321 +
322 +class Siglip2Encoder(nn.Module):
323 + """
324 + Transformer encoder consisting of `config.num_hidden_layers`
325 + self attention layers. Each layer is a [`Siglip2EncoderLayer`].
326 +
327 + Args:
328 + config: PretrainedConfig
329 + """
330 +
331 + def __init__(
332 + self,
333 + config: Siglip2VisionConfig,
334 + quant_config: QuantizationConfig | None = None,
335 + multimodal_config: MultiModalConfig | None = None,
336 + prefix: str = "",
337 + ):
338 + super().__init__()
339 + self.config = config
340 + self.layers = nn.ModuleList(
341 + [
342 + Siglip2EncoderLayer(
343 + config=config,
344 + quant_config=quant_config,
345 + multimodal_config=multimodal_config,
346 + prefix=f"{prefix}.layers.{idx}",
347 + )
348 + for idx in range(config.num_hidden_layers)
349 + ]
350 + )
351 +
352 + def forward(
353 + self,
354 + inputs_embeds: torch.Tensor,
355 + cu_seqlens: torch.Tensor,
356 + max_seqlen: int | torch.Tensor,
357 + ) -> torch.Tensor:
358 + hidden_states = inputs_embeds
359 + for encoder_layer in self.layers:
360 + layer_outputs = encoder_layer(
361 + hidden_states,
362 + cu_seqlens=cu_seqlens,
363 + max_seqlen=max_seqlen,
364 + )
365 + hidden_states = layer_outputs
366 + return hidden_states
367 +
368 +
369 +class Siglip2VisionTransformer(nn.Module):
370 + def __init__(
371 + self,
372 + config: Siglip2VisionConfig,
373 + quant_config: QuantizationConfig | None = None,
374 + multimodal_config: MultiModalConfig | None = None,
375 + prefix: str = "",
376 + ):
377 + super().__init__()
378 + embed_dim = config.hidden_size
379 + self.config = config
380 + self.embeddings = Siglip2VisionEmbeddings(config)
381 + # Keep the import local to avoid circular dependencies during model init.
382 + from vllm.compilation.backends import set_model_tag
383 +
384 + with set_model_tag("Siglip2Encoder", is_encoder=True):
385 + self.encoder = Siglip2Encoder(
386 + config,
387 + quant_config=quant_config,
388 + multimodal_config=multimodal_config,
389 + prefix=f"{prefix}.encoder",
390 + )
391 + num_hidden_layers = config.num_hidden_layers
392 + if len(self.encoder.layers) > config.num_hidden_layers:
393 + raise ValueError(
394 + f"The original encoder only has {num_hidden_layers} "
395 + f"layers, but you requested {len(self.encoder.layers)} layers."
396 + )
397 +
398 + self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
399 +
400 + def get_input_embeddings(self):
401 + return self.embeddings
402 +
403 + def forward(
404 + self,
405 + pixel_values: torch.FloatTensor,
406 + spatial_shapes: torch.LongTensor,
407 + packed_mask: torch.Tensor,
408 + cu_seqlens: torch.Tensor,
409 + max_seqlen: int | torch.Tensor,
410 + ) -> torch.Tensor:
411 + r"""
412 + spatial_shapes (`torch.LongTensor` of shape `(batch_size, 2)`):
413 + Tensor containing the spatial dimensions (height, width)
414 + of the input images.
415 + """
416 + hidden_states = self.embeddings(pixel_values, spatial_shapes)
417 + flat_mask = packed_mask.view(-1)
418 + packed_indices = flat_mask.nonzero(as_tuple=True)[0]
419 + flat_hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
420 + hidden_states = flat_hidden_states.index_select(0, packed_indices).unsqueeze(0)
421 + encoder_outputs = self.encoder(
422 + inputs_embeds=hidden_states,
423 + cu_seqlens=cu_seqlens,
424 + max_seqlen=max_seqlen,
425 + )
426 + unpacked = encoder_outputs.new_zeros(
427 + packed_mask.numel(), encoder_outputs.shape[-1]
428 + )
429 + unpacked.index_copy_(0, packed_indices, encoder_outputs.squeeze(0))
430 + encoder_outputs = unpacked.view(
431 + packed_mask.shape + (encoder_outputs.shape[-1],)
432 + )
433 + last_hidden_state = self.post_layernorm(encoder_outputs)
434 + return last_hidden_state
435 +
436 +
437 +class Siglip2Model(torch.nn.Module):
438 + def __init__(
439 + self,
440 + config: Siglip2VisionConfig,
441 + quant_config: QuantizationConfig | None = None,
442 + multimodal_config: MultiModalConfig | None = None,
443 + prefix: str = "",
444 + ):
445 + super().__init__()
446 +
447 + self.vision_model = Siglip2VisionTransformer(
448 + config,
449 + quant_config=quant_config,
450 + multimodal_config=multimodal_config,
451 + prefix=f"{prefix}.vision_model",
452 + )
453 +
454 + def forward(
455 + self,
456 + pixel_values: torch.FloatTensor,
457 + spatial_shapes: torch.LongTensor,
458 + packed_mask: torch.Tensor,
459 + cu_seqlens: torch.Tensor,
460 + max_seqlen: int | torch.Tensor,
461 + ) -> torch.Tensor:
462 + return self.vision_model(
463 + pixel_values=pixel_values,
464 + spatial_shapes=spatial_shapes,
465 + packed_mask=packed_mask,
466 + cu_seqlens=cu_seqlens,
467 + max_seqlen=max_seqlen,
468 + )
469 +
470 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
471 + stacked_params_mapping = [
472 + # (param_name, shard_name, shard_id)
473 + ("qkv_proj", "q_proj", "q"),
474 + ("qkv_proj", "k_proj", "k"),
475 + ("qkv_proj", "v_proj", "v"),
476 + ]
477 + params_dict = dict(self.named_parameters())
478 + loaded_params: set[str] = set()
479 +
480 + for name, loaded_weight in weights:
481 + for param_name, weight_name, shard_id in stacked_params_mapping:
482 + if weight_name not in name:
483 + continue
484 + name = name.replace(weight_name, param_name)
485 +
486 + param = params_dict[name]
487 + weight_loader = param.weight_loader
488 + weight_loader(param, loaded_weight, shard_id)
489 + break
490 + else:
491 + param = params_dict[name]
492 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
493 + weight_loader(param, loaded_weight)
494 + loaded_params.add(name)
495 + return loaded_params