| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Shared Kimi-K3 multimodal preprocessing.""" |
| 4 | |
| 5 | import math |
| 6 | from collections.abc import Mapping, Sequence |
| 7 | from typing import Any, cast |
| 8 | |
| 9 | import torch |
| 10 | from transformers import BatchFeature |
| 11 | |
| 12 | from vllm.config.multimodal import ImageDummyOptions, MultiModalDummyOptions |
| 13 | from vllm.inputs import MultiModalDataDict |
| 14 | from vllm.logger import init_logger |
| 15 | from vllm.multimodal.inputs import ( |
| 16 | MultiModalFieldConfig, |
| 17 | MultiModalKwargsItems, |
| 18 | ) |
| 19 | from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems |
| 20 | from vllm.multimodal.processing import ( |
| 21 | BaseDummyInputsBuilder, |
| 22 | BaseMultiModalProcessor, |
| 23 | BaseProcessingInfo, |
| 24 | InputProcessingContext, |
| 25 | PromptReplacement, |
| 26 | PromptUpdate, |
| 27 | PromptUpdateDetails, |
| 28 | cached_encode, |
| 29 | ) |
| 30 | from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config |
| 31 | from vllm.transformers_utils.processor import cached_get_image_processor |
| 32 | from vllm.transformers_utils.processors.kimi_k3 import KimiK3Processor |
| 33 | |
| 34 | logger = init_logger(__name__) |
| 35 | |
| 36 | |
| 37 | def navit_resize_image( |
| 38 | width: int, |
| 39 | height: int, |
| 40 | patch_size: int, |
| 41 | merge_kernel_size: int, |
| 42 | in_patch_limit: int, |
| 43 | patch_limit_on_one_side: int, |
| 44 | fixed_output_tokens: int | None, |
| 45 | ): |
| 46 | # Apply the patch limits. |
| 47 | s1 = math.sqrt( |
| 48 | in_patch_limit |
| 49 | / (max(1.0, width // patch_size) * max(1.0, height // patch_size)) |
| 50 | ) |
| 51 | s2 = patch_limit_on_one_side * patch_size / width |
| 52 | s3 = patch_limit_on_one_side * patch_size / height |
| 53 | scale = min(1.0, s1, s2, s3) |
| 54 | new_w, new_h = max(1, int(width * scale)), max(1, int(height * scale)) |
| 55 | new_w = min(new_w, patch_limit_on_one_side * patch_size) |
| 56 | new_h = min(new_h, patch_limit_on_one_side * patch_size) |
| 57 | |
| 58 | factor = merge_kernel_size * patch_size |
| 59 | |
| 60 | pad_height = (factor - new_h % factor) % factor |
| 61 | pad_width = (factor - new_w % factor) % factor |
| 62 | |
| 63 | if fixed_output_tokens is not None: |
| 64 | num_tokens = fixed_output_tokens |
| 65 | else: |
| 66 | # Calculate new dimensions after padding and patching |
| 67 | token_height = (new_h + pad_height) // factor |
| 68 | token_width = (new_w + pad_width) // factor |
| 69 | |
| 70 | assert token_height * merge_kernel_size <= patch_limit_on_one_side, ( |
| 71 | f"token_height {token_height} * merge_kernel_size {merge_kernel_size} > " |
| 72 | f"patch_limit_on_one_side {patch_limit_on_one_side}" |
| 73 | ) |
| 74 | assert token_width * merge_kernel_size <= patch_limit_on_one_side, ( |
| 75 | f"token_width {token_width} * merge_kernel_size {merge_kernel_size} > " |
| 76 | f"patch_limit_on_one_side {patch_limit_on_one_side}" |
| 77 | ) |
| 78 | |
| 79 | num_tokens = token_height * token_width |
| 80 | return { |
| 81 | "num_tokens": num_tokens, |
| 82 | "new_width": new_w, |
| 83 | "new_height": new_h, |
| 84 | "pad_width": pad_width, |
| 85 | "pad_height": pad_height, |
| 86 | "sampled_nframes": 1, |
| 87 | } |
| 88 | |
| 89 | |
| 90 | class KimiK3ProcessingInfo(BaseProcessingInfo): |
| 91 | """Processing information for the image-only Kimi-K3 model. |
| 92 | |
| 93 | K3 uses the standard ``image`` modality (unlike K2.5's unified |
| 94 | ``vision_chunk``), so it builds its own ``KimiK3Processor`` wrapper around |
| 95 | the checkpoint's image processor and resolves the ``<|media_pad|>`` token |
| 96 | id the same way K2.5 does. |
| 97 | """ |
| 98 | |
| 99 | def __init__(self, ctx: InputProcessingContext) -> None: |
| 100 | super().__init__(ctx) |
| 101 | |
| 102 | self.hf_config = hf_config = self.get_hf_config() |
| 103 | |
| 104 | tokenizer = self.get_tokenizer() |
| 105 | image_processor = cached_get_image_processor( |
| 106 | self.ctx.model_config.model, |
| 107 | revision=self.ctx.model_config.revision, |
| 108 | trust_remote_code=self.ctx.model_config.trust_remote_code, |
| 109 | ) |
| 110 | |
| 111 | # Resolve token ID from the tokenizer because transformers v5 |
| 112 | # may remap token IDs vs config.json. |
| 113 | config_token_id = hf_config.media_placeholder_token_id |
| 114 | resolved_token_id = tokenizer.convert_tokens_to_ids("<|media_pad|>") |
| 115 | unk_token_id = getattr(tokenizer, "unk_token_id", None) |
| 116 | is_valid_resolved = isinstance(resolved_token_id, int) and ( |
| 117 | unk_token_id is None or resolved_token_id != unk_token_id |
| 118 | ) |
| 119 | if is_valid_resolved and resolved_token_id != config_token_id: |
| 120 | logger.warning_once( |
| 121 | "Kimi-K3 config.media_placeholder_token_id (%d) disagrees " |
| 122 | "with tokenizer mapping for <|media_pad|> (%d). " |
| 123 | "Using tokenizer value.", |
| 124 | config_token_id, |
| 125 | resolved_token_id, |
| 126 | ) |
| 127 | media_token_id = resolved_token_id |
| 128 | # Patch config so downstream code also sees the correct ID. |
| 129 | hf_config.media_placeholder_token_id = resolved_token_id |
| 130 | else: |
| 131 | media_token_id = config_token_id |
| 132 | |
| 133 | self.media_token_id = media_token_id |
| 134 | self.media_token = tokenizer.decode(media_token_id) |
| 135 | |
| 136 | self.image_processor = image_processor |
| 137 | self.hf_processor = KimiK3Processor( |
| 138 | tokenizer=tokenizer, |
| 139 | image_processor=image_processor, |
| 140 | ) |
| 141 | self.media_tokens_calculator = image_processor.media_tokens_calculator |
| 142 | |
| 143 | def get_hf_processor(self, **kwargs: object) -> KimiK3Processor: |
| 144 | return self.hf_processor |
| 145 | |
| 146 | def get_hf_config(self) -> KimiK3Config: |
| 147 | return self.ctx.get_hf_config(KimiK3Config) |
| 148 | |
| 149 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 150 | # None means unlimited |
| 151 | return {"image": None} |
| 152 | |
| 153 | @classmethod |
| 154 | def get_max_image_size( |
| 155 | cls, |
| 156 | patch_size: int, |
| 157 | merge_kernel_size: int, |
| 158 | in_patch_limit: int, |
| 159 | patch_limit_on_one_side: int, |
| 160 | fixed_output_tokens: int | None, |
| 161 | ) -> ImageSize: |
| 162 | max_side = patch_limit_on_one_side * patch_size |
| 163 | best_score = (-1, -1) |
| 164 | best_size = (max_side, max_side) |
| 165 | |
| 166 | for width_patches in range(patch_limit_on_one_side + 1): |
| 167 | width = min((width_patches + 1) * patch_size - 1, max_side) |
| 168 | for height_patches in range(width_patches, patch_limit_on_one_side + 1): |
| 169 | height = min((height_patches + 1) * patch_size - 1, max_side) |
| 170 | resize_config = navit_resize_image( |
| 171 | width, |
| 172 | height, |
| 173 | patch_size, |
| 174 | merge_kernel_size, |
| 175 | in_patch_limit, |
| 176 | patch_limit_on_one_side, |
| 177 | fixed_output_tokens, |
| 178 | ) |
| 179 | padded_width = resize_config["new_width"] + resize_config["pad_width"] |
| 180 | padded_height = ( |
| 181 | resize_config["new_height"] + resize_config["pad_height"] |
| 182 | ) |
| 183 | num_patches = padded_width // patch_size * (padded_height // patch_size) |
| 184 | score = (resize_config["num_tokens"], num_patches) |
| 185 | if score > best_score: |
| 186 | best_score = score |
| 187 | best_size = (width, height) |
| 188 | return ImageSize(width=best_size[0], height=best_size[1]) |
| 189 | |
| 190 | |
| 191 | class KimiK3DummyInputsBuilder(BaseDummyInputsBuilder[KimiK3ProcessingInfo]): |
| 192 | """Builds image-based dummy inputs for K3 profiling. |
| 193 | |
| 194 | The dummy text is made of ``<|kimi_image_placeholder|>`` tokens — exactly |
| 195 | the placeholder that K3's ``_get_prompt_updates`` expands — and the dummy |
| 196 | mm data is a plain list of PIL images under the ``image`` key. |
| 197 | """ |
| 198 | |
| 199 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 200 | num_images = mm_counts.get("image", 0) |
| 201 | return self.info.get_hf_config().image_placeholder * num_images |
| 202 | |
| 203 | def get_dummy_mm_data( |
| 204 | self, |
| 205 | seq_len: int, |
| 206 | mm_counts: Mapping[str, int], |
| 207 | mm_options: MultiModalDummyOptions | None = None, |
| 208 | ) -> MultiModalDataDict: |
| 209 | media_proc_cfg = self.info.image_processor.media_proc_cfg |
| 210 | max_size = self.info.get_max_image_size( |
| 211 | media_proc_cfg["patch_size"], |
| 212 | media_proc_cfg["merge_kernel_size"], |
| 213 | media_proc_cfg["in_patch_limit"], |
| 214 | media_proc_cfg["patch_limit_on_one_side"], |
| 215 | media_proc_cfg["fixed_output_tokens"], |
| 216 | ) |
| 217 | image_overrides = cast( |
| 218 | ImageDummyOptions | None, |
| 219 | mm_options.get("image") if mm_options else None, |
| 220 | ) |
| 221 | return { |
| 222 | "image": self._get_dummy_images( |
| 223 | width=max_size.width, |
| 224 | height=max_size.height, |
| 225 | num_images=mm_counts.get("image", 0), |
| 226 | overrides=image_overrides, |
| 227 | ) |
| 228 | } |
| 229 | |
| 230 | |
| 231 | class KimiK3MultiModalProcessor(BaseMultiModalProcessor[KimiK3ProcessingInfo]): |
| 232 | """Image-only multi-modal processor for Kimi-K3.""" |
| 233 | |
| 234 | def _get_mm_fields_config( |
| 235 | self, |
| 236 | hf_inputs: BatchFeature, |
| 237 | hf_processor_mm_kwargs: Mapping[str, object], |
| 238 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 239 | """Slice the flattened patch tensor back into per-image items. |
| 240 | |
| 241 | ``pixel_values`` holds all patches from every image concatenated; each |
| 242 | image's patch count is ``prod(grid_thws[i])``. ``grid_thws`` is one |
| 243 | ``[N_t, N_h, N_w]`` row per image. |
| 244 | """ |
| 245 | grid_thws = hf_inputs.get("grid_thws", torch.empty((0, 3))) |
| 246 | grid_sizes = grid_thws.prod(-1) |
| 247 | |
| 248 | return dict( |
| 249 | pixel_values=MultiModalFieldConfig.flat_from_sizes("image", grid_sizes), |
| 250 | grid_thws=MultiModalFieldConfig.batched("image", keep_on_cpu=True), |
| 251 | ) |
| 252 | |
| 253 | def _get_prompt_updates( |
| 254 | self, |
| 255 | mm_items: MultiModalDataItems, |
| 256 | hf_processor_mm_kwargs: Mapping[str, Any], |
| 257 | out_mm_kwargs: MultiModalKwargsItems, |
| 258 | ) -> Sequence[PromptUpdate]: |
| 259 | """Expand each K3 image placeholder into a resolution-aware update. |
| 260 | |
| 261 | K3's prompt carries a single ``<|kimi_image_placeholder|>`` token per |
| 262 | image. This replaces that token with |
| 263 | ``<|media_begin|>image {w}x{h}<|media_content|>{pads}<|media_end|>``, |
| 264 | embedding the per-image resolution in the prompt and marking only the |
| 265 | ``<|media_pad|>`` positions as embedding slots (the number of pads is |
| 266 | the feature size returned by ``media_tokens_calculator``). |
| 267 | """ |
| 268 | media_token_id = self.info.media_token_id |
| 269 | media_token = self.info.media_token |
| 270 | image_placeholder = self.info.get_hf_config().image_placeholder |
| 271 | tokenizer = self.info.get_tokenizer() |
| 272 | |
| 273 | def get_replacement(item_idx: int) -> PromptUpdateDetails: |
| 274 | images = mm_items.get_items("image", ImageProcessorItems) |
| 275 | image = images.get(item_idx) |
| 276 | if image is None: |
| 277 | raise ValueError(f"Missing image data at index {item_idx}") |
| 278 | |
| 279 | # The checkpoint image processor works on media dicts, so wrap the |
| 280 | # PIL image before asking it for the token count. |
| 281 | num_media_token = self.info.media_tokens_calculator( |
| 282 | {"type": "image", "image": image} |
| 283 | ) |
| 284 | pads = media_token * num_media_token |
| 285 | |
| 286 | # NOTE: `width`/`height` are the ORIGINAL upload dimensions, not the |
| 287 | # post-preprocess (smart-resized) ones. `image` comes from the |
| 288 | # untouched parsed `mm_items`; the checkpoint image processor |
| 289 | # (`KimiK3VisionProcessor.preprocess`) only produces new tensors via |
| 290 | # `image.resize(...)` and never mutates the stored PIL. This matches |
| 291 | # the reference HF processor (`KimiK3Processor.preprocess_medias`), |
| 292 | # which also builds the prompt from the original `img.size`. The |
| 293 | # resize is reflected only in the pad count above. |
| 294 | width, height = images.get_image_size(item_idx) |
| 295 | full = ( |
| 296 | f"<|media_begin|>image {width}x{height}<|media_content|>" |
| 297 | f"{pads}<|media_end|>" |
| 298 | ) |
| 299 | |
| 300 | return PromptUpdateDetails.select_token_id( |
| 301 | cached_encode(tokenizer, full, add_special_tokens=False), |
| 302 | media_token_id, |
| 303 | ) |
| 304 | |
| 305 | return [ |
| 306 | PromptReplacement( |
| 307 | modality="image", |
| 308 | target=cached_encode( |
| 309 | tokenizer, image_placeholder, add_special_tokens=False |
| 310 | ), |
| 311 | replacement=get_replacement, |
| 312 | ), |
| 313 | ] |