diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index 1e913091b..2fe03f8b9 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -58,6 +58,11 @@ from sglang.srt.model_executor.model_runner_components.load_model_utils import ( maybe_precompile_model_kernels_after_loading, ) from sglang.srt.model_loader import get_model +from sglang.srt.multimodal.encoder_preprocessing import ( + EncoderPreprocessOutput, + get_encoder_preprocessed_items, + invoke_encoder_preprocessor, +) from sglang.srt.multimodal.processors.qwen_vl import preprocess_video from sglang.srt.observability.metrics_collector import EncoderMetricsCollector from sglang.srt.observability.req_time_stats import EncoderReqTimeStats @@ -65,7 +70,13 @@ from sglang.srt.observability.trace import ( process_tracing_init, trace_set_thread_info, ) -from sglang.srt.runtime_context import get_disagg, get_exec, get_mm, publish +from sglang.srt.runtime_context import ( + get_disagg, + get_exec, + get_mm, + get_parallel, + publish, +) from sglang.srt.server_args import ( PortArgs, ServerArgs, @@ -862,9 +873,24 @@ class MMEncoder: return slices def _calculate_hashes_from_features( - self, mm_feature, grid_thw: List, modality: Modality + self, mm_feature, grid_thw: List, modality: Modality, mm_inputs=None ) -> List[int]: """CPU Task: Compute hashes based on processed feature patches.""" + preprocessed_items = ( + get_encoder_preprocessed_items(mm_inputs) if mm_inputs is not None else None + ) + if preprocessed_items is not None: + if len(preprocessed_items) != len(grid_thw): + raise ValueError( + "Encoder preprocess item/grid mismatch: " + f"{len(preprocessed_items)} items != {len(grid_thw)} grids" + ) + hashes = [] + for item in preprocessed_items: + item.set_pad_value() + hashes.append(item.hash) + return hashes + hashes = [] if modality == Modality.AUDIO and isinstance(mm_feature, list): for feature in mm_feature: @@ -884,21 +910,37 @@ class MMEncoder: offset += num_patches return hashes - def _encode_missing( + def _build_mm_data_items( self, mm_feature, mm_inputs: dict, indices: List[int], - modality: Modality = Modality.IMAGE, - get_feature_fn=None, + modality: Modality, grid_thw: Optional[List] = None, - keep_on_gpu: bool = False, - ) -> List[torch.Tensor]: - """ - GPU Task: Run ViT inference ONLY on the subset of mm items missing from the cache. + ) -> List[MultimodalDataItem]: + """Build the model-facing items selected for one encoder forward. + + A model preprocessor can preserve an item-wise representation with + ``EncoderPreprocessOutput``. This path avoids concatenating and then + re-slicing features before encoder-DP knows which rank owns each item. + Legacy Hugging Face processor outputs retain their existing aggregate + tensor behavior. """ if grid_thw is None: grid_thw = _get_mm_grid_dim(mm_inputs, modality, self.model_type) + + preprocessed_items = get_encoder_preprocessed_items(mm_inputs) + if preprocessed_items is not None: + if len(preprocessed_items) != len(grid_thw): + raise ValueError( + "Encoder preprocess item/grid mismatch: " + f"{len(preprocessed_items)} items != {len(grid_thw)} grids" + ) + selected = [preprocessed_items[index] for index in indices] + if any(item.modality != modality for item in selected): + raise ValueError("Encoder preprocess output contains wrong modality") + return selected + split_kimi_k3_images = ( self.model_type == "kimi_k3" and modality == Modality.IMAGE ) @@ -916,11 +958,11 @@ class MMEncoder: sub_feature_list = [] offsets = [0] curr = 0 - for g in grid_thw: - curr += self.get_num_patches(g, modality) + for grid in grid_thw: + curr += self.get_num_patches(grid, modality) offsets.append(curr) - for idx in indices: - sub_feature_list.append(mm_feature[offsets[idx] : offsets[idx + 1]]) + for index in indices: + sub_feature_list.append(mm_feature[offsets[index] : offsets[index + 1]]) if not split_kimi_k3_images: sub_feature = torch.cat(sub_feature_list, dim=0) @@ -948,19 +990,39 @@ class MMEncoder: ) ] - for k, v in mm_inputs.items(): - if k in _mm_feature_attrs.get(modality, []): + for key, value in mm_inputs.items(): + if key in _mm_feature_attrs.get(modality, []): continue - val = _convert(v) - if k in _mm_grid_attrs.get(modality, []): + value = _convert(value) + if key in _mm_grid_attrs.get(modality, []): if split_kimi_k3_images: - for mm_item, idx in zip(mm_items, indices): - mm_item.set(k, val[idx : idx + 1]) + for mm_item, index in zip(mm_items, indices): + mm_item.set(key, value[index : index + 1]) else: - mm_items[0].set(k, val[indices]) + mm_items[0].set(key, value[indices]) else: for mm_item in mm_items: - mm_item.set(k, val) + mm_item.set(key, value) + return mm_items + + def _encode_missing( + self, + mm_feature, + mm_inputs: dict, + indices: List[int], + modality: Modality = Modality.IMAGE, + get_feature_fn=None, + grid_thw: Optional[List] = None, + keep_on_gpu: bool = False, + ) -> List[torch.Tensor]: + """ + GPU Task: Run ViT inference ONLY on the subset of mm items missing from the cache. + """ + if grid_thw is None: + grid_thw = _get_mm_grid_dim(mm_inputs, modality, self.model_type) + mm_items = self._build_mm_data_items( + mm_feature, mm_inputs, indices, modality, grid_thw + ) forward_start = time.perf_counter() with torch.inference_mode(): @@ -1002,7 +1064,7 @@ class MMEncoder: if self.rank == 0: if hashes is None: mm_hashes = self._calculate_hashes_from_features( - mm_feature, grid_thw, modality + mm_feature, grid_thw, modality, mm_inputs ) else: mm_hashes = hashes @@ -1651,12 +1713,31 @@ class MMEncoder: if not (self.image_processor or model_preprocessor): raise ValueError("No image processor available") images = await self._flatten_and_load_images(mm_items) - if model_preprocessor: - return model_preprocessor(images, Modality.IMAGE, self.vision_config) - image_config = self.vision_config.get("image", {}) - original_image_sizes = [_get_original_image_size(item) for item in images] if self.model_type in ["kimi_k25", "kimi_k3", "kimi_vl"]: images = self._normalize_kimi_encoder_images(images) + original_image_sizes = [_get_original_image_size(item) for item in images] + if model_preprocessor: + processor_output = invoke_encoder_preprocessor( + model_preprocessor, + images, + Modality.IMAGE, + self.vision_config, + image_processor=self.image_processor, + use_gpu_preprocessing=self.use_image_processor_gpu, + ) + if ( + isinstance(processor_output, EncoderPreprocessOutput) + and processor_output.materialize_local_items is not None + ): + parallel = get_parallel() + await asyncio.get_running_loop().run_in_executor( + self.preproc_executor, + processor_output.materialize_for_rank, + parallel.attn_tp_rank, + parallel.attn_tp_size, + ) + return processor_output + image_config = self.vision_config.get("image", {}) processor_input = await asyncio.get_running_loop().run_in_executor( self.preproc_executor, functools.partial(self.image_processor, images=images, **image_config), @@ -1775,25 +1856,25 @@ class MMEncoder: # support mm_cache mm_embedding = None mm_hash = None - - mm_item = MultimodalDataItem.from_dict( - { - "modality": modality, - "feature": _convert(_get_mm_feature(mm_inputs, modality)), - } + mm_feature = _convert(_get_mm_feature(mm_inputs, modality)) + grid_thw = _get_mm_grid_dim(mm_inputs, modality, self.model_type) + model_mm_items = self._build_mm_data_items( + mm_feature, + mm_inputs, + list(range(len(grid_thw))), + modality, + grid_thw, ) - for k, v in mm_inputs.items(): - if k in _mm_feature_attrs[modality]: - continue - mm_item.set(k, _convert(v)) cache_hit = False use_mm_cache = get_mm().enable_prefix_mm_cache and log_metrics if use_mm_cache: - mm_item.set_pad_value() - mm_hash = MultiModalStaticCache.combine_hashes([mm_item.hash]) + for item in model_mm_items: + item.set_pad_value() + item_hashes = [item.hash for item in model_mm_items] + mm_hash = MultiModalStaticCache.combine_hashes(item_hashes) async with self.mm_cache_lock: - mm_cache = self.mm_cache.get([mm_item.hash]) + mm_cache = self.mm_cache.get(item_hashes) if mm_cache is not None: mm_embedding = mm_cache.embedding cache_hit = True @@ -1801,7 +1882,7 @@ class MMEncoder: if mm_embedding is None: forward_start = time.perf_counter() with torch.inference_mode(): - mm_embedding: torch.Tensor = get_feature_fn([mm_item]) + mm_embedding: torch.Tensor = get_feature_fn(model_mm_items) mm_embedding = mm_embedding.cpu() if len(mm_embedding.shape) != 2: mm_embedding = mm_embedding.reshape(-1, mm_embedding.shape[-1]) @@ -1810,7 +1891,8 @@ class MMEncoder: time.perf_counter() - forward_start, modality=modality_str ) - # Per-request cache hit metrics: tokens = embedding rows, files = 1 item. + # Per-request cache hit metrics: tokens = embedding rows, files = + # logical multimodal items (not the legacy aggregate tensor count). if use_mm_cache and encoder_metrics_collector is not None: total_tokens = int(mm_embedding.shape[0]) hit_tokens = total_tokens if cache_hit else 0 @@ -1818,7 +1900,9 @@ class MMEncoder: hit_tokens, total_tokens, modality=modality_str ) encoder_metrics_collector.record_cache_files( - 1 if cache_hit else 0, 1, modality=modality_str + len(model_mm_items) if cache_hit else 0, + len(model_mm_items), + modality=modality_str, ) if use_mm_cache: @@ -1865,7 +1949,7 @@ class MMEncoder: ) return ( - _get_mm_grid_dim(mm_inputs, modality, self.model_type), + grid_thw, mm_embedding, aux_data, ) @@ -2223,23 +2307,21 @@ class MMEncoder: ) ) - # Build mm_item (all ranks) - mm_item = MultimodalDataItem.from_dict( - { - "modality": modality, - "feature": _convert(_get_mm_feature(mm_inputs, modality)), - } + # Build model-facing items on all ranks. Owner-deferred processor + # outputs stay per-item until encoder-DP assigns them. + mm_feature = _convert(_get_mm_feature(mm_inputs, modality)) + model_mm_items = self._build_mm_data_items( + mm_feature, + mm_inputs, + list(range(len(grid_thw))), + modality, + grid_thw, ) - for k, v in mm_inputs.items(): - if k in _mm_feature_attrs.get(modality, []): - continue - val = _convert(v) - mm_item.set(k, val) async def _run_forward(): try: with torch.inference_mode(): - emb = get_feature_fn([mm_item]) + emb = get_feature_fn(model_mm_items) if len(emb.shape) != 2: emb = emb.reshape(-1, emb.shape[-1]) # mooncake's transfer_sync is a host-side diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index 9c9817c42..2c82e986a 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -3064,6 +3064,32 @@ class KimiK3ForConditionalGeneration(nn.Module): ) self.language_model.set_dspark_layers_to_capture(layer_ids) + def preprocess_mm_for_encoder( + self, + mm_data, + modality, + config, + *, + image_processor=None, + use_gpu_preprocessing=False, + ): + """Prepare per-image raw inputs for owner-side EPD preprocessing.""" + if modality != Modality.IMAGE: + raise ValueError("Kimi-K3 encoder mode supports image input only") + if image_processor is None: + raise ValueError("Kimi-K3 encoder preprocessing needs an image processor") + + from sglang.srt.multimodal.kimi_k3_image_processing import ( + prepare_kimi_k3_encoder_inputs, + ) + + self._encoder_image_processor = image_processor + return prepare_kimi_k3_encoder_inputs( + mm_data, + image_processor, + use_gpu_preprocessing=use_gpu_preprocessing, + ) + def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: device = self.vision_tower.device target_dtype = self.vision_tower.patch_embed.proj.weight.dtype @@ -3087,6 +3113,10 @@ class KimiK3ForConditionalGeneration(nn.Module): def materialize_item_features(image_indices: List[int]) -> torch.Tensor: """Materialize only the images assigned to this vision-DP rank.""" + from sglang.srt.multimodal.encoder_preprocessing import ( + LOCAL_PREPROCESSED_KEY, + ) + # Match the configured TP consumer count captured when the # tokenizer creates MmItemMemoryPool. A live attention subgroup # size could leave acknowledgements missing and strand the lease. @@ -3104,6 +3134,21 @@ class KimiK3ForConditionalGeneration(nn.Module): ) selected_items.append(item) + locally_preprocessed = [ + item.model_specific_data.get(LOCAL_PREPROCESSED_KEY, False) + for item in selected_items + ] + if any(locally_preprocessed): + if not all(locally_preprocessed): + raise ValueError( + "Kimi-K3 cannot mix local preprocessed and deferred images" + ) + return materialize_multimodal_features( + [item.feature for item in selected_items], + device=device, + dtype=target_dtype, + ) + deferred = [ item.model_specific_data.get(DEFERRED_PREPROCESSING_KEY) for item in selected_items @@ -3113,25 +3158,44 @@ class KimiK3ForConditionalGeneration(nn.Module): raise ValueError( "Kimi-K3 cannot mix deferred and preprocessed image features" ) - from sglang.srt.multimodal.processors.kimi_k25 import ( - _gpu_preprocess_images, - ) - first_config = deferred[0] - image_scale, image_bias = normalization_tensors( - first_config["image_mean"], first_config["image_std"], device - ) - pixel_values, _ = _gpu_preprocess_images( - [item.feature for item in selected_items], - [config["resize_config"] for config in deferred], - image_scale, - image_bias, - self.vision_tower.patch_size, - to_chw=lambda image: to_chw_uint8(image, device=device), - post_resize=lambda x: fill_transparent_bg( - x, first_config["transparent_bg_config"] - ), - ) + backend = first_config["backend"] + if any(config["backend"] != backend for config in deferred): + raise ValueError( + "Kimi-K3 cannot mix deferred preprocessing backends" + ) + if backend == "gpu": + from sglang.srt.multimodal.processors.kimi_k25 import ( + _gpu_preprocess_images, + ) + + image_scale, image_bias = normalization_tensors( + first_config["image_mean"], first_config["image_std"], device + ) + pixel_values, _ = _gpu_preprocess_images( + [item.feature for item in selected_items], + [config["resize_config"] for config in deferred], + image_scale, + image_bias, + self.vision_tower.patch_size, + to_chw=lambda image: to_chw_uint8(image, device=device), + post_resize=lambda x: fill_transparent_bg( + x, first_config["transparent_bg_config"] + ), + ) + elif backend == "cpu": + from sglang.srt.multimodal.kimi_k3_image_processing import ( + materialize_kimi_k3_cpu_features, + ) + + pixel_values = materialize_kimi_k3_cpu_features( + selected_items, self._encoder_image_processor + ) + pixel_values = pixel_values.to(device, non_blocking=True) + else: + raise ValueError( + f"Unsupported Kimi-K3 deferred preprocessing backend: {backend}" + ) return pixel_values.to(dtype=target_dtype) features = [] diff --git a/python/sglang/srt/multimodal/encoder_preprocessing.py b/python/sglang/srt/multimodal/encoder_preprocessing.py new file mode 100644 index 000000000..79958b119 --- /dev/null +++ b/python/sglang/srt/multimodal/encoder_preprocessing.py @@ -0,0 +1,140 @@ +import hashlib +import inspect +from collections.abc import Mapping +from typing import Any, Callable, Iterable, Sequence + +import numpy as np +import torch + +from sglang.srt.managers.schedule_batch import MultimodalDataItem + +LOCAL_PREPROCESSED_KEY = "encoder_local_preprocessed" + + +def hash_raw_encoder_item(value: Any) -> int: + """Hash raw CPU media including layout metadata, before owner materialization.""" + if isinstance(value, torch.Tensor): + value = value.detach().cpu().contiguous().numpy() + elif not isinstance(value, np.ndarray): + from PIL import Image + + if not isinstance(value, Image.Image): + raise TypeError(f"Unsupported raw encoder item: {type(value)}") + value = np.asarray(value) + + value = np.ascontiguousarray(value) + hasher = hashlib.sha256() + hasher.update(value.dtype.str.encode()) + hasher.update(repr(value.shape).encode()) + hasher.update(memoryview(value)) + return int.from_bytes(hasher.digest()[:8], byteorder="big", signed=False) + + +class EncoderPreprocessOutput(dict): + """Processor output that preserves one feature object per multimodal item. + + Most Hugging Face processors concatenate every image into one tensor before + the vision model decides its encoder-DP assignment. A model-specific + ``preprocess_mm_for_encoder`` hook can return this mapping instead, allowing + the encoder to carry raw or partially processed items to the model. The + model can then materialize only the items owned by its local vision rank. + + The mapping remains compatible with existing encoder metadata helpers; + ``mm_items`` is an out-of-band, per-item representation used only by the + encoder forward/cache paths. + """ + + def __init__( + self, + values: Mapping[str, Any] | None = None, + *, + mm_items: Iterable[MultimodalDataItem], + item_sizes: Sequence[int] | None = None, + materialize_local_items: ( + Callable[[list[MultimodalDataItem]], Sequence[torch.Tensor]] | None + ) = None, + ) -> None: + super().__init__(values or {}) + self.mm_items = list(mm_items) + if not self.mm_items: + raise ValueError("EncoderPreprocessOutput requires at least one item") + self.item_sizes = list(item_sizes) if item_sizes is not None else None + if self.item_sizes is not None and len(self.item_sizes) != len(self.mm_items): + raise ValueError("Encoder preprocess item_sizes must match mm_items") + self.materialize_local_items = materialize_local_items + + def local_item_indices(self, rank: int, world_size: int) -> list[int]: + """Return the same size-balanced owner assignment used by vision DP.""" + if self.materialize_local_items is None: + return [] + if world_size < 1 or not 0 <= rank < world_size: + raise ValueError( + f"Invalid encoder preprocess rank {rank} for world size {world_size}" + ) + if world_size == 1: + return list(range(len(self.mm_items))) + if self.item_sizes is None: + raise ValueError( + "Owner-side encoder preprocessing requires per-item load sizes" + ) + + from sglang.srt.multimodal.mm_utils import get_dp_encoder_lb_assignment + + shuffled, counts, _ = get_dp_encoder_lb_assignment(self.item_sizes, world_size) + start = sum(counts[:rank]) + return shuffled[start : start + counts[rank]] + + def materialize_for_rank(self, rank: int, world_size: int) -> None: + """Materialize only this vision-DP rank's items in-place.""" + indices = self.local_item_indices(rank, world_size) + if not indices: + return + items = [self.mm_items[index] for index in indices] + materialize = self.materialize_local_items + assert materialize is not None + features = list(materialize(items)) + if len(features) != len(items): + raise ValueError( + "Encoder local materializer must return one feature per item" + ) + for item, feature in zip(items, features): + item.feature = feature + item.model_specific_data[LOCAL_PREPROCESSED_KEY] = True + + +def get_encoder_preprocessed_items( + processor_output: Mapping[str, Any], +) -> list[MultimodalDataItem] | None: + if isinstance(processor_output, EncoderPreprocessOutput): + return processor_output.mm_items + return None + + +def invoke_encoder_preprocessor( + preprocessor, + mm_data, + modality, + config, + **available_context, +): + """Call a model hook with only the optional context it declares. + + Existing hooks keep their three-argument contract. New model integrations + can request shared processor state or backend policy by adding named + keyword-only parameters, without requiring encode-server model branches. + """ + parameters = inspect.signature(preprocessor).parameters + accepts_kwargs = any( + parameter.kind == inspect.Parameter.VAR_KEYWORD + for parameter in parameters.values() + ) + context = ( + available_context + if accepts_kwargs + else { + name: value + for name, value in available_context.items() + if name in parameters + } + ) + return preprocessor(mm_data, modality, config, **context) diff --git a/python/sglang/srt/multimodal/kimi_k3_image_processing.py b/python/sglang/srt/multimodal/kimi_k3_image_processing.py index e4d959f79..ab959ef87 100644 --- a/python/sglang/srt/multimodal/kimi_k3_image_processing.py +++ b/python/sglang/srt/multimodal/kimi_k3_image_processing.py @@ -1,3 +1,5 @@ +import functools +import math from typing import Union import numpy as np @@ -7,6 +9,195 @@ from PIL import Image DEFERRED_PREPROCESSING_KEY = "kimi_k3_deferred_preprocessing" +def prepare_kimi_k3_encoder_inputs( + images, image_processor, *, use_gpu_preprocessing=False +): + """Keep K3 EPD images raw until the vision-DP owner is known. + + The lightweight NaViT shape calculation runs on every encoder rank. The + expensive resize, normalization, and patchification remain deferred in each + item and are executed by ``KimiK3ForConditionalGeneration`` only for images + assigned to the local vision rank. + """ + from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem + from sglang.srt.multimodal.encoder_preprocessing import ( + EncoderPreprocessOutput, + hash_raw_encoder_item, + ) + from sglang.srt.multimodal.processors.kimi_k25 import ( + _grid_thw_from_resize_config, + navit_resize_config, + ) + + media_proc_cfg = getattr(image_processor, "media_proc_cfg", None) + if not isinstance(media_proc_cfg, dict): + raise ValueError( + "Kimi-K3 EPD owner-side preprocessing requires " + "image_processor.media_proc_cfg" + ) + + required = ( + "patch_size", + "merge_kernel_size", + "in_patch_limit", + "patch_limit_on_one_side", + "image_mean", + "image_std", + ) + missing = [name for name in required if name not in media_proc_cfg] + if missing: + raise ValueError( + "Kimi-K3 image processor is missing deferred-preprocessing config: " + + ", ".join(missing) + ) + + concrete_images = [] + for image in images: + if isinstance(image, dict): + if image.get("type") != "image" or "image" not in image: + raise ValueError(f"Unsupported Kimi-K3 encoder media item: {image}") + image = image["image"] + concrete_images.append(image) + + patch_size = int(media_proc_cfg["patch_size"]) + merge_kernel_size = int(media_proc_cfg["merge_kernel_size"]) + common_deferred_config = { + "backend": "gpu" if use_gpu_preprocessing else "cpu", + "image_mean": list(media_proc_cfg["image_mean"]), + "image_std": list(media_proc_cfg["image_std"]), + "transparent_bg_config": media_proc_cfg.get("transparent_bg_config"), + } + + items = [] + grids = [] + original_image_sizes = [] + for image in concrete_images: + width, height = ( + (int(image.shape[-1]), int(image.shape[-2])) + if isinstance(image, torch.Tensor) + else image.size + ) + resize_config = navit_resize_config( + width, + height, + patch_size, + merge_kernel_size, + int(media_proc_cfg["in_patch_limit"]), + int(media_proc_cfg["patch_limit_on_one_side"]), + media_proc_cfg.get("fixed_output_tokens"), + ) + grid_thw = _grid_thw_from_resize_config(resize_config, patch_size) + grid_tensor = torch.tensor([grid_thw], dtype=torch.int64) + item = MultimodalDataItem( + modality=Modality.IMAGE, + feature=to_chw_uint8(image) if use_gpu_preprocessing else image, + model_specific_data={ + "grid_thws": grid_tensor, + DEFERRED_PREPROCESSING_KEY: { + **common_deferred_config, + "feature_layout": "chw" if use_gpu_preprocessing else "raw", + "resize_config": resize_config, + }, + }, + ) + if not use_gpu_preprocessing: + item.set_hash(hash_raw_encoder_item(image)) + items.append(item) + grids.append(grid_thw) + original_image_sizes.append([width, height]) + + grid_thws = torch.tensor(grids, dtype=torch.int64) + return EncoderPreprocessOutput( + { + # Preserve the conventional feature key for cache/accounting code; + # EncoderPreprocessOutput.mm_items is the authoritative per-item view. + "pixel_values": [item.feature for item in items], + "grid_thws": grid_thws, + "original_image_sizes": original_image_sizes, + }, + mm_items=items, + item_sizes=[math.prod(grid) for grid in grids], + materialize_local_items=( + None + if use_gpu_preprocessing + else functools.partial( + materialize_kimi_k3_cpu_item_features, + image_processor=image_processor, + ) + ), + ) + + +def materialize_kimi_k3_cpu_features(items, image_processor) -> torch.Tensor: + """Run the checkpoint's exact processor only on locally owned images.""" + medias = [] + for item in items: + image = item.feature + config = item.model_specific_data[DEFERRED_PREPROCESSING_KEY] + if config["feature_layout"] != "raw": + raise ValueError("Kimi-K3 deferred CPU preprocessing expects raw inputs") + if not isinstance(image, Image.Image): + if not isinstance(image, torch.Tensor) or image.dtype != torch.uint8: + raise TypeError( + "Kimi-K3 deferred CPU preprocessing expects PIL or uint8 tensors" + ) + image = to_hwc_uint8(image).numpy() + channels = image.shape[-1] + if channels == 1: + image = Image.fromarray(image[..., 0], mode="L") + elif channels == 3: + image = Image.fromarray(image, mode="RGB") + elif channels == 4: + image = Image.fromarray(image, mode="RGBA") + else: + raise ValueError(f"Unsupported Kimi-K3 image channel count: {channels}") + medias.append({"type": "image", "image": image}) + + output = image_processor.preprocess(medias, return_tensors="pt") + expected_grids = torch.cat( + [item.model_specific_data["grid_thws"] for item in items], dim=0 + ) + if not torch.equal(output["grid_thws"].cpu(), expected_grids.cpu()): + raise ValueError("Kimi-K3 deferred CPU preprocessing produced wrong grids") + return output["pixel_values"] + + +def materialize_kimi_k3_cpu_item_features(items, image_processor) -> list[torch.Tensor]: + """Return exact checkpoint-processor features split by logical image.""" + pixel_values = materialize_kimi_k3_cpu_features(items, image_processor) + patch_counts = [ + math.prod(item.model_specific_data["grid_thws"][0].tolist()) for item in items + ] + if sum(patch_counts) != pixel_values.shape[0]: + raise ValueError( + "Kimi-K3 processor feature length does not match image grids: " + f"{pixel_values.shape[0]} != {sum(patch_counts)}" + ) + return list(pixel_values.split(patch_counts)) + + +def to_hwc_uint8(image: Union[torch.Tensor, Image.Image]) -> torch.Tensor: + """Stage an exact CPU image without doing resize/normalize work.""" + if isinstance(image, Image.Image): + has_alpha = image.mode != "RGB" and ( + "A" in image.getbands() or "transparency" in image.info + ) + array = np.array(image.convert("RGBA" if has_alpha else "RGB"), copy=True) + return torch.from_numpy(array) + + if image.dtype != torch.uint8: + raise ValueError( + f"Kimi-K3 preprocessing expects raw uint8 pixels, got {image.dtype}" + ) + if image.dim() == 2: + image = image.unsqueeze(-1) + elif image.dim() == 3 and image.shape[0] in (1, 3, 4): + image = image.permute(1, 2, 0) + if image.shape[-1] == 1: + image = image.repeat(1, 1, 3) + return image.cpu().contiguous() + + def to_chw_uint8( image: Union[torch.Tensor, Image.Image], device: torch.device | str | None = None, diff --git a/test/registered/unit/disaggregation/test_kimi_k3_encoder_mode.py b/test/registered/unit/disaggregation/test_kimi_k3_encoder_mode.py index 751b35aba..b25362fba 100644 --- a/test/registered/unit/disaggregation/test_kimi_k3_encoder_mode.py +++ b/test/registered/unit/disaggregation/test_kimi_k3_encoder_mode.py @@ -22,11 +22,23 @@ from sglang.srt.disaggregation.encode_receiver import ( _select_mm_processor_prompt, ) from sglang.srt.disaggregation.encode_server import MMEncoder, _get_mm_grid_dim -from sglang.srt.managers.schedule_batch import Modality +from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem from sglang.srt.managers.tokenizer_manager import ( _reject_missing_dispatched_encoder_embedding, ) from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration +from sglang.srt.multimodal.encoder_preprocessing import ( + LOCAL_PREPROCESSED_KEY, + EncoderPreprocessOutput, + get_encoder_preprocessed_items, + hash_raw_encoder_item, + invoke_encoder_preprocessor, +) +from sglang.srt.multimodal.kimi_k3_image_processing import ( + DEFERRED_PREPROCESSING_KEY, + materialize_kimi_k3_cpu_features, + prepare_kimi_k3_encoder_inputs, +) from sglang.srt.runtime_context import get_context from sglang.srt.server_args import resolve_encoder_transfer_backend from sglang.test.ci.ci_register import register_cpu_ci @@ -138,6 +150,189 @@ def test_kimi_k3_encoder_passes_media_dicts_to_image_processor(): assert kwargs == {"return_tensors": "pt"} +def _kimi_k3_image_processor(): + return SimpleNamespace( + media_proc_cfg={ + "patch_size": 2, + "merge_kernel_size": 2, + "in_patch_limit": 1024, + "patch_limit_on_one_side": 64, + "fixed_output_tokens": None, + "image_mean": [0.5, 0.5, 0.5], + "image_std": [0.5, 0.5, 0.5], + "transparent_bg_config": {"type": "white"}, + } + ) + + +def test_kimi_k3_epd_preprocess_preserves_raw_per_image_items(): + first = Image.new("RGB", (8, 6), color=(1, 2, 3)) + second = Image.new("RGB", (5, 9), color=(4, 5, 6)) + + output = prepare_kimi_k3_encoder_inputs( + [ + {"type": "image", "image": first}, + {"type": "image", "image": second}, + ], + _kimi_k3_image_processor(), + ) + + items = get_encoder_preprocessed_items(output) + assert isinstance(output, EncoderPreprocessOutput) + assert len(items) == 2 + assert output["original_image_sizes"] == [[8, 6], [5, 9]] + assert output["grid_thws"].tolist() == [[1, 4, 4], [1, 6, 4]] + for item, image in zip(items, (first, second)): + assert item.modality == Modality.IMAGE + assert item.feature is image + assert item.hash is not None + assert item.pad_value is not None + deferred = item.model_specific_data[DEFERRED_PREPROCESSING_KEY] + assert deferred["image_mean"] == [0.5, 0.5, 0.5] + assert deferred["image_std"] == [0.5, 0.5, 0.5] + + +def test_kimi_k3_epd_model_preprocessor_receives_image_processor(): + image = Image.new("RGB", (8, 6), color=(1, 2, 3)) + image_processor = _kimi_k3_image_processor() + image_processor.preprocess = lambda medias, return_tensors: { + "pixel_values": torch.zeros(16, 12), + "grid_thws": torch.tensor([[1, 4, 4]]), + } + calls = [] + + def model_preprocessor( + mm_data, + modality, + config, + *, + image_processor=None, + use_gpu_preprocessing=False, + ): + calls.append( + (mm_data, modality, config, image_processor, use_gpu_preprocessing) + ) + return prepare_kimi_k3_encoder_inputs(mm_data, image_processor) + + encoder = _encoder() + encoder.image_processor = image_processor + encoder.use_image_processor_gpu = False + encoder.vision_config = {"image": {"return_tensors": "pt"}} + encoder._flatten_and_load_images = AsyncMock(return_value=[image]) + encoder.preproc_executor = ThreadPoolExecutor(max_workers=1) + try: + with patch( + "sglang.srt.disaggregation.encode_server.get_parallel", + return_value=SimpleNamespace(attn_tp_rank=0, attn_tp_size=1), + ): + output = asyncio.run( + encoder._process_image_items([image], model_preprocessor) + ) + finally: + encoder.preproc_executor.shutdown() + + assert len(calls) == 1 + assert calls[0][0][0] == {"type": "image", "image": image} + assert calls[0][1:] == ( + Modality.IMAGE, + encoder.vision_config, + image_processor, + False, + ) + assert len(get_encoder_preprocessed_items(output)) == 1 + + +def test_encoder_preprocessor_context_keeps_legacy_hooks_compatible(): + calls = [] + + def legacy_hook(mm_data, modality, config): + calls.append((mm_data, modality, config)) + return {"ok": True} + + result = invoke_encoder_preprocessor( + legacy_hook, + ["image"], + Modality.IMAGE, + {"image": {}}, + image_processor=object(), + use_gpu_preprocessing=True, + ) + + assert result == {"ok": True} + assert calls == [(["image"], Modality.IMAGE, {"image": {}})] + + +def test_kimi_k3_epd_default_cpu_materialization_is_owner_only_and_exact(): + class RecordingImageProcessor: + media_proc_cfg = _kimi_k3_image_processor().media_proc_cfg + + def __init__(self): + self.calls = [] + + def preprocess(self, medias, return_tensors): + self.calls.append(medias) + features = [ + torch.full((4, 3, 2, 2), media["image"].getpixel((0, 0))[0]) + for media in medias + ] + grids = torch.tensor([[1, 2, 2]] * len(medias)) + return {"pixel_values": torch.cat(features), "grid_thws": grids} + + processor = RecordingImageProcessor() + images = [Image.new("RGB", (4, 4), color=(value, 0, 0)) for value in (7, 11)] + output = prepare_kimi_k3_encoder_inputs(images, processor) + items = get_encoder_preprocessed_items(output) + + materialized = materialize_kimi_k3_cpu_features([items[1]], processor) + + assert len(processor.calls) == 1 + assert len(processor.calls[0]) == 1 + assert processor.calls[0][0]["image"].getpixel((0, 0)) == (11, 0, 0) + assert torch.all(materialized == 11) + assert items[0].model_specific_data[DEFERRED_PREPROCESSING_KEY]["backend"] == "cpu" + + +def test_encoder_preprocess_materializes_only_local_size_balanced_items(): + items = [ + MultimodalDataItem( + modality=Modality.IMAGE, + feature=torch.tensor([value], dtype=torch.uint8), + ) + for value in (3, 5, 7) + ] + calls = [] + + def materialize(selected): + calls.append(selected) + return [item.feature.float() + 10 for item in selected] + + output = EncoderPreprocessOutput( + {"pixel_values": [item.feature for item in items]}, + mm_items=items, + item_sizes=[8, 5, 3], + materialize_local_items=materialize, + ) + + output.materialize_for_rank(rank=1, world_size=2) + + assert calls == [[items[1], items[2]]] + assert items[0].feature.tolist() == [3] + assert items[1].feature.tolist() == [15.0] + assert items[2].feature.tolist() == [17.0] + assert LOCAL_PREPROCESSED_KEY not in items[0].model_specific_data + assert items[1].model_specific_data[LOCAL_PREPROCESSED_KEY] + assert items[2].model_specific_data[LOCAL_PREPROCESSED_KEY] + + +def test_raw_encoder_hash_includes_shape_and_dtype(): + flat = torch.arange(12, dtype=torch.uint8) + + assert hash_raw_encoder_item(flat.reshape(2, 2, 3)) != hash_raw_encoder_item( + flat.reshape(3, 2, 2) + ) + assert hash_raw_encoder_item(flat) != hash_raw_encoder_item(flat.to(torch.int16)) + + @pytest.mark.parametrize( ("use_image_processor_gpu", "expected_decode_mode"), [(False, False), (True, "nvjpeg_fancy")], @@ -233,6 +428,67 @@ def test_kimi_k3_encoder_splits_cross_request_batch_into_single_grid_items(): torch.testing.assert_close(torch.cat(output), embeddings) +def test_encoder_preprocessed_items_follow_dp_owner_selection_order(): + encoder = _encoder() + grid_thws = torch.tensor([[1, 2, 2], [1, 2, 4], [1, 4, 2]]) + items = [ + MultimodalDataItem( + modality=Modality.IMAGE, + feature=torch.full((3, i + 2, i + 3), i, dtype=torch.uint8), + model_specific_data={"grid_thws": grid_thws[i : i + 1]}, + ) + for i in range(3) + ] + mm_inputs = EncoderPreprocessOutput( + {"pixel_values": [item.feature for item in items], "grid_thws": grid_thws}, + mm_items=items, + ) + embeddings = torch.arange(4, dtype=torch.float32).reshape(4, 1) + captured = {} + + def get_feature_fn(selected_items): + captured["items"] = selected_items + return embeddings + + output = encoder._encode_missing( + mm_inputs["pixel_values"], + mm_inputs, + indices=[2, 0], + modality=Modality.IMAGE, + get_feature_fn=get_feature_fn, + grid_thw=grid_thws, + keep_on_gpu=True, + ) + + assert captured["items"] == [items[2], items[0]] + assert [part.shape[0] for part in output] == [2, 1] + torch.testing.assert_close(torch.cat(output), embeddings[:3]) + + +def test_encoder_preprocessed_items_hash_individually(): + encoder = _encoder() + grid_thws = torch.tensor([[1, 2, 2], [1, 2, 4]]) + items = [ + MultimodalDataItem( + modality=Modality.IMAGE, + feature=torch.full((3, 2, 2), value, dtype=torch.uint8), + model_specific_data={"grid_thws": grid_thws[i : i + 1]}, + ) + for i, value in enumerate((17, 29)) + ] + mm_inputs = EncoderPreprocessOutput( + {"pixel_values": [item.feature for item in items], "grid_thws": grid_thws}, + mm_items=items, + ) + + hashes = encoder._calculate_hashes_from_features( + mm_inputs["pixel_values"], grid_thws, Modality.IMAGE, mm_inputs + ) + + assert hashes == [item.hash for item in items] + assert hashes[0] != hashes[1] + + def test_kimi_k3_encoder_only_wrapper_guards_language_tower_hooks(): model = SimpleNamespace(language_model=None) diff --git a/test/registered/unit/models/test_kimi_k3_vision.py b/test/registered/unit/models/test_kimi_k3_vision.py index d80508048..ddc3b6026 100644 --- a/test/registered/unit/models/test_kimi_k3_vision.py +++ b/test/registered/unit/models/test_kimi_k3_vision.py @@ -518,6 +518,8 @@ def test_kimi_k3_preprocesses_only_dp_owner_images(monkeypatch): model.mm_projector = lambda image_embeds: image_embeds deferred_config = { + "backend": "gpu", + "feature_layout": "chw", "image_mean": [0.5, 0.5, 0.5], "image_std": [0.5, 0.5, 0.5], "transparent_bg_config": None,