@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7c8271cd1e

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Model]: support KeyeVL-1_5-8B (#23838)

Signed-off-by: wangruitao <wangruitao@kuaishou.com> Co-authored-by: wangruitao <wangruitao@kuaishou.com>

Kwai-Keye committed Sep 1, 2025 at 18:50 UTC 7c8271cd1e85670bd60c058441e1a224421dea01

9 files changed +980 -135

docs/models/supported_models.md

+2 -1

@@ -634,7 +634,8 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen

634 | `InternS1ForConditionalGeneration` | Intern-S1 | T + IE+ + VE+ | `internlm/Intern-S1`, etc. | ✅︎ | ✅︎ | ✅︎ |

635 | `InternVLChatModel` | InternVL 3.5, InternVL 3.0, InternVideo 2.5, InternVL 2.5, Mono-InternVL, InternVL 2.0 | T + IE+ + (VE+) | `OpenGVLab/InternVL3_5-14B`, `OpenGVLab/InternVL3-9B`, `OpenGVLab/InternVideo2_5_Chat_8B`, `OpenGVLab/InternVL2_5-4B`, `OpenGVLab/Mono-InternVL-2B`, `OpenGVLab/InternVL2-4B`, etc. | ✅︎ | ✅︎ | ✅︎ |

636 | `InternVLForConditionalGeneration` | InternVL 3.0 (HF format) | T + IE+ + VE+ | `OpenGVLab/InternVL3-1B-hf`, etc. | ✅︎ | ✅︎ | ✅︎ |

637 +| `KeyeForConditionalGeneration` | Keye-VL-8B-Preview | T + IE+ + VE+ | `Kwai-Keye/Keye-VL-8B-Preview` | ✅︎ | ✅︎ | ✅︎ |

638 +| `KeyeVL1_5ForConditionalGeneration` | Keye-VL-1_5-8B | T + IE+ + VE+ | `Kwai-Keye/Keye-VL-1_5-8B` | ✅︎ | ✅︎ | ✅︎ |

639 | `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I+ | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | ✅︎ | ✅︎ |

640 | `Llama4ForConditionalGeneration` | Llama 4 | T + I+ | `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. | | ✅︎ | ✅︎ |

641 | `Llama_Nemotron_Nano_VL` | Llama Nemotron Nano VL | T + IE+ | `nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1` | ✅︎ | ✅︎ | ✅︎ |

examples/offline_inference/vision_language.py

+32

@@ -683,6 +683,37 @@ def run_keye_vl(questions: list[str], modality: str) -> ModelRequestData:

683 )

684

685

686 +# Keye-VL-1.5

687 +def run_keye_vl1_5(questions: list[str], modality: str) -> ModelRequestData:

688 + model_name = "Kwai-Keye/Keye-VL-1.5-8B"

689 +

690 + engine_args = EngineArgs(

691 + model=model_name,

692 + max_model_len=8192,

693 + trust_remote_code=True,

694 + limit_mm_per_prompt={modality: 1},

695 + )

696 +

697 + if modality == "image":

698 + placeholder = "<|image_pad|>"

699 + elif modality == "video":

700 + placeholder = "<|video_pad|>"

701 +

702 + prompts = [

703 + (

704 + f"<|im_start|>user\n<|vision_start|>{placeholder}<|vision_end|>"

705 + f"{question}<|im_end|>\n"

706 + "<|im_start|>assistant\n"

707 + )

708 + for question in questions

709 + ]

710 +

711 + return ModelRequestData(

712 + engine_args=engine_args,

713 + prompts=prompts,

714 + )

715 +

716 +

717 # Kimi-VL

718 def run_kimi_vl(questions: list[str], modality: str) -> ModelRequestData:

719 assert modality == "image"

@@ -1648,6 +1679,7 @@ model_example_map = {

1679 "interns1": run_interns1,

1680 "internvl_chat": run_internvl,

1681 "keye_vl": run_keye_vl,

1682 + "keye_vl1_5": run_keye_vl1_5,

1683 "kimi_vl": run_kimi_vl,

1684 "llama4": run_llama4,

1685 "llava": run_llava,

examples/offline_inference/vision_language_multi_image.py

+38

@@ -542,6 +542,43 @@ def load_keye_vl(question: str, image_urls: list[str]) -> ModelRequestData:

542 )

543

544

545 +def load_keye_vl1_5(question: str, image_urls: list[str]) -> ModelRequestData:

546 + model_name = "Kwai-Keye/Keye-VL-1_5-8B"

547 +

548 + engine_args = EngineArgs(

549 + model=model_name,

550 + trust_remote_code=True,

551 + max_model_len=8192,

552 + max_num_seqs=5,

553 + limit_mm_per_prompt={"image": len(image_urls)},

554 + )

555 +

556 + placeholders = [{"type": "image", "image": url} for url in image_urls]

557 + messages = [

558 + {

559 + "role": "user",

560 + "content": [

561 + *placeholders,

562 + {"type": "text", "text": question},

563 + ],

564 + },

565 + ]

566 +

567 + processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)

568 +

569 + prompt = processor.apply_chat_template(

570 + messages, tokenize=False, add_generation_prompt=True

571 + )

572 +

573 + image_data = [fetch_image(url) for url in image_urls]

574 +

575 + return ModelRequestData(

576 + engine_args=engine_args,

577 + prompt=prompt,

578 + image_data=image_data,

579 + )

580 +

581 +

582 def load_kimi_vl(question: str, image_urls: list[str]) -> ModelRequestData:

583 model_name = "moonshotai/Kimi-VL-A3B-Instruct"

584

@@ -1209,6 +1246,7 @@ model_example_map = {

1246 "interns1": load_interns1,

1247 "internvl_chat": load_internvl,

1248 "keye_vl": load_keye_vl,

1249 + "keye_vl1_5": load_keye_vl1_5,

1250 "kimi_vl": load_kimi_vl,

1251 "llama4": load_llama4,

1252 "llava": load_llava,

tests/models/multimodal/processing/test_common.py

@@ -293,6 +293,7 @@ def _test_processing_correctness_one(

293 "OpenGVLab/InternVL3_5-GPT-OSS-20B-A4B-Preview",

294 "OpenGVLab/InternVL3_5-30B-A3B",

295 "Kwai-Keye/Keye-VL-8B-Preview",

296 + "Kwai-Keye/Keye-VL-1_5-8B",

297 "moonshotai/Kimi-VL-A3B-Instruct",

298 "meta-llama/Llama-4-Scout-17B-16E-Instruct",

299 "llava-hf/llava-1.5-7b-hf",

tests/models/registry.py

@@ -438,6 +438,8 @@ _MULTIMODAL_EXAMPLE_MODELS = {

438 "InternVLForConditionalGeneration": _HfExamplesInfo("OpenGVLab/InternVL3-1B-hf"), # noqa: E501

439 "KeyeForConditionalGeneration": _HfExamplesInfo("Kwai-Keye/Keye-VL-8B-Preview", # noqa: E501

440 trust_remote_code=True),

441 + "KeyeVL1_5ForConditionalGeneration": _HfExamplesInfo("Kwai-Keye/Keye-VL-1_5-8B", # noqa: E501

442 + trust_remote_code=True),

443 "KimiVLForConditionalGeneration": _HfExamplesInfo("moonshotai/Kimi-VL-A3B-Instruct", # noqa: E501

444 extras={"thinking": "moonshotai/Kimi-VL-A3B-Thinking"}, # noqa: E501

445 trust_remote_code=True),

vllm/model_executor/layers/rotary_embedding/mrope.py

+129

@@ -402,6 +402,15 @@ class MRotaryEmbedding(RotaryEmbedding):

402 context_len=context_len,

403 seq_len=seq_len,

404 )

405 + elif "KeyeVL1_5" in hf_config.model_type:

406 + return cls._keye_get_input_positions_tensor(

407 + input_tokens=input_tokens,

408 + hf_config=hf_config,

409 + image_grid_thw=image_grid_thw,

410 + video_grid_thw=video_grid_thw,

411 + context_len=context_len,

412 + seq_len=seq_len,

413 + )

414 else:

415 return cls._vl_get_input_positions_tensor(

416 input_tokens=input_tokens,

@@ -636,6 +645,126 @@ class MRotaryEmbedding(RotaryEmbedding):

645 len(input_tokens)).item()

646 return llm_positions, mrope_position_delta

647

648 + @classmethod

649 + def _keye_get_input_positions_tensor(

650 + cls,

651 + input_tokens: list[int],

652 + hf_config: PretrainedConfig,

653 + image_grid_thw: Union[list[list[int]], torch.Tensor],

654 + video_grid_thw: Union[list[list[int]], torch.Tensor],

655 + context_len: int = 0,

656 + seq_len: Optional[int] = None,

657 + ) -> tuple[torch.Tensor, int]:

658 + if isinstance(video_grid_thw, list) and len(video_grid_thw) > 0:

659 + video_grid_thw = video_grid_thw[0]

660 + """Get mrope input positions and delta value (Keye series)."""

661 +

662 + def split_thw(

663 + grid_thw: Union[torch.Tensor, list[int]]) -> list[list[int]]:

664 + """

665 + Split grid_thw along the t dimension.

666 +

667 + Args:

668 + grid_thw: shape [N, 3] tensor or nested list of [t, h, w].

669 +

670 + Returns:

671 + List of [1, h, w] rows, repeated t times for each original row.

672 + """

673 +

674 + if isinstance(grid_thw, list):

675 + grid_thw = torch.tensor(grid_thw, dtype=torch.long)

676 +

677 + if grid_thw.numel() == 0:

678 + return []

679 +

680 + t, hw = grid_thw[:, 0], grid_thw[:, 1:]

681 + ones = torch.ones_like(hw[:, :1]) # [N,1]

682 + out = torch.cat([ones, hw], dim=1).repeat_interleave(t, dim=0)

683 + return out.tolist()

684 +

685 + video_grid_thw = split_thw(video_grid_thw)

686 +

687 + image_token_id = hf_config.image_token_id

688 + video_token_id = hf_config.video_token_id

689 + spatial_merge_size = hf_config.vision_config.spatial_merge_size

690 +

691 + image_nums = len(image_grid_thw)

692 + frame_nums = len(video_grid_thw)

693 + llm_pos_ids_list: list = []

694 +

695 + st = 0

696 + remain_images, remain_frames = image_nums, frame_nums

697 +

698 + image_index, video_index = 0, 0

699 + for _ in range(image_nums + frame_nums):

700 + if remain_images > 0:

701 + try:

702 + ed_image = input_tokens.index(image_token_id, st)

703 + except ValueError:

704 + ed_image = len(input_tokens) + 1

705 + else:

706 + ed_image = len(input_tokens) + 1

707 + if remain_frames > 0:

708 + try:

709 + ed_video = input_tokens.index(video_token_id, st)

710 + except ValueError:

711 + ed_video = len(input_tokens) + 1

712 + else:

713 + ed_video = len(input_tokens) + 1

714 +

715 + if ed_image < ed_video:

716 + t, h, w = (

717 + image_grid_thw[image_index][0],

718 + image_grid_thw[image_index][1],

719 + image_grid_thw[image_index][2],

720 + )

721 + image_index += 1

722 + remain_images -= 1

723 + ed = ed_image

724 + else:

725 + t, h, w = (

726 + video_grid_thw[video_index][0],

727 + video_grid_thw[video_index][1],

728 + video_grid_thw[video_index][2],

729 + )

730 + video_index += 1

731 + remain_frames -= 1

732 + ed = ed_video

733 +

734 + llm_grid_t, llm_grid_h, llm_grid_w = \

735 + t, h // spatial_merge_size, w // spatial_merge_size

736 + text_len = ed - st

737 +

738 + st_idx = llm_pos_ids_list[-1].max() + 1 if len(

739 + llm_pos_ids_list) > 0 else 0

740 + llm_pos_ids_list.append(

741 + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)

742 +

743 + t_index = (torch.arange(llm_grid_t).view(-1, 1).expand(

744 + -1, llm_grid_h * llm_grid_w)).long().flatten()

745 +

746 + h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(

747 + llm_grid_t, -1, llm_grid_w).flatten()

748 + w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(

749 + llm_grid_t, llm_grid_h, -1).flatten()

750 + llm_pos_ids_list.append(

751 + torch.stack([t_index, h_index, w_index]) + text_len + st_idx)

752 + st = ed + llm_grid_t * llm_grid_h * llm_grid_w

753 +

754 + if st < len(input_tokens):

755 + st_idx = llm_pos_ids_list[-1].max() + 1 if len(

756 + llm_pos_ids_list) > 0 else 0

757 + text_len = len(input_tokens) - st

758 + llm_pos_ids_list.append(

759 + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)

760 +

761 + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)

762 + mrope_position_delta = (llm_positions.max() + 1 -

763 + len(input_tokens)).item()

764 + llm_positions = llm_positions[:, context_len:seq_len]

765 +

766 + return llm_positions, mrope_position_delta

767 +

768 @classmethod

769 def _vl_get_input_positions_tensor(

770 cls,

vllm/model_executor/models/keye.py

+174 -134

@@ -1,9 +1,10 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import math

4 +from abc import abstractmethod

5 from collections.abc import Iterable, Mapping, Sequence

6 from functools import partial

6 -from typing import Annotated, Any, Literal, Optional, Union

7 +from typing import Annotated, Any, Literal, Optional, TypeVar, Union

9 import numpy as np

10 import torch

@@ -57,16 +58,13 @@ from .vision import get_vit_attn_backend

59 logger = init_logger(__name__)

60 -_MAX_FRAMES_PER_VIDEO = 16

61 -_MAX_IMAGE_SIZE = 9999999

62 -

62 def smart_resize(

63 height: int,

64 width: int,

67 - factor: int = 28,

68 - min_pixels: int = 28 * 28 * 130,

69 - max_pixels: int = 28 * 28 * 1280,

65 + factor: int,

66 + min_pixels: int,

67 + max_pixels: int,

68 ):

69 if height < factor:

70 logger.warning(

@@ -887,9 +885,9 @@ class Projector(nn.Module):

885

886 def forward(

887 self,

890 - image_features: torch.Tensor,

888 + image_features: Union[torch.Tensor, list[torch.Tensor]],

889 image_grid_thw: list[tuple[int, int, int]],

892 - ) -> torch.Tensor:

890 + ) -> Union[torch.Tensor, list[torch.Tensor]]:

891 m1, m2 = self.merge_kernel_size

892 if isinstance(image_features, (list, tuple)):

893 processed_features = list()

@@ -986,6 +984,12 @@ class KeyeMultiModalDataParser(MultiModalDataParser):

984

985 class KeyeProcessingInfo(BaseProcessingInfo):

986

987 + def get_max_image_size(self) -> int:

988 + return 9999999 #_MAX_IMAGE_SIZE

989 +

990 + def get_max_frame_per_video(self) -> int:

991 + return 16 #_MAX_FRAMES_PER_VIDEO

992 +

993 def get_image_processor(self, **kwargs: object):

994 return self.get_hf_processor(**kwargs).image_processor

995

@@ -1077,8 +1081,8 @@ class KeyeProcessingInfo(BaseProcessingInfo):

1081

1082 def get_image_size_with_most_features(self, ) -> ImageSize:

1083 max_image_size, _ = self._get_vision_info(

1080 - image_width=_MAX_IMAGE_SIZE,

1081 - image_height=_MAX_IMAGE_SIZE,

1084 + image_width=self.get_max_image_size(),

1085 + image_height=self.get_max_image_size(),

1086 image_processor=None,

1087 )

1088 return max_image_size

@@ -1123,7 +1127,7 @@ class KeyeProcessingInfo(BaseProcessingInfo):

1127 max_image_tokens)

1128 max_frames_per_video = min(

1129 max_total_frames // max(max_videos, 1),

1126 - _MAX_FRAMES_PER_VIDEO,

1130 + self.get_max_frame_per_video(),

1131 )

1132

1133 return max(max_frames_per_video, 1)

@@ -1139,7 +1143,10 @@ class KeyeProcessingInfo(BaseProcessingInfo):

1143 )

1144

1145

1142 -class KeyeDummyInputsBuilder(BaseDummyInputsBuilder[KeyeProcessingInfo]):

1146 +_I = TypeVar("_I", bound=KeyeProcessingInfo)

1147 +

1148 +

1149 +class KeyeBaseDummyInputsBuilder(BaseDummyInputsBuilder[_I]):

1150

1151 def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:

1152 num_images = mm_counts.get("image", 0)

@@ -1183,6 +1190,10 @@ class KeyeDummyInputsBuilder(BaseDummyInputsBuilder[KeyeProcessingInfo]):

1190 return mm_data

1191

1192

1193 +class KeyeDummyInputsBuilder(KeyeBaseDummyInputsBuilder[KeyeProcessingInfo]):

1194 + ...

1195 +

1196 +

1197 class KeyeMultiModalProcessor(BaseMultiModalProcessor[KeyeProcessingInfo]):

1198

1199 def _get_data_parser(self) -> MultiModalDataParser:

@@ -1231,13 +1242,7 @@ class KeyeMultiModalProcessor(BaseMultiModalProcessor[KeyeProcessingInfo]):

1242 return _keye_field_config(hf_inputs)

1243

1244

1234 -@MULTIMODAL_REGISTRY.register_processor(

1235 - KeyeMultiModalProcessor,

1236 - info=KeyeProcessingInfo,

1237 - dummy_inputs=KeyeDummyInputsBuilder,

1238 -)

1239 -class KeyeForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA,

1240 - SupportsPP):

1245 +class BaseKeyeModule(nn.Module):

1246 packed_modules_mapping = {

1247 "qkv_proj": [

1248 "q_proj",

@@ -1264,6 +1269,11 @@ class KeyeForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA,

1269

1270 raise ValueError("Only image or video modality is supported")

1271

1272 + def _maybe_ignore_quant_config(self, quant_config: QuantizationConfig):

1273 + if isinstance(quant_config, (GPTQConfig, GPTQMarlinConfig)):

1274 + return None

1275 + return quant_config

1276 +

1277 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

1278 super().__init__()

1279 config: PretrainedConfig = vllm_config.model_config.hf_config

@@ -1278,7 +1288,8 @@ class KeyeForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA,

1288 quant_config=self._maybe_ignore_quant_config(quant_config),

1289 prefix=maybe_prefix(prefix, "visual"),

1290 )

1281 - self.mlp_AR = Projector(

1291 +

1292 + self.mlp_AR = self._build_projector(

1293 config,

1294 config.vision_config,

1295 quant_config=self._maybe_ignore_quant_config(quant_config),

@@ -1294,102 +1305,16 @@ class KeyeForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA,

1305 self.make_empty_intermediate_tensors = (

1306 self.language_model.make_empty_intermediate_tensors)

1307

1297 - def _maybe_ignore_quant_config(self, quant_config: QuantizationConfig):

1298 - if isinstance(quant_config, (GPTQConfig, GPTQMarlinConfig)):

1299 - return None

1300 - return quant_config

1301 -

1302 - def _validate_and_reshape_mm_tensor(self, mm_input: NestedTensors,

1303 - name: str) -> torch.Tensor:

1304 - if not isinstance(mm_input, (torch.Tensor, list)):

1305 - raise ValueError(f"Incorrect type of {name}. "

1306 - f"Got type: {type(mm_input)}")

1307 - if isinstance(mm_input, torch.Tensor):

1308 - if mm_input.ndim == 2:

1309 - return mm_input

1310 - if mm_input.ndim == 5:

1311 - return mm_input

1312 - if mm_input.ndim != 3:

1313 - raise ValueError(f"{name} should be 2D or batched 3D tensor. "

1314 - f"Got ndim: {mm_input.ndim} "

1315 - f"(shape={mm_input.shape})")

1316 - return torch.concat(list(mm_input))

1317 - elif is_list_of(mm_input, torch.Tensor):

1318 - if all(p.dim() == 4 for p in mm_input) or all(p.dim() == 2

1319 - for p in mm_input):

1320 - return mm_input

1321 - return torch.concat(list(mm_input))

1322 -

1323 - def _parse_and_validate_image_input(

1324 - self, **kwargs: object) -> Optional[KeyeImageInputs]:

1325 - pixel_values = kwargs.pop("pixel_values", None)

1326 - image_embeds = kwargs.pop("image_embeds", None)

1327 - image_grid_thw = kwargs.pop("image_grid_thw", None)

1328 -

1329 - if pixel_values is None and image_embeds is None:

1330 - return None

1331 -

1332 - if pixel_values is not None:

1333 - pixel_values = self._validate_and_reshape_mm_tensor(

1334 - pixel_values, "image pixel values")

1335 - image_grid_thw = self._validate_and_reshape_mm_tensor(

1336 - image_grid_thw, "image grid_thw")

1337 -

1338 - return KeyeImagePixelInputs(

1339 - type="pixel_values",

1340 - pixel_values=pixel_values,

1341 - image_grid_thw=image_grid_thw,

1342 - )

1343 -

1344 - if image_embeds is not None:

1345 - image_embeds = self._validate_and_reshape_mm_tensor(

1346 - image_embeds, "image embeds")

1347 - image_grid_thw = self._validate_and_reshape_mm_tensor(

1348 - image_grid_thw, "image grid_thw")

1349 -

1350 - return KeyeImageEmbeddingInputs(

1351 - type="image_embeds",

1352 - image_embeds=image_embeds,

1353 - image_grid_thw=image_grid_thw,

1354 - )

1355 -

1356 - def _parse_and_validate_video_input(

1357 - self, **kwargs: object) -> Optional[KeyeVideoInputs]:

1358 - pixel_values_videos = kwargs.pop("pixel_values_videos", None)

1359 - video_embeds = kwargs.pop("video_embeds", None)

1360 - video_grid_thw = kwargs.pop("video_grid_thw", None)

1361 -

1362 - if pixel_values_videos is None and video_embeds is None:

1363 - return None

1364 -

1365 - if pixel_values_videos is not None:

1366 - pixel_values_videos = self._validate_and_reshape_mm_tensor(

1367 - pixel_values_videos,

1368 - "video pixel values",

1369 - )

1370 - video_grid_thw = self._validate_and_reshape_mm_tensor(

1371 - video_grid_thw, "video grid_thw")

1372 -

1373 - return KeyeVideoPixelInputs(

1374 - type="pixel_values_videos",

1375 - pixel_values_videos=pixel_values_videos,

1376 - video_grid_thw=video_grid_thw,

1377 - )

1378 -

1379 - if video_embeds is not None:

1380 - video_embeds = self._validate_and_reshape_mm_tensor(

1381 - video_embeds, "video embeds")

1382 - video_grid_thw = self._validate_and_reshape_mm_tensor(

1383 - video_grid_thw, "video grid_thw")

1384 -

1385 - return KeyeVideoEmbeddingInputs(

1386 - type="video_embeds",

1387 - video_embeds=video_embeds,

1388 - video_grid_thw=video_grid_thw,

1389 - )

1308 + @abstractmethod

1309 + def _build_projector(self,

1310 + text_config: PretrainedConfig,

1311 + vision_config: PretrainedConfig,

1312 + quant_config: Optional[QuantizationConfig] = None,

1313 + prefix: str = "") -> nn.Module:

1314 + raise ValueError("Need projector")

1315

1391 - def _process_image_input(

1392 - self, image_input: KeyeImageInputs) -> tuple[torch.Tensor, ...]:

1316 + def _process_image_input(self,

1317 + image_input: Any) -> tuple[torch.Tensor, ...]:

1318 siglip_position_ids = list()

1319 image_grid_hws = list()

1320 sample_indices = list()

@@ -1434,18 +1359,20 @@ class KeyeForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA,

1359 image_embeds = tuple(self.mlp_AR(image_embeds, image_grid_thw))

1360 return image_embeds

1361

1437 - def _process_video_input(

1438 - self, video_input: KeyeVideoInputs) -> tuple[torch.Tensor, ...]:

1362 + def _process_video_embeds(

1363 + self,

1364 + video_type: Literal["video_embeds", "pixel_values_videos"],

1365 + video_grid_thw: list[torch.Tensor],

1366 + pixel_values_videos: Optional[torch.Tensor] = None

1367 + ) -> Union[torch.Tensor, list[torch.Tensor]]:

1368 siglip_position_ids = list()

1369 video_grid_hws = list()

1370 sample_indices = list()

1371 cu_seqlens = [0]

1372

1444 - video_grid_thw = video_input["video_grid_thw"]

1373 assert video_grid_thw.ndim == 2

1446 -

1447 - for idx, thaw in enumerate(video_grid_thw):

1448 - thw_tuple = tuple(thaw.detach().cpu().numpy().tolist())

1374 + for idx, sub_thw in enumerate(video_grid_thw):

1375 + thw_tuple = tuple(sub_thw.detach().cpu().numpy().tolist())

1376 numel = np.prod(thw_tuple)

1377

1378 video_grid_hws.append(thw_tuple)

@@ -1455,12 +1382,11 @@ class KeyeForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA,

1382 dtype=torch.int64))

1383 cu_seqlens.append(cu_seqlens[-1] + numel)

1384

1458 - if video_input["type"] == "video_embeds":

1385 + if video_type == "video_embeds":

1386 raise ValueError(

1387 "Video embeddings are not supported for this processing path.")

1388 else:

1462 - pixel_values_videos = video_input["pixel_values_videos"].type(

1463 - self.visual.dtype)

1389 + pixel_values_videos = pixel_values_videos.type(self.visual.dtype)

1390 siglip_position_ids = torch.concat(siglip_position_ids, dim=0).to(

1391 pixel_values_videos.device)

1392 cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32).to(

@@ -1479,7 +1405,7 @@ class KeyeForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA,

1405 use_rope=True,

1406 window_size=-1,

1407 )

1482 - video_embeds = tuple(self.mlp_AR(video_embeds, video_grid_thw))

1408 + video_embeds = self.mlp_AR(video_embeds, video_grid_thw)

1409 return video_embeds

1410

1411 def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:

@@ -1541,8 +1467,8 @@ class KeyeForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA,

1467 def get_input_embeddings_v0(

1468 self,

1469 input_ids: torch.Tensor,

1544 - image_input: Optional[KeyeImagePixelInputs] = None,

1545 - video_input: Optional[KeyeVideoPixelInputs] = None,

1470 + image_input: Optional[Any] = None,

1471 + video_input: Optional[Any] = None,

1472 ) -> torch.Tensor:

1473 inputs_embeds = self.get_input_embeddings(input_ids)

1474 if image_input is not None:

@@ -1572,7 +1498,7 @@ class KeyeForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA,

1498 inputs_embeds: Optional[torch.Tensor] = None,

1499 **kwargs: object,

1500 ) -> Union[torch.Tensor, IntermediateTensors]:

1575 - """Run forward pass for Qwen2-VL.

1501 + """Run forward pass for Keye-VL.

1502

1503 Args:

1504 input_ids: Flattened (concatenated) input_ids corresponding to a

@@ -1591,14 +1517,12 @@ class KeyeForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA,

1517 video_grid_thw: Tensor `(n_videos, 3)` of video 3D grid in LLM.

1518 `None` if no videos are passed.

1519 """

1594 -

1520 if intermediate_tensors is not None:

1521 inputs_embeds = None

1522

1523 elif inputs_embeds is None:

1524 image_input = self._parse_and_validate_image_input(**kwargs)

1525 video_input = self._parse_and_validate_video_input(**kwargs)

1601 -

1526 if image_input is None and video_input is None:

1527 inputs_embeds = None

1528 else:

@@ -1619,6 +1543,7 @@ class KeyeForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA,

1543 intermediate_tensors=intermediate_tensors,

1544 inputs_embeds=inputs_embeds,

1545 )

1546 +

1547 return hidden_states

1548

1549 def compute_logits(

@@ -1631,7 +1556,6 @@ class KeyeForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA,

1556

1557 def load_weights(self, weights: Iterable[tuple[str,

1558 torch.Tensor]]) -> set[str]:

1634 -

1559 loader = AutoWeightsLoader(self)

1560 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

1561

@@ -1639,6 +1563,122 @@ class KeyeForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA,

1563 """Get the module prefix in multimodal models."""

1564 return MultiModelKeys.from_string_field(

1565 language_model="language_model",

1642 - connector="visual.",

1643 - tower_model="mlp_AR.",

1566 + connector="mlp_AR.",

1567 + tower_model="visual.",

1568 )

1569 +

1570 +

1571 +@MULTIMODAL_REGISTRY.register_processor(

1572 + KeyeMultiModalProcessor,

1573 + info=KeyeProcessingInfo,

1574 + dummy_inputs=KeyeDummyInputsBuilder,

1575 +)

1576 +class KeyeForConditionalGeneration(BaseKeyeModule, SupportsMultiModal,

1577 + SupportsLoRA, SupportsPP):

1578 +

1579 + def _build_projector(self,

1580 + text_config: PretrainedConfig,

1581 + vision_config: PretrainedConfig,

1582 + quant_config: Optional[QuantizationConfig] = None,

1583 + prefix: str = "") -> nn.Module:

1584 + return Projector(text_config, vision_config, quant_config, prefix)

1585 +

1586 + def _validate_and_reshape_mm_tensor(

1587 + self, mm_input: NestedTensors,

1588 + name: str) -> Union[torch.Tensor, list[torch.Tensor]]:

1589 + if not isinstance(mm_input, (torch.Tensor, list)):

1590 + raise ValueError(f"Incorrect type of {name}. "

1591 + f"Got type: {type(mm_input)}")

1592 + if isinstance(mm_input, torch.Tensor):

1593 + if mm_input.ndim == 2:

1594 + return mm_input

1595 + if mm_input.ndim == 5:

1596 + return mm_input

1597 + if mm_input.ndim != 3:

1598 + raise ValueError(f"{name} should be 2D or batched 3D tensor. "

1599 + f"Got ndim: {mm_input.ndim} "

1600 + f"(shape={mm_input.shape})")

1601 + return torch.concat(list(mm_input))

1602 + elif is_list_of(mm_input, torch.Tensor):

1603 + if all(p.dim() == 4 for p in mm_input) or all(p.dim() == 2

1604 + for p in mm_input):

1605 + return mm_input

1606 + return torch.concat(list(mm_input))

1607 +

1608 + def _parse_and_validate_image_input(

1609 + self, **kwargs: object) -> Optional[KeyeImageInputs]:

1610 + pixel_values = kwargs.pop("pixel_values", None)

1611 + image_embeds = kwargs.pop("image_embeds", None)

1612 + image_grid_thw = kwargs.pop("image_grid_thw", None)

1613 +

1614 + if pixel_values is None and image_embeds is None:

1615 + return None

1616 +

1617 + if pixel_values is not None:

1618 + pixel_values = self._validate_and_reshape_mm_tensor(

1619 + pixel_values, "image pixel values")

1620 + image_grid_thw = self._validate_and_reshape_mm_tensor(

1621 + image_grid_thw, "image grid_thw")

1622 +

1623 + return KeyeImagePixelInputs(

1624 + type="pixel_values",

1625 + pixel_values=pixel_values,

1626 + image_grid_thw=image_grid_thw,

1627 + )

1628 +

1629 + if image_embeds is not None:

1630 + image_embeds = self._validate_and_reshape_mm_tensor(

1631 + image_embeds, "image embeds")

1632 + image_grid_thw = self._validate_and_reshape_mm_tensor(

1633 + image_grid_thw, "image grid_thw")

1634 +

1635 + return KeyeImageEmbeddingInputs(

1636 + type="image_embeds",

1637 + image_embeds=image_embeds,

1638 + image_grid_thw=image_grid_thw,

1639 + )

1640 +

1641 + def _parse_and_validate_video_input(

1642 + self, **kwargs: object) -> Optional[KeyeVideoInputs]:

1643 + pixel_values_videos = kwargs.pop("pixel_values_videos", None)

1644 + video_embeds = kwargs.pop("video_embeds", None)

1645 + video_grid_thw = kwargs.pop("video_grid_thw", None)

1646 +

1647 + if pixel_values_videos is None and video_embeds is None:

1648 + return None

1649 +

1650 + if pixel_values_videos is not None:

1651 + pixel_values_videos = self._validate_and_reshape_mm_tensor(

1652 + pixel_values_videos,

1653 + "video pixel values",

1654 + )

1655 + video_grid_thw = self._validate_and_reshape_mm_tensor(

1656 + video_grid_thw, "video grid_thw")

1657 +

1658 + return KeyeVideoPixelInputs(

1659 + type="pixel_values_videos",

1660 + pixel_values_videos=pixel_values_videos,

1661 + video_grid_thw=video_grid_thw,

1662 + )

1663 +

1664 + if video_embeds is not None:

1665 + video_embeds = self._validate_and_reshape_mm_tensor(

1666 + video_embeds, "video embeds")

1667 + video_grid_thw = self._validate_and_reshape_mm_tensor(

1668 + video_grid_thw, "video grid_thw")

1669 +

1670 + return KeyeVideoEmbeddingInputs(

1671 + type="video_embeds",

1672 + video_embeds=video_embeds,

1673 + video_grid_thw=video_grid_thw,

1674 + )

1675 +

1676 + def _process_video_input(

1677 + self, video_input: KeyeVideoInputs) -> tuple[torch.Tensor, ...]:

1678 + video_type = video_input["type"]

1679 + video_grid_thw = video_input["video_grid_thw"]

1680 + pixel_values_videos = video_input.get("pixel_values_videos", None)

1681 +

1682 + return tuple(

1683 + self._process_video_embeds(video_type, video_grid_thw,

1684 + pixel_values_videos))

vllm/model_executor/models/keye_vl1_5.py new

+601

@@ -0,0 +1,601 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import itertools

4 +from collections.abc import Mapping, Sequence

5 +from functools import partial

6 +from typing import Annotated, Any, Literal, Optional, 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 +

16 +from vllm.config import VllmConfig

17 +from vllm.logger import init_logger

18 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,

19 + RowParallelLinear)

20 +from vllm.model_executor.layers.quantization import QuantizationConfig

21 +from vllm.multimodal import MULTIMODAL_REGISTRY, NestedTensors

22 +from vllm.multimodal.inputs import (ImageItem, ModalityData,

23 + MultiModalFieldConfig,

24 + MultiModalKwargsItems, VideoItem)

25 +from vllm.multimodal.parse import (DictEmbeddingItems, ModalityDataItems,

26 + MultiModalDataItems, MultiModalDataParser)

27 +from vllm.multimodal.processing import (PromptReplacement, PromptUpdate,

28 + PromptUpdateDetails)

29 +from vllm.utils.tensor_schema import TensorSchema, TensorShape

30 +

31 +from .interfaces import SupportsLoRA, SupportsMultiModal, SupportsPP

32 +from .keye import (BaseKeyeModule, BaseMultiModalProcessor,

33 + KeyeBaseDummyInputsBuilder, KeyeProcessingInfo)

34 +

35 +logger = init_logger(__name__)

36 +

37 +

38 +def split_thw(grid_thw: torch.Tensor) -> torch.Tensor:

39 + """

40 + Split grid_thw in t dimension.

41 +

42 + Args:

43 + grid_thw: [N, 3] tensor of [t, h, w]

44 +

45 + Returns:

46 + [Σt, 3] tensor where each row is [1, h, w]

47 +

48 + Example:

49 + >>> grid_thw = torch.tensor([[2, 3, 4], [1, 5, 6]])

50 + >>> split_thw(grid_thw)

51 + tensor([[1, 3, 4],

52 + [1, 3, 4],

53 + [1, 5, 6]])

54 + """

55 + t = grid_thw[:, 0]

56 + h_w = grid_thw[:, 1:]

57 + ones = torch.ones_like(h_w[:, :1])

58 + return torch.cat([ones, h_w], dim=1).repeat_interleave(t, dim=0)

59 +

60 +

61 +def get_num_patches(grid_thw: torch.Tensor, num_frames: Union[list[int],

62 + torch.Tensor]):

63 + """

64 + Return num_patches per video.

65 +

66 + Args:

67 + t: tensor with shape [N, ...] where each item is a list/tensor

68 + cu_seqlens: list indicating the boundaries of groups

69 +

70 + Returns:

71 + list of ints representing the sum of products for each group

72 +

73 + Examples:

74 + >>> # Suppose there are 2 videos with a total of 3 grids

75 + >>> grid_thw = torch.tensor([[2, 2, 2], # grid 0: 2*2*2=8 patches

76 + ... [2, 2, 2], # grid 1: 2*2*2=8 patches

77 + ... [1, 1, 1]]) # grid 2: 1*1*1=1 patches

78 + >>> num_frames = [2, 1] # The first video contains 2 grids,

79 + the second contains 1 grid.

80 + >>> get_num_patches(grid_thw, num_frames)

81 + tensor([16, 1]) # Total patches for first video: 8+8=16,

82 + second video: 1.

83 + """

84 +

85 + assert len(grid_thw.shape) == 2

86 + if isinstance(num_frames, torch.Tensor):

87 + num_frames = num_frames.clone().tolist()

88 +

89 + num_grids_per_frame = grid_thw.prod(dim=1)

90 + start_idx_per_video = [0, *itertools.accumulate(num_frames)]

91 + num_patches = [

92 + num_grids_per_frame[start_idx_per_video[i]:start_idx_per_video[i + 1]].

93 + sum() for i in range(len(num_frames))

94 + ]

95 + return torch.stack(num_patches) if num_patches else torch.zeros(

96 + 0, dtype=grid_thw.dtype, device=grid_thw.device)

97 +

98 +

99 +class KeyeVL1_5ImagePixelInputs(TensorSchema):

100 + """

101 + Dimensions:

102 + - b: Batch size

103 + - np: Number of patches

104 + - c: Number of channels

105 + - ps: Patch size

106 + - ni: Number of images

107 + - g: Grid dimensions (3 for t, h, w)

108 + """

109 + type: Literal["pixel_values"]

110 +

111 + pixel_values: Annotated[

112 + torch.Tensor,

113 + TensorShape("np", 3, "ps", "ps", dynamic_dims={"np"})]

114 +

115 + image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]

116 +

117 +

118 +class KeyeVL1_5ImageEmbeddingInputs(TensorSchema):

119 + """

120 + Dimensions:

121 + - nf: Number of image features

122 + - hs: Hidden size (must match the hidden size of language model

123 + backbone)

124 + - ni: Number of images

125 + - g: Grid dimensions (3 for t, h, w)

126 + """

127 + type: Literal["image_embeds"]

128 + image_embeds: Annotated[torch.Tensor, TensorShape("nf", "hs")]

129 + image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]

130 +

131 +

132 +KeyeVL1_5ImageInputs = Union[KeyeVL1_5ImagePixelInputs,

133 + KeyeVL1_5ImageEmbeddingInputs]

134 +

135 +

136 +class KeyeVL1_5VideoPixelInputs(TensorSchema):

137 + """

138 + Dimensions:

139 + - b: Batch size

140 + - np: Number of patches

141 + - c: Number of channels

142 + - ps: Patch size

143 + - ni: Number of images

144 + - g: Grid dimensions (3 for t, h, w)

145 + """

146 + type: Literal["pixel_values_videos"]

147 + pixel_values_videos: Annotated[

148 + torch.Tensor,

149 + TensorShape("np", 3, "ps", "ps", dynamic_dims={"np"})]

150 + video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)]

151 +

152 + num_frames: torch.Tensor

153 +

154 +

155 +class KeyeVL1_5VideoEmbeddingInputs(TensorSchema):

156 + """

157 + Dimensions:

158 + - nf: Number of video features

159 + - hs: Hidden size (must match the hidden size of language model

160 + backbone)

161 + - nv: Number of videos

162 + - g: Grid dimensions (3 for t, h, w)

163 + """

164 + type: Literal["video_embeds"]

165 + video_embeds: Annotated[torch.Tensor, TensorShape("nf", "hs")]

166 + video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)]

167 + num_frames: torch.Tensor

168 +

169 +

170 +KeyeVL1_5VideoInputs = Union[KeyeVL1_5VideoPixelInputs,

171 + KeyeVL1_5VideoEmbeddingInputs]

172 +

173 +

174 +class KeyeVL1_5Projector(nn.Module):

175 +

176 + def __init__(

177 + self,

178 + text_config: PretrainedConfig,

179 + vision_config: PretrainedConfig,

180 + quant_config: Optional[QuantizationConfig] = None,

181 + prefix: str = "",

182 + ):

183 + super().__init__()

184 + self.text_config = text_config

185 + self.vision_config = vision_config

186 + self.merge_kernel_size = (2, 2)

187 +

188 + self.hidden_size = (self.vision_config.hidden_size *

189 + self.merge_kernel_size[0] *

190 + self.merge_kernel_size[1])

191 +

192 + self.pre_norm = torch.nn.LayerNorm(self.hidden_size, eps=1e-05)

193 + self.act = GELUActivation()

194 +

195 + self.linear_1 = ColumnParallelLinear(

196 + self.hidden_size,

197 + self.hidden_size,

198 + bias=True,

199 + quant_config=quant_config,

200 + prefix=f"{prefix}.linear_1",

201 + )

202 + self.linear_2 = RowParallelLinear(

203 + self.hidden_size,

204 + self.text_config.hidden_size,

205 + bias=True,

206 + quant_config=quant_config,

207 + prefix=f"{prefix}.linear_2",

208 + )

209 +

210 + def forward(

211 + self,

212 + image_features: Union[torch.Tensor, tuple[torch.Tensor],

213 + list[torch.Tensor]],

214 + image_grid_thw: list[tuple[int, int, int]],

215 + ) -> Union[torch.Tensor, list[torch.Tensor]]:

216 + m1, m2 = self.merge_kernel_size

217 + if isinstance(image_features, (list, tuple)):

218 + processed_features = list()

219 + for image_feature, image_grid in zip(image_features,

220 + image_grid_thw):

221 + t, h, w = image_grid

222 + image_feature = rearrange(

223 + image_feature,

224 + "(t h p1 w p2) d -> (t h w) (p1 p2 d)",

225 + t=t,

226 + h=h // m1,

227 + p1=m1,

228 + w=w // m2,

229 + p2=m2,

230 + )

231 + image_feature = self.pre_norm(image_feature)

232 + hidden_states, _ = self.linear_1(image_feature)

233 + hidden_states = self.act(hidden_states)

234 + hidden_states, _ = self.linear_2(hidden_states)

235 + processed_features.append(hidden_states)

236 +

237 + return processed_features

238 +

239 + dims = image_features.shape[:-1]

240 + dim = image_features.shape[-1]

241 + image_features = image_features.view(np.prod(dims), dim)

242 + hidden_states = self.pre_norm(image_features.view(

243 + -1, self.hidden_size))

244 + hidden_states = self.linear_1(hidden_states)

245 + hidden_states = self.act(hidden_states)

246 + hidden_states = self.linear_2(hidden_states)

247 +

248 + return hidden_states.view(*dims, -1)

249 +

250 +

251 +class KeyeVL1_5ProcessingInfo(KeyeProcessingInfo):

252 +

253 + def get_max_frame_per_video(self) -> int:

254 + return 2048

255 +

256 + def get_supported_mm_limits(self, ) -> Mapping[str, Optional[int]]:

257 + return {"image": None, "video": 1}

258 +

259 +

260 +def _keye_field_config(hf_inputs: Mapping[str, torch.Tensor], ):

261 + image_grid_thw = hf_inputs.get("image_grid_thw",

262 + torch.empty((0, 3), dtype=torch.int64))

263 + image_grid_sizes = image_grid_thw.prod(-1)

264 +

265 + video_grid_thw = hf_inputs.get("video_grid_thw",

266 + torch.empty((0, 3), dtype=torch.int64))

267 + video_grid_thw = split_thw(video_grid_thw)

268 + num_frames = hf_inputs.get("num_frames",

269 + video_grid_thw[:, 0]).clone().tolist()

270 +

271 + video_num_patches = get_num_patches(video_grid_thw, num_frames)

272 +

273 + video_num_grids = []

274 + if len(num_frames) > 0:

275 + i = 0

276 + j = 1

277 + cur_frames = num_frames[i]

278 + for t, _, _ in video_grid_thw.tolist():

279 + cur_frames -= t

280 + if cur_frames == 0:

281 + video_num_grids.append(j)

282 + i += 1

283 + if i < len(num_frames):

284 + cur_frames = num_frames[i]

285 + j = 1

286 + else:

287 + j += 1

288 + video_num_grids = torch.tensor(video_num_grids)

289 + return dict(pixel_values=MultiModalFieldConfig.flat_from_sizes(

290 + "image", image_grid_sizes),

291 + image_embeds=MultiModalFieldConfig.flat_from_sizes(

292 + "image", image_grid_sizes),

293 + image_grid_thw=MultiModalFieldConfig.batched("image"),

294 + pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(

295 + "video", video_num_patches),

296 + video_embeds=MultiModalFieldConfig.flat_from_sizes(

297 + "video", video_num_patches),

298 + video_grid_thw=MultiModalFieldConfig.flat_from_sizes(

299 + "video", video_num_grids),

300 + num_frames=MultiModalFieldConfig.batched("video"))

301 +

302 +

303 +class KeyeVL1_5MultiModalDataParser(MultiModalDataParser):

304 +

305 + def _parse_image_data(

306 + self,

307 + data: Union[dict[str, torch.Tensor], ModalityData[ImageItem]],

308 + ) -> ModalityDataItems[Any, Any]:

309 + if isinstance(data, dict):

310 + return DictEmbeddingItems(

311 + data,

312 + modality="image",

313 + required_fields={

314 + "image_embeds",

315 + "image_grid_thw",

316 + },

317 + fields_factory=_keye_field_config,

318 + )

319 +

320 + return super()._parse_image_data(data)

321 +

322 + def _parse_video_data(

323 + self,

324 + data: Union[dict[str, torch.Tensor], ModalityData[VideoItem]],

325 + ) -> ModalityDataItems[Any, Any]:

326 + if isinstance(data, dict):

327 + return DictEmbeddingItems(

328 + data,

329 + modality="video",

330 + required_fields={

331 + "video_embeds",

332 + "video_grid_thw",

333 + },

334 + fields_factory=_keye_field_config,

335 + )

336 +

337 + return super()._parse_video_data(data)

338 +

339 +

340 +class KeyeVL1_5MultiModalProcessor(

341 + BaseMultiModalProcessor[KeyeVL1_5ProcessingInfo]):

342 +

343 + def _get_data_parser(self) -> MultiModalDataParser:

344 + return KeyeVL1_5MultiModalDataParser()

345 +

346 + def _get_prompt_updates(

347 + self,

348 + mm_items: MultiModalDataItems,

349 + hf_processor_mm_kwargs: Mapping[str, Any],

350 + out_mm_kwargs: MultiModalKwargsItems,

351 + ) -> Sequence[PromptUpdate]:

352 + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

353 + image_processor = self.info.get_image_processor(

354 + **hf_processor_mm_kwargs)

355 + tokenizer = self.info.get_tokenizer()

356 + vocab = tokenizer.get_vocab()

357 + image_token_id = vocab[hf_processor.image_token]

358 + video_token_id = vocab[hf_processor.video_token]

359 + placeholder = {"image": image_token_id, "video": video_token_id}

360 + merge_length = image_processor.merge_size**2

361 +

362 + out_mm_kwargs_data = out_mm_kwargs.get_data()

363 + frame_types: list[torch.Tensor] = \

364 + hf_processor_mm_kwargs.get("frame_types", None)

365 + timestamps: list[torch.Tensor] = \

366 + hf_processor_mm_kwargs.get("timestamps", None)

367 + num_videos = mm_items.get_count("video", strict=False)

368 +

369 + if frame_types is None:

370 + frame_types = [None] * num_videos

371 + assert len(frame_types) == num_videos, \

372 + f"Number of frame_types={len(frame_types)} " \

373 + f"doesn't equal to number of videos={num_videos}"

374 + if timestamps is None:

375 + timestamps = [None] * num_videos

376 + assert len(timestamps) == num_videos, \

377 + f"Number of timestamps={len(timestamps)} " \

378 + f"doesn't equal to number of videos={num_videos}"

379 +

380 + video_grid_thw = out_mm_kwargs_data.get(

381 + 'video_grid_thw', torch.empty((0, 3), dtype=torch.int64))

382 + num_frames = out_mm_kwargs_data.get(

383 + 'num_frames', torch.tensor([], dtype=torch.int64))

384 +

385 + assert len(num_frames) == num_videos, \

386 + f"Size of num_frames={len(num_frames)} " \

387 + f"doesn't equal to number of videos={num_videos}"

388 +

389 + video_grid_hws = split_thw(video_grid_thw)

390 + assert int(num_frames.sum().tolist()) == video_grid_hws.shape[0], (

391 + f"The first dimension of `video_grid_hws`={video_grid_hws.shape[0]}"

392 + f"doesn't equal to num of frames.")

393 +

394 + cu_seqlens = torch.cumsum(torch.tensor([0] + num_frames.tolist()),

395 + dim=-1)

396 +

397 + def get_replacement_keye(item_idx: int, modality: str):

398 + """

399 + Args:

400 + item_idx(int): The item index of modality to replace

401 + modality(str): The modality

402 + """

403 + if modality == "image":

404 + out_item = out_mm_kwargs[modality][item_idx]

405 + grid_thw = out_item[f"{modality}_grid_thw"].data

406 + assert isinstance(grid_thw, torch.Tensor)

407 +

408 + num_tokens = int(grid_thw.prod()) // merge_length

409 + return [image_token_id] * num_tokens

410 + elif modality == "video":

411 + placeholders = []

412 + video_timestamps = timestamps[item_idx]

413 + video_frame_types = frame_types[item_idx]

414 + grid_thw = video_grid_hws[

415 + cu_seqlens[item_idx]:cu_seqlens[item_idx + 1]]

416 +

417 + nframes = grid_thw.shape[0]

418 +

419 + if video_timestamps is None:

420 + video_timestamps = [""] * nframes

421 + else:

422 + video_timestamps = [

423 + format(ts, ".1f") for ts in video_timestamps

424 + ]

425 +

426 + if video_frame_types is None:

427 + video_frame_types = [0] * nframes

428 + for i, sub_thw in enumerate(grid_thw):

429 + s = f"{hf_processor.frame_token}{video_timestamps[i]}"

430 + if video_frame_types[i] == 1:

431 + s += hf_processor.fast_start

432 + placeholders.extend(tokenizer.encode(s))

433 + num_frame_tokens = int(sub_thw.prod()) // merge_length

434 + placeholders.extend([video_token_id] * num_frame_tokens)

435 + if video_frame_types[i] == 1:

436 + placeholders.append(vocab[hf_processor.fast_end])

437 +

438 + return PromptUpdateDetails.select_token_id(

439 + placeholders, embed_token_id=video_token_id)

440 + else:

441 + raise ValueError(f"Unsupported modality {modality}")

442 +

443 + return [

444 + PromptReplacement(

445 + modality=modality,

446 + target=[placeholder[modality]],

447 + replacement=partial(get_replacement_keye, modality=modality),

448 + ) for modality in ("image", "video")

449 + ]

450 +

451 + def _get_mm_fields_config(

452 + self,

453 + hf_inputs: BatchFeature,

454 + hf_processor_mm_kwargs: Mapping[str, object],

455 + ) -> Mapping[str, MultiModalFieldConfig]:

456 + return _keye_field_config(hf_inputs)

457 +

458 +

459 +class KeyeVL1_5DummyInputsBuilder(

460 + KeyeBaseDummyInputsBuilder[KeyeVL1_5ProcessingInfo]):

461 + ...

462 +

463 +

464 +@MULTIMODAL_REGISTRY.register_processor(

465 + KeyeVL1_5MultiModalProcessor,

466 + info=KeyeVL1_5ProcessingInfo,

467 + dummy_inputs=KeyeVL1_5DummyInputsBuilder,

468 +)

469 +class KeyeVL1_5ForConditionalGeneration(BaseKeyeModule, SupportsMultiModal,

470 + SupportsLoRA, SupportsPP):

471 +

472 + def _build_projector(self,

473 + text_config: PretrainedConfig,

474 + vision_config: PretrainedConfig,

475 + quant_config: Optional[QuantizationConfig] = None,

476 + prefix: str = "") -> nn.Module:

477 + return KeyeVL1_5Projector(text_config, vision_config, quant_config,

478 + prefix)

479 +

480 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

481 + config: PretrainedConfig = vllm_config.model_config.hf_config

482 + self.merge_size = config.vision_config.spatial_merge_size

483 + super().__init__(vllm_config=vllm_config, prefix=prefix)

484 +

485 + def _validate_and_reshape_mm_tensor(self, mm_input: NestedTensors,

486 + expected_dim: int, name: str):

487 + if not isinstance(mm_input, (torch.Tensor, list)):

488 + raise ValueError(f"Incorrect type of {name}. "

489 + f"Got type: {type(mm_input)}")

490 + if isinstance(mm_input, torch.Tensor):

491 + if mm_input.ndim == expected_dim:

492 + return mm_input

493 + elif mm_input.ndim == expected_dim + 1:

494 + return torch.concat(list(mm_input))

495 + else:

496 + raise ValueError(

497 + f"{name} should be {expected_dim}D or "

498 + f"batched {expected_dim}D tensor."

499 + f"Got ndim: {mm_input.ndim} (shape={mm_input.shape})")

500 + else:

501 + return torch.concat(list(mm_input))

502 +

503 + def _parse_and_validate_image_input(

504 + self, **kwargs: object) -> Optional[KeyeVL1_5ImageInputs]:

505 + pixel_values = kwargs.pop("pixel_values", None)

506 + image_embeds = kwargs.pop("image_embeds", None)

507 + image_grid_thw = kwargs.pop("image_grid_thw", None)

508 +

509 + if pixel_values is None and image_embeds is None:

510 + return None

511 +

512 + if pixel_values is not None:

513 + pixel_values = self._validate_and_reshape_mm_tensor(

514 + pixel_values, expected_dim=4, name="image pixel values")

515 + image_grid_thw = self._validate_and_reshape_mm_tensor(

516 + image_grid_thw, expected_dim=2, name="image grid_thw")

517 +

518 + return KeyeVL1_5ImagePixelInputs(

519 + type="pixel_values",

520 + pixel_values=pixel_values,

521 + image_grid_thw=image_grid_thw,

522 + )

523 +

524 + if image_embeds is not None:

525 + image_embeds = self._validate_and_reshape_mm_tensor(

526 + image_embeds, expected_dim=2, name="image embeds")

527 + image_grid_thw = self._validate_and_reshape_mm_tensor(

528 + image_grid_thw, expected_dim=2, name="image grid_thw")

529 +

530 + return KeyeVL1_5ImageEmbeddingInputs(

531 + type="image_embeds",

532 + image_embeds=image_embeds,

533 + image_grid_thw=image_grid_thw,

534 + )

535 +

536 + def _parse_and_validate_video_input(

537 + self, **kwargs: object) -> Optional[KeyeVL1_5VideoInputs]:

538 + pixel_values_videos = kwargs.pop("pixel_values_videos", None)

539 + video_embeds = kwargs.pop("video_embeds", None)

540 + video_grid_thw = kwargs.pop("video_grid_thw", None)

541 + num_frames = kwargs.pop("num_frames", None)

542 +

543 + if pixel_values_videos is None and video_embeds is None:

544 + return None

545 +

546 + if pixel_values_videos is not None:

547 + pixel_values_videos = self._validate_and_reshape_mm_tensor(

548 + pixel_values_videos,

549 + expected_dim=4,

550 + name="video pixel values",

551 + )

552 + video_grid_thw = self._validate_and_reshape_mm_tensor(

553 + video_grid_thw, expected_dim=2, name="video grid_thw")

554 +

555 + num_frames = self._validate_and_reshape_mm_tensor(

556 + num_frames, expected_dim=1, name="video num frames")

557 +

558 + return KeyeVL1_5VideoPixelInputs(

559 + type="pixel_values_videos",

560 + pixel_values_videos=pixel_values_videos,

561 + video_grid_thw=video_grid_thw,

562 + num_frames=num_frames)

563 +

564 + if video_embeds is not None:

565 + video_embeds = self._validate_and_reshape_mm_tensor(

566 + video_embeds, expected_dim=2, name="video embeds")

567 + video_grid_thw = self._validate_and_reshape_mm_tensor(

568 + video_grid_thw, expected_dim=2, name="video grid_thw")

569 +

570 + return KeyeVL1_5VideoEmbeddingInputs(type="video_embeds",

571 + video_embeds=video_embeds,

572 + video_grid_thw=video_grid_thw,

573 + num_frames=num_frames)

574 +

575 + def _process_video_input(

576 + self,

577 + video_input: KeyeVL1_5VideoInputs) -> tuple[torch.Tensor, ...]:

578 + video_type = video_input["type"]

579 + video_grid_thw = split_thw(video_input["video_grid_thw"])

580 + pixel_values_videos = video_input.get("pixel_values_videos", None)

581 +

582 + video_embeds = self._process_video_embeds(video_type, video_grid_thw,

583 + pixel_values_videos)

584 + video_embeds = torch.concat(video_embeds, dim=0)

585 +

586 + num_frames = video_input["num_frames"].clone().tolist()

587 +

588 + num_patches = get_num_patches(video_grid_thw, num_frames).tolist()

589 +

590 + patch_cu_seqlens = torch.cumsum(

591 + torch.tensor([0] + num_patches).detach().clone(), dim=-1)

592 + patch_cu_seqlens = torch.div(patch_cu_seqlens,

593 + self.merge_size**2,

594 + rounding_mode="floor")

595 +

596 + new_video_embeds = []

597 + for idx in range(patch_cu_seqlens.shape[0] - 1):

598 + start = patch_cu_seqlens[idx]

599 + end = patch_cu_seqlens[idx + 1]

600 + new_video_embeds.append(video_embeds[start:end])

601 + return tuple(new_video_embeds)

vllm/model_executor/models/registry.py

@@ -227,6 +227,7 @@ _MULTIMODAL_MODELS = {

227 "Idefics3ForConditionalGeneration":("idefics3","Idefics3ForConditionalGeneration"),

228 "SmolVLMForConditionalGeneration": ("smolvlm","SmolVLMForConditionalGeneration"), # noqa: E501

229 "KeyeForConditionalGeneration": ("keye", "KeyeForConditionalGeneration"),

230 + "KeyeVL1_5ForConditionalGeneration": ("keye_vl1_5", "KeyeVL1_5ForConditionalGeneration"), # noqa: E501

231 "RForConditionalGeneration": ("rvl", "RForConditionalGeneration"),

232 "KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"), # noqa: E501

233 "Llama_Nemotron_Nano_VL": ("nemotron_vl", "LlamaNemotronVLChatModel"),