diff --git a/vllm/model_executor/models/kimi_k25.py b/vllm/model_executor/models/kimi_k25.py index 7321b913605..e2b74744b41 100644 --- a/vllm/model_executor/models/kimi_k25.py +++ b/vllm/model_executor/models/kimi_k25.py @@ -56,6 +56,10 @@ from vllm.sequence import IntermediateTensors from vllm.transformers_utils.configs.kimi_k25 import KimiK25Config from vllm.transformers_utils.processor import cached_get_image_processor from vllm.transformers_utils.processors.kimi_k25 import KimiK25Processor +from vllm.transformers_utils.processors.kimi_k25_vision_fused import ( + KimiK25FusedVisionProcessor, +) +from vllm.utils.import_utils import is_numba_available from vllm.utils.tensor_schema import TensorSchema, TensorShape from .utils import ( @@ -108,10 +112,16 @@ class KimiK25ProcessingInfo(BaseProcessingInfo): self.hf_config = hf_config = self.get_hf_config() tokenizer = self.get_tokenizer() + processor_cls = KimiK25FusedVisionProcessor if is_numba_available() else None + logger.info_once( + "Using %s image preprocessing for Kimi-K2.5/K2.6 vision chunks.", + "fused CPU" if processor_cls is not None else "remote HF", + ) image_processor = cached_get_image_processor( self.ctx.model_config.model, revision=self.ctx.model_config.revision, trust_remote_code=self.ctx.model_config.trust_remote_code, + processor_cls_overrides=processor_cls, ) # Resolve token ID from the tokenizer because transformers v5 diff --git a/vllm/transformers_utils/processor.py b/vllm/transformers_utils/processor.py index fa4c558a739..aa33faed916 100644 --- a/vllm/transformers_utils/processor.py +++ b/vllm/transformers_utils/processor.py @@ -81,6 +81,7 @@ _transformers_v4_compatibility_import() _transformers_v4_compatibility_init() _P = TypeVar("_P", bound=ProcessorMixin, default=ProcessorMixin) +_I = TypeVar("_I", bound=BaseImageProcessor, default=BaseImageProcessor) _V = TypeVar("_V", bound=BaseVideoProcessor, default=BaseVideoProcessor) @@ -440,12 +441,14 @@ def get_image_processor( *args: Any, revision: str | None = None, trust_remote_code: bool = False, + processor_cls_overrides: type[_I] | None = None, **kwargs: Any, ): """Load an image processor for the given model name via HuggingFace.""" try: processor_name = convert_model_repo_to_path(processor_name) - processor = AutoImageProcessor.from_pretrained( + processor_cls = processor_cls_overrides or AutoImageProcessor + processor = processor_cls.from_pretrained( processor_name, *args, revision=revision, diff --git a/vllm/transformers_utils/processors/kimi_k25_vision_fused.py b/vllm/transformers_utils/processors/kimi_k25_vision_fused.py new file mode 100644 index 00000000000..63907898175 --- /dev/null +++ b/vllm/transformers_utils/processors/kimi_k25_vision_fused.py @@ -0,0 +1,352 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Optimized CPU image processor for Kimi-K2.5/K2.6 vision chunks.""" + +import io +import json +import math +from typing import Any + +import numpy as np +import pybase64 as base64 +import torch +from PIL import Image +from transformers.image_processing_utils import BaseImageProcessor, BatchFeature +from transformers.utils import TensorType + +from vllm.utils.import_utils import is_numba_available +from vllm.utils.jit_monitor import numba_workqueue_threading_layer + +if is_numba_available(): + from numba import njit, prange + + @njit(parallel=True, cache=True) + def _write_fused_patches( + frames: np.ndarray, + out: np.ndarray, + out_offset: int, + new_h: int, + new_w: int, + padded_h: int, + padded_w: int, + patch_size: int, + normalize_lut: np.ndarray, + ) -> None: + # frames: [T, new_h, new_w, 3] uint8, without padding. + # out: [total_patches, 3, patch_size, patch_size] float32. + t_size = frames.shape[0] + patch_h = padded_h // patch_size + patch_w = padded_w // patch_size + total = t_size * padded_h * padded_w * 3 + hwc = padded_h * padded_w * 3 + wc = padded_w * 3 + + for linear in prange(total): + t = linear // hwc + rem = linear - t * hwc + y = rem // wc + rem = rem - y * wc + x = rem // 3 + c = rem - x * 3 + + value = frames[t, y, x, c] if y < new_h and x < new_w else 0 + + patch_idx = ( + out_offset + + t * patch_h * patch_w + + (y // patch_size) * patch_w + + (x // patch_size) + ) + out[patch_idx, c, y % patch_size, x % patch_size] = normalize_lut[value, c] + +else: + + def _write_fused_patches(*args: Any, **kwargs: Any) -> None: + raise RuntimeError("numba is required for fused Kimi image preprocessing") + + +def navit_resize_image( + width: int, + height: int, + patch_size: int, + merge_kernel_size: int, + in_patch_limit: int, + patch_limit_on_one_side: int, + fixed_output_tokens: int | None, +) -> dict[str, int]: + s1 = math.sqrt( + in_patch_limit + / (max(1.0, width // patch_size) * max(1.0, height // patch_size)) + ) + s2 = patch_limit_on_one_side * patch_size / width + s3 = patch_limit_on_one_side * patch_size / height + scale = min(1.0, s1, s2, s3) + new_w = min(max(1, int(width * scale)), patch_limit_on_one_side * patch_size) + new_h = min(max(1, int(height * scale)), patch_limit_on_one_side * patch_size) + + factor = merge_kernel_size * patch_size + pad_height = (factor - new_h % factor) % factor + pad_width = (factor - new_w % factor) % factor + + if fixed_output_tokens is not None: + num_tokens = fixed_output_tokens + else: + token_height = (new_h + pad_height) // factor + token_width = (new_w + pad_width) // factor + num_tokens = token_height * token_width + + return { + "num_tokens": num_tokens, + "new_width": new_w, + "new_height": new_h, + "pad_width": pad_width, + "pad_height": pad_height, + "sampled_nframes": 1, + } + + +def navit_resize_video( + width: int, + height: int, + nframes: int, + avg_fps: float, + sample_fps: float, + patch_size: int, + merge_kernel_size: int, + in_patch_limit_each_frame: int, + patch_limit_on_one_side: int, + in_patch_limit_total: int | None, + max_num_frames_each_video: int | None, + fixed_output_tokens_each_frame: int | None, +) -> dict[str, int]: + sample_fps = min(sample_fps, avg_fps) + sampled_nframes = max(round(nframes * sample_fps / avg_fps), 1) + if max_num_frames_each_video is not None: + sampled_nframes = min(sampled_nframes, max_num_frames_each_video) + + if in_patch_limit_total is not None: + in_patch_limit_each_frame = min( + round(in_patch_limit_total / sampled_nframes), + in_patch_limit_each_frame, + ) + + ret = navit_resize_image( + width, + height, + patch_size, + merge_kernel_size, + in_patch_limit_each_frame, + patch_limit_on_one_side, + fixed_output_tokens_each_frame, + ) + ret["sampled_nframes"] = sampled_nframes + return ret + + +def _to_pil(data: Any) -> Image.Image: + if hasattr(data, "media") and hasattr(data, "original_bytes"): + data = data.media + if isinstance(data, Image.Image): + return data if data.mode == "RGB" else data.convert("RGB") + if isinstance(data, str): + if data.startswith("data:"): + raw_base64 = data.split(",", 1)[1] + return Image.open(io.BytesIO(base64.b64decode(raw_base64))).convert("RGB") + return Image.open(data).convert("RGB") + if isinstance(data, bytes): + return Image.open(io.BytesIO(data)).convert("RGB") + raise ValueError(f"Unsupported data type: {type(data)}") + + +def _ensure_media_type(media: dict[str, Any]) -> dict[str, Any]: + if media["type"] == "image": + media["image"] = _to_pil(media["image"]) + return media + if media["type"] == "video_chunk": + media["video_chunk"] = [_to_pil(frame) for frame in media["video_chunk"]] + return media + raise ValueError(f"Unsupported media type: {media['type']}") + + +class KimiK25FusedVisionProcessor(BaseImageProcessor): + model_type = "kimi_k25" + + def __init__(self, media_proc_cfg: dict[str, Any], **kwargs: Any) -> None: + super().__init__(**kwargs) + media_proc_cfg = dict(media_proc_cfg) + merge_kernel_size = media_proc_cfg["merge_kernel_size"] + if isinstance(merge_kernel_size, (list, tuple)): + media_proc_cfg["merge_kernel_size"] = int(merge_kernel_size[0]) + self.media_proc_cfg = media_proc_cfg + self.num_frames_per_chunk = media_proc_cfg["temporal_merge_kernel_size"] + values = np.arange(256, dtype=np.float32)[:, None] + image_mean = np.asarray(media_proc_cfg["image_mean"], dtype=np.float32) + image_std_inv = 1.0 / np.asarray(media_proc_cfg["image_std"], dtype=np.float32) + self.normalize_lut = (values / 255.0 - image_mean[None, :]) * image_std_inv[ + None, : + ] + + def media_tokens_calculator(self, media: dict[str, Any]) -> int: + media = _ensure_media_type(media) + ret = self.get_resize_config(media) + return ret["num_tokens"] + + def get_resize_config(self, media_input: dict[str, Any]) -> dict[str, int]: + if media_input["type"] == "image": + width, height = media_input["image"].size + return navit_resize_image( + width, + height, + self.media_proc_cfg["patch_size"], + self.media_proc_cfg["merge_kernel_size"], + self.media_proc_cfg["in_patch_limit"], + self.media_proc_cfg["patch_limit_on_one_side"], + self.media_proc_cfg["fixed_output_tokens"], + ) + + if media_input["type"] == "video_chunk": + frame = media_input["video_chunk"][0] + width, height = frame.size + num_frames = len(media_input["video_chunk"]) + in_patch_limit_each_frame = self.media_proc_cfg["in_patch_limit_each_frame"] + if in_patch_limit_each_frame is None: + in_patch_limit_each_frame = self.media_proc_cfg["in_patch_limit"] + + return navit_resize_video( + width, + height, + num_frames, + 1.0, + math.inf, + self.media_proc_cfg["patch_size"], + self.media_proc_cfg["merge_kernel_size"], + in_patch_limit_each_frame, + self.media_proc_cfg["patch_limit_on_one_side"], + self.media_proc_cfg["in_patch_limit_video"], + None, + self.media_proc_cfg["fixed_output_tokens"], + ) + + raise ValueError(f"Unsupported type: {media_input['type']}") + + @staticmethod + def resize_image(image: Image.Image, new_width: int, new_height: int) -> np.ndarray: + image = image.resize((new_width, new_height), resample=Image.Resampling.BICUBIC) + return np.asarray(image) + + def preprocess( + self, + medias: list[dict[str, Any]], + return_tensors: str | TensorType | None = None, + ) -> BatchFeature: + if not isinstance(medias, list): + medias = [medias] + if not medias: + return BatchFeature(data={}, tensor_type=return_tensors) + + if njit is None: + raise RuntimeError("numba is required for fused Kimi image preprocessing") + + patch_size = int(self.media_proc_cfg["patch_size"]) + prepared = [] + grid_thws_np = np.empty((len(medias), 3), dtype=np.int64) + total_patches = 0 + + for idx, item in enumerate(medias): + item = _ensure_media_type(item) + resize_config = self.get_resize_config(item) + new_width = resize_config["new_width"] + new_height = resize_config["new_height"] + pad_width = resize_config["pad_width"] + pad_height = resize_config["pad_height"] + padded_width = new_width + pad_width + padded_height = new_height + pad_height + + if item["type"] == "image": + image_np = self.resize_image(item["image"], new_width, new_height) + frames = image_np[np.newaxis, ...] + elif item["type"] == "video_chunk": + frames = np.stack( + [ + self.resize_image(frame, new_width, new_height) + for frame in item["video_chunk"] + ], + axis=0, + ) + else: + raise ValueError(f"Unsupported type: {item['type']}") + + t_size = frames.shape[0] + grid_h = padded_height // patch_size + grid_w = padded_width // patch_size + grid_thws_np[idx, 0] = t_size + grid_thws_np[idx, 1] = grid_h + grid_thws_np[idx, 2] = grid_w + + num_patches = t_size * grid_h * grid_w + prepared.append( + ( + frames, + new_height, + new_width, + padded_height, + padded_width, + num_patches, + ) + ) + total_patches += num_patches + + pixel_values_np = np.empty( + (total_patches, 3, patch_size, patch_size), dtype=np.float32 + ) + out_offset = 0 + with numba_workqueue_threading_layer(): + for ( + frames, + new_height, + new_width, + padded_height, + padded_width, + num_patches, + ) in prepared: + _write_fused_patches( + frames, + pixel_values_np, + out_offset, + new_height, + new_width, + padded_height, + padded_width, + patch_size, + self.normalize_lut, + ) + out_offset += num_patches + + data = { + "pixel_values": torch.from_numpy(pixel_values_np), + "grid_thws": torch.from_numpy(grid_thws_np), + } + return BatchFeature(data=data, tensor_type=return_tensors) + + def __repr__(self): + return f"KimiK25FusedVisionProcessor(media_proc_cfg={self.media_proc_cfg})" + + def to_dict(self) -> dict[str, Any]: + output = super().to_dict() + output["media_proc_cfg"] = self.media_proc_cfg + if "media_processor" in output: + del output["media_processor"] + return output + + @classmethod + def from_dict(cls, config_dict: dict[str, Any], **kwargs): + config = config_dict.copy() + media_proc_cfg = config.pop("media_proc_cfg", {}) + return cls(media_proc_cfg=media_proc_cfg, **config, **kwargs) + + def to_json_string(self): + dictionary = self.to_dict() + for key, value in dictionary.items(): + if hasattr(value, "tolist"): + dictionary[key] = value.tolist() + return json.dumps(dictionary, indent=2, sort_keys=True) + "\n" diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index 043798a584b..78dcde6f061 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -487,6 +487,11 @@ def has_nixl_ep() -> bool: return _has_module("nixl_ep") +def is_numba_available() -> bool: + """Whether the optional `numba` package is available.""" + return _has_module("numba") + + def has_triton_kernels() -> bool: """Whether the optional `triton_kernels` package is available.""" is_available = _has_module("triton_kernels") or _has_module( diff --git a/vllm/utils/jit_monitor.py b/vllm/utils/jit_monitor.py index 4565ffdae06..7ba8ecde653 100644 --- a/vllm/utils/jit_monitor.py +++ b/vllm/utils/jit_monitor.py @@ -20,9 +20,10 @@ Currently monitors: (via ``knobs.runtime.jit_post_compile_hook``) """ +import contextlib import functools import os -from collections.abc import Mapping +from collections.abc import Iterator, Mapping from typing import Literal from vllm.logger import init_logger @@ -287,3 +288,32 @@ def _setup_cutedsl_jit_hook() -> None: cute.compile = _compile_with_monitor _cutedsl_hook_installed = True + + +@contextlib.contextmanager +def numba_workqueue_threading_layer() -> Iterator[None]: + """Force numba's fork-safe `workqueue` threading layer for this block. + + GNU OpenMP (numba's default `omp` threading layer) aborts the process + if a forked child re-enters an OpenMP-active runtime. vLLM forks the + EngineCore subprocess from a process that may already have launched + numba's parallel accelerator, so the first call to any + `@njit(parallel=True)` function must happen under `workqueue` instead. + The threading layer choice is sticky for the life of the process once + launched, so restoring the config on exit does not undo the effect. + """ + import numba + + key = "NUMBA_THREADING_LAYER" + previous_env = os.environ.get(key) + previous_config = numba.config.THREADING_LAYER + os.environ[key] = "workqueue" + numba.config.THREADING_LAYER = "workqueue" + try: + yield + finally: + if previous_env is None: + os.environ.pop(key, None) + else: + os.environ[key] = previous_env + numba.config.THREADING_LAYER = previous_config