@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8452946c06
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 15, 03:30 UTC

[Model][VLM] Support Keye-VL-8B-Preview (#20126)

Signed-off-by: Kwai-Keye <Keye@kuaishou.com>

Kwai-Keye committed Jul 2, 2025 at 14:35 UTC 8452946c06a3b8a76233d2b390d886a5a8c78182
7 files changed +1801 -2
docs/models/supported_models.md
+1
@@ -559,6 +559,7 @@ Specified using `--task generate`.
559 | `H2OVLChatModel` | H2OVL | T + I<sup>E+</sup> | `h2oai/h2ovl-mississippi-800m`, `h2oai/h2ovl-mississippi-2b`, etc. | | ✅︎ | ✅︎\* |
560 | `Idefics3ForConditionalGeneration` | Idefics3 | T + I | `HuggingFaceM4/Idefics3-8B-Llama3` etc. | ✅︎ | | ✅︎ |
561 | `InternVLChatModel` | InternVL 3.0, InternVideo 2.5, InternVL 2.5, Mono-InternVL, InternVL 2.0 | T + I<sup>E+</sup> + (V<sup>E+</sup>) | `OpenGVLab/InternVL3-9B`, `OpenGVLab/InternVideo2_5_Chat_8B`, `OpenGVLab/InternVL2_5-4B`, `OpenGVLab/Mono-InternVL-2B`, `OpenGVLab/InternVL2-4B`, etc. | ✅︎ | ✅︎ | ✅︎ |
562 +| `KeyeForConditionalGeneration` | Keye-VL-8B-Preview | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-8B-Preview` | | | ✅︎ |
563 | `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I<sup>+</sup> | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | | ✅︎ |
564 | `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. | | ✅︎ | ✅︎ |
565 | `LlavaForConditionalGeneration` | LLaVA-1.5 | T + I<sup>E+</sup> | `llava-hf/llava-1.5-7b-hf`, `TIGER-Lab/Mantis-8B-siglip-llama3` (see note), etc. | | ✅︎ | ✅︎ |
examples/offline_inference/vision_language.py
+32
@@ -429,6 +429,37 @@ def run_internvl(questions: list[str], modality: str) -> ModelRequestData:
429 )
430
431
432 +# Keye-VL
433 +def run_keye_vl(questions: list[str], modality: str) -> ModelRequestData:
434 + model_name = "Kwai-Keye/Keye-VL-8B-Preview"
435 +
436 + engine_args = EngineArgs(
437 + model=model_name,
438 + max_model_len=8192,
439 + trust_remote_code=True,
440 + limit_mm_per_prompt={modality: 1},
441 + )
442 +
443 + if modality == "image":
444 + placeholder = "<|image_pad|>"
445 + elif modality == "video":
446 + placeholder = "<|video_pad|>"
447 +
448 + prompts = [
449 + (
450 + f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"
451 + f"{question}<|im_end|>\n"
452 + "<|im_start|>assistant\n"
453 + )
454 + for question in questions
455 + ]
456 +
457 + return ModelRequestData(
458 + engine_args=engine_args,
459 + prompts=prompts,
460 + )
461 +
462 +
463 # Kimi-VL
464 def run_kimi_vl(questions: list[str], modality: str) -> ModelRequestData:
465 assert modality == "image"
@@ -1154,6 +1185,7 @@ model_example_map = {
1185 "h2ovl_chat": run_h2ovl,
1186 "idefics3": run_idefics3,
1187 "internvl_chat": run_internvl,
1188 + "keye_vl": run_keye_vl,
1189 "kimi_vl": run_kimi_vl,
1190 "llava": run_llava,
1191 "llava-next": run_llava_next,
examples/offline_inference/vision_language_multi_image.py
+38
@@ -423,6 +423,43 @@ def load_llama4(question: str, image_urls: list[str]) -> ModelRequestData:
423 )
424
425
426 +def load_keye_vl(question: str, image_urls: list[str]) -> ModelRequestData:
427 + model_name = "Kwai-Keye/Keye-VL-8B-Preview"
428 +
429 + engine_args = EngineArgs(
430 + model=model_name,
431 + trust_remote_code=True,
432 + max_model_len=8192,
433 + max_num_seqs=5,
434 + limit_mm_per_prompt={"image": len(image_urls)},
435 + )
436 +
437 + placeholders = [{"type": "image", "image": url} for url in image_urls]
438 + messages = [
439 + {
440 + "role": "user",
441 + "content": [
442 + *placeholders,
443 + {"type": "text", "text": question},
444 + ],
445 + },
446 + ]
447 +
448 + processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
449 +
450 + prompt = processor.apply_chat_template(
451 + messages, tokenize=False, add_generation_prompt=True
452 + )
453 +
454 + image_data = [fetch_image(url) for url in image_urls]
455 +
456 + return ModelRequestData(
457 + engine_args=engine_args,
458 + prompt=prompt,
459 + image_data=image_data,
460 + )
461 +
462 +
463 def load_kimi_vl(question: str, image_urls: list[str]) -> ModelRequestData:
464 model_name = "moonshotai/Kimi-VL-A3B-Instruct"
465
@@ -862,6 +899,7 @@ model_example_map = {
899 "h2ovl_chat": load_h2ovl,
900 "idefics3": load_idefics3,
901 "internvl_chat": load_internvl,
902 + "keye_vl": load_keye_vl,
903 "kimi_vl": load_kimi_vl,
904 "llava": load_llava,
905 "llava-next": load_llava_next,
tests/models/registry.py
+2
@@ -351,6 +351,8 @@ _MULTIMODAL_EXAMPLE_MODELS = {
351 trust_remote_code=True),
352 "Idefics3ForConditionalGeneration": _HfExamplesInfo("HuggingFaceM4/Idefics3-8B-Llama3", # noqa: E501
353 {"tiny": "HuggingFaceTB/SmolVLM-256M-Instruct"}), # noqa: E501
354 + "KeyeForConditionalGeneration": _HfExamplesInfo("Kwai-Keye/Keye-VL-8B-Preview", # noqa: E501
355 + trust_remote_code=True),
356 "KimiVLForConditionalGeneration": _HfExamplesInfo("moonshotai/Kimi-VL-A3B-Instruct", # noqa: E501
357 extras={"thinking": "moonshotai/Kimi-VL-A3B-Thinking"}, # noqa: E501
358 trust_remote_code=True,
vllm/entrypoints/chat_utils.py
+2 -2
@@ -540,7 +540,7 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):
540 return "<image>"
541 if model_type in ("mllama", "llama4"):
542 return "<|image|>"
543 - if model_type in ("qwen2_vl", "qwen2_5_vl"):
543 + if model_type in ("qwen2_vl", "qwen2_5_vl", "keye", "Keye"):
544 return "<|vision_start|><|image_pad|><|vision_end|>"
545 if model_type == "qwen2_5_omni":
546 return "<|vision_start|><|IMAGE|><|vision_end|>"
@@ -570,7 +570,7 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):
570 return "<video>"
571 if model_type == "glm4v":
572 return "<|begin_of_video|><|video|><|end_of_video|>"
573 - if model_type in ("qwen2_vl", "qwen2_5_vl"):
573 + if model_type in ("qwen2_vl", "qwen2_5_vl", "keye", "Keye"):
574 return "<|vision_start|><|video_pad|><|vision_end|>"
575 if model_type == "qwen2_5_omni":
576 return "<|vision_start|><|VIDEO|><|vision_end|>"
vllm/model_executor/models/keye.py new
+1725
@@ -0,0 +1,1725 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import math
4 +from collections.abc import Iterable, Mapping, Sequence
5 +from functools import partial
6 +from typing import Any, Literal, Optional, TypedDict, Union
7 +
8 +import numpy as np
9 +import torch
10 +import torch.nn as nn
11 +from einops import rearrange
12 +from transformers import PretrainedConfig
13 +from transformers.activations import GELUActivation
14 +from transformers.feature_extraction_utils import BatchFeature
15 +from transformers.modeling_outputs import (BaseModelOutput,
16 + BaseModelOutputWithPooling)
17 +from transformers.utils import torch_int
18 +
19 +from vllm.config import VllmConfig
20 +from vllm.distributed import get_tensor_model_parallel_world_size
21 +from vllm.logger import init_logger
22 +from vllm.model_executor import SamplingMetadata
23 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,
24 + QKVParallelLinear,
25 + RowParallelLinear)
26 +from vllm.model_executor.layers.quantization import QuantizationConfig
27 +from vllm.model_executor.layers.quantization.gptq import GPTQConfig
28 +from vllm.model_executor.layers.quantization.gptq_marlin import (
29 + GPTQMarlinConfig)
30 +from vllm.model_executor.model_loader.weight_utils import (
31 + default_weight_loader, maybe_remap_kv_scale_name)
32 +from vllm.model_executor.models.module_mapping import MultiModelKeys
33 +from vllm.multimodal import MULTIMODAL_REGISTRY
34 +from vllm.multimodal.inputs import (ImageItem, ModalityData,
35 + MultiModalDataDict, MultiModalFieldConfig,
36 + MultiModalKwargs, VideoItem)
37 +from vllm.multimodal.parse import (DictEmbeddingItems, ImageSize,
38 + ModalityDataItems, MultiModalDataItems,
39 + MultiModalDataParser)
40 +from vllm.multimodal.processing import (BaseMultiModalProcessor,
41 + BaseProcessingInfo, PromptReplacement,
42 + PromptUpdate)
43 +from vllm.multimodal.profiling import BaseDummyInputsBuilder
44 +from vllm.platforms import _Backend
45 +from vllm.sequence import IntermediateTensors
46 +from vllm.transformers_utils.config import uses_mrope
47 +from vllm.transformers_utils.processor import (
48 + cached_image_processor_from_config)
49 +
50 +from .interfaces import (MultiModalEmbeddings, SupportsLoRA,
51 + SupportsMultiModal, SupportsPP)
52 +from .siglip import SiglipMLP
53 +from .utils import (AutoWeightsLoader, WeightsMapper,
54 + init_vllm_registered_model, is_pp_missing_parameter,
55 + maybe_prefix, merge_multimodal_embeddings)
56 +from .vision import get_vit_attn_backend
57 +
58 +logger = init_logger(__name__)
59 +
60 +_MAX_FRAMES_PER_VIDEO = 16
61 +_MAX_IMAGE_SIZE = 9999999
62 +
63 +
64 +def smart_resize(
65 + height: int,
66 + width: int,
67 + factor: int = 28,
68 + min_pixels: int = 28 * 28 * 130,
69 + max_pixels: int = 28 * 28 * 1280,
70 +):
71 + if height < factor:
72 + logger.warning(
73 + "smart_resize: height=%s < factor=%s, reset height=factor",
74 + height,
75 + factor,
76 + )
77 + width = round((width * factor) / height)
78 + height = factor
79 +
80 + if width < factor:
81 + logger.warning(
82 + "smart_resize: width=%s < factor=%s, reset width=factor",
83 + width,
84 + factor,
85 + )
86 + height = round((height * factor) / width)
87 + width = factor
88 +
89 + if max(height, width) / min(height, width) > 200:
90 + raise ValueError("absolute aspect ratio must be smaller than 200, got "
91 + "{max(height, width) / min(height, width)}")
92 + h_bar = round(height / factor) * factor
93 + w_bar = round(width / factor) * factor
94 + if h_bar * w_bar > max_pixels:
95 + beta = math.sqrt((height * width) / max_pixels)
96 + h_bar = math.floor(height / beta / factor) * factor
97 + w_bar = math.floor(width / beta / factor) * factor
98 + elif h_bar * w_bar < min_pixels:
99 + beta = math.sqrt(min_pixels / (height * width))
100 + h_bar = math.ceil(height * beta / factor) * factor
101 + w_bar = math.ceil(width * beta / factor) * factor
102 + return h_bar, w_bar
103 +
104 +
105 +class KeyeImagePixelInputs(TypedDict):
106 + type: Literal["pixel_values"]
107 + pixel_values: torch.Tensor
108 + """Shape:
109 + `(num_patches, num_channels * patch_size * patch_size)`
110 + """
111 +
112 + image_grid_thw: torch.Tensor
113 + """Shape: `(num_images, 3)`
114 + This should be in `(grid_t, grid_h, grid_w)` format.
115 + """
116 +
117 +
118 +class KeyeImageEmbeddingInputs(TypedDict):
119 + type: Literal["image_embeds"]
120 + image_embeds: torch.Tensor
121 + """Supported types:
122 + - list[`torch.Tensor`]: A list of tensors holding all images' features.
123 + Each tensor holds an image's features.
124 + - `torch.Tensor`: A tensor holding all images' features
125 + (concatenation of all images' feature tensors).
126 +
127 + Tensor shape: `(num_image_features, hidden_size)`
128 + - `num_image_features` varies based on
129 + the number and resolution of the images.
130 + - `hidden_size` must match the hidden size of language model backbone.
131 + """
132 +
133 + image_grid_thw: torch.Tensor
134 + """Shape: `(num_images, 3)`
135 + This should be in `(grid_t, grid_h, grid_w)` format.
136 + """
137 +
138 +
139 +KeyeImageInputs = Union[KeyeImagePixelInputs, KeyeImageEmbeddingInputs]
140 +
141 +
142 +class KeyeVideoPixelInputs(TypedDict):
143 + type: Literal["pixel_values_videos"]
144 + pixel_values_videos: torch.Tensor
145 + """Shape:
146 + `(num_patches,
147 + num_channels * temporal_patch_size * patch_size * patch_size)`
148 + """
149 +
150 + video_grid_thw: torch.Tensor
151 + """Shape: `(num_videos, 3)`
152 +
153 + This should be in `(grid_t, grid_h, grid_w)` format.
154 + """
155 +
156 +
157 +class KeyeVideoEmbeddingInputs(TypedDict):
158 + type: Literal["video_embeds"]
159 + video_embeds: torch.Tensor
160 + """Supported types:
161 + - list[`torch.Tensor`]: A list of tensors holding all videos' features.
162 + Each tensor holds an video's features.
163 + - `torch.Tensor`: A tensor holding all videos' features
164 + (concatenation of all videos' feature tensors).
165 +
166 + Tensor shape: `(num_image_features, hidden_size)`
167 + - `num_image_features` varies based on
168 + the number and resolution of the videos.
169 + - `hidden_size` must match the hidden size of language model backbone.
170 + """
171 +
172 + video_grid_thw: torch.Tensor
173 + """Shape: `(num_videos, 3)`
174 + This should be in `(grid_t, grid_h, grid_w)` format.
175 + """
176 +
177 +
178 +KeyeVideoInputs = Union[KeyeVideoPixelInputs, KeyeVideoEmbeddingInputs]
179 +
180 +
181 +class KeyeVisionEmbeddings(nn.Module):
182 +
183 + def __init__(self, config: PretrainedConfig):
184 + super().__init__()
185 + self.config = config
186 + self.embed_dim = config.hidden_size
187 + self.image_size = config.image_size
188 + self.patch_size = config.patch_size
189 +
190 + self.patch_embedding = nn.Conv2d(
191 + in_channels=config.num_channels,
192 + out_channels=self.embed_dim,
193 + kernel_size=self.patch_size,
194 + stride=self.patch_size,
195 + padding="valid",
196 + )
197 +
198 + self.num_patches = (self.image_size // self.patch_size)**2
199 + self.num_positions = self.num_patches
200 + self.cache_position_embedding = dict()
201 + self.cache_position_count = dict()
202 + self.position_embedding = nn.Embedding(self.num_positions,
203 + self.embed_dim)
204 + self.packing_position_embedding = nn.Embedding(32768, self.embed_dim)
205 +
206 + self.register_buffer(
207 + "position_ids",
208 + torch.arange(self.num_positions).expand((1, -1)),
209 + persistent=False,
210 + )
211 +
212 + def interpolate_pos_encoding(
213 + self,
214 + embeddings: torch.Tensor,
215 + height: int,
216 + width: int,
217 + is_after_patchify: bool = False,
218 + ) -> torch.Tensor:
219 +
220 + num_positions = self.position_embedding.weight.shape[0]
221 +
222 + patch_pos_embed = self.position_embedding.weight.unsqueeze(0)
223 +
224 + dim = embeddings.shape[-1]
225 +
226 + if is_after_patchify:
227 + new_height = height
228 + new_width = width
229 + else:
230 + new_height = height // self.patch_size
231 + new_width = width // self.patch_size
232 +
233 + sqrt_num_positions = torch_int(num_positions**0.5)
234 + patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions,
235 + sqrt_num_positions, dim)
236 + patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
237 +
238 + patch_pos_embed = nn.functional.interpolate(
239 + patch_pos_embed,
240 + size=(new_height, new_width),
241 + mode="bilinear",
242 + align_corners=False,
243 + )
244 +
245 + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
246 + return patch_pos_embed
247 +
248 + def fetch_position_embedding_lfu_cache(self,
249 + embeddings,
250 + h,
251 + w,
252 + max_cache: int = 20):
253 + grid = (h, w)
254 + if grid in self.cache_position_embedding:
255 + self.cache_position_count[grid] += 1
256 + return self.cache_position_embedding[grid]
257 +
258 + if len(self.cache_position_embedding) >= max_cache:
259 + min_hit_grid = min(
260 + self.cache_position_count,
261 + key=self.cache_position_count.get,
262 + )
263 + self.cache_position_count.pop(min_hit_grid)
264 + self.cache_position_embedding.pop(min_hit_grid)
265 +
266 + position_embedding = self.interpolate_pos_encoding(
267 + embeddings, h, w, True)
268 + self.cache_position_count[grid] = 1
269 + self.cache_position_embedding[grid] = position_embedding
270 + return position_embedding
271 +
272 + def forward(
273 + self,
274 + pixel_values: torch.FloatTensor,
275 + position_ids: Optional[torch.Tensor] = None,
276 + image_grid_thw: Optional[list[Union[
277 + tuple[int, int, int],
278 + list[tuple[int, int, int]],
279 + ]]] = None,
280 + interpolate_pos_encoding=False,
281 + ) -> torch.Tensor:
282 + if pixel_values.dim() == 4:
283 + pixel_values = pixel_values.unsqueeze(0)
284 + if pixel_values.dim() == 5:
285 + if position_ids is None:
286 + raise ValueError(
287 + "position_ids cannot be None when pixel_values.dim() is 5."
288 + )
289 + (
290 + batch_size,
291 + squence_len,
292 + channel,
293 + height,
294 + width,
295 + ) = pixel_values.shape
296 + target_dtype = self.patch_embedding.weight.dtype
297 + pixel_values = rearrange(pixel_values, "b l c h w -> (b l) c h w")
298 + patch_embeds = self.patch_embedding(
299 + pixel_values.to(dtype=target_dtype))
300 + embeddings = patch_embeds.flatten(-2).squeeze(-1)
301 +
302 + if interpolate_pos_encoding and image_grid_thw is not None:
303 + start = 0
304 + tmp_embeddings = list()
305 + for image_grid in image_grid_thw:
306 + t, h, w = image_grid
307 + end = start + t * h * w
308 + image_embeddings = embeddings[start:end, :]
309 + position_embedding = (self.interpolate_pos_encoding(
310 + image_embeddings, h, w, True).squeeze(0).repeat(t, 1))
311 + image_embeddings = image_embeddings + position_embedding
312 + tmp_embeddings.append(image_embeddings)
313 + start = end
314 + embeddings = torch.concat(tmp_embeddings, dim=0).unsqueeze(0)
315 + else:
316 + embeddings = embeddings + self.packing_position_embedding(
317 + position_ids)
318 + return embeddings
319 + else:
320 + raise ValueError("Unsupported pixel_values dimension:"
321 + f" {pixel_values.dim()}. Expected 4 or 5.")
322 +
323 +
324 +def apply_rotary_pos_emb_flashatt(
325 + q: torch.Tensor,
326 + k: torch.Tensor,
327 + cos: torch.Tensor,
328 + sin: torch.Tensor,
329 +) -> tuple[torch.Tensor, torch.Tensor]:
330 + cos = cos.chunk(2, dim=-1)[0].contiguous()
331 + sin = sin.chunk(2, dim=-1)[0].contiguous()
332 +
333 + from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
334 +
335 + q_embed = apply_rotary_emb(q.float(), cos.float(), sin.float()).type_as(q)
336 + k_embed = apply_rotary_emb(k.float(), cos.float(), sin.float()).type_as(k)
337 + return q_embed, k_embed
338 +
339 +
340 +class KeyeSiglipAttention(nn.Module):
341 + """Multi-headed attention from 'Attention Is All You
342 + Need' paper."""
343 +
344 + def __init__(
345 + self,
346 + config: PretrainedConfig,
347 + quant_config: Optional[QuantizationConfig] = None,
348 + prefix: str = "",
349 + ):
350 + super().__init__()
351 + self.config = config
352 +
353 + hidden_size = config.hidden_size
354 + self.hidden_size = config.hidden_size
355 + tp_size = get_tensor_model_parallel_world_size()
356 + self.total_num_heads = config.num_attention_heads
357 + assert self.total_num_heads % tp_size == 0
358 + self.num_heads = self.total_num_heads // tp_size
359 + self.total_num_kv_heads = config.num_attention_heads
360 + if self.total_num_kv_heads >= tp_size:
361 + assert self.total_num_kv_heads % tp_size == 0
362 + else:
363 + assert tp_size % self.total_num_kv_heads == 0
364 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
365 + self.head_dim = config.hidden_size // self.total_num_heads
366 + self.q_size = self.num_heads * self.head_dim
367 + self.kv_size = self.num_kv_heads * self.head_dim
368 + self.scale = self.head_dim**-0.5
369 +
370 + self.qkv_proj = QKVParallelLinear(
371 + hidden_size,
372 + self.head_dim,
373 + self.total_num_heads,
374 + self.total_num_kv_heads,
375 + bias=True,
376 + quant_config=quant_config,
377 + prefix=f"{prefix}.qkv_proj",
378 + )
379 + self.out_proj = RowParallelLinear(
380 + input_size=hidden_size,
381 + output_size=hidden_size,
382 + quant_config=quant_config,
383 + prefix=f"{prefix}.out_proj",
384 + )
385 +
386 + # Detect attention implementation.
387 + self.attn_backend: _Backend = get_vit_attn_backend(support_fa=True)
388 + if self.attn_backend not in {_Backend.FLASH_ATTN, _Backend.XFORMERS}:
389 + raise RuntimeError(
390 + f"Keye-VL does not support {self.attn_backend} backend now.")
391 +
392 + def forward(
393 + self,
394 + hidden_states: torch.Tensor,
395 + attention_mask: Optional[torch.Tensor] = None,
396 + output_attentions: Optional[bool] = False,
397 + cu_seqlens: Optional[list[torch.Tensor]] = None,
398 + rope_emb: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
399 + ) -> torch.Tensor:
400 + qkv, _ = self.qkv_proj(hidden_states)
401 + q, k, v = qkv.split(
402 + [self.q_size, self.kv_size, self.kv_size],
403 + dim=-1,
404 + )
405 +
406 + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
407 + seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
408 + batch_size = q.shape[0]
409 +
410 + if rope_emb is None:
411 + q = q.view(*q.shape[:-1], self.num_heads, self.head_dim)
412 + k = k.view(
413 + *k.shape[:-1],
414 + self.num_kv_heads,
415 + self.head_dim,
416 + )
417 + v = v.view(
418 + *v.shape[:-1],
419 + self.num_kv_heads,
420 + self.head_dim,
421 + )
422 + else:
423 + if cu_seqlens is None:
424 + raise ValueError(
425 + "cu_seqlens cannot be None when rope_emb is not None.")
426 + cos, sin = rope_emb
427 + q = q.view(*q.shape[:-1], self.num_heads, self.head_dim)
428 + k = k.view(
429 + *k.shape[:-1],
430 + self.num_kv_heads,
431 + self.head_dim,
432 + )
433 + q, k = apply_rotary_pos_emb_flashatt(q, k, cos, sin)
434 + v = v.view(
435 + *v.shape[:-1],
436 + self.num_kv_heads,
437 + self.head_dim,
438 + )
439 +
440 + if self.attn_backend == _Backend.FLASH_ATTN:
441 + from flash_attn import flash_attn_varlen_func
442 +
443 + q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
444 +
445 + output = flash_attn_varlen_func(
446 + q,
447 + k,
448 + v,
449 + cu_seqlens_q=cu_seqlens,
450 + cu_seqlens_k=cu_seqlens,
451 + max_seqlen_q=max_seqlen,
452 + max_seqlen_k=max_seqlen,
453 + causal=False,
454 + softmax_scale=self.scale,
455 + )
456 + context_layer = rearrange(output,
457 + "(b s) ... -> b s ...",
458 + b=batch_size)
459 + elif self.attn_backend == _Backend.XFORMERS:
460 + from xformers import ops as xops
461 + from xformers.ops.fmha.attn_bias import BlockDiagonalMask
462 +
463 + attn_bias = BlockDiagonalMask.from_seqlens(q_seqlen=seqlens,
464 + kv_seqlen=None,
465 + device=q.device)
466 +
467 + context_layer = xops.memory_efficient_attention_forward(
468 + q, k, v, attn_bias=attn_bias, p=0, scale=None)
469 +
470 + context_layer = rearrange(context_layer,
471 + "b s h d -> b s (h d)").contiguous()
472 +
473 + output, _ = self.out_proj(context_layer)
474 + return output
475 +
476 +
477 +class SigLIPRotaryEmbedding(nn.Module):
478 +
479 + def __init__(self, dim: int, theta: float = 10000.0) -> None:
480 + super().__init__()
481 + self.dim = dim
482 + self.theta = theta
483 + self.rope_init()
484 +
485 + def rope_init(self):
486 + inv_freq = 1.0 / (self.theta**(
487 + torch.arange(0, self.dim, 2, dtype=torch.float) / self.dim))
488 + self.register_buffer("inv_freq", inv_freq, persistent=False)
489 +
490 + def forward(self, seqlen: int) -> torch.Tensor:
491 + seq = torch.arange(
492 + seqlen,
493 + device=self.inv_freq.device,
494 + dtype=self.inv_freq.dtype,
495 + )
496 + freqs = torch.outer(seq, self.inv_freq)
497 + return freqs
498 +
499 +
500 +class KeyeSiglipEncoderLayer(nn.Module):
501 +
502 + def __init__(
503 + self,
504 + config: Union[PretrainedConfig],
505 + quant_config: Optional[QuantizationConfig] = None,
506 + prefix: str = "",
507 + ):
508 + super().__init__()
509 + self.embed_dim = config.hidden_size
510 + self.layer_norm1 = nn.LayerNorm(self.embed_dim,
511 + eps=config.layer_norm_eps)
512 + self.self_attn = KeyeSiglipAttention(
513 + config,
514 + quant_config=quant_config,
515 + prefix=f"{prefix}.self_attn",
516 + )
517 + self.layer_norm2 = nn.LayerNorm(self.embed_dim,
518 + eps=config.layer_norm_eps)
519 + self.mlp = SiglipMLP(
520 + config,
521 + quant_config=quant_config,
522 + prefix=f"{prefix}.mlp",
523 + )
524 +
525 + def forward(
526 + self,
527 + hidden_states: torch.Tensor,
528 + attention_mask: torch.Tensor,
529 + output_attentions: Optional[bool] = False,
530 + cu_seqlens: Optional[list[torch.Tensor]] = None,
531 + rope_emb: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
532 + ) -> tuple[torch.FloatTensor]:
533 +
534 + residual = hidden_states
535 +
536 + hidden_states = self.layer_norm1(hidden_states)
537 + hidden_states = self.self_attn(
538 + hidden_states=hidden_states,
539 + attention_mask=attention_mask,
540 + output_attentions=output_attentions,
541 + cu_seqlens=cu_seqlens,
542 + rope_emb=rope_emb,
543 + )
544 +
545 + hidden_states = residual + hidden_states
546 +
547 + residual = hidden_states
548 + hidden_states = self.layer_norm2(hidden_states)
549 + hidden_states = self.mlp(hidden_states)
550 +
551 + hidden_states = residual + hidden_states
552 +
553 + return hidden_states
554 +
555 +
556 +class KeyeSiglipEncoder(nn.Module):
557 +
558 + def __init__(
559 + self,
560 + config: PretrainedConfig,
561 + quant_config: Optional[QuantizationConfig] = None,
562 + prefix: str = "",
563 + ):
564 + super().__init__()
565 + self.config = config
566 + embed_dim = config.hidden_size
567 + num_heads = config.num_attention_heads
568 + head_dim = embed_dim // num_heads
569 + self.layers = nn.ModuleList([
570 + KeyeSiglipEncoderLayer(
571 + config,
572 + quant_config=quant_config,
573 + prefix=f"{prefix}.layers.{layer_idx}",
574 + ) for layer_idx in range(config.num_hidden_layers)
575 + ])
576 + self.rotary_pos_emb = SigLIPRotaryEmbedding(head_dim // 2)
577 +
578 + @staticmethod
579 + def flatten_list(image_grid_thw):
580 + tmp_image_grid_thw = list()
581 + for image_grid in image_grid_thw:
582 + if isinstance(image_grid, list):
583 + tmp_image_grid_thw.extend(image_grid)
584 + else:
585 + tmp_image_grid_thw.append(image_grid)
586 + return tmp_image_grid_thw
587 +
588 + def forward(
589 + self,
590 + inputs_embeds,
591 + attention_mask: Optional[torch.Tensor] = None,
592 + output_attentions: Optional[bool] = None,
593 + output_hidden_states: Optional[bool] = None,
594 + cu_seqlens: Optional[list[torch.Tensor]] = None,
595 + image_grid_thw: Optional[list[Union[
596 + tuple[int, int, int],
597 + list[tuple[int, int, int]],
598 + ]]] = None,
599 + height_position_ids: Optional[torch.Tensor] = None,
600 + width_position_ids: Optional[torch.Tensor] = None,
601 + use_rope: Optional[bool] = False,
602 + window_size: Optional[bool] = -1,
603 + vision_or_text: str = "vision",
604 + ) -> BaseModelOutput:
605 + device = inputs_embeds.device
606 + hidden_states = inputs_embeds
607 + if use_rope is True:
608 + flatten_image_grid_thw = self.flatten_list(image_grid_thw)
609 +
610 + if width_position_ids is None or height_position_ids is None:
611 + split_hids = list()
612 + split_wids = list()
613 + for t, h, w in flatten_image_grid_thw:
614 + image_pids = torch.arange(t * h * w,
615 + device=device) % (h * w)
616 + sample_hids = image_pids // w
617 + sample_wids = image_pids % w
618 + split_hids.append(sample_hids)
619 + split_wids.append(sample_wids)
620 + width_position_ids = torch.concat(split_wids, dim=0)
621 + height_position_ids = torch.concat(split_hids, dim=0)
622 +
623 + pids = torch.stack(
624 + [height_position_ids, width_position_ids],
625 + dim=-1,
626 + )
627 + max_grid_size = pids.max() + 1
628 + rope_emb_max_grid = self.rotary_pos_emb(max_grid_size)
629 + rope_emb = rope_emb_max_grid[pids].flatten(1)
630 + rope_emb = rope_emb.repeat(1, 2)
631 + rope_emb = (rope_emb.cos(), rope_emb.sin())
632 + else:
633 + rope_emb = None
634 +
635 + attn_cu_seqlens = cu_seqlens
636 + hidden_states = inputs_embeds
637 + assert attention_mask is None
638 +
639 + for encoder_layer in self.layers:
640 + hidden_states = encoder_layer(
641 + hidden_states,
642 + attention_mask,
643 + output_attentions=output_attentions,
644 + cu_seqlens=attn_cu_seqlens,
645 + rope_emb=rope_emb,
646 + )
647 + return hidden_states
648 +
649 +
650 +class KeyeSiglipVisionTransformer(nn.Module):
651 +
652 + def __init__(
653 + self,
654 + config: PretrainedConfig,
655 + quant_config: Optional[QuantizationConfig] = None,
656 + prefix: str = "",
657 + ):
658 + super().__init__()
659 + self.config = config
660 + embed_dim = config.hidden_size
661 +
662 + self.embeddings = KeyeVisionEmbeddings(config)
663 + self.encoder = KeyeSiglipEncoder(
664 + config,
665 + quant_config=quant_config,
666 + prefix=f"{prefix}.encoder",
667 + )
668 + self.post_layernorm = nn.LayerNorm(embed_dim,
669 + eps=config.layer_norm_eps)
670 +
671 + def forward(
672 + self,
673 + pixel_values,
674 + output_attentions: Optional[bool] = None,
675 + output_hidden_states: Optional[bool] = None,
676 + interpolate_pos_encoding: Optional[bool] = False,
677 + attention_mask: Optional[torch.Tensor] = None,
678 + sample_indices: Optional[torch.Tensor] = None,
679 + image_indices: Optional[torch.Tensor] = None,
680 + position_ids: Optional[torch.Tensor] = None,
681 + height_position_ids: Optional[torch.Tensor] = None,
682 + width_position_ids: Optional[torch.Tensor] = None,
683 + cu_seqlens: Optional[list[torch.Tensor]] = None,
684 + padding_mask: Optional[torch.Tensor] = None,
685 + vision_return_embed_list: Optional[bool] = False,
686 + image_grid_thw: Optional[list[Union[
687 + tuple[int, int, int],
688 + list[tuple[int, int, int]],
689 + ]]] = None,
690 + return_pooler_output: Optional[bool] = True,
691 + use_rope: Optional[bool] = False,
692 + window_size: Optional[bool] = -1,
693 + ) -> BaseModelOutputWithPooling:
694 +
695 + hidden_states = self.embeddings(
696 + pixel_values,
697 + interpolate_pos_encoding=interpolate_pos_encoding,
698 + position_ids=position_ids,
699 + image_grid_thw=image_grid_thw,
700 + )
701 +
702 + last_hidden_state = self.encoder(
703 + inputs_embeds=hidden_states,
704 + output_attentions=output_attentions,
705 + output_hidden_states=output_hidden_states,
706 + attention_mask=attention_mask,
707 + cu_seqlens=cu_seqlens,
708 + image_grid_thw=image_grid_thw,
709 + use_rope=use_rope,
710 + height_position_ids=height_position_ids,
711 + width_position_ids=width_position_ids,
712 + window_size=window_size,
713 + vision_or_text="vision",
714 + )
715 +
716 + last_hidden_state = self.post_layernorm(last_hidden_state)
717 +
718 + sample_hidden_state = list()
719 + if cu_seqlens is None:
720 + raise ValueError("cu_seqlens cannot be None for "
721 + "SiglipVisionTransformer output processing.")
722 + for i in range(cu_seqlens.shape[0] - 1):
723 + start = cu_seqlens[i]
724 + end = cu_seqlens[i + 1]
725 + tensor = last_hidden_state[:, start:end, :].squeeze(0)
726 + sample_hidden_state.append(tensor)
727 +
728 + return sample_hidden_state
729 +
730 +
731 +class KeyeSiglipVisionModel(nn.Module):
732 + config_class = PretrainedConfig
733 + main_input_name = "pixel_values"
734 +
735 + def __init__(
736 + self,
737 + config: PretrainedConfig,
738 + quant_config: Optional[QuantizationConfig] = None,
739 + prefix: str = "",
740 + ):
741 + super().__init__()
742 +
743 + self.vision_model = KeyeSiglipVisionTransformer(
744 + config,
745 + quant_config=quant_config,
746 + prefix=f"{prefix}.vision_model",
747 + )
748 + self.quant_config = quant_config
749 +
750 + @property
751 + def dtype(self) -> torch.dtype:
752 + return self.vision_model.embeddings.patch_embedding.weight.dtype
753 +
754 + @property
755 + def device(self) -> torch.device:
756 + return self.vision_model.embeddings.patch_embedding.weight.device
757 +
758 + def get_input_embeddings(self) -> nn.Module:
759 + return self.vision_model.embeddings.patch_embedding
760 +
761 + def forward(
762 + self,
763 + pixel_values,
764 + sample_indices: Optional[torch.Tensor] = None,
765 + output_attentions: Optional[bool] = None,
766 + output_hidden_states: Optional[bool] = None,
767 + interpolate_pos_encoding: bool = False,
768 + position_ids: Optional[torch.Tensor] = None,
769 + vision_return_embed_list: Optional[bool] = False,
770 + image_grid_thw: Optional[list[Union[
771 + tuple[int, int, int],
772 + list[tuple[int, int, int]],
773 + ]]] = None,
774 + cu_seqlens: Optional[list[torch.Tensor]] = None,
775 + return_pooler_output: Optional[bool] = True,
776 + use_rope: Optional[bool] = False,
777 + window_size: Optional[bool] = -1,
778 + ) -> BaseModelOutputWithPooling:
779 +
780 + return self.vision_model(
781 + pixel_values=pixel_values,
782 + output_attentions=output_attentions,
783 + output_hidden_states=output_hidden_states,
784 + interpolate_pos_encoding=interpolate_pos_encoding,
785 + position_ids=position_ids,
786 + vision_return_embed_list=vision_return_embed_list,
787 + image_grid_thw=image_grid_thw,
788 + sample_indices=sample_indices,
789 + cu_seqlens=cu_seqlens,
790 + return_pooler_output=return_pooler_output,
791 + use_rope=use_rope,
792 + window_size=window_size,
793 + )
794 +
795 + def load_weights(self, weights: Iterable[tuple[str,
796 + torch.Tensor]]) -> set[str]:
797 + stacked_params_mapping = [
798 + ("qkv_proj", "q_proj", "q"),
799 + ("qkv_proj", "k_proj", "k"),
800 + ("qkv_proj", "v_proj", "v"),
801 + ]
802 + params_dict = dict(self.named_parameters(remove_duplicate=False))
803 + loaded_params: set[str] = set()
804 + for name, loaded_weight in weights:
805 + if "rotary_emb.inv_freq" in name:
806 + continue
807 + if "head.attention" in name or "head.layernorm" in name:
808 + continue
809 + if "head.mlp" in name or "head.probe" in name:
810 + continue
811 + if self.quant_config is not None and (
812 + scale_name := self.quant_config.get_cache_scale(name)):
813 + param = params_dict[scale_name]
814 + weight_loader = getattr(
815 + param,
816 + "weight_loader",
817 + default_weight_loader,
818 + )
819 + loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
820 + loaded_weight[0])
821 + weight_loader(param, loaded_weight)
822 + loaded_params.add(scale_name)
823 + continue
824 + for (
825 + param_name,
826 + weight_name,
827 + shard_id,
828 + ) in stacked_params_mapping:
829 + if weight_name not in name:
830 + continue
831 + name = name.replace(weight_name, param_name)
832 + if name.endswith(".bias") and name not in params_dict:
833 + continue
834 + if is_pp_missing_parameter(name, self):
835 + continue
836 + param = params_dict[name]
837 + weight_loader = param.weight_loader
838 + weight_loader(param, loaded_weight, shard_id)
839 + break
840 + else:
841 + if name.endswith(".bias") and name not in params_dict:
842 + continue
843 + name = maybe_remap_kv_scale_name(name, params_dict)
844 + if name is None:
845 + continue
846 + if is_pp_missing_parameter(name, self):
847 + continue
848 + param = params_dict[name]
849 + weight_loader = getattr(
850 + param,
851 + "weight_loader",
852 + default_weight_loader,
853 + )
854 + weight_loader(param, loaded_weight)
855 + loaded_params.add(name)
856 + return loaded_params
857 +
858 +
859 +class Projector(nn.Module):
860 +
861 + def __init__(
862 + self,
863 + text_config: PretrainedConfig,
864 + vision_config: PretrainedConfig,
865 + quant_config: Optional[QuantizationConfig] = None,
866 + prefix: str = "",
867 + ):
868 + super().__init__()
869 + self.text_config = text_config
870 + self.vision_config = vision_config
871 + self.merge_kernel_size = (2, 2)
872 +
873 + self.hidden_size = (self.vision_config.hidden_size *
874 + self.merge_kernel_size[0] *
875 + self.merge_kernel_size[1])
876 +
877 + self.pre_norm = torch.nn.LayerNorm(self.vision_config.hidden_size,
878 + eps=1e-05)
879 + self.act = GELUActivation()
880 +
881 + self.linear_1 = ColumnParallelLinear(
882 + self.hidden_size,
883 + self.hidden_size,
884 + bias=True,
885 + quant_config=quant_config,
886 + prefix=f"{prefix}.linear_1",
887 + )
888 + self.linear_2 = RowParallelLinear(
889 + self.hidden_size,
890 + self.text_config.hidden_size,
891 + bias=True,
892 + quant_config=quant_config,
893 + prefix=f"{prefix}.linear_2",
894 + )
895 +
896 + def forward(
897 + self,
898 + image_features: torch.Tensor,
899 + image_grid_thw: list[tuple[int, int, int]],
900 + ) -> torch.Tensor:
901 + m1, m2 = self.merge_kernel_size
902 + if isinstance(image_features, (list, tuple)):
903 + processed_features = list()
904 + for image_feature, image_grid in zip(image_features,
905 + image_grid_thw):
906 + image_feature = self.pre_norm(image_feature)
907 + t, h, w = image_grid
908 +
909 + image_feature = rearrange(
910 + image_feature,
911 + "(t h p1 w p2) d -> (t h w) (p1 p2 d)",
912 + t=t,
913 + h=h // m1,
914 + p1=m1,
915 + w=w // m2,
916 + p2=m2,
917 + )
918 + hidden_states, _ = self.linear_1(image_feature)
919 + hidden_states = self.act(hidden_states)
920 + hidden_states, _ = self.linear_2(hidden_states)
921 + processed_features.append(hidden_states)
922 +
923 + return processed_features
924 +
925 + dims = image_features.shape[:-1]
926 + dim = image_features.shape[-1]
927 + image_features = image_features.view(np.prod(dims), dim)
928 + hidden_states = self.pre_norm(image_features).view(
929 + -1, self.hidden_size)
930 + hidden_states = self.linear_1(hidden_states)
931 + hidden_states = self.act(hidden_states)
932 + hidden_states = self.linear_2(hidden_states)
933 +
934 + return hidden_states.view(*dims, -1)
935 +
936 +
937 +def _keye_field_config(hf_inputs: Mapping[str, torch.Tensor], ):
938 + image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3)))
939 + image_grid_sizes = image_grid_thw.prod(-1)
940 +
941 + video_grid_thw = hf_inputs.get("video_grid_thw", torch.empty((0, 3)))
942 + video_grid_sizes = video_grid_thw.prod(-1)
943 +
944 + return dict(
945 + pixel_values=MultiModalFieldConfig.flat_from_sizes(
946 + "image", image_grid_sizes),
947 + image_embeds=MultiModalFieldConfig.flat_from_sizes(
948 + "image", image_grid_sizes),
949 + image_grid_thw=MultiModalFieldConfig.batched("image"),
950 + pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
951 + "video", video_grid_sizes),
952 + video_embeds=MultiModalFieldConfig.flat_from_sizes(
953 + "video", video_grid_sizes),
954 + video_grid_thw=MultiModalFieldConfig.batched("video"),
955 + )
956 +
957 +
958 +class KeyeMultiModalDataParser(MultiModalDataParser):
959 +
960 + def _parse_image_data(
961 + self,
962 + data: Union[dict[str, torch.Tensor], ModalityData[ImageItem]],
963 + ) -> ModalityDataItems[Any, Any]:
964 + if isinstance(data, dict):
965 + return DictEmbeddingItems(
966 + data,
967 + modality="image",
968 + required_fields={
969 + "image_embeds",
970 + "image_grid_thw",
971 + },
972 + fields_factory=_keye_field_config,
973 + )
974 +
975 + return super()._parse_image_data(data)
976 +
977 + def _parse_video_data(
978 + self,
979 + data: Union[dict[str, torch.Tensor], ModalityData[VideoItem]],
980 + ) -> ModalityDataItems[Any, Any]:
981 + if isinstance(data, dict):
982 + return DictEmbeddingItems(
983 + data,
984 + modality="video",
985 + required_fields={
986 + "video_embeds",
987 + "video_grid_thw",
988 + },
989 + fields_factory=_keye_field_config,
990 + )
991 +
992 + return super()._parse_video_data(data)
993 +
994 +
995 +class KeyeProcessingInfo(BaseProcessingInfo):
996 +
997 + def get_hf_config(self):
998 + return self.ctx.get_hf_config(PretrainedConfig)
999 +
1000 + def get_hf_processor(
1001 + self,
1002 + *,
1003 + min_pixels: Optional[int] = None,
1004 + max_pixels: Optional[int] = None,
1005 + size: Optional[dict[str, int]] = None,
1006 + **kwargs: object,
1007 + ):
1008 + return self.ctx.get_hf_processor(
1009 + image_processor=self.get_image_processor(
1010 + min_pixels=min_pixels,
1011 + max_pixels=max_pixels,
1012 + size=size,
1013 + ),
1014 + **kwargs,
1015 + )
1016 +
1017 + def _get_image_processor_kwargs(
1018 + self,
1019 + *,
1020 + min_pixels: Optional[int] = None,
1021 + max_pixels: Optional[int] = None,
1022 + size: Optional[dict[str, int]] = None,
1023 + **kwargs: object,
1024 + ):
1025 + if self.ctx.model_config.mm_processor_kwargs:
1026 + kwargs.update(self.ctx.model_config.mm_processor_kwargs)
1027 +
1028 + if min_pixels is not None:
1029 + kwargs["min_pixels"] = min_pixels
1030 +
1031 + if size is None:
1032 + size = {"shortest_edge": min_pixels}
1033 + else:
1034 + size["shortest_edge"] = min_pixels
1035 +
1036 + if max_pixels is not None:
1037 + kwargs["max_pixels"] = max_pixels
1038 +
1039 + if size is None:
1040 + size = {"longest_edge": max_pixels}
1041 + else:
1042 + size["longest_edge"] = max_pixels
1043 +
1044 + if size is not None:
1045 + kwargs["size"] = size
1046 +
1047 + return kwargs
1048 +
1049 + def get_image_processor(
1050 + self,
1051 + *,
1052 + min_pixels: Optional[int] = None,
1053 + max_pixels: Optional[int] = None,
1054 + size: Optional[dict[str, int]] = None,
1055 + **kwargs: object,
1056 + ):
1057 + return cached_image_processor_from_config(
1058 + self.ctx.model_config,
1059 + **self._get_image_processor_kwargs(
1060 + min_pixels=min_pixels,
1061 + max_pixels=max_pixels,
1062 + size=size,
1063 + **kwargs,
1064 + ),
1065 + )
1066 +
1067 + def get_supported_mm_limits(self, ) -> Mapping[str, Optional[int]]:
1068 + return {"image": None, "video": None}
1069 +
1070 + def get_mm_max_tokens_per_item(
1071 + self,
1072 + seq_len: int,
1073 + mm_counts: Mapping[str, int],
1074 + ) -> Mapping[str, int]:
1075 + return {
1076 + "image": self.get_max_image_tokens(),
1077 + "video": self.get_max_video_tokens(seq_len),
1078 + }
1079 +
1080 + def _get_vision_info(
1081 + self,
1082 + *,
1083 + image_width: int,
1084 + image_height: int,
1085 + num_frames: int = 1,
1086 + do_resize: bool = True,
1087 + image_processor,
1088 + ) -> tuple[ImageSize, int]:
1089 + if image_processor is None:
1090 + image_processor = self.get_image_processor()
1091 +
1092 + hf_config = self.get_hf_config()
1093 + vision_config = hf_config.vision_config
1094 + patch_size = vision_config.patch_size
1095 + merge_size = vision_config.spatial_merge_size
1096 + temporal_patch_size = 1
1097 +
1098 + if do_resize:
1099 + resized_height, resized_width = smart_resize(
1100 + height=image_height,
1101 + width=image_width,
1102 + factor=patch_size * merge_size,
1103 + min_pixels=image_processor.min_pixels,
1104 + max_pixels=image_processor.max_pixels,
1105 + )
1106 + preprocessed_size = ImageSize(width=resized_width,
1107 + height=resized_height)
1108 + else:
1109 + preprocessed_size = ImageSize(width=image_width,
1110 + height=image_height)
1111 +
1112 + padded_num_frames = num_frames + num_frames % temporal_patch_size
1113 +
1114 + grid_t = max(padded_num_frames // temporal_patch_size, 1)
1115 + grid_h = preprocessed_size.height // patch_size
1116 + grid_w = preprocessed_size.width // patch_size
1117 +
1118 + num_patches = grid_t * grid_h * grid_w
1119 + num_vision_tokens = num_patches // (merge_size**2)
1120 +
1121 + return preprocessed_size, num_vision_tokens
1122 +
1123 + def get_num_image_tokens(
1124 + self,
1125 + *,
1126 + image_width: int,
1127 + image_height: int,
1128 + image_processor,
1129 + ) -> int:
1130 + _, num_image_tokens = self._get_vision_info(
1131 + image_width=image_width,
1132 + image_height=image_height,
1133 + image_processor=image_processor,
1134 + )
1135 + return num_image_tokens
1136 +
1137 + def get_num_video_tokens(
1138 + self,
1139 + *,
1140 + image_width: int,
1141 + image_height: int,
1142 + num_frames: int,
1143 + image_processor,
1144 + ) -> int:
1145 + _, num_video_tokens = self._get_vision_info(
1146 + image_width=image_width,
1147 + image_height=image_height,
1148 + num_frames=num_frames,
1149 + image_processor=image_processor,
1150 + )
1151 + return num_video_tokens
1152 +
1153 + def get_image_size_with_most_features(self, ) -> ImageSize:
1154 + max_image_size, _ = self._get_vision_info(
1155 + image_width=_MAX_IMAGE_SIZE,
1156 + image_height=_MAX_IMAGE_SIZE,
1157 + image_processor=None,
1158 + )
1159 + return max_image_size
1160 +
1161 + def get_max_image_tokens(self) -> int:
1162 + target_width, target_height = self.get_image_size_with_most_features()
1163 +
1164 + return self.get_num_image_tokens(
1165 + image_width=target_width,
1166 + image_height=target_height,
1167 + image_processor=None,
1168 + )
1169 +
1170 + def _get_max_video_frames(self, max_tokens: int) -> int:
1171 + target_width, target_height = self.get_image_size_with_most_features()
1172 +
1173 + num_frames = 0
1174 +
1175 + while True:
1176 + next_num_frames = num_frames + 1
1177 + next_max_tokens = self.get_num_video_tokens(
1178 + image_width=target_width,
1179 + image_height=target_height,
1180 + num_frames=next_num_frames,
1181 + image_processor=None,
1182 + )
1183 +
1184 + if next_max_tokens > max_tokens:
1185 + break
1186 +
1187 + num_frames = next_num_frames
1188 +
1189 + return num_frames
1190 +
1191 + def get_num_frames_with_most_features(self, seq_len: int) -> int:
1192 + mm_config = self.ctx.get_mm_config()
1193 + max_images = mm_config.get_limit_per_prompt("image")
1194 + max_videos = mm_config.get_limit_per_prompt("video")
1195 +
1196 + max_image_tokens = self.get_max_image_tokens() * max_images
1197 + max_total_frames = self._get_max_video_frames(seq_len -
1198 + max_image_tokens)
1199 + max_frames_per_video = min(
1200 + max_total_frames // max(max_videos, 1),
1201 + _MAX_FRAMES_PER_VIDEO,
1202 + )
1203 +
1204 + return max(max_frames_per_video, 1)
1205 +
1206 + def get_max_video_tokens(self, seq_len: int) -> int:
1207 + target_width, target_height = self.get_image_size_with_most_features()
1208 +
1209 + return self.get_num_video_tokens(
1210 + image_width=target_width,
1211 + image_height=target_height,
1212 + num_frames=self.get_num_frames_with_most_features(seq_len),
1213 + image_processor=None,
1214 + )
1215 +
1216 +
1217 +class KeyeDummyInputsBuilder(BaseDummyInputsBuilder[KeyeProcessingInfo]):
1218 +
1219 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
1220 + num_images = mm_counts.get("image", 0)
1221 + num_videos = mm_counts.get("video", 0)
1222 +
1223 + hf_processor = self.info.get_hf_processor()
1224 + image_token: str = hf_processor.image_token
1225 + video_token: str = hf_processor.video_token
1226 +
1227 + return image_token * num_images + video_token * num_videos
1228 +
1229 + def get_dummy_mm_data(
1230 + self,
1231 + seq_len: int,
1232 + mm_counts: Mapping[str, int],
1233 + ) -> MultiModalDataDict:
1234 + num_images = mm_counts.get("image", 0)
1235 + num_videos = mm_counts.get("video", 0)
1236 +
1237 + target_width, target_height = (
1238 + self.info.get_image_size_with_most_features())
1239 + target_num_frames = self.info.get_num_frames_with_most_features(
1240 + seq_len)
1241 +
1242 + mm_data = {
1243 + "image":
1244 + self._get_dummy_images(
1245 + width=target_width,
1246 + height=target_height,
1247 + num_images=num_images,
1248 + ),
1249 + "video":
1250 + self._get_dummy_videos(
1251 + width=target_width,
1252 + height=target_height,
1253 + num_frames=target_num_frames,
1254 + num_videos=num_videos,
1255 + ),
1256 + }
1257 +
1258 + return mm_data
1259 +
1260 +
1261 +class KeyeMultiModalProcessor(BaseMultiModalProcessor[KeyeProcessingInfo]):
1262 +
1263 + def _get_data_parser(self) -> MultiModalDataParser:
1264 + return KeyeMultiModalDataParser()
1265 +
1266 + def _call_hf_processor(
1267 + self,
1268 + prompt: str,
1269 + mm_data: Mapping[str, object],
1270 + mm_kwargs: Mapping[str, object],
1271 + ) -> BatchFeature:
1272 + return self.info.ctx.call_hf_processor(
1273 + self.info.get_hf_processor(**mm_kwargs),
1274 + dict(text=prompt, **mm_data),
1275 + self.info._get_image_processor_kwargs(**mm_kwargs),
1276 + )
1277 +
1278 + def _get_prompt_updates(
1279 + self,
1280 + mm_items: MultiModalDataItems,
1281 + hf_processor_mm_kwargs: Mapping[str, Any],
1282 + out_mm_kwargs: MultiModalKwargs,
1283 + ) -> Sequence[PromptUpdate]:
1284 + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
1285 + image_processor = self.info.get_image_processor(
1286 + **hf_processor_mm_kwargs)
1287 + tokenizer = self.info.get_tokenizer()
1288 + vocab = tokenizer.get_vocab()
1289 +
1290 + placeholder = {
1291 + "image": vocab[hf_processor.image_token],
1292 + "video": vocab[hf_processor.video_token],
1293 + }
1294 +
1295 + merge_length = image_processor.merge_size**2
1296 +
1297 + def get_replacement_keye(item_idx: int, modality: str):
1298 + grid_thw = out_mm_kwargs[f"{modality}_grid_thw"][item_idx]
1299 + assert isinstance(grid_thw, torch.Tensor)
1300 +
1301 + num_tokens = int(grid_thw.prod()) // merge_length
1302 + return [placeholder[modality]] * num_tokens
1303 +
1304 + return [
1305 + PromptReplacement(
1306 + modality=modality,
1307 + target=[placeholder[modality]],
1308 + replacement=partial(get_replacement_keye, modality=modality),
1309 + ) for modality in ("image", "video")
1310 + ]
1311 +
1312 + def _get_mm_fields_config(
1313 + self,
1314 + hf_inputs: BatchFeature,
1315 + hf_processor_mm_kwargs: Mapping[str, object],
1316 + ) -> Mapping[str, MultiModalFieldConfig]:
1317 + return _keye_field_config(hf_inputs)
1318 +
1319 +
1320 +@MULTIMODAL_REGISTRY.register_processor(
1321 + KeyeMultiModalProcessor,
1322 + info=KeyeProcessingInfo,
1323 + dummy_inputs=KeyeDummyInputsBuilder,
1324 +)
1325 +class KeyeForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA,
1326 + SupportsPP):
1327 + packed_modules_mapping = {
1328 + "qkv_proj": [
1329 + "q_proj",
1330 + "k_proj",
1331 + "v_proj",
1332 + ],
1333 + "gate_up_proj": [
1334 + "gate_proj",
1335 + "up_proj",
1336 + ],
1337 + }
1338 +
1339 + hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={
1340 + "lm_head.": "language_model.lm_head.",
1341 + "model.": "language_model.model.",
1342 + })
1343 +
1344 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1345 + super().__init__()
1346 + config: PretrainedConfig = vllm_config.model_config.hf_config
1347 + quant_config = vllm_config.quant_config
1348 + multimodal_config = vllm_config.model_config.multimodal_config
1349 +
1350 + self.config = config
1351 + self.multimodal_config = multimodal_config
1352 +
1353 + self.visual = KeyeSiglipVisionModel(
1354 + config.vision_config,
1355 + quant_config=self._maybe_ignore_quant_config(quant_config),
1356 + prefix=maybe_prefix(prefix, "visual"),
1357 + )
1358 + self.mlp_AR = Projector(
1359 + config,
1360 + config.vision_config,
1361 + quant_config=self._maybe_ignore_quant_config(quant_config),
1362 + prefix=maybe_prefix(prefix, "mlp_AR"),
1363 + )
1364 +
1365 + self.language_model = init_vllm_registered_model(
1366 + vllm_config=vllm_config,
1367 + prefix=maybe_prefix(prefix, "language_model"),
1368 + architectures=["Qwen3ForCausalLM"],
1369 + )
1370 +
1371 + self.make_empty_intermediate_tensors = (
1372 + self.language_model.make_empty_intermediate_tensors)
1373 +
1374 + def _maybe_ignore_quant_config(self, quant_config: QuantizationConfig):
1375 + if isinstance(quant_config, (GPTQConfig, GPTQMarlinConfig)):
1376 + return None
1377 + return quant_config
1378 +
1379 + def _validate_and_reshape_mm_tensor(self, mm_input: object,
1380 + name: str) -> torch.Tensor:
1381 + if not isinstance(mm_input, (torch.Tensor, list)):
1382 + raise ValueError(f"Incorrect type of {name}. "
1383 + f"Got type: {type(mm_input)}")
1384 + if isinstance(mm_input, torch.Tensor):
1385 + if mm_input.ndim == 2:
1386 + return mm_input
1387 + if mm_input.ndim == 5:
1388 + return mm_input
1389 + if mm_input.ndim != 3:
1390 + raise ValueError(f"{name} should be 2D or batched 3D tensor. "
1391 + f"Got ndim: {mm_input.ndim} "
1392 + f"(shape={mm_input.shape})")
1393 + return torch.concat(list(mm_input))
1394 + else:
1395 + return torch.concat(mm_input)
1396 +
1397 + def _parse_and_validate_image_input(
1398 + self, **kwargs: object) -> Optional[KeyeImageInputs]:
1399 + pixel_values = kwargs.pop("pixel_values", None)
1400 + image_embeds = kwargs.pop("image_embeds", None)
1401 + image_grid_thw = kwargs.pop("image_grid_thw", None)
1402 +
1403 + if pixel_values is None and image_embeds is None:
1404 + return None
1405 +
1406 + if pixel_values is not None:
1407 + pixel_values = self._validate_and_reshape_mm_tensor(
1408 + pixel_values, "image pixel values")
1409 + image_grid_thw = self._validate_and_reshape_mm_tensor(
1410 + image_grid_thw, "image grid_thw")
1411 +
1412 + if not isinstance(pixel_values, (torch.Tensor, list)):
1413 + raise ValueError("Incorrect type of image pixel values. "
1414 + f"Got type: {type(pixel_values)}")
1415 +
1416 + return KeyeImagePixelInputs(
1417 + type="pixel_values",
1418 + pixel_values=pixel_values,
1419 + image_grid_thw=image_grid_thw,
1420 + )
1421 +
1422 + if image_embeds is not None:
1423 + image_embeds = self._validate_and_reshape_mm_tensor(
1424 + image_embeds, "image embeds")
1425 + image_grid_thw = self._validate_and_reshape_mm_tensor(
1426 + image_grid_thw, "image grid_thw")
1427 +
1428 + if not isinstance(image_embeds, torch.Tensor):
1429 + raise ValueError("Incorrect type of image embeddings. "
1430 + f"Got type: {type(image_embeds)}")
1431 + return KeyeImageEmbeddingInputs(
1432 + type="image_embeds",
1433 + image_embeds=image_embeds,
1434 + image_grid_thw=image_grid_thw,
1435 + )
1436 +
1437 + def _parse_and_validate_video_input(
1438 + self, **kwargs: object) -> Optional[KeyeVideoInputs]:
1439 + pixel_values_videos = kwargs.pop("pixel_values_videos", None)
1440 + video_embeds = kwargs.pop("video_embeds", None)
1441 + video_grid_thw = kwargs.pop("video_grid_thw", None)
1442 +
1443 + if pixel_values_videos is None and video_embeds is None:
1444 + return None
1445 +
1446 + if pixel_values_videos is not None:
1447 + pixel_values_videos = self._validate_and_reshape_mm_tensor(
1448 + pixel_values_videos,
1449 + "video pixel values",
1450 + )
1451 + video_grid_thw = self._validate_and_reshape_mm_tensor(
1452 + video_grid_thw, "video grid_thw")
1453 +
1454 + return KeyeVideoPixelInputs(
1455 + type="pixel_values_videos",
1456 + pixel_values_videos=pixel_values_videos,
1457 + video_grid_thw=video_grid_thw,
1458 + )
1459 +
1460 + if video_embeds is not None:
1461 + video_embeds = self._validate_and_reshape_mm_tensor(
1462 + video_embeds, "video embeds")
1463 + video_grid_thw = self._validate_and_reshape_mm_tensor(
1464 + video_grid_thw, "video grid_thw")
1465 +
1466 + if not isinstance(video_embeds, torch.Tensor):
1467 + raise ValueError("Incorrect type of video embeddings. "
1468 + f"Got type: {type(video_embeds)}")
1469 + return KeyeVideoEmbeddingInputs(
1470 + type="video_embeds",
1471 + video_embeds=video_embeds,
1472 + video_grid_thw=video_grid_thw,
1473 + )
1474 +
1475 + def _process_image_input(
1476 + self, image_input: KeyeImageInputs) -> tuple[torch.Tensor, ...]:
1477 + siglip_position_ids = list()
1478 + image_grid_hws = list()
1479 + sample_indices = list()
1480 + cu_seqlens = [0]
1481 +
1482 + image_grid_thw = image_input["image_grid_thw"]
1483 + assert image_grid_thw.ndim == 2
1484 +
1485 + for idx, thaw in enumerate(image_grid_thw):
1486 + thw_tuple = tuple(thaw.detach().cpu().numpy().tolist())
1487 + numel = np.prod(thw_tuple)
1488 + image_grid_hws.append(thw_tuple)
1489 + image_position_ids = torch.arange(numel) % np.prod(thw_tuple[1:])
1490 + siglip_position_ids.append(image_position_ids)
1491 + sample_indices.append(torch.full((numel, ), idx,
1492 + dtype=torch.int64))
1493 + cu_seqlens.append(cu_seqlens[-1] + numel)
1494 +
1495 + if image_input["type"] == "image_embeds":
1496 + raise ValueError(
1497 + "Image embeddings are not supported for this processing path.")
1498 + else:
1499 + pixel_values = image_input["pixel_values"].type(self.visual.dtype)
1500 + siglip_position_ids = torch.concat(siglip_position_ids,
1501 + dim=0).to(pixel_values.device)
1502 + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32).to(
1503 + pixel_values.device)
1504 + sample_indices = torch.concat(sample_indices,
1505 + dim=0).to(pixel_values.device)
1506 +
1507 + image_embeds = self.visual(
1508 + pixel_values=pixel_values,
1509 + image_grid_thw=image_grid_hws,
1510 + position_ids=siglip_position_ids,
1511 + vision_return_embed_list=False,
1512 + interpolate_pos_encoding=True,
1513 + sample_indices=sample_indices,
1514 + cu_seqlens=cu_seqlens,
1515 + use_rope=True,
1516 + window_size=-1,
1517 + )
1518 + image_embeds = tuple(self.mlp_AR(image_embeds, image_grid_thw))
1519 + return image_embeds
1520 +
1521 + def _process_video_input(
1522 + self, video_input: KeyeVideoInputs) -> tuple[torch.Tensor, ...]:
1523 + siglip_position_ids = list()
1524 + video_grid_hws = list()
1525 + sample_indices = list()
1526 + cu_seqlens = [0]
1527 +
1528 + video_grid_thw = video_input["video_grid_thw"]
1529 + assert video_grid_thw.ndim == 2
1530 +
1531 + for idx, thaw in enumerate(video_grid_thw):
1532 + thw_tuple = tuple(thaw.detach().cpu().numpy().tolist())
1533 + numel = np.prod(thw_tuple)
1534 +
1535 + video_grid_hws.append(thw_tuple)
1536 + video_position_ids = torch.arange(numel) % np.prod(thw_tuple[1:])
1537 + siglip_position_ids.append(video_position_ids)
1538 + sample_indices.append(torch.full((numel, ), idx,
1539 + dtype=torch.int64))
1540 + cu_seqlens.append(cu_seqlens[-1] + numel)
1541 +
1542 + if video_input["type"] == "video_embeds":
1543 + raise ValueError(
1544 + "Video embeddings are not supported for this processing path.")
1545 + else:
1546 + pixel_values_videos = video_input["pixel_values_videos"].type(
1547 + self.visual.dtype)
1548 + siglip_position_ids = torch.concat(siglip_position_ids, dim=0).to(
1549 + pixel_values_videos.device)
1550 + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32).to(
1551 + pixel_values_videos.device)
1552 + sample_indices = torch.concat(sample_indices,
1553 + dim=0).to(pixel_values_videos.device)
1554 +
1555 + video_embeds = self.visual(
1556 + pixel_values=pixel_values_videos,
1557 + image_grid_thw=video_grid_hws,
1558 + position_ids=siglip_position_ids,
1559 + vision_return_embed_list=True,
1560 + interpolate_pos_encoding=True,
1561 + sample_indices=sample_indices,
1562 + cu_seqlens=cu_seqlens,
1563 + use_rope=True,
1564 + window_size=-1,
1565 + )
1566 + video_embeds = tuple(self.mlp_AR(video_embeds, video_grid_thw))
1567 + return video_embeds
1568 +
1569 + def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
1570 + modalities = {}
1571 +
1572 + for input_key in kwargs:
1573 + if (input_key in ("pixel_values", "image_embeds")
1574 + and "images" not in modalities):
1575 + modalities["images"] = self._parse_and_validate_image_input(
1576 + **kwargs)
1577 + if (input_key in ("pixel_values_videos", "video_embeds")
1578 + and "videos" not in modalities):
1579 + modalities["videos"] = self._parse_and_validate_video_input(
1580 + **kwargs)
1581 +
1582 + return modalities
1583 +
1584 + def get_multimodal_embeddings(
1585 + self, **kwargs: object) -> Optional[MultiModalEmbeddings]:
1586 +
1587 + modalities = self._parse_and_validate_multimodal_inputs(**kwargs)
1588 + if not modalities:
1589 + return None
1590 +
1591 + multimodal_embeddings: tuple[torch.Tensor, ...] = ()
1592 +
1593 + for modality in modalities:
1594 + if modality == "images":
1595 + image_input = modalities["images"]
1596 + vision_embeddings = self._process_image_input(image_input)
1597 + multimodal_embeddings += vision_embeddings
1598 + if modality == "videos":
1599 + video_input = modalities["videos"]
1600 + video_embeddings = self._process_video_input(video_input)
1601 + multimodal_embeddings += video_embeddings
1602 + return multimodal_embeddings
1603 +
1604 + def get_input_embeddings(
1605 + self,
1606 + input_ids: torch.Tensor,
1607 + multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
1608 + ) -> torch.Tensor:
1609 + inputs_embeds = self.language_model.get_input_embeddings(input_ids)
1610 + if multimodal_embeddings is not None:
1611 + inputs_embeds = merge_multimodal_embeddings(
1612 + input_ids,
1613 + inputs_embeds,
1614 + multimodal_embeddings,
1615 + [
1616 + self.config.image_token_id,
1617 + self.config.video_token_id,
1618 + ],
1619 + )
1620 + return inputs_embeds
1621 +
1622 + def get_input_embeddings_v0(
1623 + self,
1624 + input_ids: torch.Tensor,
1625 + image_input: Optional[KeyeImagePixelInputs] = None,
1626 + video_input: Optional[KeyeVideoPixelInputs] = None,
1627 + ) -> torch.Tensor:
1628 + inputs_embeds = self.get_input_embeddings(input_ids)
1629 + if image_input is not None:
1630 + image_embeds = self._process_image_input(image_input)
1631 + inputs_embeds = merge_multimodal_embeddings(
1632 + input_ids,
1633 + inputs_embeds,
1634 + image_embeds,
1635 + placeholder_token_id=self.config.image_token_id,
1636 + )
1637 +
1638 + if video_input is not None:
1639 + video_embeds = self._process_video_input(video_input)
1640 + inputs_embeds = merge_multimodal_embeddings(
1641 + input_ids,
1642 + inputs_embeds,
1643 + video_embeds,
1644 + placeholder_token_id=self.config.video_token_id,
1645 + )
1646 + return inputs_embeds
1647 +
1648 + def forward(
1649 + self,
1650 + input_ids: torch.Tensor,
1651 + positions: torch.Tensor,
1652 + intermediate_tensors: Optional[IntermediateTensors] = None,
1653 + inputs_embeds: Optional[torch.Tensor] = None,
1654 + **kwargs: object,
1655 + ) -> Union[torch.Tensor, IntermediateTensors]:
1656 + """Run forward pass for Qwen2-VL.
1657 +
1658 + Args:
1659 + input_ids: Flattened (concatenated) input_ids corresponding to a
1660 + batch.
1661 + positions: Flattened (concatenated) position ids corresponding to a
1662 + batch.
1663 + **NOTE**: If mrope is enabled (default setting for Qwen2-VL
1664 + opensource models), the shape will be `(3, seq_len)`,
1665 + otherwise it will be `(seq_len,).
1666 + pixel_values: Pixel values to be fed to a model.
1667 + `None` if no images are passed.
1668 + image_grid_thw: Tensor `(n_images, 3)` of image 3D grid in LLM.
1669 + `None` if no images are passed.
1670 + pixel_values_videos: Pixel values of videos to be fed to a model.
1671 + `None` if no videos are passed.
1672 + video_grid_thw: Tensor `(n_videos, 3)` of video 3D grid in LLM.
1673 + `None` if no videos are passed.
1674 + """
1675 +
1676 + if intermediate_tensors is not None:
1677 + inputs_embeds = None
1678 +
1679 + elif inputs_embeds is None:
1680 + image_input = self._parse_and_validate_image_input(**kwargs)
1681 + video_input = self._parse_and_validate_video_input(**kwargs)
1682 +
1683 + if image_input is None and video_input is None:
1684 + inputs_embeds = None
1685 + else:
1686 + if uses_mrope(self.config):
1687 + assert positions.ndim == 2 and positions.size(0) == 3, (
1688 + "multimodal section rotary embedding requires "
1689 + f"(3, seq_len) positions, but got {positions.size()}")
1690 + inputs_embeds = self.get_input_embeddings_v0(
1691 + input_ids,
1692 + image_input=image_input,
1693 + video_input=video_input,
1694 + )
1695 + input_ids = None
1696 +
1697 + hidden_states = self.language_model.model(
1698 + input_ids=input_ids,
1699 + positions=positions,
1700 + intermediate_tensors=intermediate_tensors,
1701 + inputs_embeds=inputs_embeds,
1702 + )
1703 + return hidden_states
1704 +
1705 + def compute_logits(
1706 + self,
1707 + hidden_states: torch.Tensor,
1708 + sampling_metadata: SamplingMetadata,
1709 + ) -> Optional[torch.Tensor]:
1710 + return self.language_model.compute_logits(hidden_states,
1711 + sampling_metadata)
1712 +
1713 + def load_weights(self, weights: Iterable[tuple[str,
1714 + torch.Tensor]]) -> set[str]:
1715 +
1716 + loader = AutoWeightsLoader(self)
1717 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1718 +
1719 + def get_mm_mapping(self) -> MultiModelKeys:
1720 + """Get the module prefix in multimodal models."""
1721 + return MultiModelKeys.from_string_field(
1722 + language_model="language_model",
1723 + connector="visual.",
1724 + tower_model="mlp_AR.",
1725 + )
vllm/model_executor/models/registry.py
+1
@@ -197,6 +197,7 @@ _MULTIMODAL_MODELS = {
197 "InternVLChatModel": ("internvl", "InternVLChatModel"),
198 "Idefics3ForConditionalGeneration":("idefics3","Idefics3ForConditionalGeneration"),
199 "SmolVLMForConditionalGeneration": ("smolvlm","SmolVLMForConditionalGeneration"), # noqa: E501
200 + "KeyeForConditionalGeneration": ("keye", "KeyeForConditionalGeneration"),
201 "KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"), # noqa: E501
202 "LlavaForConditionalGeneration": ("llava", "LlavaForConditionalGeneration"),
203 "LlavaNextForConditionalGeneration": ("llava_next", "LlavaNextForConditionalGeneration"), # noqa: E501