@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6bc9c8473e
Fork 0

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

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

Synced Nov 10, 06:30 UTC

[MODEL] New model support for kakaocorp/kanana-1.5-v-3b-instruct (#29384)

Signed-off-by: Jaehyun An <steve.ai@kakaocorp.com>

Jaehyun An committed Jan 13, 2026 at 01:39 UTC 6bc9c8473e23d908d5ca1ae6a305704135597b42
5 files changed +790
docs/models/supported_models.md
+1
@@ -680,6 +680,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
680 | `InternS1ForConditionalGeneration` | Intern-S1 | T + I<sup>E+</sup> + V<sup>E+</sup> | `internlm/Intern-S1`, `internlm/Intern-S1-mini`, etc. | ✅︎ | ✅︎ |
681 | `InternVLChatModel` | InternVL 3.5, InternVL 3.0, InternVideo 2.5, InternVL 2.5, Mono-InternVL, InternVL 2.0 | T + I<sup>E+</sup> + (V<sup>E+</sup>) | `OpenGVLab/InternVL3_5-14B`, `OpenGVLab/InternVL3-9B`, `OpenGVLab/InternVideo2_5_Chat_8B`, `OpenGVLab/InternVL2_5-4B`, `OpenGVLab/Mono-InternVL-2B`, `OpenGVLab/InternVL2-4B`, etc. | ✅︎ | ✅︎ |
682 | `InternVLForConditionalGeneration` | InternVL 3.0 (HF format) | T + I<sup>E+</sup> + V<sup>E+</sup> | `OpenGVLab/InternVL3-1B-hf`, etc. | ✅︎ | ✅︎ |
683 +| `KananaVForConditionalGeneration` | Kanana-V | T + I<sup>+</sup> | `kakaocorp/kanana-1.5-v-3b-instruct`, etc. | | ✅︎ |
684 | `KeyeForConditionalGeneration` | Keye-VL-8B-Preview | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-8B-Preview` | ✅︎ | ✅︎ |
685 | `KeyeVL1_5ForConditionalGeneration` | Keye-VL-1_5-8B | T + I<sup>E+</sup> + V<sup>E+</sup> | `Kwai-Keye/Keye-VL-1_5-8B` | ✅︎ | ✅︎ |
686 | `KimiVLForConditionalGeneration` | Kimi-VL-A3B-Instruct, Kimi-VL-A3B-Thinking | T + I<sup>+</sup> | `moonshotai/Kimi-VL-A3B-Instruct`, `moonshotai/Kimi-VL-A3B-Thinking` | | ✅︎ |
examples/offline_inference/vision_language.py
+28
@@ -769,6 +769,33 @@ def run_internvl(questions: list[str], modality: str) -> ModelRequestData:
769 )
770
771
772 +# Kanana-V
773 +def run_kanana_v(questions: list[str], modality: str) -> ModelRequestData:
774 + assert modality == "image"
775 +
776 + model_name = "kakaocorp/kanana-1.5-v-3b-instruct"
777 +
778 + engine_args = EngineArgs(
779 + model=model_name,
780 + max_model_len=8192,
781 + trust_remote_code=True,
782 + limit_mm_per_prompt={modality: 1},
783 + )
784 +
785 + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
786 + messages = [
787 + [{"role": "user", "content": f"<image>\n{question}"}] for question in questions
788 + ]
789 + prompts = tokenizer.apply_chat_template(
790 + messages, tokenize=False, add_generation_prompt=True
791 + )
792 +
793 + return ModelRequestData(
794 + engine_args=engine_args,
795 + prompts=prompts,
796 + )
797 +
798 +
799 # Keye-VL
800 def run_keye_vl(questions: list[str], modality: str) -> ModelRequestData:
801 model_name = "Kwai-Keye/Keye-VL-8B-Preview"
@@ -1876,6 +1903,7 @@ model_example_map = {
1903 "idefics3": run_idefics3,
1904 "interns1": run_interns1,
1905 "internvl_chat": run_internvl,
1906 + "kanana_v": run_kanana_v,
1907 "keye_vl": run_keye_vl,
1908 "keye_vl1_5": run_keye_vl1_5,
1909 "kimi_vl": run_kimi_vl,
tests/models/registry.py
+4
@@ -728,6 +728,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {
728 trust_remote_code=True,
729 ),
730 "InternVLForConditionalGeneration": _HfExamplesInfo("OpenGVLab/InternVL3-1B-hf"),
731 + "KananaVForConditionalGeneration": _HfExamplesInfo(
732 + "kakaocorp/kanana-1.5-v-3b-instruct",
733 + trust_remote_code=True,
734 + ),
735 "KeyeForConditionalGeneration": _HfExamplesInfo(
736 "Kwai-Keye/Keye-VL-8B-Preview",
737 trust_remote_code=True,
vllm/model_executor/models/kanana_v.py new
+756
@@ -0,0 +1,756 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Iterable, Mapping, Sequence
4 +from functools import partial
5 +from typing import Annotated, Literal, TypeAlias
6 +
7 +import numpy as np
8 +import regex as re
9 +import torch
10 +from einops import rearrange
11 +from PIL import Image
12 +from timm.layers import LayerNorm2d
13 +from timm.layers.pos_embed import resample_abs_pos_embed
14 +from timm.models.regnet import RegStage
15 +from torch import nn
16 +from transformers import BatchFeature
17 +from transformers.modeling_outputs import BaseModelOutput
18 +from transformers.models.qwen2_vl.configuration_qwen2_vl import Qwen2VLVisionConfig
19 +
20 +from vllm.config import VllmConfig
21 +from vllm.config.multimodal import BaseDummyOptions
22 +from vllm.logger import init_logger
23 +from vllm.multimodal import MULTIMODAL_REGISTRY
24 +from vllm.multimodal.inputs import (
25 + MultiModalDataDict,
26 + MultiModalFieldConfig,
27 + MultiModalKwargsItems,
28 +)
29 +from vllm.multimodal.parse import ImageSize, MultiModalDataItems
30 +from vllm.multimodal.processing import (
31 + BaseMultiModalProcessor,
32 + BaseProcessingInfo,
33 + PromptReplacement,
34 + PromptUpdate,
35 +)
36 +from vllm.multimodal.profiling import BaseDummyInputsBuilder
37 +from vllm.sequence import IntermediateTensors
38 +from vllm.utils.import_utils import resolve_obj_by_qualname
39 +from vllm.utils.tensor_schema import TensorSchema, TensorShape
40 +
41 +from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
42 +from .qwen2_vl import Qwen2VisionTransformer
43 +from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
44 +
45 +logger = init_logger(__name__)
46 +
47 +
48 +class KananaVImagePixelInputs(TensorSchema):
49 + """
50 + Dimensions:
51 + - np: The total number of patches over all images in the batch
52 + - cps: Number of channels * patch_size * patch_size
53 + - ni: Number of images
54 + """
55 +
56 + type: Literal["pixel_values"]
57 +
58 + pixel_values: Annotated[
59 + torch.Tensor,
60 + TensorShape("np", "cps"),
61 + ]
62 +
63 + vision_grid_thw: Annotated[
64 + torch.Tensor,
65 + TensorShape("ni", 3),
66 + ]
67 +
68 +
69 +KananaVImageInputs: TypeAlias = KananaVImagePixelInputs
70 +
71 +
72 +def build_pos_embeds(
73 + config: Qwen2VLVisionConfig,
74 + num_input_tokens: int,
75 + vision_hidden_size: int,
76 +) -> nn.Parameter | None:
77 + """Build positional embeddings for the visual encoder output."""
78 + if config.pos_emb:
79 + pos_emb = nn.Parameter(torch.zeros(1, num_input_tokens, vision_hidden_size))
80 + nn.init.trunc_normal_(pos_emb, mean=0.0, std=0.02)
81 + else:
82 + pos_emb = None
83 +
84 + return pos_emb
85 +
86 +
87 +def build_mlp(
88 + depth: int,
89 + hidden_size: int,
90 + output_hidden_size: int,
91 +) -> nn.Sequential:
92 + """Simple SiLU-activated MLP used as a projector readout."""
93 + layers = [nn.Linear(hidden_size, output_hidden_size)]
94 + for _ in range(1, depth):
95 + layers.append(nn.SiLU())
96 + layers.append(nn.Linear(output_hidden_size, output_hidden_size))
97 + return nn.Sequential(*layers)
98 +
99 +
100 +class PatchMerge(nn.Module):
101 + """Merge neighboring patches spatially to reduce resolution."""
102 +
103 + def __init__(self, merge_size: int) -> None:
104 + super().__init__()
105 + self.merge_size = merge_size
106 +
107 + def forward(
108 + self,
109 + x: torch.Tensor,
110 + channel_last: bool = False,
111 + ) -> torch.Tensor:
112 + """Merge patches by `merge_size x merge_size`."""
113 + if channel_last:
114 + x = rearrange(x, "B H W D -> B D H W")
115 + _, _, H, W = x.shape
116 + merged_x = rearrange(
117 + x,
118 + "B D (H h2) (W w2) -> B (D h2 w2) H W",
119 + h2=self.merge_size,
120 + w2=self.merge_size,
121 + )
122 + return merged_x
123 +
124 +
125 +class DynamicCAbstractor(nn.Module):
126 + """Dynamic C-Abstractor based on RegNet blocks."""
127 +
128 + def __init__(
129 + self,
130 + config: Qwen2VLVisionConfig,
131 + num_input_tokens: int,
132 + ) -> None:
133 + super().__init__()
134 + assert hasattr(config, "merge_size"), "merge_size must be provided."
135 + self.config = config
136 + self.merge_size = config.merge_size
137 + self.pos_emb_size = config.pos_emb_size
138 + if num_input_tokens == -1:
139 + num_input_tokens = config.pos_emb_size
140 + self.num_input_tokens = num_input_tokens
141 + self.pos_emb = build_pos_embeds(
142 + config, num_input_tokens, config.encoder_hidden_size
143 + )
144 + self.build_net()
145 +
146 + def _load_from_state_dict(self, state_dict, *args, **kwargs) -> None:
147 + if not state_dict:
148 + return
149 +
150 + if self.pos_emb is not None:
151 + key_re = re.compile(r"[\w,.]*abstractor[\w,.]*pos_emb")
152 + pos_emb_key = None
153 + for key in state_dict:
154 + if key_re.match(key):
155 + pos_emb_key = key
156 + break
157 +
158 + assert pos_emb_key is not None
159 + # update old ckpt compatible with current code
160 + pos_emb = state_dict[pos_emb_key]
161 + if pos_emb.size(1) == self.pos_emb.size(1) + 1:
162 + # remove obsolete first pos emb (for cls token originally)
163 + state_dict[pos_emb_key] = pos_emb[:, 1:]
164 +
165 + super()._load_from_state_dict(state_dict, *args, **kwargs)
166 +
167 + def build_net(self) -> None:
168 + encoder_hidden_size = self.config.encoder_hidden_size
169 + hidden_size = self.config.hidden_size
170 + output_hidden_size = self.config.output_hidden_size
171 + depth = self.config.depth
172 + mlp_depth = self.config.mlp_depth
173 +
174 + RegBlock = partial(
175 + RegStage,
176 + stride=1,
177 + dilation=1,
178 + act_layer=nn.SiLU,
179 + norm_layer=LayerNorm2d,
180 + )
181 +
182 + s1 = RegBlock(
183 + depth,
184 + encoder_hidden_size,
185 + hidden_size,
186 + )
187 + sampler = PatchMerge(merge_size=self.merge_size)
188 + s2 = RegBlock(
189 + depth,
190 + self.merge_size**2 * hidden_size,
191 + hidden_size,
192 + )
193 +
194 + if depth:
195 + self.net = nn.ModuleList([s1, sampler, s2])
196 + self.readout = build_mlp(mlp_depth, hidden_size, output_hidden_size)
197 + else:
198 + self.net = sampler
199 + self.readout = build_mlp(mlp_depth, encoder_hidden_size, output_hidden_size)
200 +
201 + def forward(
202 + self,
203 + flattened_visual_embeds: torch.Tensor,
204 + grid_thw: torch.Tensor,
205 + **unused_kwargs: object,
206 + ) -> BaseModelOutput:
207 + """Apply the dynamic abstractor over flattened visual embeddings."""
208 + n_token_loc = torch.prod(grid_thw, dim=1)
209 + split_visual_embeds = torch.split(flattened_visual_embeds, n_token_loc.tolist())
210 +
211 + flattened_visual_embeds = []
212 + for _visual_embeds, _grid_thw in zip(split_visual_embeds, grid_thw):
213 + T, H, W = _grid_thw
214 + assert T == 1, "T must be 1. Video is not supported yet."
215 + reshaped_visual_embeds = rearrange(
216 + _visual_embeds, "(t h w) d -> 1 t h w d", t=T, h=H, w=W
217 + )
218 + # remove temporal dim
219 + reshaped_visual_embeds = reshaped_visual_embeds[:, 0]
220 +
221 + if self.pos_emb is not None:
222 + # interpolate pos emb and add to visual embeds
223 + _local_pos_emb = resample_abs_pos_embed(
224 + posemb=self.pos_emb,
225 + old_size=tuple([int(self.pos_emb_size**0.5)] * 2),
226 + new_size=(H, W),
227 + num_prefix_tokens=0,
228 + )
229 + _local_pos_emb = rearrange(
230 + _local_pos_emb,
231 + "1 (h w) d -> 1 h w d",
232 + h=H,
233 + w=W,
234 + )
235 + reshaped_visual_embeds = reshaped_visual_embeds + _local_pos_emb
236 +
237 + reshaped_visual_embeds = self._forward(
238 + reshaped_visual_embeds,
239 + input_size=(H, W),
240 + )
241 + flattened_visual_embeds.append(reshaped_visual_embeds)
242 + reshaped_visual_embeds = torch.cat(flattened_visual_embeds, dim=0)
243 + return BaseModelOutput(last_hidden_state=reshaped_visual_embeds)
244 +
245 + def _forward(
246 + self,
247 + x: torch.Tensor,
248 + input_size: tuple[int, int],
249 + ) -> torch.Tensor:
250 + h, w = input_size
251 + x = rearrange(x, "1 h w d -> 1 d h w", h=h, w=w)
252 + if self.config.depth:
253 + x = self.net[0](x)
254 + x = self.net[1](x)
255 + x = self.net[2](x)
256 + else:
257 + # When depth=0, self.net is a single PatchMerge module
258 + x = self.net(x)
259 + x = rearrange(x, "1 d h w -> (h w) d")
260 + x = self.readout(x)
261 + return x
262 +
263 +
264 +class CustomQwen2VLVE(Qwen2VisionTransformer):
265 + """Thin wrapper around the Qwen2-VL used as a vision encoder.
266 +
267 + This mirrors the original HF-based vision encoder used in Kanana-V, but
268 + reuses vLLM's optimized `Qwen2VisionTransformer` building blocks.
269 + """
270 +
271 + def __init__(self, config: Qwen2VLVisionConfig) -> None:
272 + super().__init__(
273 + vision_config=config,
274 + norm_eps=getattr(config, "rms_norm_eps", 1e-6),
275 + quant_config=None,
276 + prefix="",
277 + )
278 +
279 + # Kanana-V uses its own projector/abstractor instead of the Qwen2
280 + # built-in patch merger, so we drop the merger module to keep the
281 + # parameter set compatible with the original checkpoint.
282 + if hasattr(self, "merger"):
283 + del self.merger
284 +
285 + @classmethod
286 + def _from_config(cls, config: Qwen2VLVisionConfig) -> "CustomQwen2VLVE":
287 + """Drop-in replacement for the HF `_from_config` constructor."""
288 + return cls(config)
289 +
290 + def forward(
291 + self,
292 + pixel_values: torch.Tensor,
293 + grid_thw: torch.Tensor,
294 + output_hidden_states: bool | None = None,
295 + return_dict: bool | None = None,
296 + ) -> tuple | BaseModelOutput:
297 + """Run the vision transformer and optionally return intermediate states.
298 +
299 + Unlike the base `Qwen2VisionTransformer`, this wrapper exposes the
300 + pre-merger patch-level representations and a HF-style `BaseModelOutput`
301 + so that the existing projector / abstractor code can be reused.
302 + """
303 + assert return_dict, "Only return_dict=True is supported."
304 +
305 + # Patchify
306 + x = pixel_values.to(device=self.device, dtype=self.dtype)
307 + x = self.patch_embed(x) # (num_patches, embed_dim)
308 +
309 + # Prepare grid and rotary embeddings – mirror base implementation.
310 + if isinstance(grid_thw, list):
311 + grid_thw_list = grid_thw
312 + grid_thw_np = np.array(grid_thw, dtype=np.int32)
313 + else:
314 + grid_thw_list = grid_thw.tolist()
315 + grid_thw_np = grid_thw.cpu().numpy()
316 +
317 + rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw_list)
318 +
319 + # Compute cu_seqlens in numpy then move to device, same as base model.
320 + cu_seqlens = np.repeat(
321 + grid_thw_np[:, 1] * grid_thw_np[:, 2],
322 + grid_thw_np[:, 0],
323 + ).cumsum(axis=0, dtype=np.int32)
324 + cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])
325 + cu_seqlens = torch.from_numpy(cu_seqlens).to(
326 + self.device,
327 + non_blocking=True,
328 + )
329 +
330 + # Shape to (S, B, D) with batch dimension 1 as expected by the blocks.
331 + x = x.unsqueeze(1)
332 +
333 + # Pre-compute seqlens for attention backend.
334 + max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)
335 +
336 + encoder_states = () if output_hidden_states else None
337 +
338 + for blk in self.blocks:
339 + if output_hidden_states:
340 + # Store patch-level states (S, D).
341 + encoder_states = encoder_states + (x.squeeze(1),)
342 +
343 + x = blk(
344 + x,
345 + cu_seqlens=cu_seqlens,
346 + rotary_pos_emb_cos=rotary_pos_emb_cos,
347 + rotary_pos_emb_sin=rotary_pos_emb_sin,
348 + max_seqlen=max_seqlen,
349 + )
350 +
351 + # Final hidden state at patch level (S, D).
352 + hidden_states = x.squeeze(1)
353 + if output_hidden_states:
354 + encoder_states = encoder_states + (hidden_states,)
355 +
356 + if not return_dict:
357 + return tuple(v for v in [hidden_states, encoder_states] if v is not None)
358 + return BaseModelOutput(
359 + last_hidden_state=hidden_states,
360 + hidden_states=encoder_states,
361 + )
362 +
363 + def get_num_tokens(self) -> int:
364 + # Not used in the current Kanana-V pipeline, kept for API compatibility.
365 + return -1
366 +
367 +
368 +class KananaVProcessingInfo(BaseProcessingInfo):
369 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:
370 + return {"image": None}
371 +
372 + def get_image_size_with_most_features(self) -> ImageSize:
373 + max_image_size, _ = self._get_vision_info(
374 + image_width=9999,
375 + image_height=9999,
376 + num_frames=1,
377 + )
378 + return max_image_size
379 +
380 + def _get_vision_info(
381 + self,
382 + *,
383 + image_width: int,
384 + image_height: int,
385 + num_frames: int = 1,
386 + do_resize: bool = True,
387 + ) -> tuple[ImageSize, int]:
388 + image_processor = self.ctx.get_hf_processor().image_processor
389 + smart_resize = resolve_obj_by_qualname(
390 + f"{type(image_processor).__module__}.smart_resize"
391 + )
392 +
393 + hf_config = self.get_hf_config()
394 + vision_config = hf_config.vision_config
395 + patch_size = vision_config.patch_size
396 + merge_size = vision_config.spatial_merge_size
397 + temporal_patch_size = vision_config.temporal_patch_size
398 +
399 + if do_resize:
400 + resized_height, resized_width = smart_resize(
401 + height=image_height,
402 + width=image_width,
403 + factor=patch_size * merge_size,
404 + min_pixels=image_processor.min_pixels,
405 + max_pixels=image_processor.max_pixels,
406 + )
407 + preprocessed_size = ImageSize(width=resized_width, height=resized_height)
408 + else:
409 + preprocessed_size = ImageSize(width=image_width, height=image_height)
410 +
411 + # NOTE: Frames are padded to be divisible by `temporal_patch_size`
412 + # https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/qwen2_vl/image_processing_qwen2_vl.py#L294
413 + padded_num_frames = num_frames + num_frames % temporal_patch_size
414 +
415 + grid_t = max(padded_num_frames // temporal_patch_size, 1)
416 + grid_h = preprocessed_size.height // patch_size
417 + grid_w = preprocessed_size.width // patch_size
418 +
419 + num_patches = grid_t * grid_h * grid_w
420 + num_vision_tokens = num_patches // (merge_size**2)
421 +
422 + return preprocessed_size, num_vision_tokens
423 +
424 + def get_mm_max_tokens_per_item(
425 + self,
426 + seq_len: int,
427 + mm_counts: Mapping[str, int],
428 + ) -> Mapping[str, int]:
429 + target_width, target_height = self.get_image_size_with_most_features()
430 + num_vision_tokens = self._get_vision_info(
431 + image_width=target_width,
432 + image_height=target_height,
433 + num_frames=1,
434 + )[1]
435 + return {"image": num_vision_tokens}
436 +
437 +
438 +class KananaVDummyInputsBuilder(BaseDummyInputsBuilder[KananaVProcessingInfo]):
439 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
440 + num_images = mm_counts.get("image", 0)
441 + return "<image>" * num_images
442 +
443 + def get_dummy_mm_data(
444 + self,
445 + seq_len: int,
446 + mm_counts: Mapping[str, int],
447 + mm_options: Mapping[str, BaseDummyOptions] | None = None,
448 + ) -> MultiModalDataDict:
449 + num_images = mm_counts.get("image", 0)
450 + return {
451 + "image": self._get_dummy_images(
452 + width=9999, height=9999, num_images=num_images
453 + ),
454 + }
455 +
456 +
457 +class KananaVMultiModalProcessor(BaseMultiModalProcessor[KananaVProcessingInfo]):
458 + """vLLM multimodal processor for Kanana-V (text + image)."""
459 +
460 + @property
461 + def media_token_id(self) -> int:
462 + return self.info.get_hf_config().text_config.eos_token_id + 1
463 +
464 + def _call_hf_processor(
465 + self,
466 + prompt: str,
467 + mm_data: Mapping[str, object],
468 + mm_kwargs: Mapping[str, object],
469 + tok_kwargs: Mapping[str, object],
470 + ) -> BatchFeature:
471 + """Run the underlying HF processor on text and image data."""
472 + # Text-only input is handled as a special case here.
473 + if not mm_data or not mm_data.get("images", []):
474 + prompt_ids = self.info.get_tokenizer().encode(prompt)
475 + return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
476 +
477 + # Images
478 + image_inputs = mm_data.get("images", [])
479 + pixel_sizes = []
480 + if not isinstance(image_inputs[0], Image.Image):
481 + image_inputs = [Image.fromarray(image) for image in image_inputs]
482 +
483 + image_processor = self.info.get_hf_processor().image_processor
484 + processor_output = [image_processor(image) for image in image_inputs]
485 + pixel_values = [o["pixel_values"] for o in processor_output]
486 + image_meta = [o["image_meta"] for o in processor_output]
487 + # list of dict -> dict of list
488 + image_meta = {k: [d[k] for d in image_meta] for k in image_meta[0]}
489 +
490 + for pixel_value in pixel_values:
491 + pixel_sizes.append(pixel_value.shape[0])
492 + # flattened pixel_values for single example (already includes batch dim)
493 + pixel_values = torch.concat(pixel_values, dim=0)
494 +
495 + tokenizer = self.info.get_tokenizer()
496 + media_token = tokenizer.convert_ids_to_tokens([self.media_token_id])[0]
497 + prompt_replaced = prompt.replace("<image>", media_token)
498 + input_ids = tokenizer.encode(prompt_replaced)
499 + input_ids = torch.tensor(input_ids)
500 +
501 + # Ensure HF output is consistent with vLLM prompt-update expectations:
502 + # if the HF tokenizer emits exactly 1 placeholder token per image, expand
503 + # it to `T*H*W` placeholder tokens per image so placeholder detection works.
504 + num_images = len(image_inputs)
505 + image_token_thw = torch.tensor(image_meta["image_token_thw"])
506 + per_image_token_counts = image_token_thw.prod(dim=1).tolist()
507 + expected_total = int(sum(int(x) for x in per_image_token_counts))
508 +
509 + n_placeholders = int((input_ids == self.media_token_id).sum().item())
510 + if n_placeholders == num_images and expected_total != num_images:
511 + expanded: list[int] = []
512 + img_i = 0
513 + for tok in input_ids.tolist():
514 + if tok == self.media_token_id and img_i < num_images:
515 + expanded.extend(
516 + [self.media_token_id] * int(per_image_token_counts[img_i])
517 + )
518 + img_i += 1
519 + else:
520 + expanded.append(tok)
521 + input_ids = input_ids.new_tensor(expanded)
522 +
523 + combined_outputs = dict(
524 + # Add batch dimension to input_ids.
525 + input_ids=input_ids.unsqueeze(0),
526 + pixel_values=pixel_values,
527 + vision_grid_thw=torch.tensor(image_meta["vision_grid_thw"]),
528 + image_token_thw=torch.tensor(image_meta["image_token_thw"]),
529 + pixel_sizes=torch.tensor(pixel_sizes),
530 + )
531 + return BatchFeature(combined_outputs, tensor_type="pt")
532 +
533 + def _get_prompt_updates(
534 + self,
535 + mm_items: MultiModalDataItems,
536 + hf_processor_mm_kwargs: Mapping[str, object],
537 + out_mm_kwargs: MultiModalKwargsItems,
538 + ) -> Sequence[PromptUpdate]:
539 + def get_replacement(idx: int) -> Sequence[int]:
540 + out_item = out_mm_kwargs["image"][idx]
541 + image_token_thw = out_item["image_token_thw"].data
542 + assert isinstance(image_token_thw, torch.Tensor)
543 +
544 + num_tokens = int(image_token_thw.prod().item())
545 + return [self.media_token_id] * num_tokens
546 +
547 + return [
548 + PromptReplacement(
549 + modality="image",
550 + target="<image>",
551 + replacement=get_replacement,
552 + ),
553 + ]
554 +
555 + def _get_mm_fields_config(
556 + self,
557 + hf_inputs: BatchFeature,
558 + hf_processor_mm_kwargs: Mapping[str, object],
559 + ) -> Mapping[str, MultiModalFieldConfig]:
560 + pixel_sizes = hf_inputs.get("pixel_sizes", torch.empty(0))
561 +
562 + mm_fields_config = dict(
563 + pixel_values=MultiModalFieldConfig.flat_from_sizes("image", pixel_sizes),
564 + vision_grid_thw=MultiModalFieldConfig.batched("image"),
565 + image_token_thw=MultiModalFieldConfig.batched("image"),
566 + )
567 + return mm_fields_config
568 +
569 +
570 +@MULTIMODAL_REGISTRY.register_processor(
571 + KananaVMultiModalProcessor,
572 + info=KananaVProcessingInfo,
573 + dummy_inputs=KananaVDummyInputsBuilder,
574 +)
575 +class KananaVForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
576 + @classmethod
577 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:
578 + if modality.startswith("image"):
579 + return "<image>"
580 + else:
581 + raise ValueError(f"Unsupported modality: {modality}")
582 +
583 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
584 + super().__init__()
585 +
586 + config = vllm_config.model_config.hf_config
587 + self.config = config
588 +
589 + self.vision_model = CustomQwen2VLVE._from_config(config.vision_config)
590 + self.abstractor = DynamicCAbstractor(
591 + config.projector_config, num_input_tokens=self.vision_model.get_num_tokens()
592 + )
593 + self.language_model = init_vllm_registered_model(
594 + vllm_config=vllm_config,
595 + hf_config=config.text_config,
596 + prefix=maybe_prefix(prefix, "model"),
597 + architectures=["LlamaForCausalLM"],
598 + )
599 + self.make_empty_intermediate_tensors = (
600 + self.language_model.make_empty_intermediate_tensors
601 + )
602 +
603 + def _parse_and_validate_image_input(
604 + self, **kwargs: object
605 + ) -> KananaVImageInputs | None:
606 + pixel_values = kwargs.pop("pixel_values", None)
607 + vision_grid_thw = kwargs.pop("vision_grid_thw", None)
608 +
609 + if pixel_values is None:
610 + return None
611 +
612 + if vision_grid_thw is None:
613 + raise ValueError(
614 + "vision_grid_thw is required when pixel_values is provided"
615 + )
616 +
617 + # Normalize pixel_values to 2D tensor (num_patches, channels*patch*patch)
618 + if isinstance(pixel_values, torch.Tensor):
619 + if pixel_values.ndim == 2:
620 + pass # Already in expected shape
621 + elif pixel_values.ndim == 3:
622 + pixel_values = pixel_values.flatten(0, 1)
623 + else:
624 + raise ValueError(
625 + f"pixel_values should be 2D or batched 3D tensor. "
626 + f"Got ndim: {pixel_values.ndim} "
627 + f"(shape={pixel_values.shape})"
628 + )
629 + else:
630 + pixel_values = torch.concat(pixel_values)
631 +
632 + # Normalize vision_grid_thw to 2D tensor (num_images, 3)
633 + if isinstance(vision_grid_thw, torch.Tensor):
634 + if vision_grid_thw.ndim == 3:
635 + vision_grid_thw = vision_grid_thw.flatten(0, 1)
636 + else:
637 + vision_grid_thw = torch.concat(vision_grid_thw)
638 +
639 + return KananaVImagePixelInputs(
640 + type="pixel_values",
641 + pixel_values=pixel_values,
642 + vision_grid_thw=vision_grid_thw,
643 + )
644 +
645 + def _process_image_input(self, image_input: KananaVImageInputs) -> torch.Tensor:
646 + pixel_values = image_input["pixel_values"]
647 + vision_grid_thw = image_input["vision_grid_thw"]
648 +
649 + image_metas = {"vision_grid_thw": vision_grid_thw}
650 + visual_embeds = self.forward_and_project_vision(pixel_values, image_metas)
651 +
652 + merge_size = self.abstractor.merge_size
653 + batch_size = vision_grid_thw.size(0)
654 + multi_modal_embeddings: tuple[torch.Tensor, ...] = ()
655 + sample_index = 0
656 + for i in range(batch_size):
657 + t, h, w = (
658 + vision_grid_thw[i][0],
659 + vision_grid_thw[i][1] // merge_size,
660 + vision_grid_thw[i][2] // merge_size,
661 + )
662 + num_tokens = t * h * w
663 + visual_embed = visual_embeds[sample_index : sample_index + num_tokens]
664 + multi_modal_embeddings += (visual_embed,)
665 + sample_index += num_tokens
666 +
667 + return multi_modal_embeddings
668 +
669 + def _get_visual_feature_at(
670 + self,
671 + v_output: Sequence[torch.Tensor],
672 + layer_index: int | Sequence[int],
673 + ) -> torch.Tensor:
674 + if isinstance(layer_index, (list, tuple)):
675 + visual_features = torch.stack(v_output, dim=1)[
676 + :, layer_index
677 + ] # [B, n_scales, L, dim]
678 + else:
679 + visual_features = v_output[layer_index] # [B, L, dim]
680 + return visual_features
681 +
682 + def forward_vision(
683 + self,
684 + pixel_values: torch.Tensor,
685 + image_metas: dict | None = None,
686 + ) -> torch.Tensor:
687 + vision_model_args = {
688 + "pixel_values": pixel_values,
689 + "return_dict": True,
690 + "output_hidden_states": True,
691 + "grid_thw": image_metas["vision_grid_thw"],
692 + }
693 + v_outputs = self.vision_model(**vision_model_args)
694 + layer_index = self.config.projector_config.feature_layer_index
695 + visual_features = self._get_visual_feature_at(
696 + v_outputs.hidden_states, layer_index
697 + )
698 + return visual_features
699 +
700 + def forward_projector(
701 + self,
702 + visual_features: torch.Tensor,
703 + image_metas: dict | None = None,
704 + ) -> torch.Tensor:
705 + visual_embeds = self.abstractor(
706 + visual_features,
707 + grid_thw=image_metas["vision_grid_thw"],
708 + )["last_hidden_state"]
709 + return visual_embeds
710 +
711 + def forward_and_project_vision(
712 + self,
713 + pixel_values: torch.Tensor,
714 + image_metas: dict | None = None,
715 + ) -> torch.Tensor:
716 + assert pixel_values is not None
717 + visual_features = self.forward_vision(pixel_values, image_metas=image_metas)
718 + visual_embeds = self.forward_projector(visual_features, image_metas=image_metas)
719 + return visual_embeds
720 +
721 + def get_language_model(self) -> torch.nn.Module:
722 + return self.language_model
723 +
724 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
725 + image_input = self._parse_and_validate_image_input(**kwargs)
726 + if image_input is None:
727 + return []
728 +
729 + return self._process_image_input(image_input)
730 +
731 + def forward(
732 + self,
733 + input_ids: torch.Tensor,
734 + positions: torch.Tensor,
735 + intermediate_tensors: IntermediateTensors | None = None,
736 + inputs_embeds: torch.Tensor | None = None,
737 + **kwargs,
738 + ):
739 + if intermediate_tensors is not None:
740 + inputs_embeds = None
741 +
742 + hidden_states = self.language_model(
743 + input_ids=input_ids,
744 + positions=positions,
745 + intermediate_tensors=intermediate_tensors,
746 + inputs_embeds=inputs_embeds,
747 + )
748 +
749 + return hidden_states
750 +
751 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
752 + return self.language_model.compute_logits(hidden_states)
753 +
754 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
755 + loader = AutoWeightsLoader(self)
756 + return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+1
@@ -343,6 +343,7 @@ _MULTIMODAL_MODELS = {
343 ),
344 "IsaacForConditionalGeneration": ("isaac", "IsaacForConditionalGeneration"),
345 "SmolVLMForConditionalGeneration": ("smolvlm", "SmolVLMForConditionalGeneration"), # noqa: E501
346 + "KananaVForConditionalGeneration": ("kanana_v", "KananaVForConditionalGeneration"),
347 "KeyeForConditionalGeneration": ("keye", "KeyeForConditionalGeneration"),
348 "KeyeVL1_5ForConditionalGeneration": (
349 "keye_vl1_5",