[VLM] GPU Image Preprocessing for Kimi-K2.5 (#22368)

.codespellrc

@@ -1,3 +1,3 @@
 [codespell]
-ignore-words-list = ans, als, hel, boostrap, childs, te, vas, hsa, ment, cann, thi, makro, wil, rouge, PRIS, ather, MIS
+ignore-words-list = ans, als, hel, boostrap, childs, te, vas, hsa, ment, cann, thi, makro, wil, rouge, PRIS, ather, MIS, medias
 skip = *.json,*.jsonl,*.patch,*.txt

python/sglang/benchmark/datasets/image.py

@@ -118,6 +118,14 @@ def create_mm_data_row(
         prompt_str = f"<image>{text_prompt}"
 
     # Calculate total tokens (text + vision)
+    if type(processor).__name__ == "KimiK25Processor":
+        medias = [{"type": "image", "image": img} for img in images]
+        prompt_len = processor(
+            text=prompt_str,
+            medias=medias,
+            return_tensors="pt",
+        )["input_ids"].numel()
+    else:
         prompt_len = processor(
             text=[prompt_str],
             images=images,

python/sglang/srt/multimodal/processors/kimi_k25.py

@@ -1,7 +1,12 @@
+import math
 import re
-from typing import Dict, List, Tuple, Union
+from collections import defaultdict
+from typing import Dict, List, Union
 
+import numpy as np
 import torch
+import torch.nn.functional as F
+from PIL import Image
 
 from sglang.srt.managers.schedule_batch import (
     Modality,
@@ -16,11 +21,317 @@ from sglang.srt.multimodal.processors.base_processor import (
     MultimodalSpecialTokens,
 )
 
+# ---------------------------------------------------------------------------
+# GPU image preprocessing utilities (resize, pad, normalize, patchify on CUDA)
+# ---------------------------------------------------------------------------
+
+
+def navit_resize_config(
+    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 = None,
+) -> dict:
+    """Compute NaViT resize target dimensions and token count.
+
+    Pure math -- no image data needed, only (width, height).
+    """
+    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,
+    }
+
+
+def _get_image_dimensions(image: Union[torch.Tensor, Image.Image]) -> tuple[int, int]:
+    """Get (width, height) from a CUDA tensor or PIL Image."""
+    if isinstance(image, torch.Tensor):
+        # nvJPEG returns (C, H, W) uint8
+        return image.shape[2], image.shape[1]
+    return image.size  # PIL returns (width, height)
+
+
+def _pil_to_cuda_chw(image: Image.Image) -> torch.Tensor:
+    """Convert PIL Image to (C, H, W) uint8 CUDA tensor."""
+    arr = np.asarray(image.convert("RGB"))
+    return torch.from_numpy(arr).permute(2, 0, 1).cuda()
+
+
+def _process_single_image(
+    image: Union[torch.Tensor, Image.Image],
+    config: dict,
+    image_mean: torch.Tensor,
+    image_std_inv: torch.Tensor,
+    patch_size: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Process a single image on GPU: resize -> pad -> normalize -> patchify."""
+    if isinstance(image, Image.Image):
+        image = _pil_to_cuda_chw(image)
+
+    new_h, new_w = config["new_height"], config["new_width"]
+    pad_h, pad_w = config["pad_height"], config["pad_width"]
+
+    x = image.unsqueeze(0).float()
+    x = F.interpolate(x, size=(new_h, new_w), mode="bicubic", align_corners=False)
+
+    if pad_h > 0 or pad_w > 0:
+        x = F.pad(x, (0, pad_w, 0, pad_h), value=0.0)
+
+    x = x / 255.0
+    x = (x - image_mean) * image_std_inv
+
+    _, C, H, W = x.shape
+    T = 1
+    gh, gw = H // patch_size, W // patch_size
+    x = x.view(T, C, gh, patch_size, gw, patch_size)
+    x = x.permute(0, 2, 4, 1, 3, 5).reshape(-1, C, patch_size, patch_size)
+
+    grid_thw = torch.tensor([T, gh, gw], dtype=torch.int64, device=x.device)
+    return x, grid_thw
+
+
+def _gpu_preprocess_images(
+    images: list[Union[torch.Tensor, Image.Image]],
+    resize_configs: list[dict],
+    image_mean: torch.Tensor,
+    image_std_inv: torch.Tensor,
+    patch_size: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """GPU preprocessing pipeline for a batch of images.
+
+    Groups images with the same target padded size for batch processing.
+    """
+    n = len(images)
+    if n == 0:
+        device = image_mean.device
+        return (
+            torch.empty(0, 3, patch_size, patch_size, device=device),
+            torch.empty(0, 3, dtype=torch.int64, device=device),
+        )
+
+    groups = defaultdict(list)
+    for idx, (image, config) in enumerate(zip(images, resize_configs)):
+        padded_h = config["new_height"] + config["pad_height"]
+        padded_w = config["new_width"] + config["pad_width"]
+        target_h = config["new_height"]
+        target_w = config["new_width"]
+        groups[(target_h, target_w, padded_h, padded_w)].append((idx, image, config))
+
+    all_patches = [None] * n
+    all_grids = [None] * n
+
+    for (target_h, target_w, padded_h, padded_w), group in groups.items():
+        if len(group) == 1:
+            idx, image, config = group[0]
+            patches, grid = _process_single_image(
+                image, config, image_mean, image_std_inv, patch_size
+            )
+            all_patches[idx] = patches
+            all_grids[idx] = grid
+        else:
+            tensors = []
+            for _, image, _ in group:
+                if isinstance(image, Image.Image):
+                    image = _pil_to_cuda_chw(image)
+                tensors.append(image.unsqueeze(0).float())
+
+            resized = []
+            for t in tensors:
+                r = F.interpolate(
+                    t, size=(target_h, target_w), mode="bicubic", align_corners=False
+                )
+                resized.append(r)
+            batch = torch.cat(resized, dim=0)
+
+            pad_h = padded_h - target_h
+            pad_w = padded_w - target_w
+            if pad_h > 0 or pad_w > 0:
+                batch = F.pad(batch, (0, pad_w, 0, pad_h), value=0.0)
+
+            batch = batch / 255.0
+            batch = (batch - image_mean) * image_std_inv
+
+            B, C, H, W = batch.shape
+            T = 1
+            gh, gw = H // patch_size, W // patch_size
+            batch = batch.view(B, C, gh, patch_size, gw, patch_size)
+            batch = batch.permute(0, 2, 4, 1, 3, 5).reshape(
+                B, -1, C, patch_size, patch_size
+            )
+
+            grid = torch.tensor([T, gh, gw], dtype=torch.int64, device=batch.device)
+            for i, (idx, _, _) in enumerate(group):
+                all_patches[idx] = batch[i]
+                all_grids[idx] = grid
+
+    pixel_values = torch.cat(all_patches, dim=0)
+    grid_thws = torch.stack(all_grids, dim=0)
+    return pixel_values, grid_thws
+
+
+# ---------------------------------------------------------------------------
+# Kimi K2.5 GPU processor wrapper
+# ---------------------------------------------------------------------------
+
+
+class KimiGPUProcessorWrapper:
+    """Wraps Kimi's HF processor to do GPU image preprocessing.
+
+    GPU path: nvJPEG CUDA tensor / PIL -> _gpu_preprocess_images()
+    CPU fallback: PIL -> medias kwarg -> original HF KimiK25Processor.__call__
+
+    Exposes attributes that base class's process_mm_data needs so it behaves
+    like a normal HF processor from the outside.
+    """
+
+    def __init__(
+        self,
+        hf_processor,
+        image_token,
+        patch_size,
+        merge_kernel_size,
+        in_patch_limit,
+        patch_limit_on_one_side,
+        fixed_output_tokens,
+        image_mean,
+        image_std,
+    ):
+        self._hf_processor = hf_processor
+        self._image_token = image_token
+        self._patch_size = patch_size
+        self._merge_kernel_size = merge_kernel_size
+        self._in_patch_limit = in_patch_limit
+        self._patch_limit_on_one_side = patch_limit_on_one_side
+        self._fixed_output_tokens = fixed_output_tokens
+        self._image_mean = image_mean
+        self._image_std = image_std
+        self._gpu_norm_tensors = None
+
+        # Explicitly expose attributes that base class process_mm_data needs:
+        # - image_processor: checked via isinstance(..., BaseImageProcessorFast)
+        # - tokenizer: used for tokenization
+        # - media_processor: used by CPU fallback path
+        self.image_processor = hf_processor.image_processor
+        self.tokenizer = hf_processor.tokenizer
+        self.media_processor = hf_processor.media_processor
+
+    def __call__(self, text=None, images=None, **kwargs):
+        # process_mm_data passes images via kwargs["images"]
+        images = images or kwargs.pop("images", None)
+
+        if images and torch.cuda.is_available():
+            return self._gpu_call(text, images)
+        return self._cpu_call(text, images, **kwargs)
+
+    def _gpu_call(self, text, images):
+        """Bypass HF KimiK25VisionProcessor.preprocess entirely -- use GPU ops."""
+        input_text = text[0] if isinstance(text, list) else text
+
+        # 1. Compute resize configs (CPU math)
+        resize_configs = []
+        for image in images:
+            w, h = _get_image_dimensions(image)
+            resize_configs.append(
+                navit_resize_config(
+                    w,
+                    h,
+                    self._patch_size,
+                    self._merge_kernel_size,
+                    self._in_patch_limit,
+                    self._patch_limit_on_one_side,
+                    self._fixed_output_tokens,
+                )
+            )
+
+        # 2. Expand image tokens
+        parts = input_text.split(self._image_token)
+        result = [parts[0]]
+        for config, part in zip(resize_configs, parts[1:]):
+            result.append(self._image_token * config["num_tokens"] + part)
+        input_text = "".join(result)
+
+        # 3. Tokenize
+        text_inputs = self._hf_processor.tokenizer(input_text, return_tensors="pt")
+
+        # 4. GPU image preprocessing
+        image_mean, image_std_inv = self._get_gpu_norm_tensors()
+        pixel_values, grid_thws = _gpu_preprocess_images(
+            images, resize_configs, image_mean, image_std_inv, self._patch_size
+        )
+
+        return {
+            "input_ids": text_inputs["input_ids"],
+            "pixel_values": pixel_values,
+            "grid_thws": grid_thws,
+        }
+
+    def _cpu_call(self, text, images, **kwargs):
+        """Fallback: token expansion + medias kwarg -> original HF processor."""
+        input_text = text[0] if isinstance(text, list) else text
+
+        if images:
+            # Token expansion via media_tokens_calculator
+            parts = input_text.split(self._image_token)
+            result = [parts[0]]
+            for image, part in zip(images, parts[1:]):
+                num_tokens = self._hf_processor.media_processor.media_tokens_calculator(
+                    {"type": "image", "image": image}
+                )
+                result.append(self._image_token * num_tokens + part)
+            input_text = "".join(result)
+
+            # Convert to medias format for Kimi's HF processor
+            kwargs["medias"] = [{"type": "image", "image": img} for img in images]
+
+        return self._hf_processor(text=[input_text], **kwargs)
+
+    def _get_gpu_norm_tensors(self, device="cuda"):
+        if self._gpu_norm_tensors is None:
+            image_mean = torch.tensor(
+                self._image_mean, device=device, dtype=torch.float32
+            ).view(1, 3, 1, 1)
+            image_std_inv = (
+                1.0 / torch.tensor(self._image_std, device=device, dtype=torch.float32)
+            ).view(1, 3, 1, 1)
+            self._gpu_norm_tensors = (image_mean, image_std_inv)
+        return self._gpu_norm_tensors
+
+
+# ---------------------------------------------------------------------------
+# Kimi K2.5 SGLang multimodal processor
+# ---------------------------------------------------------------------------
+
 
 # Compatible with KimiVLForConditionalGeneration
 class KimiK2_5VLImageProcessor(SGLangBaseProcessor):
     models = [KimiK25ForConditionalGeneration]
-    gpu_image_decode = False  # KimiK2.5VL HF processor does not support tensor inputs
+    gpu_image_decode = True  # nvJPEG for JPEG, PIL fallback for others
 
     def __init__(self, hf_config, server_args, _processor, *args, **kwargs):
         super().__init__(hf_config, server_args, _processor, *args, **kwargs)
@@ -31,6 +342,22 @@ class KimiK2_5VLImageProcessor(SGLangBaseProcessor):
             image_token_regex=re.compile(r"(?:<\|media_pad\|>)+"),
         ).build(_processor)
 
+        # Extract media processing config from HF processor
+        media_proc_cfg = _processor.media_processor.media_proc_cfg
+
+        # Replace with GPU-capable wrapper
+        self._processor = KimiGPUProcessorWrapper(
+            _processor,
+            image_token=self.mm_tokens.image_token,
+            patch_size=media_proc_cfg["patch_size"],
+            merge_kernel_size=media_proc_cfg["merge_kernel_size"],
+            in_patch_limit=media_proc_cfg["in_patch_limit"],
+            patch_limit_on_one_side=media_proc_cfg["patch_limit_on_one_side"],
+            fixed_output_tokens=media_proc_cfg.get("fixed_output_tokens"),
+            image_mean=media_proc_cfg["image_mean"],
+            image_std=media_proc_cfg["image_std"],
+        )
+
     async def process_mm_data_async(
         self,
         image_data: List[Union[str, bytes, Dict]],
@@ -44,7 +371,6 @@ class KimiK2_5VLImageProcessor(SGLangBaseProcessor):
             image_data=image_data,
             multimodal_tokens=self.mm_tokens,
         )
-        prompt = base_output.input_text
 
         mm_items, input_ids, _ = self.process_and_combine_mm_data(
             base_output, self.mm_tokens
@@ -119,41 +445,3 @@ class KimiK2_5VLImageProcessor(SGLangBaseProcessor):
             mm_items=mm_items,
             im_token_id=image_token_id,
         )
-
-    def _process_and_collect_mm_items(
-        self, input_text: str, images=None, audios=None, videos=None, **kwargs
-    ) -> Tuple[List[MultimodalDataItem], torch.Tensor, dict]:
-        """
-        Helper method to process multimodal data and create mm_items in one step.
-
-        Returns:
-            Tuple of (created mm_items, input_ids)
-        """
-
-        parts = input_text.split(self.mm_tokens.image_token)
-
-        result = [parts[0]]
-        for image, part in zip(images, parts[1:]):
-            num_tokens = self._processor.media_processor.media_tokens_calculator(
-                {"type": "image", "image": image}
-            )
-            result.append(self.mm_tokens.image_token * num_tokens + part)
-
-        input_text = "".join(result)
-
-        if images:  # for kimi k2 vl
-            mediums = []
-            for image in images:
-                mediums.append({"type": "image", "image": image})
-            key = "_medias"[1:]  # bypass lint
-            kwargs[key] = mediums
-            images = None
-
-        ret = self.process_mm_data(
-            input_text=input_text, images=images, audios=audios, videos=videos, **kwargs
-        )
-
-        input_ids = ret["input_ids"].flatten()
-        collected_items = self.collect_mm_items_from_processor_output(ret)
-
-        return collected_items, input_ids, ret