From afc868517b7fbd27f6bb654ee6b88a808a788691 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Mon, 3 Aug 2026 21:57:49 -0700 Subject: [PATCH] [Perf] Speed up the Kimi-K2.5 vision path and match PIL bicubic in the GPU resize (#33349) Co-authored-by: Mick --- python/sglang/srt/layers/attention/vision.py | 4 +- python/sglang/srt/models/kimi_k25.py | 111 +++--- python/sglang/srt/models/kimi_vl_moonvit.py | 31 +- python/sglang/srt/multimodal/mm_utils.py | 14 +- .../multimodal/processors/base_processor.py | 4 + .../srt/multimodal/processors/kimi_common.py | 18 +- .../srt/multimodal/processors/kimi_k25.py | 331 ++++++++++++------ .../srt/multimodal/processors/kimi_vl.py | 6 + test/manual/vlm/verify_k25_equivalence.py | 145 ++++++++ .../ops/test_kimi_k3_prerequisite_ops.py | 29 ++ test/registered/unit/models/test_kimi_k25.py | 249 ++++++++++++- 11 files changed, 761 insertions(+), 181 deletions(-) create mode 100644 test/manual/vlm/verify_k25_equivalence.py diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index a6fcc0e37..7f50be156 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/python/sglang/srt/layers/attention/vision.py @@ -148,6 +148,7 @@ def prepare_vision_attention_metadata( cu_seqlens: torch.Tensor, device: torch.device, *, + max_seqlen: Optional[int] = None, packed_indptrs: Optional[torch.Tensor] = None, sequence_lengths: Optional[torch.Tensor] = None, flashinfer_max_seqlen: Optional[int] = None, @@ -156,7 +157,8 @@ def prepare_vision_attention_metadata( cu_seqlens = cu_seqlens.to(device=device, dtype=torch.int32, non_blocking=True) seq_lens = cu_seqlens[1:] - cu_seqlens[:-1] - max_seqlen = int(seq_lens.max().item()) + if max_seqlen is None: + max_seqlen = int(seq_lens.max().item()) return VisionAttentionMetadata( cu_seqlens=cu_seqlens, seq_lens=seq_lens, diff --git a/python/sglang/srt/models/kimi_k25.py b/python/sglang/srt/models/kimi_k25.py index 1e920ace0..a4806a884 100644 --- a/python/sglang/srt/models/kimi_k25.py +++ b/python/sglang/srt/models/kimi_k25.py @@ -8,6 +8,11 @@ import torch.nn.functional as F from torch import nn from transformers.activations import PytorchGELUTanh +from sglang.kernels.ops.attention.vision_rope import ( + PreparedInplaceComplexRoPE, + apply_fused_qk_complex_rope_inplace, + prepare_fused_qk_complex_rope_inplace, +) from sglang.srt.configs.kimi_k25 import KimiK25Config, KimiK25VisionConfig from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation from sglang.srt.layers.attention.vision import ( @@ -33,32 +38,46 @@ from sglang.srt.managers.schedule_batch import ( from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.deepseek_v2 import DeepseekV3ForCausalLM -from sglang.srt.models.kimi_vl_moonvit import MLP2 +from sglang.srt.models.kimi_vl_moonvit import MLP2, tpool_patch_merger from sglang.srt.models.utils import WeightsMapper from sglang.srt.multimodal.mm_utils import ( + concat_or_single, materialize_multimodal_features, run_dp_sharded_mrope_vision_model, ) -from sglang.srt.runtime_context import get_mm, get_parallel, get_server_args -from sglang.srt.utils import add_prefix, is_npu +from sglang.srt.runtime_context import ( + get_exec, + get_mm, + get_parallel, + get_server_args, +) +from sglang.srt.utils import add_prefix, is_cuda, is_npu logger = logging.getLogger(__name__) _is_npu = is_npu() +_is_cuda = is_cuda() def apply_rope( - xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor, x_shape=None + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cis: torch.Tensor | PreparedInplaceComplexRoPE, + x_shape=None, ) -> tuple[torch.Tensor, torch.Tensor]: """ Args: (The leading dimensions of all inputs should be the same) xq: query, tensor of shape (..., num_heads, head_dim) xk: key, tensor of shape (..., num_heads, head_dim) - freqs_cis: tensor of shape (..., head_dim/2), dtype=torch.complex64. It contains the precomputed cis(freqs) for each position in the 2D grid. + freqs_cis: Complex frequencies for the portable path, or inputs + prepared once for the contiguous in-place CUDA kernel. Returns: xq_out, xk_out: tensors of shape (..., num_heads, head_dim) """ + if isinstance(freqs_cis, tuple): + return apply_fused_qk_complex_rope_inplace(xq, xk, freqs_cis) + freqs_cis = freqs_cis.unsqueeze(-2) # ..., 1, head_dim/2 # ..., num_heads, head_dim/2 xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2)) @@ -68,36 +87,6 @@ def apply_rope( return xq_out.type_as(xq), xk_out.type_as(xk) -def tpool_patch_merger( - x: torch.Tensor, - grid_thws: torch.Tensor, - merge_kernel_size: tuple[int, int] = (2, 2), -) -> list[torch.Tensor]: - d_model = x.size(-1) - - outputs = [] - pre_sum = 0 - for t, h, w in grid_thws.tolist(): - # Get the current sequence - seq = x[pre_sum : pre_sum + t * h * w] - # Reshape along self.merge_kernel_size and concat to the last dimension - kernel_height, kernel_width = merge_kernel_size - new_height, new_width = h // kernel_height, w // kernel_width - reshaped_seq = seq.view( - t, new_height, kernel_height, new_width, kernel_width, d_model - ) - reshaped_seq = ( - reshaped_seq.permute(0, 1, 3, 2, 4, 5).contiguous().mean(dim=0) - ) # temporal pooling - padded_seq = reshaped_seq.view( - new_height * new_width, kernel_height * kernel_width, -1 - ) - outputs.append(padded_seq) - pre_sum += t * h * w - - return outputs - - class MoonViTEncoderLayer(nn.Module): def __init__( @@ -433,6 +422,9 @@ class MoonVision3dPatchEmbed(nn.Module): class MoonViT3dEncoder(nn.Module): + # Class-level default so forward() stays usable on instances built with + # __new__ (unit tests skip __init__). + use_fused_rope = False def __init__( self, @@ -452,6 +444,9 @@ class MoonViT3dEncoder(nn.Module): self.rope_2d = Rope2DPosEmbRepeated( block_cfg["hidden_dim"] // block_cfg["num_heads"], 512, 512 ) + self.use_fused_rope = ( + _is_cuda and get_exec().deterministic.rl_on_policy_target is None + ) self.blocks = nn.ModuleList( [ MoonViTEncoderLayer( @@ -472,8 +467,18 @@ class MoonViT3dEncoder(nn.Module): rope_freqs_cis = self.rope_2d.get_freqs_cis( grid_thws=grid_thws, device=hidden_states.device ) + # The in-place kernel is a JIT template on the q/k dtype, and only + # fp16/bf16 are exercised by test_vision_rope_inplace. Leave other + # dtypes on the portable path rather than ship an untested one. + if self.use_fused_rope and hidden_states.dtype in ( + torch.float16, + torch.bfloat16, + ): + rope_freqs_cis = prepare_fused_qk_complex_rope_inplace(rope_freqs_cis) - sequence_lengths = (grid_thws[:, 0] * grid_thws[:, 1] * grid_thws[:, 2]).to( + sequence_lengths = grid_thws[:, 0] * grid_thws[:, 1] * grid_thws[:, 2] + max_seqlen = int(sequence_lengths.max().item()) + sequence_lengths = sequence_lengths.to( device=hidden_states.device, dtype=torch.int32 ) lengths = torch.cat( @@ -483,14 +488,12 @@ class MoonViT3dEncoder(nn.Module): ) ) - # FlashAttention needs a host integer. Compute it once per MoonViT - # forward and pass it to every encoder block instead of synchronizing - # once per block inside the attention backend. - max_seqlen = int(lengths.max().item()) cu_seqlens = lengths.to(hidden_states.device).cumsum(dim=0, dtype=torch.int32) forward_metadata = prepare_vision_attention_metadata( - cu_seqlens, device=hidden_states.device + cu_seqlens, + device=hidden_states.device, + max_seqlen=max_seqlen, ) for block in self.blocks: @@ -625,18 +628,12 @@ class K2VLMultiModalProjector(nn.Module): @torch.inference_mode() def mm_projection_auto( - mm_projector: torch.nn.Module | None, vt_output: list[torch.Tensor] -): - """Apply MM projector to vision tower outputs.""" - if mm_projector is None: - return vt_output - - num_embedding_list = [x.shape[0] for x in vt_output] - batched = torch.cat(vt_output, dim=0) - proj_out = mm_projector(batched) if mm_projector else batched - proj_out = proj_out.reshape(-1, proj_out.shape[-1]) - proj_out = torch.split(proj_out, num_embedding_list) - return proj_out + mm_projector: torch.nn.Module, + vt_output: Sequence[torch.Tensor], +) -> torch.Tensor: + """Project MoonViT's per-image outputs into one flattened (tokens, dim) feature.""" + projected = mm_projector(concat_or_single(vt_output, dim=0)) + return projected.reshape(-1, projected.shape[-1]) class KimiK25ForConditionalGeneration(nn.Module): @@ -769,9 +766,11 @@ class KimiK25ForConditionalGeneration(nn.Module): return image_features pixel_values = materialize_item_features(list(range(len(items)))) - image_embeds = self.vision_tower(pixel_values, grid_thws.to(device)) - proj_out = mm_projection_auto(self.mm_projector, image_embeds) - return torch.cat(proj_out, dim=0) + # grid_thws stays on the host: MoonViT3d only reads it as shape metadata + # (.tolist() in the pos-emb, RoPE and merger), so a device copy would + # buy one sync per read. Same contract the encoder-DP path relies on. + image_embeds = self.vision_tower(pixel_values, grid_thws) + return mm_projection_auto(self.mm_projector, image_embeds) def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs): pattern = MultiModalityDataPaddingPatternMultimodalTokens() diff --git a/python/sglang/srt/models/kimi_vl_moonvit.py b/python/sglang/srt/models/kimi_vl_moonvit.py index d1bea46ac..c4adf0a07 100644 --- a/python/sglang/srt/models/kimi_vl_moonvit.py +++ b/python/sglang/srt/models/kimi_vl_moonvit.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # ruff: noqa: E501 # Adapted from https://huggingface.co/moonshotai/Kimi-VL-A3B-Instruct/blob/main/modeling_kimi_vl.py -# This file is meant to be used in kimi_vl.py only +# Shared MoonViT building blocks for kimi_vl.py and kimi_k25.py # Copyright 2025 The Moonshot AI Team, DeepSeek-AI, and HuggingFace Inc. team. All rights reserved. # # The code is based on llava (llava/modeling_llava.py) and DeepSeek-V3 (DeepSeek-V3/modeling_deepseek.py), but modified for KimiVL. @@ -566,6 +566,35 @@ def patch_merger( return outputs +def tpool_patch_merger( + x: torch.Tensor, + grid_thws: torch.Tensor, + merge_kernel_size: tuple[int, int] = (2, 2), +) -> List[torch.Tensor]: + """Group spatial patches and average only across real video frames.""" + + d_model = x.size(-1) + outputs = [] + pre_sum = 0 + for t, h, w in grid_thws.tolist(): + seq = x[pre_sum : pre_sum + t * h * w] + kernel_height, kernel_width = merge_kernel_size + new_height, new_width = h // kernel_height, w // kernel_width + reshaped_seq = seq.view( + t, new_height, kernel_height, new_width, kernel_width, d_model + ) + reshaped_seq = reshaped_seq.permute(0, 1, 3, 2, 4, 5).contiguous() + reshaped_seq = reshaped_seq.squeeze(0) if t == 1 else reshaped_seq.mean(dim=0) + outputs.append( + reshaped_seq.view( + new_height * new_width, kernel_height * kernel_width, d_model + ) + ) + pre_sum += t * h * w + + return outputs + + class MoonVitVLProjector(nn.Module): def __init__( diff --git a/python/sglang/srt/multimodal/mm_utils.py b/python/sglang/srt/multimodal/mm_utils.py index 4065092df..3c4b50fa1 100644 --- a/python/sglang/srt/multimodal/mm_utils.py +++ b/python/sglang/srt/multimodal/mm_utils.py @@ -544,6 +544,12 @@ def run_dp_sharded_vision_model( return vision_embeddings +def concat_or_single(tensors: Sequence[torch.Tensor], dim: int = 0) -> torch.Tensor: + """Concatenate multiple tensors without copying a singleton input.""" + + return tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=dim) + + # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/vision.py def run_dp_sharded_mrope_vision_model( vision_model: torch.nn.Module, @@ -617,12 +623,12 @@ def run_dp_sharded_mrope_vision_model( # already concatenates these tensors before returning, so keep the # TP=1 DP-encoder path on the same projector-facing contract. if isinstance(image_embeds, list): - return torch.cat(image_embeds, dim=0) + return concat_or_single(image_embeds, dim=0) return image_embeds if rope_type == "rope_2d_packed": image_embeds = vision_model(pixel_values, grid_thw) if isinstance(image_embeds, list): - return torch.cat(image_embeds, dim=0) + return concat_or_single(image_embeds, dim=0) return image_embeds return vision_model(pixel_values, grid_thw=grid_thw) @@ -719,7 +725,7 @@ def run_dp_sharded_mrope_vision_model( else: image_embeds_local = vision_model(pixel_values_local, local_grid_thw) if isinstance(image_embeds_local, list): - image_embeds_local = torch.cat(image_embeds_local, dim=0) + image_embeds_local = concat_or_single(image_embeds_local, dim=0) else: out_dim = getattr(vision_model.config, "hidden_size", None) image_embeds_local = torch.empty( @@ -734,7 +740,7 @@ def run_dp_sharded_mrope_vision_model( pixel_values_local, torch.tensor(local_grid_thw_list) ) if isinstance(image_embeds_local, list): - image_embeds_local = torch.cat(image_embeds_local, dim=0) + image_embeds_local = concat_or_single(image_embeds_local, dim=0) else: # Handle empty case out_dim = getattr(vision_model, "out_hidden_size", None) diff --git a/python/sglang/srt/multimodal/processors/base_processor.py b/python/sglang/srt/multimodal/processors/base_processor.py index 2f91b40dd..b4e37c721 100644 --- a/python/sglang/srt/multimodal/processors/base_processor.py +++ b/python/sglang/srt/multimodal/processors/base_processor.py @@ -181,6 +181,9 @@ class BaseMultimodalProcessor(ABC): gpu_image_decode = True # Enable GPU decoding by default prefer_tokenized_input = False precompute_hash_before_cpu_transfer = False + # Set by processors that already build input_ids from the request's own + # tokens, so the retokenize-avoidance rebuild below has nothing to add. + preserve_processor_input_ids = False auto_mm_processor_worker_num = 1 auto_mm_io_worker_num = 4 supports_mm_processor_concurrency = False @@ -1510,6 +1513,7 @@ class BaseMultimodalProcessor(ABC): # Drift happens when Retokenization is not identity: Decode(X) => String => Re-tokenize => Y, X != Y. if ( envs.SGLANG_MM_AVOID_RETOKENIZE.get() + and not self.preserve_processor_input_ids and base_output.input_ids is not None and input_ids is not None and raw_images diff --git a/python/sglang/srt/multimodal/processors/kimi_common.py b/python/sglang/srt/multimodal/processors/kimi_common.py index 5371fee71..f59c8f3fe 100644 --- a/python/sglang/srt/multimodal/processors/kimi_common.py +++ b/python/sglang/srt/multimodal/processors/kimi_common.py @@ -3,7 +3,7 @@ Shared by KimiVLImageProcessor and KimiK2_5VLImageProcessor. """ -from typing import Union +from typing import Optional, Union import numpy as np import torch @@ -38,6 +38,22 @@ class KimiGridMMDataMixin: for image in images ] + @staticmethod + def count_image_placeholders(input_ids, image_token_id: int) -> Optional[int]: + """Structural image tokens in a pre-tokenized prompt, None if it is text.""" + if not isinstance(input_ids, (list, torch.Tensor)): + return None + + token_ids = np.asarray( + ( + input_ids.detach().flatten().cpu() + if isinstance(input_ids, torch.Tensor) + else input_ids + ), + dtype=np.int64, + ) + return int(np.count_nonzero(token_ids == image_token_id)) + def _num_image_tokens_from_grid( self, grid_thw: Union[torch.Tensor, np.ndarray, list, tuple] ) -> int: diff --git a/python/sglang/srt/multimodal/processors/kimi_k25.py b/python/sglang/srt/multimodal/processors/kimi_k25.py index 299a0ef1a..4fb0ab375 100644 --- a/python/sglang/srt/multimodal/processors/kimi_k25.py +++ b/python/sglang/srt/multimodal/processors/kimi_k25.py @@ -8,6 +8,7 @@ import torch import torch.nn.functional as F from PIL import Image +from sglang.kernels.ops.mm.process import normalize_and_patchify from sglang.srt.managers.schedule_batch import ( MultimodalProcessorOutput, ) @@ -79,6 +80,34 @@ def _get_image_dimensions(image: Union[torch.Tensor, Image.Image]) -> tuple[int, return image.size # PIL returns (width, height) +def _expand_image_token_ids( + input_ids: Union[List[int], torch.Tensor], + image_token_id: int, + image_token_counts: List[int], +) -> torch.Tensor: + """Expand one placeholder per image without tokenizing the media string again. + + Same rebuild as ``BaseMultimodalProcessor._expand_input_ids``, but staying in + the array domain skips a list round trip on the way to the output tensor. + test_kimi_k25.py pins the two together. + """ + if isinstance(input_ids, torch.Tensor): + input_ids = input_ids.detach().flatten().cpu().numpy() + input_ids = np.asarray(input_ids, dtype=np.int64) + + placeholder_mask = input_ids == image_token_id + placeholder_count = np.count_nonzero(placeholder_mask) + if placeholder_count != len(image_token_counts): + raise ValueError( + f"Expected {len(image_token_counts)} image placeholder token(s), " + f"found {placeholder_count}." + ) + + repeats = np.ones(input_ids.shape, dtype=np.int64) + repeats[placeholder_mask] = image_token_counts + return torch.from_numpy(np.repeat(input_ids, repeats)).unsqueeze(0) + + 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")) @@ -95,10 +124,16 @@ def _ensure_chw_rgb(image: torch.Tensor) -> torch.Tensor: consistent channel dimension. Normalize every tensor to 3 channels. Also move the tensor to the GPU (matching _pil_to_cuda_chw) so a CPU - input does not trip a device mismatch against the CUDA image_mean / - image_std_inv normalization constants downstream. No-op if already on - the device. + input does not trip a device mismatch against the CUDA normalization + constants downstream. No-op if already on the device. """ + if image.dtype != torch.uint8: + # Raw 0-255 is load-bearing downstream: the resize rounds to integers + # and the normalization folds in a 1/255 scale, so a normalized float + # image would collapse to 0/1 and then be rescaled. + raise ValueError( + f"Kimi GPU preprocessing expects raw uint8 pixels, got {image.dtype}" + ) image = image.cuda() if image.dim() == 2: # (H, W) grayscale -> (1, H, W) image = image.unsqueeze(0) @@ -111,39 +146,67 @@ def _ensure_chw_rgb(image: torch.Tensor) -> torch.Tensor: return image[:3] +def _resize_bicubic_if_needed( + image: torch.Tensor, target_height: int, target_width: int +) -> torch.Tensor: + """Track the checkpoint processor's ``PIL.Image.resize(..., BICUBIC)``. + + NaViT only ever downscales, and PIL's bicubic widens its kernel support by + the scale factor -- it always antialiases, which ``F.interpolate`` only does + under ``antialias=True``. PIL also returns uint8, so round and clip back to + integer pixels; the bicubic overshoot would otherwise survive normalization. + + Close but not exact: PIL evaluates uint8 resizes in fixed point, so a few + 8-bit levels of residual remain -- against PIL's float path we agree to 5e-3, + i.e. the kernel matches and only the arithmetic differs. + """ + image = image.float() + if image.shape[-2:] == (target_height, target_width): + return image + return ( + F.interpolate( + image, + size=(target_height, target_width), + mode="bicubic", + align_corners=False, + antialias=True, + ) + .round_() + .clamp_(0.0, 255.0) + ) + + +def _grid_thw_from_resize_config(config: dict, patch_size: int) -> tuple[int, int, int]: + height = config["new_height"] + config["pad_height"] + width = config["new_width"] + config["pad_width"] + return 1, height // patch_size, width // patch_size + + +def _to_cuda_chw(image: Union[torch.Tensor, Image.Image]) -> torch.Tensor: + if isinstance(image, Image.Image): + return _pil_to_cuda_chw(image) + return _ensure_chw_rgb(image) + + def _process_single_image( image: Union[torch.Tensor, Image.Image], config: dict, - image_mean: torch.Tensor, - image_std_inv: torch.Tensor, + image_scale: torch.Tensor, + image_bias: torch.Tensor, patch_size: int, -) -> tuple[torch.Tensor, 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) - else: - image = _ensure_chw_rgb(image) + image = _to_cuda_chw(image) new_h, new_w = config["new_height"], config["new_width"] - pad_h, pad_w = config["pad_height"], config["pad_width"] + padded_h = new_h + config["pad_height"] + padded_w = new_w + config["pad_width"] - x = image.unsqueeze(0).float() - x = F.interpolate(x, size=(new_h, new_w), mode="bicubic", align_corners=False) + x = _resize_bicubic_if_needed(image.unsqueeze(0), new_h, new_w) - 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 + return normalize_and_patchify( + x, image_scale, image_bias, patch_size, padded_h, padded_w + ).squeeze(0) def _resize_images_by_source_shape( @@ -166,22 +229,14 @@ def _resize_images_by_source_shape( for images in by_source_shape.values(): if len(images) == 1: index, image = images[0] - resized_by_index[index] = F.interpolate( - image.unsqueeze(0).float(), - size=(target_height, target_width), - mode="bicubic", - align_corners=False, + resized_by_index[index] = _resize_bicubic_if_needed( + image.unsqueeze(0), target_height, target_width ) continue - source_batch = torch.cat( - [image.unsqueeze(0) for _, image in images], dim=0 - ).float() - resized_batch = F.interpolate( - source_batch, - size=(target_height, target_width), - mode="bicubic", - align_corners=False, + source_batch = torch.cat([image.unsqueeze(0) for _, image in images], dim=0) + resized_batch = _resize_bicubic_if_needed( + source_batch, target_height, target_width ) for local_index, (index, _) in enumerate(images): resized_by_index[index] = resized_batch[local_index : local_index + 1] @@ -192,8 +247,8 @@ def _resize_images_by_source_shape( def _gpu_preprocess_images( images: list[Union[torch.Tensor, Image.Image]], resize_configs: list[dict], - image_mean: torch.Tensor, - image_std_inv: torch.Tensor, + image_scale: torch.Tensor, + image_bias: torch.Tensor, patch_size: int, ) -> tuple[torch.Tensor, torch.Tensor]: """GPU preprocessing pipeline for a batch of images. @@ -202,10 +257,10 @@ def _gpu_preprocess_images( """ n = len(images) if n == 0: - device = image_mean.device + device = image_scale.device return ( torch.empty(0, 3, patch_size, patch_size, device=device), - torch.empty(0, 3, dtype=torch.int64, device=device), + torch.empty(0, 3, dtype=torch.int64), ) groups = defaultdict(list) @@ -222,19 +277,13 @@ def _gpu_preprocess_images( 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 + patches = _process_single_image( + image, config, image_scale, image_bias, patch_size ) all_patches[idx] = patches - all_grids[idx] = grid + all_grids[idx] = _grid_thw_from_resize_config(config, patch_size) else: - indexed_images = [] - for idx, image, _ in group: - if isinstance(image, Image.Image): - image = _pil_to_cuda_chw(image) - else: - image = _ensure_chw_rgb(image) - indexed_images.append((idx, image)) + indexed_images = [(idx, _to_cuda_chw(image)) for idx, image, _ in group] # One NaViT target group can include several original resolutions. # Batch only source-compatible images, which removes redundant @@ -245,29 +294,24 @@ def _gpu_preprocess_images( 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 + gh, gw = padded_h // patch_size, padded_w // patch_size + batch = normalize_and_patchify( + batch, + image_scale, + image_bias, + patch_size, + padded_h, + padded_w, ) - grid = torch.tensor([T, gh, gw], dtype=torch.int64, device=batch.device) + grid = (T, gh, gw) 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) + grid_thws = torch.tensor(all_grids, dtype=torch.int64) return pixel_values, grid_thws @@ -290,6 +334,7 @@ class KimiGPUProcessorWrapper: self, hf_processor, image_token, + image_token_id, patch_size, merge_kernel_size, in_patch_limit, @@ -300,6 +345,7 @@ class KimiGPUProcessorWrapper: ): self._hf_processor = hf_processor self._image_token = image_token + self._image_token_id = image_token_id self._patch_size = patch_size self._merge_kernel_size = merge_kernel_size self._in_patch_limit = in_patch_limit @@ -320,12 +366,30 @@ class KimiGPUProcessorWrapper: def __call__(self, text=None, images=None, **kwargs): # process_mm_data passes images via kwargs["images"] images = images or kwargs.pop("images", None) + original_input_ids = kwargs.pop("sglang_original_input_ids", None) if images and torch.cuda.is_available(): - return self._gpu_call(text, images) - return self._cpu_call(text, images, **kwargs) + return self._gpu_call(text, images, original_input_ids) + return self._cpu_call(text, images, original_input_ids, **kwargs) - def _gpu_call(self, text, images): + def _prepare_input_ids(self, input_text, resize_configs, original_input_ids): + if original_input_ids is not None: + return _expand_image_token_ids( + original_input_ids, + self._image_token_id, + [config["num_tokens"] for config in resize_configs], + ) + + 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) + expanded_text = "".join(result) + return self._hf_processor.tokenizer(expanded_text, return_tensors="pt")[ + "input_ids" + ] + + def _gpu_call(self, text, images, original_input_ids=None): """Bypass HF KimiK25VisionProcessor.preprocess entirely -- use GPU ops.""" input_text = text[0] if isinstance(text, list) else text @@ -345,44 +409,44 @@ class KimiGPUProcessorWrapper: ) ) - # 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 + # 2. Reuse the request's tokenization when available: expanding the + # placeholders is exact, and skips tokenizing thousands of repeated + # ``<|media_pad|>`` strings. + input_ids = self._prepare_input_ids( + input_text, resize_configs, original_input_ids ) - grid_thws = grid_thws.cpu() + # 3. GPU image preprocessing + image_scale, image_bias = self._get_gpu_norm_tensors() + pixel_values, grid_thws = _gpu_preprocess_images( + images, resize_configs, image_scale, image_bias, self._patch_size + ) return { - "input_ids": text_inputs["input_ids"], + "input_ids": input_ids, "pixel_values": pixel_values, # Use SGL-standard key so get_new_expanded_mm_items() can split # per-image for cache granularity (it looks up 'image_grid_thw'). "image_grid_thw": grid_thws, } - def _cpu_call(self, text, images, **kwargs): + def _cpu_call(self, text, images, original_input_ids=None, **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 + image_token_counts = [ + int( + self._hf_processor.media_processor.media_tokens_calculator( + {"type": "image", "image": image} + ) + ) + for image in images + ] 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} - ) + for num_tokens, part in zip(image_token_counts, parts[1:]): result.append(self._image_token * num_tokens + part) input_text = "".join(result) @@ -390,6 +454,12 @@ class KimiGPUProcessorWrapper: kwargs["medias"] = [{"type": "image", "image": img} for img in images] out = self._hf_processor(text=[input_text], **kwargs) + if images and original_input_ids is not None: + # preserve_processor_input_ids turns off the base class rebuild, so + # this path has to keep the request's own tokens itself. + out["input_ids"] = _expand_image_token_ids( + original_input_ids, self._image_token_id, image_token_counts + ) grid_thws = out.pop("grid_thws", None) if grid_thws is not None: out["image_grid_thw"] = grid_thws @@ -397,13 +467,17 @@ class KimiGPUProcessorWrapper: 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 + image_scale = torch.tensor( + [1.0 / (255.0 * std) for std in self._image_std], + 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) + image_bias = torch.tensor( + [-mean / std for mean, std in zip(self._image_mean, self._image_std)], + device=device, + dtype=torch.float32, ).view(1, 3, 1, 1) - self._gpu_norm_tensors = (image_mean, image_std_inv) + self._gpu_norm_tensors = (image_scale, image_bias) return self._gpu_norm_tensors @@ -418,23 +492,25 @@ class KimiK2_5VLImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor): gpu_image_decode = True # nvJPEG for JPEG, PIL fallback for others prefer_tokenized_input = True precompute_hash_before_cpu_transfer = True + # The GPU wrapper expands placeholders from the request's own token IDs. + preserve_processor_input_ids = True + auto_mm_processor_worker_num = 2 + auto_mm_io_worker_num = 16 + supports_mm_processor_concurrency = True def __init__(self, hf_config, server_args, _processor, *args, **kwargs): - super().__init__(hf_config, server_args, _processor, *args, **kwargs) - self.mm_tokens = MultimodalSpecialTokens( + mm_tokens = MultimodalSpecialTokens( image_token="<|media_pad|>", # TODO: could we convert in MultimodalSpecialTokens? image_token_id=hf_config.media_placeholder_token_id, 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 = KimiGPUProcessorWrapper( _processor, - image_token=self.mm_tokens.image_token, + image_token=mm_tokens.image_token, + image_token_id=mm_tokens.image_token_id, 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"], @@ -443,6 +519,10 @@ class KimiK2_5VLImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor): image_mean=media_proc_cfg["image_mean"], image_std=media_proc_cfg["image_std"], ) + # Initialize the executor from the final GPU wrapper. Cloning the raw + # HF processor here would silently bypass Kimi's GPU preprocessing. + super().__init__(hf_config, server_args, processor, *args, **kwargs) + self.mm_tokens = mm_tokens async def process_mm_data_async( self, @@ -452,14 +532,43 @@ class KimiK2_5VLImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor): *args, **kwargs, ): - base_output = await self.load_mm_data( - prompt=input_text, - image_data=image_data, - multimodal_tokens=self.mm_tokens, + expected_image_count = len(image_data or []) + placeholder_count = self.count_image_placeholders( + input_text, self.mm_tokens.image_token_id ) + if placeholder_count is not None: + if placeholder_count != expected_image_count: + raise ValueError( + "Kimi image placeholders must map one-to-one to image data: " + f"expected {expected_image_count}, found {placeholder_count} token(s)" + ) + base_output = await self.fast_load_mm_data( + prompt=input_text, + image_data=image_data, + multimodal_tokens=self.mm_tokens, + # fast_load_mm_data, unlike load_mm_data, does not derive + # input_ids from the prompt; without this the wrapper falls back + # to re-tokenizing the expanded string. + input_ids=input_text, + ) + else: + base_output = await self.load_mm_data( + prompt=input_text, + image_data=image_data, + multimodal_tokens=self.mm_tokens, + ) + # Only the text-scanning loader can come back with a different + # count; fast_load_mm_data fills one slot per image_data entry. + if len(base_output.images) != expected_image_count: + raise ValueError( + "Kimi image placeholders must map one-to-one to image data: " + f"expected {expected_image_count}, loaded {len(base_output.images)}" + ) - mm_items, input_ids, _ = self.process_and_combine_mm_data( - base_output, self.mm_tokens + mm_items, input_ids, _ = await self.process_and_combine_mm_data_async( + base_output, + self.mm_tokens, + sglang_original_input_ids=base_output.input_ids, ) # K2.5/K2.7 encoder-DP assigns an image to exactly one TP rank. Keep diff --git a/python/sglang/srt/multimodal/processors/kimi_vl.py b/python/sglang/srt/multimodal/processors/kimi_vl.py index 7bcf2885c..e54011021 100644 --- a/python/sglang/srt/multimodal/processors/kimi_vl.py +++ b/python/sglang/srt/multimodal/processors/kimi_vl.py @@ -39,6 +39,12 @@ class KimiVLImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor): image_data=image_data, multimodal_tokens=self.mm_tokens, ) + expected_image_count = len(image_data or []) + if len(base_output.images) != expected_image_count: + raise ValueError( + "Kimi image placeholders must map one-to-one to image data: " + f"expected {expected_image_count}, loaded {len(base_output.images)}" + ) mm_items, input_ids, _ = self.process_and_combine_mm_data( base_output, self.mm_tokens diff --git a/test/manual/vlm/verify_k25_equivalence.py b/test/manual/vlm/verify_k25_equivalence.py new file mode 100644 index 000000000..97a9d95b4 --- /dev/null +++ b/test/manual/vlm/verify_k25_equivalence.py @@ -0,0 +1,145 @@ +"""Prove the two GPU-only rewrites in the K2.5 port are equivalent to main. + +1. normalize_and_patchify(scale/bias) == pad -> /255 -> (x-mean)*inv_std -> patchify +2. apply_fused_qk_complex_rope_inplace == the torch complex reference +""" + +import torch +import torch.nn.functional as F + +from sglang.kernels.ops.attention.vision_rope import ( + apply_fused_qk_complex_rope_inplace, + prepare_fused_qk_complex_rope_inplace, +) +from sglang.kernels.ops.mm.process import normalize_and_patchify + +MEAN = [0.5, 0.5, 0.5] +STD = [0.5, 0.5, 0.5] +ASYM_MEAN = [0.481, 0.457, 0.408] +ASYM_STD = [0.268, 0.261, 0.275] + + +def reference_preprocess(batch_u8, mean, std, patch_size, padded_h, padded_w): + """Exactly what main does, in main's order.""" + image_mean = torch.tensor(mean, device="cuda", dtype=torch.float32).view(1, 3, 1, 1) + image_std_inv = (1.0 / torch.tensor(std, device="cuda", dtype=torch.float32)).view( + 1, 3, 1, 1 + ) + x = batch_u8.float() + pad_h = padded_h - x.shape[-2] + pad_w = padded_w - x.shape[-1] + 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 + B, C, H, W = x.shape + gh, gw = H // patch_size, W // patch_size + x = x.view(B, C, gh, patch_size, gw, patch_size) + return x.permute(0, 2, 4, 1, 3, 5).reshape(B, -1, C, patch_size, patch_size) + + +def check_patchify(): + print("== normalize_and_patchify vs main's pad/normalize/patchify ==") + torch.manual_seed(0) + cases = [ + # (H, W, padded_h, padded_w, patch, mean, std, label) + (32, 24, 32, 24, 8, MEAN, STD, "no padding, symmetric norm"), + (30, 22, 32, 24, 8, MEAN, STD, "padded, symmetric norm"), + (30, 22, 32, 24, 8, ASYM_MEAN, ASYM_STD, "padded, per-channel norm"), + (64, 64, 64, 64, 16, ASYM_MEAN, ASYM_STD, "large patch"), + ] + ok = True + for h, w, ph, pw, patch, mean, std, label in cases: + raw = torch.randint(0, 256, (3, 3, h, w), dtype=torch.uint8, device="cuda") + ref = reference_preprocess(raw, mean, std, patch, ph, pw) + + scale = torch.tensor( + [1.0 / (255.0 * s) for s in std], device="cuda", dtype=torch.float32 + ).view(1, 3, 1, 1) + bias = torch.tensor( + [-m / s for m, s in zip(mean, std)], device="cuda", dtype=torch.float32 + ).view(1, 3, 1, 1) + got = normalize_and_patchify(raw.float(), scale, bias, patch, ph, pw) + + max_abs = (got - ref).abs().max().item() + # The padded rows must carry -mean/std, not zero. + pad_ok = True + if ph > h or pw > w: + pad_ok = torch.allclose( + got.flatten()[(got - ref).abs().argmax()], + ref.flatten()[(got - ref).abs().argmax()], + atol=1e-5, + ) + good = max_abs < 1e-5 and pad_ok + ok &= good + print(f" {'PASS' if good else 'FAIL'} {label:32s} max|d|={max_abs:.3e}") + return ok + + +def check_padded_value_is_not_zero(): + """The old pipeline padded in raw space, so pad cells become -mean/std.""" + print("== padded cells carry -mean/std, not 0 ==") + raw = torch.full((1, 3, 8, 8), 128, dtype=torch.uint8, device="cuda") + scale = torch.tensor( + [1.0 / (255.0 * s) for s in ASYM_STD], device="cuda", dtype=torch.float32 + ).view(1, 3, 1, 1) + bias = torch.tensor( + [-m / s for m, s in zip(ASYM_MEAN, ASYM_STD)], + device="cuda", + dtype=torch.float32, + ).view(1, 3, 1, 1) + got = normalize_and_patchify(raw.float(), scale, bias, 8, 16, 16) + # patch index 1 is the (row 0, col 1) patch -- entirely padding. + pad_patch = got[0, 1] + expected = bias.view(3, 1, 1).expand(3, 8, 8) + good = torch.allclose(pad_patch, expected, atol=1e-6) + print( + f" {'PASS' if good else 'FAIL'} pad cell = {pad_patch[0, 0, 0].item():.6f}, " + f"expected -mean/std = {expected[0, 0, 0].item():.6f}" + ) + return good + + +def reference_rope(xq, xk, freqs_cis): + freqs_cis = freqs_cis.unsqueeze(-2) + xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2)) + xk_ = torch.view_as_complex(xk.float().view(*xk.shape[:-1], -1, 2)) + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(-2) + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(-2) + return xq_out.type_as(xq), xk_out.type_as(xk) + + +def check_rope(): + print("== fused vision RoPE vs the torch complex reference ==") + torch.manual_seed(0) + ok = True + for dtype, tol in ((torch.bfloat16, 8e-3), (torch.float16, 2e-3)): + for tokens, heads, head_dim in ((1024, 16, 72), (4096, 8, 128), (37, 4, 64)): + xq = torch.randn(tokens, heads, head_dim, device="cuda", dtype=dtype) + xk = torch.randn(tokens, heads, head_dim, device="cuda", dtype=dtype) + angle = torch.randn(tokens, head_dim // 2, device="cuda") + freqs_cis = torch.polar(torch.ones_like(angle), angle) + + ref_q, ref_k = reference_rope(xq, xk, freqs_cis) + prepared = prepare_fused_qk_complex_rope_inplace(freqs_cis) + got_q, got_k = apply_fused_qk_complex_rope_inplace( + xq.clone(), xk.clone(), prepared + ) + + dq = (got_q.float() - ref_q.float()).abs().max().item() + dk = (got_k.float() - ref_k.float()).abs().max().item() + good = dq < tol and dk < tol + ok &= good + print( + f" {'PASS' if good else 'FAIL'} {str(dtype):16s} " + f"t={tokens:5d} h={heads:2d} d={head_dim:3d} " + f"max|dq|={dq:.2e} max|dk|={dk:.2e}" + ) + return ok + + +if __name__ == "__main__": + results = [check_patchify(), check_padded_value_is_not_zero(), check_rope()] + print() + print("ALL PASS" if all(results) else "SOME CHECKS FAILED") + raise SystemExit(0 if all(results) else 1) diff --git a/test/registered/kernels/ops/test_kimi_k3_prerequisite_ops.py b/test/registered/kernels/ops/test_kimi_k3_prerequisite_ops.py index 51324a7e3..cb26a235d 100644 --- a/test/registered/kernels/ops/test_kimi_k3_prerequisite_ops.py +++ b/test/registered/kernels/ops/test_kimi_k3_prerequisite_ops.py @@ -20,6 +20,8 @@ from sglang.kernels.ops.attention.set_mla_kv_concat_q import ( from sglang.kernels.ops.attention.utils import concat_mla_absorb_q_general from sglang.kernels.ops.attention.vision_rope import ( apply_fused_qk_complex_rope, + apply_fused_qk_complex_rope_inplace, + prepare_fused_qk_complex_rope_inplace, ) from sglang.kernels.ops.elementwise import add3 from sglang.kernels.ops.gemm.tiny_gemm import ( @@ -427,6 +429,33 @@ class TestKimiK3PrerequisiteOps(CustomTestCase): torch.testing.assert_close(actual_q, reference(q), rtol=0, atol=atol) torch.testing.assert_close(actual_k, reference(k), rtol=0, atol=atol) + def test_vision_rope_inplace(self): + # VisionAttention hands the applier contiguous q/k, which is what the + # in-place kernel requires; mirror that rather than qkv.unbind views. + for dtype in (torch.bfloat16, torch.float16): + torch.manual_seed(4) + q = torch.randn(480, 12, 128, device="cuda", dtype=dtype) + k = torch.randn(480, 12, 128, device="cuda", dtype=dtype) + angles = torch.randn(480, 64, device="cuda") + freqs = torch.polar(torch.ones_like(angles), angles) + freqs_expanded = freqs.unsqueeze(-2) + + def reference(x): + value = torch.view_as_complex(x.float().view(*x.shape[:-1], -1, 2)) + return torch.view_as_real(value * freqs_expanded).flatten(-2).type_as(x) + + expected_q, expected_k = reference(q), reference(k) + prepared = prepare_fused_qk_complex_rope_inplace(freqs) + actual_q, actual_k = apply_fused_qk_complex_rope_inplace(q, k, prepared) + + atol = 2 * torch.finfo(dtype).eps + torch.testing.assert_close(actual_q, expected_q, rtol=0, atol=atol) + torch.testing.assert_close(actual_k, expected_k, rtol=0, atol=atol) + + def test_vision_rope_inplace_rejects_non_complex_frequencies(self): + with self.assertRaises(ValueError): + prepare_fused_qk_complex_rope_inplace(torch.randn(8, 64, device="cuda")) + def test_normalize_and_patchify(self): torch.manual_seed(5) image = torch.randn(2, 3, 17, 19, device="cuda") diff --git a/test/registered/unit/models/test_kimi_k25.py b/test/registered/unit/models/test_kimi_k25.py index 7f6c18dec..51ed152bc 100644 --- a/test/registered/unit/models/test_kimi_k25.py +++ b/test/registered/unit/models/test_kimi_k25.py @@ -3,10 +3,12 @@ from types import SimpleNamespace from unittest.mock import Mock, patch +import numpy as np import pytest import torch import torch.nn as nn import torch.nn.functional as F +from PIL import Image from sglang.srt.managers.schedule_batch import ( Modality, @@ -14,9 +16,19 @@ from sglang.srt.managers.schedule_batch import ( MultimodalInputs, MultimodalProcessorOutput, ) -from sglang.srt.models.kimi_k25 import KimiK25ForConditionalGeneration +from sglang.srt.models.kimi_k25 import ( + KimiK25ForConditionalGeneration, + mm_projection_auto, +) +from sglang.srt.models.kimi_vl_moonvit import tpool_patch_merger from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model +from sglang.srt.multimodal.processors.base_processor import BaseMultimodalProcessor +from sglang.srt.multimodal.processors.kimi_common import KimiGridMMDataMixin from sglang.srt.multimodal.processors.kimi_k25 import ( + KimiGPUProcessorWrapper, + _ensure_chw_rgb, + _expand_image_token_ids, + _resize_bicubic_if_needed, _resize_images_by_source_shape, ) from sglang.srt.runtime_context import get_parallel @@ -63,14 +75,12 @@ def _image_item(feature, grid_thw): def test_kimi_gpu_preprocess_batches_only_source_compatible_images(): torch.manual_seed(0) indexed_images = [ - (0, torch.randn(3, 32, 24)), - (1, torch.randn(3, 32, 24)), - (2, torch.randn(3, 28, 20)), + (0, torch.randint(0, 256, (3, 32, 24), dtype=torch.uint8)), + (1, torch.randint(0, 256, (3, 32, 24), dtype=torch.uint8)), + (2, torch.randint(0, 256, (3, 28, 20), dtype=torch.uint8)), ] expected = [ - F.interpolate( - image.unsqueeze(0), size=(16, 12), mode="bicubic", align_corners=False - ) + _resize_bicubic_if_needed(image.unsqueeze(0), 16, 12) for _, image in indexed_images ] real_interpolate = F.interpolate @@ -92,6 +102,209 @@ def test_kimi_gpu_preprocess_batches_only_source_compatible_images(): torch.testing.assert_close(result, reference) +def test_kimi_resize_tracks_the_checkpoint_processors_pil_bicubic(): + # Plain F.interpolate skips PIL's implicit antialiasing on downscale and + # drifts far outside 8-bit rounding; photo-like content, not pure noise. + rng = np.random.default_rng(0) + yy, xx = np.mgrid[0:512, 0:512].astype(np.float32) + plane = np.clip( + 128 + + 90 * np.sin(xx / 40) * np.cos(yy / 55) + + 40 * ((xx // 37 + yy // 41) % 2) + + rng.normal(0, 6, (512, 512)), + 0, + 255, + ) + array = np.stack([plane, np.roll(plane, 7, 0), np.roll(plane, 13, 1)], -1).astype( + np.uint8 + ) + pil = torch.from_numpy( + np.asarray(Image.fromarray(array).resize((252, 252), Image.BICUBIC)).astype( + np.float32 + ) + ).permute(2, 0, 1) + source = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0) + + resized = _resize_bicubic_if_needed(source, 252, 252) + + assert resized.shape == (1, 3, 252, 252) + torch.testing.assert_close(resized, resized.round()) + assert resized.min() >= 0.0 and resized.max() <= 255.0 + # Within a couple of 8-bit levels of PIL; the non-antialiased resize is off + # by an order of magnitude more, which is the regression this guards. + assert (resized[0] - pil).abs().max() <= 4.0 + naive = F.interpolate( + source.float(), size=(252, 252), mode="bicubic", align_corners=False + ) + assert (naive[0] - pil).abs().max() > 20.0 + + +def test_kimi_resize_is_a_dtype_only_cast_when_already_at_target(): + image = torch.randint(0, 256, (1, 3, 16, 12), dtype=torch.uint8) + + resized = _resize_bicubic_if_needed(image, 16, 12) + + assert resized.dtype == torch.float32 + torch.testing.assert_close(resized, image.float()) + + +def test_kimi_expands_one_placeholder_per_image_from_existing_ids(): + # 7 is the placeholder; the two images claim 3 and 2 tokens. + input_ids = [1, 7, 2, 7, 3] + + expanded = _expand_image_token_ids( + input_ids, image_token_id=7, image_token_counts=[3, 2] + ) + + assert expanded.tolist() == [[1, 7, 7, 7, 2, 7, 7, 3]] + + +def test_kimi_expansion_rejects_a_placeholder_count_mismatch(): + with pytest.raises(ValueError, match="placeholder"): + _expand_image_token_ids([1, 7, 2], image_token_id=7, image_token_counts=[3, 2]) + + +def test_kimi_expansion_matches_the_base_retokenize_avoidance_rebuild(): + # preserve_processor_input_ids skips the base rebuild, which is only safe + # while both produce the same sequence. Reference is the original loop. + def reference(original_ids, counts, placeholder): + rebuilt, next_image = [], 0 + for token_id in original_ids: + if token_id == placeholder: + rebuilt.extend([placeholder] * counts[next_image]) + next_image += 1 + else: + rebuilt.append(token_id) + return rebuilt + + rng = np.random.default_rng(0) + for n_images in (1, 3, 8): + # Placeholder 7 is below the random range, so only the inserted + # positions count as placeholders. + ids = rng.integers(100, 5000, 400).tolist() + for slot in range(n_images): + ids.insert(slot * 37 + 5, 7) + counts = rng.integers(1, 400, n_images).tolist() + expected = reference(ids, counts, 7) + + assert BaseMultimodalProcessor._expand_input_ids(ids, counts, 7) == expected + wrapper = _expand_image_token_ids( + ids, image_token_id=7, image_token_counts=counts + ) + assert wrapper.flatten().tolist() == expected + + +def test_kimi_cpu_fallback_keeps_the_request_tokens(): + # preserve_processor_input_ids disables the base rebuild on every path. + hf_processor = Mock() + hf_processor.media_processor.media_tokens_calculator = Mock(return_value=3) + hf_processor.return_value = {"input_ids": torch.tensor([[99, 99, 99]])} + + wrapper = KimiGPUProcessorWrapper.__new__(KimiGPUProcessorWrapper) + wrapper._hf_processor = hf_processor + wrapper._image_token = "<|media_pad|>" + wrapper._image_token_id = 7 + + out = wrapper._cpu_call( + "a<|media_pad|>b", ["img"], original_input_ids=[1, 7, 2], medias=None + ) + + # Not the [99, 99, 99] the HF processor returned. + assert out["input_ids"].flatten().tolist() == [1, 7, 7, 7, 2] + + +def test_kimi_cpu_fallback_falls_back_to_the_hf_tokens_without_request_ids(): + hf_processor = Mock() + hf_processor.media_processor.media_tokens_calculator = Mock(return_value=3) + hf_processor.return_value = {"input_ids": torch.tensor([[99, 99, 99]])} + + wrapper = KimiGPUProcessorWrapper.__new__(KimiGPUProcessorWrapper) + wrapper._hf_processor = hf_processor + wrapper._image_token = "<|media_pad|>" + wrapper._image_token_id = 7 + + out = wrapper._cpu_call("a<|media_pad|>b", ["img"], medias=None) + + assert out["input_ids"].flatten().tolist() == [99, 99, 99] + + +def test_kimi_refuses_already_normalized_float_pixels(): + with pytest.raises(ValueError, match="uint8"): + _ensure_chw_rgb(torch.rand(3, 8, 8)) + + +def test_kimi_placeholder_count_only_reads_real_token_ids(): + count = KimiGridMMDataMixin.count_image_placeholders + + assert count([1, 7, 2, 7], 7) == 2 + assert count(torch.tensor([[1, 7, 2]]), 7) == 1 + assert count([1, 2, 3], 7) == 0 + # A prompt string carries no token IDs, so the caller must not take the + # tokenized fast path. + assert count("<|media_pad|>", 7) is None + + +def test_kimi_single_frame_pool_matches_the_temporal_mean(): + torch.manual_seed(0) + x = torch.randn(1 * 4 * 4, 8) + grid_thws = torch.tensor([[1, 4, 4]]) + + (merged,) = tpool_patch_merger(x, grid_thws) + + # t == 1 skips the mean; it must stay bit-identical to averaging one frame. + reference = ( + x.view(1, 2, 2, 2, 2, 8).permute(0, 1, 3, 2, 4, 5).contiguous().mean(dim=0) + ) + assert torch.equal(merged, reference.view(4, 4, 8)) + + +def test_kimi_multi_frame_pool_still_averages_across_frames(): + torch.manual_seed(0) + x = torch.randn(3 * 4 * 4, 8) + grid_thws = torch.tensor([[3, 4, 4]]) + + (merged,) = tpool_patch_merger(x, grid_thws) + + reference = ( + x.view(3, 2, 2, 2, 2, 8).permute(0, 1, 3, 2, 4, 5).contiguous().mean(dim=0) + ) + assert merged.shape == (4, 4, 8) + torch.testing.assert_close(merged, reference.view(4, 4, 8)) + + +class _IdentityProjector(nn.Module): + """Stands in for K2VLMultiModalProjector, which is never None in production.""" + + def __init__(self): + super().__init__() + self.seen = None + + def forward(self, x): + self.seen = x + return x + + +def test_kimi_projection_returns_one_flattened_feature_tensor(): + torch.manual_seed(0) + per_image = [torch.randn(4, 2, 8), torch.randn(6, 2, 8)] + + packed = mm_projection_auto(_IdentityProjector(), per_image) + + assert packed.shape == (20, 8) + torch.testing.assert_close(packed, torch.cat(per_image, dim=0).reshape(-1, 8)) + + +def test_kimi_projection_does_not_copy_a_single_image(): + single = torch.randn(4, 2, 8) + projector = _IdentityProjector() + + packed = mm_projection_auto(projector, [single]) + + # The projector must receive the tensor itself, not a one-element cat of it. + assert projector.seen.data_ptr() == single.data_ptr() + assert packed.data_ptr() == single.data_ptr() + + def test_dp_helper_supports_moonvit3d_packed_embeddings_on_tp1(): tower = _MoonViT3dTower() pixel_values = torch.randn(4, 2) @@ -248,6 +461,28 @@ def test_kimi_k25_encoder_dp_selects_packed_moonvit_contract(): assert callable(run_dp.call_args.kwargs["load_local_pixel_values"]) +def test_kimi_non_dp_keeps_grid_thws_on_the_host(): + model = KimiK25ForConditionalGeneration.__new__(KimiK25ForConditionalGeneration) + nn.Module.__init__(model) + model.use_data_parallel = False + model.vision_tower = _MoonViT3dTower() + # Not the host, so a stray .to(tower.device) shows up without a GPU. + model.vision_tower.device = torch.device("meta") + model.mm_projector = _IdentityProjector() + items = [_image_item(torch.randn(4, 2), [[1, 2, 2]])] + + with get_parallel().override( + tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0 + ), patch( + "sglang.srt.models.kimi_k25.get_server_args", + return_value=SimpleNamespace(tp_size=1), + ): + model.get_image_feature(items) + + # A device copy would cost one sync per .tolist() inside MoonViT3d. + assert model.vision_tower.grid_thws.device.type == "cpu" + + def test_kimi_lazy_ipc_feature_skips_scheduler_reconstruction(): proxy = CudaIpcTensorTransportProxy.__new__(CudaIpcTensorTransportProxy) proxy.reconstruct_on_target_device = Mock()