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

[Model][VLM] Add Kimi-VL model support (#16387)

Signed-off-by: courage17340 <courage17340@163.com>

courage17340 committed Apr 15, 2025 at 05:41 UTC b1308b84a3a6323f33dc995923b6cffa86e7199e
18 files changed +1436 -14
docs/source/models/supported_models.md
+7
@@ -886,6 +886,13 @@ See [this page](#generative-models) for more information on how to use generativ
886 *
887 * ✅︎
888 * ✅︎
889 +- * `KimiVLForConditionalGeneration`
890 + * Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking
891 + * T + I<sup>+</sup>
892 + * `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking`
893 + *
894 + *
895 + * ✅︎
896 - * `Llama4ForConditionalGeneration`
897 * Llama 4
898 * T + I<sup>+</sup>
examples/offline_inference/vision_language.py
+24
@@ -364,6 +364,29 @@ def run_internvl(questions: list[str], modality: str) -> ModelRequestData:
364 )
365
366
367 +# Kimi-VL
368 +def run_kimi_vl(questions: list[str], modality: str) -> ModelRequestData:
369 + assert modality == "image"
370 +
371 + prompts = [
372 + "<|im_user|>user<|im_middle|><|media_start|>image<|media_content|>"
373 + f"<|media_pad|><|media_end|>{question}<|im_end|>"
374 + "<|im_assistant|>assistant<|im_middle|>" for question in questions
375 + ]
376 +
377 + engine_args = EngineArgs(
378 + model="moonshotai/Kimi-VL-A3B-Instruct",
379 + max_model_len=4096,
380 + disable_mm_preprocessor_cache=args.disable_mm_preprocessor_cache,
381 + trust_remote_code=True,
382 + )
383 +
384 + return ModelRequestData(
385 + engine_args=engine_args,
386 + prompts=prompts,
387 + )
388 +
389 +
390 # LLaVA-1.5
391 def run_llava(questions: list[str], modality: str) -> ModelRequestData:
392 assert modality == "image"
@@ -966,6 +989,7 @@ model_example_map = {
989 "h2ovl_chat": run_h2ovl,
990 "idefics3": run_idefics3,
991 "internvl_chat": run_internvl,
992 + "kimi_vl": run_kimi_vl,
993 "llava": run_llava,
994 "llava-next": run_llava_next,
995 "llava-next-video": run_llava_next_video,
examples/offline_inference/vision_language_multi_image.py
+40
@@ -326,6 +326,45 @@ def load_llama4(question: str, image_urls: list[str]) -> ModelRequestData:
326 )
327
328
329 +def load_kimi_vl(question: str, image_urls: list[str]) -> ModelRequestData:
330 + model_name = "moonshotai/Kimi-VL-A3B-Instruct"
331 +
332 + engine_args = EngineArgs(
333 + model=model_name,
334 + max_model_len=4096,
335 + max_num_seqs=4,
336 + tensor_parallel_size=1,
337 + limit_mm_per_prompt={"image": len(image_urls)},
338 + trust_remote_code=True,
339 + )
340 +
341 + placeholders = [{"type": "image", "image": url} for url in image_urls]
342 + messages = [{
343 + "role":
344 + "user",
345 + "content": [
346 + *placeholders,
347 + {
348 + "type": "text",
349 + "text": question
350 + },
351 + ],
352 + }]
353 +
354 + processor = AutoProcessor.from_pretrained(model_name,
355 + trust_remote_code=True)
356 +
357 + prompt = processor.apply_chat_template(messages,
358 + tokenize=False,
359 + add_generation_prompt=True)
360 +
361 + return ModelRequestData(
362 + engine_args=engine_args,
363 + prompt=prompt,
364 + image_data=[fetch_image(url) for url in image_urls],
365 + )
366 +
367 +
368 def load_mistral3(question: str, image_urls: list[str]) -> ModelRequestData:
369 model_name = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
370
@@ -640,6 +679,7 @@ model_example_map = {
679 "h2ovl_chat": load_h2ovl,
680 "idefics3": load_idefics3,
681 "internvl_chat": load_internvl,
682 + "kimi_vl": load_kimi_vl,
683 "llama4": load_llama4,
684 "mistral3": load_mistral3,
685 "mllama": load_mllama,
requirements/test.in
+1
@@ -10,6 +10,7 @@ pytest-timeout
10 # testing utils
11 awscli
12 backoff # required for phi4mm test
13 +blobfile # required for kimi-vl test
14 einops # required for MPT, qwen-vl and Mamba
15 httpx
16 librosa # required for audio tests
requirements/test.txt
+9 -1
@@ -39,6 +39,8 @@ bitsandbytes==0.45.3
39 # via -r requirements/test.in
40 black==24.10.0
41 # via datamodel-code-generator
42 +blobfile==3.0.0
43 + # via -r requirements/test.in
44 boto3==1.35.57
45 # via tensorizer
46 botocore==1.35.57
@@ -127,6 +129,7 @@ fastsafetensors==0.1.10
129 # via -r requirements/test.in
130 filelock==3.16.1
131 # via
132 + # blobfile
133 # datasets
134 # huggingface-hub
135 # ray
@@ -227,7 +230,9 @@ llvmlite==0.44.0
230 lm-eval==0.4.8
231 # via -r requirements/test.in
232 lxml==5.3.0
230 - # via sacrebleu
233 + # via
234 + # blobfile
235 + # sacrebleu
236 markdown-it-py==3.0.0
237 # via rich
238 markupsafe==3.0.2
@@ -426,6 +431,8 @@ pybind11==2.13.6
431 # via lm-eval
432 pycparser==2.22
433 # via cffi
434 +pycryptodomex==3.22.0
435 + # via blobfile
436 pydantic==2.9.2
437 # via
438 # datamodel-code-generator
@@ -689,6 +696,7 @@ tzdata==2024.2
696 # via pandas
697 urllib3==2.2.3
698 # via
699 + # blobfile
700 # botocore
701 # requests
702 # responses
tests/models/decoder_only/vision_language/test_models.py
+12
@@ -318,6 +318,18 @@ VLM_TEST_SETTINGS = {
318 use_tokenizer_eos=True,
319 patch_hf_runner=model_utils.internvl_patch_hf_runner,
320 ),
321 + "kimi_vl": VLMTestInfo(
322 + models=["moonshotai/Kimi-VL-A3B-Instruct"],
323 + test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
324 + prompt_formatter=lambda img_prompt: f"<|im_user|>user<|im_middle|>{img_prompt}<|im_end|><|im_assistant|>assistant<|im_middle|>", # noqa: E501
325 + img_idx_to_prompt=lambda _: "<|media_start|>image<|media_content|><|media_pad|><|media_end|>", # noqa: E501
326 + max_model_len=8192,
327 + max_num_seqs=2,
328 + dtype="bfloat16",
329 + tensor_parallel_size=1,
330 + vllm_output_post_proc=model_utils.kimiv_vl_vllm_to_hf_output,
331 + marks=[large_gpu_mark(min_gb=48)],
332 + ),
333 "llama4": VLMTestInfo(
334 models=["meta-llama/Llama-4-Scout-17B-16E-Instruct"],
335 prompt_formatter=lambda img_prompt: f"<|begin_of_text|><|header_start|>user<|header_end|>\n\n{img_prompt}<|eot|><|header_start|>assistant<|header_end|>\n\n", # noqa: E501
tests/models/decoder_only/vision_language/vlm_utils/model_utils.py
+11
@@ -68,6 +68,17 @@ def qwen2_vllm_to_hf_output(
68 return output_ids, hf_output_str, out_logprobs
69
70
71 +def kimiv_vl_vllm_to_hf_output(
72 + vllm_output: RunnerOutput,
73 + model: str) -> tuple[list[int], str, Optional[SampleLogprobs]]:
74 + """Sanitize vllm output [kimi_vl models] to be comparable with hf output."""
75 + output_ids, output_str, out_logprobs = vllm_output
76 +
77 + hf_output_str = output_str + "<|im_end|>[EOS]"
78 +
79 + return output_ids, hf_output_str, out_logprobs
80 +
81 +
82 def llava_image_vllm_to_hf_output(vllm_output: RunnerOutput,
83 model: str) -> RunnerOutput:
84 config = AutoConfig.from_pretrained(model)
tests/models/multimodal/processing/test_common.py
+1
@@ -258,6 +258,7 @@ def _test_processing_correctness_mistral(
258 "OpenGVLab/InternVL2-1B",
259 "HuggingFaceM4/Idefics3-8B-Llama3",
260 "HuggingFaceTB/SmolVLM2-2.2B-Instruct",
261 + "moonshotai/Kimi-VL-A3B-Instruct",
262 "meta-llama/Llama-4-Scout-17B-16E-Instruct",
263 "llava-hf/llava-1.5-7b-hf",
264 "llava-hf/llava-v1.6-mistral-7b-hf",
tests/models/registry.py
+3
@@ -302,6 +302,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {
302 trust_remote_code=True),
303 "Idefics3ForConditionalGeneration": _HfExamplesInfo("HuggingFaceM4/Idefics3-8B-Llama3", # noqa: E501
304 {"tiny": "HuggingFaceTB/SmolVLM-256M-Instruct"}), # noqa: E501
305 + "KimiVLForConditionalGeneration": _HfExamplesInfo("moonshotai/Kimi-VL-A3B-Instruct", # noqa: E501
306 + extras={"thinking": "moonshotai/Kimi-VL-A3B-Thinking"}, # noqa: E501
307 + trust_remote_code=True),
308 "Llama4ForConditionalGeneration": _HfExamplesInfo("meta-llama/Llama-4-Scout-17B-16E-Instruct", # noqa: E501
309 min_transformers_version="4.51"),
310 "LlavaForConditionalGeneration": _HfExamplesInfo("llava-hf/llava-1.5-7b-hf",
vllm/entrypoints/chat_utils.py
+2
@@ -512,6 +512,8 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):
512 return "<|fim_prefix|><|img|><|fim_suffix|>"
513 if model_type == "gemma3":
514 return "<start_of_image>"
515 + if model_type == "kimi_vl":
516 + return "<|media_start|>image<|media_content|><|media_pad|><|media_end|>" # noqa: E501
517
518 raise TypeError(f"Unknown {modality} model type: {model_type}")
519 elif modality == "audio":
vllm/model_executor/models/kimi_vl.py new
+608
@@ -0,0 +1,608 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# ruff: noqa: E501
3 +# Adapted from https://huggingface.co/moonshotai/Kimi-VL-A3B-Instruct/blob/main/modeling_kimi_vl.py
4 +# Copyright 2025 The Moonshot AI Team, DeepSeek-AI, and HuggingFace Inc. team. All rights reserved.
5 +#
6 +# The code is based on llava (llava/modeling_llava.py) and DeepSeek-V3 (DeepSeek-V3/modeling_deepseek.py), but modified for KimiVL.
7 +#
8 +# Licensing Information:
9 +# - Code derived from llava (llava/modeling_llava.py) and DeepSeek-V3 (DeepSeek-V3/modeling_deepseek.py) is licensed under the Apache License, Version 2.0.
10 +# - Other parts of the code are licensed under the MIT License.
11 +#
12 +# Apache License, Version 2.0:
13 +# Licensed under the Apache License, Version 2.0 (the "License");
14 +# you may not use this file except in compliance with the License.
15 +# You may obtain a copy of the License at
16 +#
17 +# http://www.apache.org/licenses/LICENSE-2.0
18 +#
19 +# Unless required by applicable law or agreed to in writing, software
20 +# distributed under the License is distributed on an "AS IS" BASIS,
21 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
22 +# See the License for the specific language governing permissions and
23 +# limitations under the License.
24 +#
25 +# MIT License:
26 +# Permission is hereby granted, free of charge, to any person obtaining a copy
27 +# of this software and associated documentation files (the "Software"), to deal
28 +# in the Software without restriction, including without limitation the rights
29 +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
30 +# copies of the Software, and to permit persons to whom the Software is
31 +# furnished to do so, subject to the following conditions:
32 +#
33 +# The above copyright notice and this permission notice shall be included in all
34 +# copies or substantial portions of the Software.
35 +#
36 +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
37 +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
38 +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
39 +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
40 +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
41 +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
42 +# SOFTWARE.
43 +
44 +import copy
45 +import math
46 +from collections.abc import Mapping
47 +from dataclasses import dataclass
48 +from typing import (Any, Iterable, List, Literal, Optional, Sequence, Tuple,
49 + TypedDict, Union)
50 +
51 +import torch
52 +from torch import nn
53 +from transformers import BatchFeature
54 +from transformers.activations import GELUActivation
55 +
56 +from vllm.config import VllmConfig
57 +from vllm.distributed import (get_tensor_model_parallel_rank,
58 + get_tensor_model_parallel_world_size)
59 +from vllm.logger import init_logger
60 +from vllm.model_executor.layers.fused_moe import FusedMoE
61 +from vllm.model_executor.layers.logits_processor import LogitsProcessor
62 +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
63 +from vllm.model_executor.layers.vocab_parallel_embedding import (
64 + DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead)
65 +from vllm.model_executor.model_loader.weight_utils import (
66 + default_weight_loader, maybe_remap_kv_scale_name)
67 +from vllm.model_executor.models.deepseek_v2 import DeepseekV2Model
68 +from vllm.model_executor.models.interfaces import SupportsMultiModal
69 +from vllm.model_executor.models.moonvit import MoonVitPretrainedModel
70 +from vllm.model_executor.models.utils import merge_multimodal_embeddings
71 +from vllm.model_executor.sampling_metadata import SamplingMetadata
72 +from vllm.multimodal import MULTIMODAL_REGISTRY
73 +from vllm.multimodal.inputs import (MultiModalFieldConfig, MultiModalKwargs,
74 + NestedTensors)
75 +from vllm.multimodal.parse import (ImageEmbeddingItems, ImageProcessorItems,
76 + MultiModalDataItems)
77 +from vllm.multimodal.processing import (BaseMultiModalProcessor,
78 + BaseProcessingInfo, PromptReplacement,
79 + PromptUpdate)
80 +from vllm.multimodal.profiling import BaseDummyInputsBuilder, ProcessorInputs
81 +from vllm.sequence import IntermediateTensors
82 +from vllm.transformers_utils.configs import KimiVLConfig, MoonViTConfig
83 +from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekV2Config
84 +
85 +from .utils import is_pp_missing_parameter, maybe_prefix
86 +
87 +logger = init_logger(__name__)
88 +
89 +
90 +# For dummy input only
91 +@dataclass
92 +class MaxImageTokenMeta:
93 + width: int = 1024
94 + height: int = 1024
95 +
96 +
97 +class KimiVLMultiModalProjector(nn.Module):
98 +
99 + def __init__(self, config: KimiVLConfig):
100 + super().__init__()
101 +
102 + self.hidden_size = (config.vision_config.hidden_size *
103 + config.vision_config.merge_kernel_size[0] *
104 + config.vision_config.merge_kernel_size[1])
105 +
106 + self.pre_norm = torch.nn.LayerNorm(config.vision_config.hidden_size,
107 + eps=1e-5)
108 + self.linear_1 = nn.Linear(self.hidden_size,
109 + self.hidden_size,
110 + bias=True)
111 + self.act = GELUActivation()
112 + self.linear_2 = nn.Linear(self.hidden_size,
113 + config.text_config.hidden_size,
114 + bias=True)
115 +
116 + def forward(self, image_features: torch.Tensor) -> torch.Tensor:
117 + hidden_states = self.pre_norm(image_features).view(
118 + -1, self.hidden_size)
119 + hidden_states = self.linear_1(hidden_states)
120 + hidden_states = self.act(hidden_states)
121 + hidden_states = self.linear_2(hidden_states)
122 + return hidden_states
123 +
124 +
125 +class KimiVLImagePixelInputs(TypedDict):
126 + type: Literal["pixel_values"]
127 + pixel_values: Union[torch.Tensor, List[torch.Tensor]]
128 + """
129 + Shape:`(num_patches, num_channels, patch_size, patch_size)`
130 + """
131 +
132 + image_grid_hws: torch.Tensor
133 + """Shape:`(num_images, 2)`"""
134 +
135 +
136 +# TODO: support embeds too
137 +# We only support pixel input for kimi-vl now
138 +KimiVLImageInputs = KimiVLImagePixelInputs
139 +
140 +
141 +class KimiVLProcessingInfo(BaseProcessingInfo):
142 +
143 + def get_hf_config(self):
144 + return self.ctx.get_hf_config(KimiVLConfig)
145 +
146 + def get_num_image_tokens(
147 + self,
148 + *,
149 + image_width: int,
150 + image_height: int,
151 + ) -> int:
152 + hf_processor = self.get_hf_processor()
153 + patch_size = hf_processor.image_processor.patch_size
154 + kernel_size = hf_processor.image_processor.merge_kernel_size
155 + in_token_limit = hf_processor.image_processor.in_token_limit
156 + height = image_height
157 + width = image_width
158 + assert isinstance(height,
159 + int), f"height must be int, current height {height}"
160 + assert isinstance(width,
161 + int), f"width must be int, current width {width}"
162 + assert kernel_size is not None, "kernel_size must be specified"
163 +
164 + if (width // patch_size) * (height // patch_size) > in_token_limit:
165 + scale = math.sqrt(in_token_limit / ((width // patch_size) *
166 + (height // patch_size)))
167 + new_w, new_h = int(width * scale), int(height * scale)
168 + width, height = new_w, new_h
169 +
170 + kernel_height, kernel_width = kernel_size
171 +
172 + pad_height = (kernel_height * patch_size - height %
173 + (kernel_height * patch_size)) % (kernel_height *
174 + patch_size)
175 + pad_width = (kernel_width * patch_size - width %
176 + (kernel_width * patch_size)) % (kernel_width * patch_size)
177 +
178 + # Calculate new dimensions after padding and patching
179 + token_height = (height + pad_height) // (kernel_size[0] * patch_size)
180 + token_width = (width + pad_width) // (kernel_size[1] * patch_size)
181 + return int(token_height * token_width)
182 +
183 + def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
184 + # None means unlimited
185 + return {"image": None}
186 +
187 + def get_mm_max_tokens_per_item(
188 + self,
189 + seq_len: int,
190 + mm_counts: Mapping[str, int],
191 + ) -> Mapping[str, int]:
192 + return {
193 + "image":
194 + self.get_num_image_tokens(
195 + image_width=MaxImageTokenMeta.width,
196 + image_height=MaxImageTokenMeta.height,
197 + ),
198 + }
199 +
200 + @property
201 + def image_token_id(self) -> int:
202 + return self.get_hf_config().media_placeholder_token_id
203 +
204 +
205 +class KimiVLDummyInputsBuilder(BaseDummyInputsBuilder[KimiVLProcessingInfo]):
206 +
207 + def __init__(self, info: KimiVLProcessingInfo) -> None:
208 + super().__init__(info)
209 +
210 + self.image_token_id = self.info.image_token_id
211 + self.image_token = self.info.get_tokenizer().decode(
212 + self.image_token_id)
213 +
214 + def get_dummy_processor_inputs(
215 + self,
216 + seq_len: int,
217 + mm_counts: Mapping[str, int],
218 + ) -> ProcessorInputs:
219 + num_images = mm_counts.get("image", 0)
220 +
221 + width = MaxImageTokenMeta.width
222 + height = MaxImageTokenMeta.height
223 + mm_data = {
224 + "image":
225 + self._get_dummy_images(width=width,
226 + height=height,
227 + num_images=num_images)
228 + }
229 +
230 + return ProcessorInputs(
231 + prompt_text=self.image_token * num_images,
232 + mm_data=mm_data,
233 + )
234 +
235 +
236 +class KimiVLMultiModalProcessor(BaseMultiModalProcessor[KimiVLProcessingInfo]):
237 +
238 + def _get_mm_fields_config(
239 + self,
240 + hf_inputs: BatchFeature,
241 + hf_processor_mm_kwargs: Mapping[str, object],
242 + ) -> Mapping[str, MultiModalFieldConfig]:
243 + image_grid_hws = hf_inputs.get("image_grid_hws", torch.empty((0, 2)))
244 + image_grid_sizes = image_grid_hws.prod(-1)
245 +
246 + # pixel_values is merged as a single large tensor
247 + # image_grid_hws is shapes for each subtensor in pixel_values
248 + return dict(
249 + pixel_values=MultiModalFieldConfig.flat_from_sizes(
250 + "image", image_grid_sizes),
251 + image_grid_hws=MultiModalFieldConfig.batched("image"),
252 + )
253 +
254 + def _get_prompt_updates(
255 + self,
256 + mm_items: MultiModalDataItems,
257 + hf_processor_mm_kwargs: Mapping[str, Any],
258 + out_mm_kwargs: MultiModalKwargs,
259 + ) -> Sequence[PromptUpdate]:
260 + image_token_id = self.info.image_token_id
261 +
262 + def get_replacement(item_idx: int):
263 + images = mm_items.get_items(
264 + "image", (ImageEmbeddingItems, ImageProcessorItems))
265 +
266 + if isinstance(images, ImageEmbeddingItems):
267 + num_image_tokens = images.get_feature_size(item_idx)
268 + else:
269 + image_size = images.get_image_size(item_idx)
270 + num_image_tokens = self.info.get_num_image_tokens(
271 + image_width=image_size.width,
272 + image_height=image_size.height,
273 + )
274 +
275 + return [image_token_id] * num_image_tokens
276 +
277 + return [
278 + PromptReplacement(
279 + modality="image",
280 + target=[image_token_id],
281 + replacement=get_replacement,
282 + ),
283 + ]
284 +
285 +
286 +@MULTIMODAL_REGISTRY.register_processor(KimiVLMultiModalProcessor,
287 + info=KimiVLProcessingInfo,
288 + dummy_inputs=KimiVLDummyInputsBuilder)
289 +class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal):
290 +
291 + def __init__(
292 + self,
293 + vllm_config: VllmConfig,
294 + prefix: str = "",
295 + ) -> None:
296 + super().__init__()
297 + model_config = vllm_config.model_config
298 + config: KimiVLConfig = model_config.hf_config
299 + self.config = config
300 + quant_config = vllm_config.quant_config
301 +
302 + assert isinstance(config.vision_config, MoonViTConfig)
303 +
304 + self.vision_tower = MoonVitPretrainedModel(config.vision_config)
305 +
306 + self.multi_modal_projector = KimiVLMultiModalProjector(config=config)
307 +
308 + self.quant_config = quant_config
309 + sub_vllm_config = copy.deepcopy(vllm_config)
310 + sub_vllm_config.model_config.hf_config = sub_vllm_config.model_config.hf_config.text_config
311 + self.language_model = DeepseekV2Model(
312 + vllm_config=sub_vllm_config,
313 + prefix=maybe_prefix(prefix, "language_model"),
314 + )
315 + self.unpadded_vocab_size = config.text_config.vocab_size
316 + self.lm_head = ParallelLMHead(
317 + self.unpadded_vocab_size,
318 + config.text_config.hidden_size,
319 + org_num_embeddings=self.config.text_config.vocab_size,
320 + padding_size=DEFAULT_VOCAB_PADDING_SIZE)
321 + logit_scale = getattr(config, "logit_scale", 1.0)
322 + self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
323 + config.vocab_size, logit_scale)
324 + self.sampler = get_sampler()
325 + self.media_placeholder: int = self.config.media_placeholder_token_id
326 + self.tp_rank = get_tensor_model_parallel_rank()
327 + self.tp_world_size = get_tensor_model_parallel_world_size()
328 +
329 + # ref: qwen2_vl.py
330 + def _validate_and_reshape_mm_tensor(self, mm_input: object,
331 + name: str) -> torch.Tensor:
332 + if not isinstance(mm_input, (torch.Tensor, list)):
333 + raise ValueError(f"Incorrect type of {name}. "
334 + f"Got type: {type(mm_input)}")
335 + if isinstance(mm_input, torch.Tensor):
336 + if mm_input.ndim == 2:
337 + return mm_input
338 + if mm_input.ndim != 3:
339 + raise ValueError(f"{name} should be 2D or batched 3D tensor. "
340 + f"Got ndim: {mm_input.ndim} "
341 + f"(shape={mm_input.shape})")
342 + return mm_input.reshape(-1, mm_input.shape[-1])
343 + else:
344 + return torch.concat(mm_input)
345 +
346 + def _parse_and_validate_image_input(
347 + self, **kwargs: object) -> Optional[KimiVLImageInputs]:
348 + # image input type must be pixel values now
349 + pixel_values = kwargs.pop("pixel_values", None)
350 + image_grid_hws = kwargs.pop("image_grid_hws", None)
351 +
352 + if pixel_values is None:
353 + return None
354 +
355 + image_grid_hws = self._validate_and_reshape_mm_tensor(
356 + image_grid_hws, "image grid hws")
357 + # pixel_values may have complex shapes
358 + num_channels = 3
359 + patch_size = self.config.vision_config.patch_size
360 + if isinstance(pixel_values, list):
361 + pixel_values = torch.cat([
362 + x.reshape(-1, num_channels, patch_size, patch_size)
363 + for x in pixel_values
364 + ])
365 + else:
366 + pixel_values = pixel_values.reshape(-1, num_channels, patch_size,
367 + patch_size)
368 + # fp32 -> bf16
369 + pixel_values = pixel_values.to(torch.bfloat16)
370 + # image_grid_hws.shape = (N, 2)
371 + assert image_grid_hws.ndim == 2, f"unexpected shape for image_grid_hws: {image_grid_hws.shape}"
372 +
373 + return KimiVLImagePixelInputs(
374 + type="pixel_values",
375 + pixel_values=pixel_values,
376 + image_grid_hws=image_grid_hws,
377 + )
378 +
379 + # perform vt on processored pixel_values
380 + @torch.inference_mode()
381 + def _process_image_pixels(self,
382 + inputs: KimiVLImagePixelInputs) -> torch.Tensor:
383 + assert self.vision_tower is not None
384 +
385 + pixel_values = inputs["pixel_values"]
386 + image_grid_hws = inputs["image_grid_hws"]
387 + return self.vision_tower(pixel_values, image_grid_hws)
388 +
389 + def _process_image_input(self,
390 + image_input: KimiVLImageInputs) -> torch.Tensor:
391 + assert image_input["type"] == "pixel_values"
392 + image_features = self._process_image_pixels(image_input)
393 + assert isinstance(image_features, list)
394 + lengths = [x.shape[0] for x in image_features]
395 + return self.multi_modal_projector(
396 + torch.cat(image_features)).split(lengths)
397 +
398 + def get_multimodal_embeddings(self,
399 + **kwargs: object) -> Optional[NestedTensors]:
400 + # Validate the multimodal input keyword arguments
401 + image_input = self._parse_and_validate_image_input(**kwargs)
402 + if image_input is None:
403 + return None
404 +
405 + # Run multimodal inputs through encoder and projector
406 + vision_embeddings = self._process_image_input(image_input)
407 + return vision_embeddings
408 +
409 + def get_input_embeddings(
410 + self,
411 + input_ids: torch.Tensor,
412 + multimodal_embeddings: Optional[NestedTensors] = None,
413 + ) -> torch.Tensor:
414 +
415 + # `get_input_embeddings` should already be implemented for the language
416 + # model as one of the requirements of basic vLLM model implementation.
417 + inputs_embeds = self.language_model.get_input_embeddings(input_ids)
418 +
419 + if multimodal_embeddings is not None:
420 + inputs_embeds = merge_multimodal_embeddings(
421 + input_ids=input_ids,
422 + inputs_embeds=inputs_embeds,
423 + multimodal_embeddings=multimodal_embeddings,
424 + placeholder_token_id=self.config.media_placeholder_token_id)
425 +
426 + return inputs_embeds
427 +
428 + def forward(
429 + self,
430 + input_ids: torch.Tensor,
431 + positions: torch.Tensor,
432 + intermediate_tensors: Optional[IntermediateTensors] = None,
433 + inputs_embeds: Optional[torch.Tensor] = None,
434 + **kwargs: object,
435 + ) -> SamplerOutput:
436 + if intermediate_tensors is not None:
437 + inputs_embeds = None
438 + # NOTE: In v1, inputs_embeds is always generated at model runner from
439 + # `get_multimodal_embeddings` and `get_input_embeddings`, this
440 + # condition is only for v0 compatibility.
441 + elif inputs_embeds is None:
442 + image_input = self._parse_and_validate_image_input(**kwargs)
443 + if image_input is None:
444 + inputs_embeds = None
445 + else:
446 + inputs_embeds = self.get_input_embeddings(input_ids)
447 + image_embeds = self._process_image_input(image_input)
448 + inputs_embeds = merge_multimodal_embeddings(
449 + input_ids,
450 + inputs_embeds,
451 + image_embeds,
452 + placeholder_token_id=self.config.
453 + media_placeholder_token_id,
454 + )
455 + input_ids = None
456 +
457 + hidden_states = self.language_model(
458 + input_ids=input_ids,
459 + positions=positions,
460 + intermediate_tensors=intermediate_tensors,
461 + inputs_embeds=inputs_embeds,
462 + )
463 +
464 + return hidden_states
465 +
466 + def compute_logits(self, hidden_states: torch.Tensor,
467 + sampling_metadata: SamplingMetadata,
468 + **kwargs) -> torch.Tensor:
469 + logits = self.logits_processor(self.lm_head, hidden_states,
470 + sampling_metadata, **kwargs)
471 + return logits
472 +
473 + def sample(
474 + self,
475 + logits: torch.Tensor,
476 + sampling_metadata: SamplingMetadata,
477 + ) -> Optional[SamplerOutput]:
478 + next_tokens = self.sampler(logits, sampling_metadata)
479 + return next_tokens
480 +
481 + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
482 + config = self.config.text_config
483 + _KEYS_TO_MODIFY_MAPPING = {
484 + "language_model.lm_head": "lm_head",
485 + "language_model.model": "language_model",
486 + }
487 + # only doing this for language model part for now.
488 + stacked_params_mapping = [
489 + # (param_name, shard_name, shard_id)
490 + (".gate_up_proj", ".gate_proj", 0),
491 + (".gate_up_proj", ".up_proj", 1),
492 + ]
493 + if not config.use_mla:
494 + stacked_params_mapping += [
495 + (".qkv_proj", ".q_proj", "q"),
496 + (".qkv_proj", ".k_proj", "k"),
497 + (".qkv_proj", ".v_proj", "v"),
498 + ]
499 + if getattr(config, "n_routed_experts", None):
500 + # Params for weights, fp8 weight scales, fp8 activation scales
501 + # (param_name, weight_name, expert_id, shard_id)
502 + expert_params_mapping = FusedMoE.make_expert_params_mapping(
503 + ckpt_gate_proj_name="gate_proj",
504 + ckpt_down_proj_name="down_proj",
505 + ckpt_up_proj_name="up_proj",
506 + num_experts=config.n_routed_experts)
507 + else:
508 + expert_params_mapping = []
509 +
510 + params_dict = dict(self.named_parameters())
511 + for args in weights:
512 + name, loaded_weight = args[:2]
513 + kwargs = args[2] if len(args) > 2 else {}
514 + if "rotary_emb.inv_freq" in name:
515 + continue
516 +
517 + spec_layer = get_spec_layer_idx_from_weight_name(config, name)
518 + if spec_layer is not None:
519 + continue # skip spec decode layers for main model
520 +
521 + if ("rotary_emb.cos_cached" in name
522 + or "rotary_emb.sin_cached" in name):
523 + # Models trained using ColossalAI may include these tensors in
524 + # the checkpoint. Skip them.
525 + continue
526 + for key_to_modify, new_key in _KEYS_TO_MODIFY_MAPPING.items():
527 + if key_to_modify in name:
528 + name = name.replace(key_to_modify, new_key)
529 + use_default_weight_loading = False
530 + if "vision" in name:
531 + if self.vision_tower is not None:
532 + # We only do sharding for language model and
533 + # not vision model for now.
534 + use_default_weight_loading = True
535 + else:
536 + for (param_name, weight_name,
537 + shard_id) in stacked_params_mapping:
538 + if weight_name not in name:
539 + continue
540 + # We have mlp.experts[0].gate_proj in the checkpoint.
541 + # Since we handle the experts below in expert_params_mapping,
542 + # we need to skip here BEFORE we update the name, otherwise
543 + # name will be updated to mlp.experts[0].gate_up_proj, which
544 + # will then be updated below in expert_params_mapping
545 + # for mlp.experts[0].gate_gate_up_proj, which breaks load.
546 + if (("mlp.experts." in name) and name not in params_dict):
547 + continue
548 + name = name.replace(weight_name, param_name)
549 + # Skip loading extra bias for GPTQ models.
550 + if name.endswith(".bias") and name not in params_dict:
551 + continue
552 +
553 + if is_pp_missing_parameter(name, self):
554 + continue
555 +
556 + param = params_dict[name]
557 + weight_loader = param.weight_loader
558 + weight_loader(param, loaded_weight, shard_id, **kwargs)
559 + break
560 + else:
561 + for idx, (param_name, weight_name, expert_id,
562 + shard_id) in enumerate(expert_params_mapping):
563 + if weight_name not in name:
564 + continue
565 + name = name.replace(weight_name, param_name)
566 +
567 + if is_pp_missing_parameter(name, self):
568 + continue
569 +
570 + param = params_dict[name]
571 + weight_loader = param.weight_loader
572 + weight_loader(param,
573 + loaded_weight,
574 + name,
575 + expert_id=expert_id,
576 + shard_id=shard_id,
577 + **kwargs)
578 + break
579 + else:
580 + use_default_weight_loading = True
581 + if use_default_weight_loading:
582 + # Skip loading extra bias for GPTQ models.
583 + if name.endswith(".bias") and name not in params_dict:
584 + continue
585 + # Remapping the name of FP8 kv-scale.
586 + name = maybe_remap_kv_scale_name(name, params_dict)
587 + if name is None:
588 + continue
589 +
590 + if is_pp_missing_parameter(name, self):
591 + continue
592 +
593 + param = params_dict[name]
594 + weight_loader = getattr(param, "weight_loader",
595 + default_weight_loader)
596 + weight_loader(param, loaded_weight, **kwargs)
597 +
598 +
599 +def get_spec_layer_idx_from_weight_name(config: DeepseekV2Config,
600 + weight_name: str) -> Optional[int]:
601 + if hasattr(config,
602 + "num_nextn_predict_layers") and (config.num_nextn_predict_layers
603 + > 0):
604 + layer_idx = config.num_hidden_layers
605 + for i in range(config.num_nextn_predict_layers):
606 + if weight_name.startswith(f"model.layers.{layer_idx+i}."):
607 + return layer_idx + i
608 + return None
vllm/model_executor/models/moonvit.py new
+628
@@ -0,0 +1,628 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# ruff: noqa: E501
3 +# Adapted from https://huggingface.co/moonshotai/Kimi-VL-A3B-Instruct/blob/main/modeling_kimi_vl.py
4 +# This file is meant to be used in kimi_vl.py only
5 +# Copyright 2025 The Moonshot AI Team, DeepSeek-AI, and HuggingFace Inc. team. All rights reserved.
6 +#
7 +# The code is based on llava (llava/modeling_llava.py) and DeepSeek-V3 (DeepSeek-V3/modeling_deepseek.py), but modified for KimiVL.
8 +#
9 +# Licensing Information:
10 +# - Code derived from llava (llava/modeling_llava.py) and DeepSeek-V3 (DeepSeek-V3/modeling_deepseek.py) is licensed under the Apache License, Version 2.0.
11 +# - Other parts of the code are licensed under the MIT License.
12 +#
13 +# Apache License, Version 2.0:
14 +# Licensed under the Apache License, Version 2.0 (the "License");
15 +# you may not use this file except in compliance with the License.
16 +# You may obtain a copy of the License at
17 +#
18 +# http://www.apache.org/licenses/LICENSE-2.0
19 +#
20 +# Unless required by applicable law or agreed to in writing, software
21 +# distributed under the License is distributed on an "AS IS" BASIS,
22 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
23 +# See the License for the specific language governing permissions and
24 +# limitations under the License.
25 +#
26 +# MIT License:
27 +# Permission is hereby granted, free of charge, to any person obtaining a copy
28 +# of this software and associated documentation files (the "Software"), to deal
29 +# in the Software without restriction, including without limitation the rights
30 +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
31 +# copies of the Software, and to permit persons to whom the Software is
32 +# furnished to do so, subject to the following conditions:
33 +#
34 +# The above copyright notice and this permission notice shall be included in all
35 +# copies or substantial portions of the Software.
36 +#
37 +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
38 +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
39 +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
40 +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
41 +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
42 +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
43 +# SOFTWARE.
44 +import math
45 +from copy import deepcopy
46 +from functools import cached_property
47 +from typing import List, Optional, Sequence, Tuple, Union
48 +
49 +import torch
50 +import torch.nn as nn
51 +import torch.nn.functional as F
52 +from transformers.activations import ACT2FN, PytorchGELUTanh
53 +from transformers.modeling_utils import PreTrainedModel
54 +from transformers.utils import is_flash_attn_2_available
55 +
56 +from vllm.transformers_utils.configs.moonvit import MoonViTConfig
57 +
58 +if is_flash_attn_2_available():
59 + from flash_attn import flash_attn_varlen_func
60 +else:
61 + flash_attn_varlen_func = None
62 +
63 +
64 +def multihead_attention(
65 + q: torch.Tensor,
66 + k: torch.Tensor,
67 + v: torch.Tensor,
68 + q_cu_seqlens: Optional[torch.Tensor] = None,
69 + k_cu_seqlens: Optional[torch.Tensor] = None,
70 +):
71 + """Multi-head attention using flash attention 2.
72 +
73 + Args:
74 + q, k, v: tensor of shape (batch_size, seqlen, num_heads, head_dim),
75 + or (tot_seqlens, num_heads, head_dim) if packing.
76 + q_cu_seqlens (torch.Tensor): cumulative sequence lengths of q.
77 + The first element should be 0 and the last element should be q.shape[0].
78 + k_cu_seqlens (torch.Tensor): cumulative sequence lengths of k.
79 + The first element should be 0 and the last element should be k.shape[0].
80 +
81 + Returns:
82 + output: shape (batch_size, seqlen, dim) or (tot_seqlens, dim) if packing,
83 + where dim = num_heads * head_dim
84 + """
85 + # Unified format legal check
86 + assert q.dim() == k.dim() == v.dim() == 3, "q, k, v must have 3 dims"
87 + assert q_cu_seqlens[-1] == q.shape[
88 + 0], "q_cu_seqlens must sum to q.shape[0]"
89 + assert (k_cu_seqlens[-1] == k.shape[0] ==
90 + v.shape[0]), "k_cu_seqlens must sum to k.shape[0]"
91 + assert q.dtype in [
92 + torch.bfloat16,
93 + torch.float16,
94 + ], f"unsupported dtype {q.dtype} for multihead attn"
95 +
96 + max_seqlen_q = (q_cu_seqlens[1:] - q_cu_seqlens[:-1]).max().item()
97 + max_seqlen_k = (k_cu_seqlens[1:] - k_cu_seqlens[:-1]).max().item()
98 + attn_out = flash_attn_varlen_func(
99 + q,
100 + k,
101 + v,
102 + q_cu_seqlens,
103 + k_cu_seqlens,
104 + max_seqlen_q,
105 + max_seqlen_k,
106 + causal=False,
107 + )
108 + attn_out = attn_out.flatten(start_dim=-2)
109 +
110 + return attn_out
111 +
112 +
113 +def sdpa_attention(
114 + q: torch.Tensor,
115 + k: torch.Tensor,
116 + v: torch.Tensor,
117 + q_cu_seqlens: Optional[torch.Tensor] = None,
118 + k_cu_seqlens: Optional[torch.Tensor] = None,
119 +) -> torch.Tensor:
120 + """SDPA attention.
121 +
122 + Args:
123 + q, k, v: tensor of shape (batch_size, seqlen, num_heads, head_dim),
124 + or (tot_seqlens, num_heads, head_dim) if packing.
125 + """
126 + seq_length = q.shape[0]
127 + attention_mask = torch.zeros([1, seq_length, seq_length],
128 + device=q.device,
129 + dtype=torch.bool)
130 + for i in range(1, len(q_cu_seqlens)):
131 + attention_mask[
132 + ...,
133 + q_cu_seqlens[i - 1]:q_cu_seqlens[i],
134 + q_cu_seqlens[i - 1]:q_cu_seqlens[i],
135 + ] = True
136 + q = q.transpose(0, 1)
137 + k = k.transpose(0, 1)
138 + v = v.transpose(0, 1)
139 + attn_output = F.scaled_dot_product_attention(q,
140 + k,
141 + v,
142 + attention_mask,
143 + dropout_p=0.0)
144 + attn_output = attn_output.transpose(0, 1)
145 + attn_output = attn_output.reshape(seq_length, -1)
146 + return attn_output
147 +
148 +
149 +VL_VISION_ATTENTION_FUNCTIONS = {
150 + "flash_attention_2": multihead_attention,
151 + "sdpa": sdpa_attention,
152 +}
153 +
154 +
155 +def _apply_rope_input_validation(x, freqs_cis):
156 + assert x.ndim == freqs_cis.ndim + 1, (x.shape, freqs_cis.shape)
157 + assert x.shape[:-2] == freqs_cis.shape[:-1], (x.shape, freqs_cis.shape)
158 + assert x.shape[-1] == 2 * freqs_cis.shape[-1], (x.shape, freqs_cis.shape)
159 + assert freqs_cis.dtype == torch.complex64, freqs_cis.dtype
160 +
161 +
162 +def apply_rope(xq: torch.Tensor, xk: torch.Tensor,
163 + freqs_cis: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
164 + """
165 + Args: (The leading dimensions of all inputs should be the same)
166 + xq: query, tensor of shape (..., num_heads, head_dim)
167 + xk: key, tensor of shape (..., num_heads, head_dim)
168 + freqs_cis: tensor of shape (..., head_dim/2), dtype=torch.complex64. It contains the precomputed cis(freqs) for each position in the 2D grid.
169 + Returns:
170 + xq_out, xk_out: tensors of shape (..., num_heads, head_dim)
171 + """
172 + _apply_rope_input_validation(xq, freqs_cis)
173 + _apply_rope_input_validation(xk, freqs_cis)
174 +
175 + freqs_cis = freqs_cis.unsqueeze(-2) # ..., 1, head_dim/2
176 + # ..., num_heads, head_dim/2
177 + xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2))
178 + xk_ = torch.view_as_complex(xk.float().view(*xq.shape[:-1], -1, 2))
179 + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(
180 + -2) # ..., num_heads, head_dim
181 + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(
182 + -2) # ..., num_heads, head_dim
183 + return xq_out.type_as(xq), xk_out.type_as(xk)
184 +
185 +
186 +class Learnable2DInterpPosEmb(nn.Module):
187 +
188 + def __init__(self,
189 + height: int,
190 + width: int,
191 + dim: int,
192 + interpolation_mode: str = "bicubic") -> None:
193 + super().__init__()
194 + self.height = height
195 + self.width = width
196 + self.interpolation_mode = interpolation_mode
197 + self.weight = nn.Parameter(torch.empty(height, width, dim))
198 + self.reset_parameters()
199 +
200 + def reset_parameters(self):
201 + nn.init.normal_(self.weight)
202 +
203 + def forward(self, x: torch.Tensor, grid_hws: torch.Tensor) -> torch.Tensor:
204 + pos_embs = []
205 + for shape in grid_hws.tolist():
206 + if shape == self.weight.shape[:-1]:
207 + pos_embs.append(self.weight.flatten(end_dim=1))
208 + else:
209 + pos_embs.append(
210 + F.interpolate(
211 + self.weight.permute((2, 0, 1)).unsqueeze(0),
212 + size=shape,
213 + mode=self.interpolation_mode,
214 + ).squeeze(0).permute((1, 2, 0)).flatten(end_dim=1))
215 + out = x + torch.cat(pos_embs)
216 + return out
217 +
218 +
219 +class MoonVisionPatchEmbed(nn.Module):
220 +
221 + def __init__(
222 + self,
223 + out_dim: int,
224 + in_dim: int = 3,
225 + patch_size: Union[int, Tuple[int, int]] = (14, 14),
226 + pos_emb_height: int = 14,
227 + pos_emb_width: int = 14,
228 + ):
229 + super().__init__()
230 + assert isinstance(
231 + patch_size,
232 + (int, Sequence)), f"Invalid patch_size type: {type(patch_size)}"
233 + if isinstance(patch_size, int):
234 + patch_size = (patch_size, patch_size)
235 + assert (len(patch_size) == 2
236 + ), f"Expected patch_size to be a tuple of 2, got {patch_size}"
237 + self.patch_size = patch_size
238 +
239 + self.proj = nn.Conv2d(in_dim,
240 + out_dim,
241 + kernel_size=patch_size,
242 + stride=patch_size)
243 +
244 + self.pos_emb = Learnable2DInterpPosEmb(height=pos_emb_height,
245 + width=pos_emb_width,
246 + dim=out_dim)
247 +
248 + def forward(self, x: torch.Tensor, grid_hw: torch.Tensor) -> torch.Tensor:
249 + """
250 + Args:
251 + x (L, Channels): input tensor
252 + grid_hw (N, 2): grid height and width
253 +
254 + Returns:
255 + (L, Cout) tensor
256 + """
257 + x = self.proj(x).view(x.size(0), -1)
258 + # apply positional embedding
259 + x = self.pos_emb(x, grid_hw)
260 + return x
261 +
262 +
263 +class Rope2DPosEmb(nn.Module):
264 + """2D rotary position embedding with multi-resolution support.
265 +
266 + This class is intended to be used in the following way:
267 + 1. Before training, create an instance of Rope2DPosEmb. This instance will hold the precomputed cis.
268 + 2. Before each forward pass, call `get_freqs_cis_by_*` to get the `freqs_cis` tensor for this iteration.
269 + 3. During the forward pass, pass the `freqs_cis` tensor to each attention layer, and call `apply` just before each attention operation.
270 + The rope is shared across all attention layers and all heads.
271 +
272 + Refs:
273 + - RoFormer: https://arxiv.org/abs/2104.09864
274 + - VisionLLaMA: https://arxiv.org/abs/2403.00522
275 + - https://github.com/Meituan-AutoML/VisionLLaMA/blob/main/dit/models.py
276 +
277 + Args:
278 + dim (int): usually the multi-head attention dimension, should be divisible by 4 (TODO: relax this constraint if needed)
279 + max_height (int): the maximum height of the 2D grid
280 + max_width (int): the maximum width of the 2D grid
281 + theta_base (float): the base of the theta
282 + device (str): the device to store the precomputed cis
283 + """
284 +
285 + def __init__(self,
286 + dim: int,
287 + max_height: int,
288 + max_width: int,
289 + theta_base=10000,
290 + device="cuda"):
291 + super().__init__()
292 + self.dim = dim
293 + assert self.dim % 4 == 0, "dim must be divisible by 4"
294 + self.max_height = max_height
295 + self.max_width = max_width
296 + self.theta_base = theta_base
297 + self.device = device
298 +
299 + def extra_repr(self):
300 + return f"dim={self.dim}, max_height={self.max_height}, max_width={self.max_width}, theta_base={self.theta_base}"
301 +
302 + @cached_property
303 + def precomputed_freqs_cis(self) -> torch.Tensor:
304 + """Calculate the cis(freqs) for each position in the 2D grid.
305 +
306 + Return: complex tensor of shape (max_height, max_width, dim//2) and value:
307 + height axis: ret[h, w, 2*i] = cis(h * theta_base**(-4*i/dim))
308 + weight axis: ret[h, w, 2*i+1] = cis(w * theta_base**(-4*i/dim)) with (i in [0, dim//4))
309 + note: `cis` is a mathematical notation defined by cis x = cos x + i sin x,
310 + """
311 + N = self.max_height * self.max_width
312 + flat_pos = torch.arange(0, N).float().to(self.device)
313 + x_pos = flat_pos % self.max_width
314 + y_pos = flat_pos // self.max_width
315 + dim_range = (torch.arange(0, self.dim,
316 + 4)[:(self.dim // 4)].float().to(self.device)
317 + ) # C/4
318 + freqs = 1.0 / (self.theta_base**(dim_range / self.dim))
319 + x_freqs = torch.outer(x_pos, freqs).float() # N, C/4
320 + y_freqs = torch.outer(y_pos, freqs).float() # N, C/4
321 + x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs) # N, C/4
322 + y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs) # N, C/4
323 + # N, C/4, 2
324 + freqs_cis = torch.cat(
325 + [x_cis.unsqueeze(dim=-1),
326 + y_cis.unsqueeze(dim=-1)], dim=-1)
327 + # max_height, max_width, C/2
328 + freqs_cis = freqs_cis.reshape(self.max_height, self.max_width, -1)
329 + return freqs_cis
330 +
331 + def get_freqs_cis_by_seqlens(self, grid_hws: torch.Tensor) -> torch.Tensor:
332 + """
333 + Args:
334 + grid_hws (torch.Tensor): containing list of (height, width) or (t, height, width) tuples.
335 + Returns:
336 + freqs_cis: tensor of shape (sum(t * height * width), dim//2)
337 + """
338 + shapes = grid_hws.tolist()
339 + assert all(1 <= h <= self.max_height and 1 <= w <= self.max_width
340 + for h, w in shapes), (
341 + shapes,
342 + self.max_height,
343 + self.max_width,
344 + )
345 + freqs_cis = torch.cat(
346 + [
347 + self.precomputed_freqs_cis[:h, :w].reshape(-1, self.dim // 2)
348 + for h, w in shapes
349 + ],
350 + dim=0,
351 + )
352 + return freqs_cis
353 +
354 + def get_freqs_cis_by_idx(self, pos_idx: torch.Tensor,
355 + pos_idx_mask: torch.Tensor) -> torch.Tensor:
356 + """
357 + Args:
358 + pos_idx: tensor of shape (..., 2), It contains the (h, w) position indices of each 2D token.
359 + pos_idx_mask: a mask of shape (...), the leading dimensions should be the same as pos_idx.
360 + Rope will only be applied to the tokens with True mask. `freqs_cis` for the tokens with False mask with be ones.
361 + Return:
362 + freqs_cis: tensor of shape (..., dim//2)
363 + """
364 + assert (pos_idx.shape[:-1] == pos_idx_mask.shape
365 + and pos_idx.shape[-1] == 2 and pos_idx.ndim
366 + == pos_idx_mask.ndim + 1), (pos_idx.shape, pos_idx_mask.shape)
367 + assert pos_idx_mask.dtype == torch.bool, pos_idx_mask.dtype
368 +
369 + shp = pos_idx_mask.shape + (self.dim // 2, ) # ..., head_dim/2
370 + freqs_cis = torch.ones(shp, dtype=torch.complex64,
371 + device=self.device) # ..., head_dim/2
372 + freqs_cis[pos_idx_mask] = self.precomputed_freqs_cis[pos_idx[
373 + ..., 0][pos_idx_mask], pos_idx[..., 1][pos_idx_mask]]
374 + return freqs_cis
375 +
376 +
377 +class MLP2(nn.Module):
378 + """
379 + Args:
380 + dims: [in_dim, hidden_dim, out_dim]
381 + bias: whether to use bias in linear layer.
382 + """
383 +
384 + def __init__(self, dims: list[int], activation, bias=True):
385 + super().__init__()
386 + assert len(dims) == 3
387 + self.fc0 = nn.Linear(dims[0], dims[1], bias=bias)
388 + self.fc1 = nn.Linear(dims[1], dims[2], bias=bias)
389 + self.activation = activation
390 + for m in [self.fc0, self.fc1]:
391 + nn.init.trunc_normal_(m.weight, std=math.sqrt(2 / m.in_features))
392 + if m.bias is not None:
393 + nn.init.zeros_(m.bias)
394 +
395 + def forward(self, x: torch.Tensor) -> torch.Tensor:
396 + x = self.fc0(x)
397 + x = self.activation(x)
398 + return self.fc1(x)
399 +
400 +
401 +class MoonVitEncoderLayer(nn.Module):
402 +
403 + def __init__(
404 + self,
405 + num_heads: int,
406 + hidden_dim: int,
407 + mlp_dim: int,
408 + *,
409 + attn_implementation: str = "sdpa",
410 + activation=F.gelu,
411 + attn_bias: bool = False,
412 + ):
413 + super().__init__()
414 + self.num_heads = num_heads
415 + self.hidden_dim = hidden_dim
416 + self.hidden_size_per_attention_head = self.hidden_dim // self.num_heads
417 + self.attn_implementation = attn_implementation
418 + # use fa2 in vllm by default
419 + if is_flash_attn_2_available():
420 + self.attn_implementation = "flash_attention_2"
421 +
422 + self.norm0 = nn.LayerNorm(hidden_dim)
423 + self.norm1 = nn.LayerNorm(hidden_dim)
424 + self.mlp = MLP2([hidden_dim, mlp_dim, hidden_dim], activation)
425 + self.wqkv = nn.Linear(hidden_dim, hidden_dim * 3, bias=attn_bias)
426 + self.wo = nn.Linear(hidden_dim, hidden_dim, bias=attn_bias)
427 +
428 + def attention_qkvpacked(
429 + self,
430 + x: torch.Tensor,
431 + cu_seqlens: torch.Tensor,
432 + rope_freqs_cis: Optional[torch.Tensor] = None,
433 + ):
434 + """
435 + Args:
436 + x (torch.Tensor): (batch_size, seqlen, hidden_dim)
437 + cu_seqlens (torch.Tensor):
438 + """
439 + xqkv = self.wqkv(x)
440 +
441 + qkv_shape = xqkv.size()[:-1] + (
442 + 3,
443 + self.num_heads,
444 + self.hidden_size_per_attention_head,
445 + )
446 + # xqkv: (batch_size, seqlen, 3, nheads, headdim)
447 + xqkv = xqkv.view(*qkv_shape)
448 + xq, xk, xv = torch.unbind(xqkv, dim=-3)
449 +
450 + xq, xk = apply_rope(xq, xk, rope_freqs_cis)
451 +
452 + attn_func = VL_VISION_ATTENTION_FUNCTIONS[self.attn_implementation]
453 + attn_out = attn_func(xq,
454 + xk,
455 + xv,
456 + q_cu_seqlens=cu_seqlens,
457 + k_cu_seqlens=cu_seqlens)
458 +
459 + attn_out = self.wo(attn_out)
460 + return attn_out
461 +
462 + def forward(
463 + self,
464 + hidden_states: torch.Tensor,
465 + cu_seqlens: torch.Tensor,
466 + rope_freqs_cis: Union[torch.Tensor, None] = None,
467 + ) -> torch.Tensor:
468 + """
469 + Args:
470 + hidden_states: non-packed (B, N, D) or packed (L, D). if non-packed, seqlens should be None, if packed, seqlens should be set
471 +
472 + Returns:
473 + output: same shape of input, non-packed (B, N, D) for non-packed input, (L, D) for packed input
474 + """
475 + residual = hidden_states
476 + hidden_states = self.norm0(hidden_states)
477 + attn_out = self.attention_qkvpacked(hidden_states,
478 + cu_seqlens,
479 + rope_freqs_cis=rope_freqs_cis)
480 + hidden_states = residual + attn_out
481 +
482 + residual = hidden_states
483 + hidden_states = self.mlp(self.norm1(hidden_states))
484 + hidden_states = residual + hidden_states
485 + return hidden_states
486 +
487 +
488 +class MoonVitEncoder(nn.Module):
489 +
490 + def __init__(
491 + self,
492 + hidden_dim: int,
493 + num_layers: int,
494 + block_cfg: dict,
495 + ) -> None:
496 + super().__init__()
497 +
498 + self.rope_2d = Rope2DPosEmb(
499 + block_cfg["hidden_dim"] // block_cfg["num_heads"], 512, 512)
500 + self.blocks = nn.ModuleList(
501 + [MoonVitEncoderLayer(**block_cfg) for _ in range(num_layers)])
502 + self.final_layernorm = nn.LayerNorm(hidden_dim)
503 +
504 + def forward(self, hidden_states: torch.Tensor,
505 + grid_hw: torch.Tensor) -> torch.Tensor:
506 + rope_freqs_cis = self.rope_2d.get_freqs_cis_by_seqlens(
507 + grid_hws=grid_hw)
508 +
509 + lengths = torch.cat((
510 + torch.zeros(1, device=hidden_states.device, dtype=grid_hw.dtype),
511 + grid_hw[:, 0] * grid_hw[:, 1],
512 + ))
513 + cu_seqlens = lengths.cumsum(dim=0, dtype=torch.int32)
514 +
515 + for _, block in enumerate(self.blocks):
516 + hidden_states = block(hidden_states,
517 + cu_seqlens,
518 + rope_freqs_cis=rope_freqs_cis)
519 +
520 + hidden_states = self.final_layernorm(hidden_states)
521 +
522 + return hidden_states
523 +
524 +
525 +def patch_merger(
526 + x: torch.Tensor,
527 + grid_hw: torch.Tensor,
528 + merge_kernel_size: list[int, int] = (2, 2),
529 +) -> List[torch.Tensor]:
530 + d_model = x.size(-1)
531 +
532 + outputs = []
533 + pre_sum = 0
534 + for x_shape in grid_hw.tolist():
535 + height, width = x_shape[0], x_shape[1]
536 + # Get the current sequence
537 + seq = x[pre_sum:pre_sum + height * width]
538 + # Reshape along self.merge_kernel_size and concat to the last dimension
539 + kernel_height, kernel_width = merge_kernel_size
540 + new_height, new_width = height // kernel_height, width // kernel_width
541 + reshaped_seq = seq.view(new_height, kernel_height, new_width,
542 + kernel_width, d_model)
543 + reshaped_seq = reshaped_seq.permute(0, 2, 1, 3, 4).contiguous()
544 + padded_seq = reshaped_seq.view(new_height * new_width,
545 + kernel_height * kernel_width, -1)
546 + outputs.append(padded_seq)
547 + pre_sum += height * width
548 +
549 + return outputs
550 +
551 +
552 +class MoonVitVLProjector(nn.Module):
553 +
554 + def __init__(
555 + self,
556 + in_channels: int,
557 + merge_kernel_size: list[int, int],
558 + hidden_act: str = "gelu",
559 + ln_eps: float = 1e-5,
560 + out_dim: int = 4096,
561 + ):
562 + super().__init__()
563 + self.hidden_size = in_channels * merge_kernel_size[
564 + 0] * merge_kernel_size[1]
565 +
566 + self.pre_norm = nn.nn.LayerNorm(in_channels, eps=ln_eps)
567 + self.linear_1 = nn.Linear(self.hidden_size,
568 + self.hidden_size,
569 + bias=True)
570 + self.act = ACT2FN[hidden_act]
571 + self.linear_2 = nn.Linear(self.hidden_size, out_dim, bias=True)
572 +
573 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
574 + hidden_states = self.pre_norm(hidden_states).view(-1, self.hidden_size)
575 + hidden_states = self.linear_1(hidden_states)
576 + hidden_states = self.act(hidden_states)
577 + hidden_states = self.linear_2(hidden_states)
578 + return hidden_states
579 +
580 +
581 +class MoonVitPretrainedModel(PreTrainedModel):
582 + config_class = MoonViTConfig
583 + model_type = "moonvit"
584 + _no_split_modules = ["PackingTransformer"]
585 + _supports_flash_attn_2 = True
586 + _supports_sdpa = True
587 +
588 + def __init__(self, config: MoonViTConfig, *inputs, **kwargs):
589 + super().__init__(config, *inputs, **kwargs)
590 + config = deepcopy(config)
591 + self.merge_kernel_size = config.merge_kernel_size
592 + self.patch_size = config.patch_size
593 + self.patch_embed = MoonVisionPatchEmbed(
594 + out_dim=config.hidden_size,
595 + patch_size=config.patch_size,
596 + pos_emb_height=config.init_pos_emb_height,
597 + pos_emb_width=config.init_pos_emb_width,
598 + )
599 +
600 + self.encoder = MoonVitEncoder(
601 + hidden_dim=config.hidden_size,
602 + num_layers=config.num_hidden_layers,
603 + block_cfg={
604 + "num_heads": config.num_attention_heads,
605 + "hidden_dim": config.hidden_size,
606 + "mlp_dim": config.intermediate_size,
607 + "activation": PytorchGELUTanh(),
608 + "attn_bias": True,
609 + "attn_implementation": config._attn_implementation,
610 + },
611 + )
612 +
613 + def forward(self, pixel_values: torch.Tensor,
614 + grid_hw: torch.Tensor) -> torch.Tensor:
615 + """
616 + Args:
617 + pixel_values (torch.Tensor): The input pixel values.
618 + grid_hw (torch.Tensor): The grid height and width.
619 +
620 + Returns:
621 + torch.Tensor: The output tokens.
622 + """
623 + hidden_states = self.patch_embed(pixel_values, grid_hw)
624 + hidden_states = self.encoder(hidden_states, grid_hw)
625 + hidden_states = patch_merger(hidden_states,
626 + grid_hw,
627 + merge_kernel_size=self.merge_kernel_size)
628 + return hidden_states
vllm/model_executor/models/registry.py
+1
@@ -177,6 +177,7 @@ _MULTIMODAL_MODELS = {
177 "InternVLChatModel": ("internvl", "InternVLChatModel"),
178 "Idefics3ForConditionalGeneration":("idefics3","Idefics3ForConditionalGeneration"),
179 "SmolVLMForConditionalGeneration": ("smolvlm","SmolVLMForConditionalGeneration"), # noqa: E501
180 + "KimiVLForConditionalGeneration": ("kimi_vl", "KimiVLForConditionalGeneration"), # noqa: E501
181 "LlavaForConditionalGeneration": ("llava", "LlavaForConditionalGeneration"),
182 "LlavaNextForConditionalGeneration": ("llava_next", "LlavaNextForConditionalGeneration"), # noqa: E501
183 "LlavaNextVideoForConditionalGeneration": ("llava_next_video", "LlavaNextVideoForConditionalGeneration"), # noqa: E501
vllm/transformers_utils/config.py
+8 -6
@@ -33,12 +33,13 @@ from vllm.transformers_utils.configs import (ChatGLMConfig, Cohere2Config,
33 EAGLEConfig, ExaoneConfig,
34 H2OVLChatConfig,
35 InternVLChatConfig, JAISConfig,
36 - MedusaConfig, MllamaConfig,
37 - MLPSpeculatorConfig, MPTConfig,
38 - NemotronConfig, NVLM_D_Config,
39 - Olmo2Config, RWConfig,
40 - SkyworkR1VChatConfig, SolarConfig,
41 - Telechat2Config, UltravoxConfig)
36 + KimiVLConfig, MedusaConfig,
37 + MllamaConfig, MLPSpeculatorConfig,
38 + MPTConfig, NemotronConfig,
39 + NVLM_D_Config, Olmo2Config,
40 + RWConfig, SkyworkR1VChatConfig,
41 + SolarConfig, Telechat2Config,
42 + UltravoxConfig)
43 # yapf: enable
44 from vllm.transformers_utils.utils import check_gguf_file
45 from vllm.utils import resolve_obj_by_qualname
@@ -62,6 +63,7 @@ _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
63 "cohere2": Cohere2Config,
64 "dbrx": DbrxConfig,
65 "deepseek_vl_v2": DeepseekVLV2Config,
66 + "kimi_vl": KimiVLConfig,
67 "mpt": MPTConfig,
68 "RefinedWeb": RWConfig, # For tiiuae/falcon-40b(-instruct)
69 "RefinedWebModel": RWConfig, # For tiiuae/falcon-7b(-instruct)
vllm/transformers_utils/configs/__init__.py
+4
@@ -13,9 +13,11 @@ from vllm.transformers_utils.configs.falcon import RWConfig
13 from vllm.transformers_utils.configs.h2ovl import H2OVLChatConfig
14 from vllm.transformers_utils.configs.internvl import InternVLChatConfig
15 from vllm.transformers_utils.configs.jais import JAISConfig
16 +from vllm.transformers_utils.configs.kimi_vl import KimiVLConfig
17 from vllm.transformers_utils.configs.medusa import MedusaConfig
18 from vllm.transformers_utils.configs.mllama import MllamaConfig
19 from vllm.transformers_utils.configs.mlp_speculator import MLPSpeculatorConfig
20 +from vllm.transformers_utils.configs.moonvit import MoonViTConfig
21 from vllm.transformers_utils.configs.mpt import MPTConfig
22 from vllm.transformers_utils.configs.nemotron import NemotronConfig
23 from vllm.transformers_utils.configs.nvlm_d import NVLM_D_Config
@@ -40,6 +42,8 @@ __all__ = [
42 "ExaoneConfig",
43 "MllamaConfig",
44 "MLPSpeculatorConfig",
45 + "MoonViTConfig",
46 + "KimiVLConfig",
47 "NemotronConfig",
48 "NVLM_D_Config",
49 "Olmo2Config",
vllm/transformers_utils/configs/kimi_vl.py new
+36
@@ -0,0 +1,36 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# Adapted from https://huggingface.co/moonshotai/Kimi-VL-A3B-Instruct/blob/main/configuration_kimi_vl.py
3 +from typing import Optional, Union
4 +
5 +from transformers.configuration_utils import PretrainedConfig
6 +
7 +from vllm.transformers_utils.configs.deepseek_vl2 import DeepseekV2Config
8 +from vllm.transformers_utils.configs.moonvit import MoonViTConfig
9 +
10 +
11 +class KimiVLConfig(PretrainedConfig):
12 + model_type = "kimi_vl"
13 +
14 + def __init__(self,
15 + vision_config: Optional[Union[dict, MoonViTConfig]] = None,
16 + text_config: Optional[Union[dict, DeepseekV2Config]] = None,
17 + ignore_index: int = -100,
18 + media_placeholder_token_id: int = 163605,
19 + pad_token_id: int = 0,
20 + **kwargs):
21 + if vision_config is None:
22 + vision_config = MoonViTConfig()
23 + elif isinstance(vision_config, dict):
24 + vision_config = MoonViTConfig(**vision_config)
25 + self.vision_config = vision_config
26 +
27 + if text_config is None:
28 + text_config = DeepseekV2Config()
29 + elif isinstance(text_config, dict):
30 + text_config = DeepseekV2Config(**text_config)
31 + self.text_config = text_config
32 +
33 + self.ignore_index = ignore_index
34 + self.media_placeholder_token_id = media_placeholder_token_id
35 +
36 + super().__init__(pad_token_id=pad_token_id, **kwargs)
vllm/transformers_utils/configs/moonvit.py new
+32
@@ -0,0 +1,32 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# Adapted from https://huggingface.co/moonshotai/Kimi-VL-A3B-Instruct/blob/main/configuration_kimi_vl.py
3 +from transformers.configuration_utils import PretrainedConfig
4 +
5 +
6 +class MoonViTConfig(PretrainedConfig):
7 + model_type = "moonvit"
8 +
9 + def __init__(
10 + self,
11 + patch_size: int = 14,
12 + init_pos_emb_height: int = 64,
13 + init_pos_emb_width: int = 64,
14 + num_attention_heads: int = 16,
15 + num_hidden_layers: int = 27,
16 + hidden_size: int = 1152,
17 + intermediate_size: int = 4304,
18 + merge_kernel_size: tuple[int, int] = (2, 2),
19 + **kwargs,
20 + ):
21 + super().__init__(**kwargs)
22 + self.patch_size = patch_size
23 + # Positional embedding config
24 + self.init_pos_emb_height = init_pos_emb_height
25 + self.init_pos_emb_width = init_pos_emb_width
26 + # Transformer config
27 + self.num_hidden_layers = num_hidden_layers
28 + self.num_attention_heads = num_attention_heads
29 + self.hidden_size = hidden_size
30 + self.intermediate_size = intermediate_size
31 + # Patch merger config
32 + self.merge_kernel_size = merge_kernel_size
vllm/v1/worker/gpu_model_runner.py
+9 -7
@@ -997,13 +997,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
997 # Return empty ModelRunnerOutput if there's no work to do.
998 return EMPTY_MODEL_RUNNER_OUTPUT
999
1000 - if self.is_multimodal_model:
1001 - # Run the multimodal encoder if any.
1002 - self._execute_mm_encoder(scheduler_output)
1003 - mm_embeds = self._gather_mm_embeddings(scheduler_output)
1004 - else:
1005 - mm_embeds = []
1006 -
1000 # Prepare the decoder inputs.
1001 attn_metadata, logits_indices, spec_decode_metadata = (
1002 self._prepare_inputs(scheduler_output))
@@ -1019,6 +1012,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1012 num_input_tokens = num_scheduled_tokens
1013 attn_metadata.num_input_tokens = num_input_tokens
1014
1015 + # _prepare_inputs may reorder the batch, so we must gather multi
1016 + # modal outputs after that to ensure the correct order
1017 + if self.is_multimodal_model:
1018 + # Run the multimodal encoder if any.
1019 + self._execute_mm_encoder(scheduler_output)
1020 + mm_embeds = self._gather_mm_embeddings(scheduler_output)
1021 + else:
1022 + mm_embeds = []
1023 +
1024 if self.is_multimodal_model:
1025 # NOTE(woosuk): To unify token ids and soft tokens (vision
1026 # embeddings), we always use embeddings (rather than token ids)