[perf]Add fused Kimi image preprocessing (#47416)
Signed-off-by: Kevin-XiongC <kevin_xiong1997@outlook.com> Signed-off-by: Kevin_Xiong <kevin_xiong1997@outlook.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: Isotr0py <2037008807@qq.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Kevin_Xiong committed
Jul 6, 2026 at 23:46 UTC
5ad11172b791104f2d9df3c7877a53a5aeb28878
5 files changed
+402
-2
vllm/model_executor/models/kimi_k25.py
+10
@@ -56,6 +56,10 @@ from vllm.sequence import IntermediateTensors
56
from vllm.transformers_utils.configs.kimi_k25 import KimiK25Config
57
from vllm.transformers_utils.processor import cached_get_image_processor
58
from vllm.transformers_utils.processors.kimi_k25 import KimiK25Processor
59
+from vllm.transformers_utils.processors.kimi_k25_vision_fused import (
60
+ KimiK25FusedVisionProcessor,
61
+)
62
+from vllm.utils.import_utils import is_numba_available
63
from vllm.utils.tensor_schema import TensorSchema, TensorShape
64
65
from .utils import (
@@ -108,10 +112,16 @@ class KimiK25ProcessingInfo(BaseProcessingInfo):
112
self.hf_config = hf_config = self.get_hf_config()
113
114
tokenizer = self.get_tokenizer()
115
+ processor_cls = KimiK25FusedVisionProcessor if is_numba_available() else None
116
+ logger.info_once(
117
+ "Using %s image preprocessing for Kimi-K2.5/K2.6 vision chunks.",
118
+ "fused CPU" if processor_cls is not None else "remote HF",
119
+ )
120
image_processor = cached_get_image_processor(
121
self.ctx.model_config.model,
122
revision=self.ctx.model_config.revision,
123
trust_remote_code=self.ctx.model_config.trust_remote_code,
124
+ processor_cls_overrides=processor_cls,
125
)
126
127
# Resolve token ID from the tokenizer because transformers v5
vllm/transformers_utils/processor.py
+4
-1
@@ -81,6 +81,7 @@ _transformers_v4_compatibility_import()
81
_transformers_v4_compatibility_init()
82
83
_P = TypeVar("_P", bound=ProcessorMixin, default=ProcessorMixin)
84
+_I = TypeVar("_I", bound=BaseImageProcessor, default=BaseImageProcessor)
85
_V = TypeVar("_V", bound=BaseVideoProcessor, default=BaseVideoProcessor)
86
87
@@ -440,12 +441,14 @@ def get_image_processor(
441
*args: Any,
442
revision: str | None = None,
443
trust_remote_code: bool = False,
444
+ processor_cls_overrides: type[_I] | None = None,
445
**kwargs: Any,
446
):
447
"""Load an image processor for the given model name via HuggingFace."""
448
try:
449
processor_name = convert_model_repo_to_path(processor_name)
448
- processor = AutoImageProcessor.from_pretrained(
450
+ processor_cls = processor_cls_overrides or AutoImageProcessor
451
+ processor = processor_cls.from_pretrained(
452
processor_name,
453
*args,
454
revision=revision,
vllm/transformers_utils/processors/kimi_k25_vision_fused.py
new
+352
@@ -0,0 +1,352 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Optimized CPU image processor for Kimi-K2.5/K2.6 vision chunks."""
4
+
5
+import io
6
+import json
7
+import math
8
+from typing import Any
9
+
10
+import numpy as np
11
+import pybase64 as base64
12
+import torch
13
+from PIL import Image
14
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
15
+from transformers.utils import TensorType
16
+
17
+from vllm.utils.import_utils import is_numba_available
18
+from vllm.utils.jit_monitor import numba_workqueue_threading_layer
19
+
20
+if is_numba_available():
21
+ from numba import njit, prange
22
+
23
+ @njit(parallel=True, cache=True)
24
+ def _write_fused_patches(
25
+ frames: np.ndarray,
26
+ out: np.ndarray,
27
+ out_offset: int,
28
+ new_h: int,
29
+ new_w: int,
30
+ padded_h: int,
31
+ padded_w: int,
32
+ patch_size: int,
33
+ normalize_lut: np.ndarray,
34
+ ) -> None:
35
+ # frames: [T, new_h, new_w, 3] uint8, without padding.
36
+ # out: [total_patches, 3, patch_size, patch_size] float32.
37
+ t_size = frames.shape[0]
38
+ patch_h = padded_h // patch_size
39
+ patch_w = padded_w // patch_size
40
+ total = t_size * padded_h * padded_w * 3
41
+ hwc = padded_h * padded_w * 3
42
+ wc = padded_w * 3
43
+
44
+ for linear in prange(total):
45
+ t = linear // hwc
46
+ rem = linear - t * hwc
47
+ y = rem // wc
48
+ rem = rem - y * wc
49
+ x = rem // 3
50
+ c = rem - x * 3
51
+
52
+ value = frames[t, y, x, c] if y < new_h and x < new_w else 0
53
+
54
+ patch_idx = (
55
+ out_offset
56
+ + t * patch_h * patch_w
57
+ + (y // patch_size) * patch_w
58
+ + (x // patch_size)
59
+ )
60
+ out[patch_idx, c, y % patch_size, x % patch_size] = normalize_lut[value, c]
61
+
62
+else:
63
+
64
+ def _write_fused_patches(*args: Any, **kwargs: Any) -> None:
65
+ raise RuntimeError("numba is required for fused Kimi image preprocessing")
66
+
67
+
68
+def navit_resize_image(
69
+ width: int,
70
+ height: int,
71
+ patch_size: int,
72
+ merge_kernel_size: int,
73
+ in_patch_limit: int,
74
+ patch_limit_on_one_side: int,
75
+ fixed_output_tokens: int | None,
76
+) -> dict[str, int]:
77
+ s1 = math.sqrt(
78
+ in_patch_limit
79
+ / (max(1.0, width // patch_size) * max(1.0, height // patch_size))
80
+ )
81
+ s2 = patch_limit_on_one_side * patch_size / width
82
+ s3 = patch_limit_on_one_side * patch_size / height
83
+ scale = min(1.0, s1, s2, s3)
84
+ new_w = min(max(1, int(width * scale)), patch_limit_on_one_side * patch_size)
85
+ new_h = min(max(1, int(height * scale)), patch_limit_on_one_side * patch_size)
86
+
87
+ factor = merge_kernel_size * patch_size
88
+ pad_height = (factor - new_h % factor) % factor
89
+ pad_width = (factor - new_w % factor) % factor
90
+
91
+ if fixed_output_tokens is not None:
92
+ num_tokens = fixed_output_tokens
93
+ else:
94
+ token_height = (new_h + pad_height) // factor
95
+ token_width = (new_w + pad_width) // factor
96
+ num_tokens = token_height * token_width
97
+
98
+ return {
99
+ "num_tokens": num_tokens,
100
+ "new_width": new_w,
101
+ "new_height": new_h,
102
+ "pad_width": pad_width,
103
+ "pad_height": pad_height,
104
+ "sampled_nframes": 1,
105
+ }
106
+
107
+
108
+def navit_resize_video(
109
+ width: int,
110
+ height: int,
111
+ nframes: int,
112
+ avg_fps: float,
113
+ sample_fps: float,
114
+ patch_size: int,
115
+ merge_kernel_size: int,
116
+ in_patch_limit_each_frame: int,
117
+ patch_limit_on_one_side: int,
118
+ in_patch_limit_total: int | None,
119
+ max_num_frames_each_video: int | None,
120
+ fixed_output_tokens_each_frame: int | None,
121
+) -> dict[str, int]:
122
+ sample_fps = min(sample_fps, avg_fps)
123
+ sampled_nframes = max(round(nframes * sample_fps / avg_fps), 1)
124
+ if max_num_frames_each_video is not None:
125
+ sampled_nframes = min(sampled_nframes, max_num_frames_each_video)
126
+
127
+ if in_patch_limit_total is not None:
128
+ in_patch_limit_each_frame = min(
129
+ round(in_patch_limit_total / sampled_nframes),
130
+ in_patch_limit_each_frame,
131
+ )
132
+
133
+ ret = navit_resize_image(
134
+ width,
135
+ height,
136
+ patch_size,
137
+ merge_kernel_size,
138
+ in_patch_limit_each_frame,
139
+ patch_limit_on_one_side,
140
+ fixed_output_tokens_each_frame,
141
+ )
142
+ ret["sampled_nframes"] = sampled_nframes
143
+ return ret
144
+
145
+
146
+def _to_pil(data: Any) -> Image.Image:
147
+ if hasattr(data, "media") and hasattr(data, "original_bytes"):
148
+ data = data.media
149
+ if isinstance(data, Image.Image):
150
+ return data if data.mode == "RGB" else data.convert("RGB")
151
+ if isinstance(data, str):
152
+ if data.startswith("data:"):
153
+ raw_base64 = data.split(",", 1)[1]
154
+ return Image.open(io.BytesIO(base64.b64decode(raw_base64))).convert("RGB")
155
+ return Image.open(data).convert("RGB")
156
+ if isinstance(data, bytes):
157
+ return Image.open(io.BytesIO(data)).convert("RGB")
158
+ raise ValueError(f"Unsupported data type: {type(data)}")
159
+
160
+
161
+def _ensure_media_type(media: dict[str, Any]) -> dict[str, Any]:
162
+ if media["type"] == "image":
163
+ media["image"] = _to_pil(media["image"])
164
+ return media
165
+ if media["type"] == "video_chunk":
166
+ media["video_chunk"] = [_to_pil(frame) for frame in media["video_chunk"]]
167
+ return media
168
+ raise ValueError(f"Unsupported media type: {media['type']}")
169
+
170
+
171
+class KimiK25FusedVisionProcessor(BaseImageProcessor):
172
+ model_type = "kimi_k25"
173
+
174
+ def __init__(self, media_proc_cfg: dict[str, Any], **kwargs: Any) -> None:
175
+ super().__init__(**kwargs)
176
+ media_proc_cfg = dict(media_proc_cfg)
177
+ merge_kernel_size = media_proc_cfg["merge_kernel_size"]
178
+ if isinstance(merge_kernel_size, (list, tuple)):
179
+ media_proc_cfg["merge_kernel_size"] = int(merge_kernel_size[0])
180
+ self.media_proc_cfg = media_proc_cfg
181
+ self.num_frames_per_chunk = media_proc_cfg["temporal_merge_kernel_size"]
182
+ values = np.arange(256, dtype=np.float32)[:, None]
183
+ image_mean = np.asarray(media_proc_cfg["image_mean"], dtype=np.float32)
184
+ image_std_inv = 1.0 / np.asarray(media_proc_cfg["image_std"], dtype=np.float32)
185
+ self.normalize_lut = (values / 255.0 - image_mean[None, :]) * image_std_inv[
186
+ None, :
187
+ ]
188
+
189
+ def media_tokens_calculator(self, media: dict[str, Any]) -> int:
190
+ media = _ensure_media_type(media)
191
+ ret = self.get_resize_config(media)
192
+ return ret["num_tokens"]
193
+
194
+ def get_resize_config(self, media_input: dict[str, Any]) -> dict[str, int]:
195
+ if media_input["type"] == "image":
196
+ width, height = media_input["image"].size
197
+ return navit_resize_image(
198
+ width,
199
+ height,
200
+ self.media_proc_cfg["patch_size"],
201
+ self.media_proc_cfg["merge_kernel_size"],
202
+ self.media_proc_cfg["in_patch_limit"],
203
+ self.media_proc_cfg["patch_limit_on_one_side"],
204
+ self.media_proc_cfg["fixed_output_tokens"],
205
+ )
206
+
207
+ if media_input["type"] == "video_chunk":
208
+ frame = media_input["video_chunk"][0]
209
+ width, height = frame.size
210
+ num_frames = len(media_input["video_chunk"])
211
+ in_patch_limit_each_frame = self.media_proc_cfg["in_patch_limit_each_frame"]
212
+ if in_patch_limit_each_frame is None:
213
+ in_patch_limit_each_frame = self.media_proc_cfg["in_patch_limit"]
214
+
215
+ return navit_resize_video(
216
+ width,
217
+ height,
218
+ num_frames,
219
+ 1.0,
220
+ math.inf,
221
+ self.media_proc_cfg["patch_size"],
222
+ self.media_proc_cfg["merge_kernel_size"],
223
+ in_patch_limit_each_frame,
224
+ self.media_proc_cfg["patch_limit_on_one_side"],
225
+ self.media_proc_cfg["in_patch_limit_video"],
226
+ None,
227
+ self.media_proc_cfg["fixed_output_tokens"],
228
+ )
229
+
230
+ raise ValueError(f"Unsupported type: {media_input['type']}")
231
+
232
+ @staticmethod
233
+ def resize_image(image: Image.Image, new_width: int, new_height: int) -> np.ndarray:
234
+ image = image.resize((new_width, new_height), resample=Image.Resampling.BICUBIC)
235
+ return np.asarray(image)
236
+
237
+ def preprocess(
238
+ self,
239
+ medias: list[dict[str, Any]],
240
+ return_tensors: str | TensorType | None = None,
241
+ ) -> BatchFeature:
242
+ if not isinstance(medias, list):
243
+ medias = [medias]
244
+ if not medias:
245
+ return BatchFeature(data={}, tensor_type=return_tensors)
246
+
247
+ if njit is None:
248
+ raise RuntimeError("numba is required for fused Kimi image preprocessing")
249
+
250
+ patch_size = int(self.media_proc_cfg["patch_size"])
251
+ prepared = []
252
+ grid_thws_np = np.empty((len(medias), 3), dtype=np.int64)
253
+ total_patches = 0
254
+
255
+ for idx, item in enumerate(medias):
256
+ item = _ensure_media_type(item)
257
+ resize_config = self.get_resize_config(item)
258
+ new_width = resize_config["new_width"]
259
+ new_height = resize_config["new_height"]
260
+ pad_width = resize_config["pad_width"]
261
+ pad_height = resize_config["pad_height"]
262
+ padded_width = new_width + pad_width
263
+ padded_height = new_height + pad_height
264
+
265
+ if item["type"] == "image":
266
+ image_np = self.resize_image(item["image"], new_width, new_height)
267
+ frames = image_np[np.newaxis, ...]
268
+ elif item["type"] == "video_chunk":
269
+ frames = np.stack(
270
+ [
271
+ self.resize_image(frame, new_width, new_height)
272
+ for frame in item["video_chunk"]
273
+ ],
274
+ axis=0,
275
+ )
276
+ else:
277
+ raise ValueError(f"Unsupported type: {item['type']}")
278
+
279
+ t_size = frames.shape[0]
280
+ grid_h = padded_height // patch_size
281
+ grid_w = padded_width // patch_size
282
+ grid_thws_np[idx, 0] = t_size
283
+ grid_thws_np[idx, 1] = grid_h
284
+ grid_thws_np[idx, 2] = grid_w
285
+
286
+ num_patches = t_size * grid_h * grid_w
287
+ prepared.append(
288
+ (
289
+ frames,
290
+ new_height,
291
+ new_width,
292
+ padded_height,
293
+ padded_width,
294
+ num_patches,
295
+ )
296
+ )
297
+ total_patches += num_patches
298
+
299
+ pixel_values_np = np.empty(
300
+ (total_patches, 3, patch_size, patch_size), dtype=np.float32
301
+ )
302
+ out_offset = 0
303
+ with numba_workqueue_threading_layer():
304
+ for (
305
+ frames,
306
+ new_height,
307
+ new_width,
308
+ padded_height,
309
+ padded_width,
310
+ num_patches,
311
+ ) in prepared:
312
+ _write_fused_patches(
313
+ frames,
314
+ pixel_values_np,
315
+ out_offset,
316
+ new_height,
317
+ new_width,
318
+ padded_height,
319
+ padded_width,
320
+ patch_size,
321
+ self.normalize_lut,
322
+ )
323
+ out_offset += num_patches
324
+
325
+ data = {
326
+ "pixel_values": torch.from_numpy(pixel_values_np),
327
+ "grid_thws": torch.from_numpy(grid_thws_np),
328
+ }
329
+ return BatchFeature(data=data, tensor_type=return_tensors)
330
+
331
+ def __repr__(self):
332
+ return f"KimiK25FusedVisionProcessor(media_proc_cfg={self.media_proc_cfg})"
333
+
334
+ def to_dict(self) -> dict[str, Any]:
335
+ output = super().to_dict()
336
+ output["media_proc_cfg"] = self.media_proc_cfg
337
+ if "media_processor" in output:
338
+ del output["media_processor"]
339
+ return output
340
+
341
+ @classmethod
342
+ def from_dict(cls, config_dict: dict[str, Any], **kwargs):
343
+ config = config_dict.copy()
344
+ media_proc_cfg = config.pop("media_proc_cfg", {})
345
+ return cls(media_proc_cfg=media_proc_cfg, **config, **kwargs)
346
+
347
+ def to_json_string(self):
348
+ dictionary = self.to_dict()
349
+ for key, value in dictionary.items():
350
+ if hasattr(value, "tolist"):
351
+ dictionary[key] = value.tolist()
352
+ return json.dumps(dictionary, indent=2, sort_keys=True) + "\n"
vllm/utils/import_utils.py
+5
@@ -487,6 +487,11 @@ def has_nixl_ep() -> bool:
487
return _has_module("nixl_ep")
488
489
490
+def is_numba_available() -> bool:
491
+ """Whether the optional `numba` package is available."""
492
+ return _has_module("numba")
493
+
494
+
495
def has_triton_kernels() -> bool:
496
"""Whether the optional `triton_kernels` package is available."""
497
is_available = _has_module("triton_kernels") or _has_module(
vllm/utils/jit_monitor.py
+31
-1
@@ -20,9 +20,10 @@ Currently monitors:
20
(via ``knobs.runtime.jit_post_compile_hook``)
21
"""
22
23
+import contextlib
24
import functools
25
import os
25
-from collections.abc import Mapping
26
+from collections.abc import Iterator, Mapping
27
from typing import Literal
28
29
from vllm.logger import init_logger
@@ -287,3 +288,32 @@ def _setup_cutedsl_jit_hook() -> None:
288
289
cute.compile = _compile_with_monitor
290
_cutedsl_hook_installed = True
291
+
292
+
293
+@contextlib.contextmanager
294
+def numba_workqueue_threading_layer() -> Iterator[None]:
295
+ """Force numba's fork-safe `workqueue` threading layer for this block.
296
+
297
+ GNU OpenMP (numba's default `omp` threading layer) aborts the process
298
+ if a forked child re-enters an OpenMP-active runtime. vLLM forks the
299
+ EngineCore subprocess from a process that may already have launched
300
+ numba's parallel accelerator, so the first call to any
301
+ `@njit(parallel=True)` function must happen under `workqueue` instead.
302
+ The threading layer choice is sticky for the life of the process once
303
+ launched, so restoring the config on exit does not undo the effect.
304
+ """
305
+ import numba
306
+
307
+ key = "NUMBA_THREADING_LAYER"
308
+ previous_env = os.environ.get(key)
309
+ previous_config = numba.config.THREADING_LAYER
310
+ os.environ[key] = "workqueue"
311
+ numba.config.THREADING_LAYER = "workqueue"
312
+ try:
313
+ yield
314
+ finally:
315
+ if previous_env is None:
316
+ os.environ.pop(key, None)
317
+ else:
318
+ os.environ[key] = previous_env
319
+ numba.config.THREADING_LAYER = previous_config