diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index 5d42150d0..9f4e4eeea 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -124,7 +124,8 @@ def _convert(data): _mm_grid_attrs = { - Modality.IMAGE: ["image_grid_thw", "image_grid_hws"], + # Kimi K2.5 HF processor uses grid_thws (see base_processor.ATTR_NAME_TO_MODALITY). + Modality.IMAGE: ["image_grid_thw", "image_grid_hws", "grid_thws"], Modality.VIDEO: ["video_grid_thw"], Modality.AUDIO: ["audio_feature_lens_raw"], } @@ -136,9 +137,14 @@ _mm_feature_attrs = { } -def _get_mm_grid_dim(mm_inputs, modality): - for attr in _mm_grid_attrs[modality]: - if attr in mm_inputs: +def _get_mm_grid_dim(mm_inputs, modality, model_type: Optional[str] = None): + # Kimi K2.5 vision processor only emits `grid_thws`; prefer it over generic keys + # so we never pick a mis-typed or stale `image_grid_hws` field from kwargs. + attrs = _mm_grid_attrs[modality] + if (model_type or "").lower() == "kimi_k25" and modality == Modality.IMAGE: + attrs = ("grid_thws", "image_grid_thw", "image_grid_hws") + for attr in attrs: + if attr in mm_inputs and mm_inputs[attr] is not None: return mm_inputs[attr] raise ValueError(f"Grid dim ({_mm_grid_attrs[modality]}) not found in {mm_inputs}") @@ -577,6 +583,18 @@ class MMEncoder: else: return int(grid[0] * grid[1] * grid[2]) + def _kimi_k25_tokens_from_patch_grid( + self, grid: Union[torch.Tensor, List[int]] + ) -> int: + """MoonViT + tpool: output len is (h//mh)*(w//mw); temporal dim is pooled (not t*h*w/merge^2).""" + if isinstance(grid, torch.Tensor): + flat = grid.flatten() + _t, h, w = (int(x) for x in flat[:3].tolist()) + else: + _t, h, w = int(grid[0]), int(grid[1]), int(grid[2]) + merge_h, merge_w = self.model_config.hf_config.vision_config.merge_kernel_size + return (h * w) // (merge_h * merge_w) + def get_num_tokens( self, grid: Union[torch.Tensor, List[int]], modality: Modality ) -> int: @@ -585,6 +603,8 @@ class MMEncoder: input_length = self.get_num_patches(grid, modality) return self._get_feat_extract_output_lengths(input_length) else: + if self.model_type == "kimi_k25" and modality == Modality.IMAGE: + return self._kimi_k25_tokens_from_patch_grid(grid) merge_size = getattr(self.image_processor, "merge_size", 2) return self.get_num_patches(grid, modality) // (merge_size**2) @@ -625,7 +645,7 @@ class MMEncoder: """ GPU Task: Run ViT inference ONLY on the subset of mm items missing from the cache. """ - grid_thw = _get_mm_grid_dim(mm_inputs, modality) + grid_thw = _get_mm_grid_dim(mm_inputs, modality, self.model_type) # 1. Slice mm_feature to get only the patches for missing mm items sub_feature_list = [] @@ -675,7 +695,7 @@ class MMEncoder: ) -> torch.Tensor: # mm_inputs: dict mm_inputs, get_feature_fn = await self._process_mm_items(mm_items, modality) - grid_thw = _get_mm_grid_dim(mm_inputs, modality) + grid_thw = _get_mm_grid_dim(mm_inputs, modality, self.model_type) mm_feature = _convert(_get_mm_feature(mm_inputs, modality)) num_items = len(grid_thw) @@ -849,10 +869,30 @@ class MMEncoder: ] return timestamps + def _normalize_kimi_k25_encoder_images(self, images): + """KimiK25VisionProcessor.preprocess expects MediaInput dicts, not raw PIL.""" + from PIL import Image as PILImage + + def wrap_one(img): + if isinstance(img, dict) and img.get("type") in ("image", "video_chunk"): + return img + if isinstance(img, PILImage.Image): + return {"type": "image", "image": img} + return img + + if not images: + return images + # Disagg may supply a nested list; Kimi preprocess expects a flat list of media. + if isinstance(images[0], (list, tuple)): + images = [x for group in images for x in group] + return [wrap_one(img) for img in images] + async def _process_mm_items(self, mm_items, modality): if modality == Modality.IMAGE and self.image_processor: images = await self._flatten_and_load_images(mm_items) image_config = self.vision_config.get("image", {}) + if self.model_type == "kimi_k25": + images = self._normalize_kimi_k25_encoder_images(images) processor_input = self.image_processor(images=images, **image_config) if hasattr(self.model, "thinker"): # for omni models get_feature_method = self.model.thinker.get_image_feature @@ -985,7 +1025,11 @@ class MMEncoder: self.profiler.step() aux_data = _build_mm_aux_data(mm_inputs) - return _get_mm_grid_dim(mm_inputs, modality), mm_embedding, aux_data + return ( + _get_mm_grid_dim(mm_inputs, modality, self.model_type), + mm_embedding, + aux_data, + ) except BadRequestError as e: raise BadRequestError(f"Bad request error: {str(e)}") except Exception as e: diff --git a/python/sglang/srt/models/kimi_k25.py b/python/sglang/srt/models/kimi_k25.py index 7fd30d248..d07bf94cf 100644 --- a/python/sglang/srt/models/kimi_k25.py +++ b/python/sglang/srt/models/kimi_k25.py @@ -708,33 +708,32 @@ class KimiK25ForConditionalGeneration(nn.Module): # Create mm projector self.mm_projector = K2VLMultiModalProjector(config.vision_config) - self.language_model = DeepseekV3ForCausalLM( - config.text_config, - quant_config, - prefix=( - "language_model" if isinstance(quant_config, ModelSlimConfig) else "" - ), - ) - - self.model = self.language_model.model + self.language_model = None + if not config.encoder_only: + self.language_model = DeepseekV3ForCausalLM( + config.text_config, + quant_config, + prefix=( + "language_model" + if isinstance(quant_config, ModelSlimConfig) + else "" + ), + ) # Ensure that the dtype of the vision_tower and mm_projector matches that of the language_model. # This solves the dtype mismatch issue when using device_map="auto" and torch_dtype. - if hasattr(self.language_model, "dtype"): + if self.language_model is not None and hasattr(self.language_model, "dtype"): target_dtype = self.language_model.dtype self.vision_tower = self.vision_tower.to(dtype=target_dtype) self.mm_projector = self.mm_projector.to(dtype=target_dtype) def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: - pixel_values = torch.cat([item.feature for item in items], dim=0).type( - self.vision_tower.dtype - ) - grid_thws = torch.concat([item.grid_thws for item in items], dim=0).to( - self.vision_tower.device - ) - + device = self.vision_tower.device target_dtype = self.vision_tower.patch_embed.proj.weight.dtype - pixel_values = pixel_values.to(target_dtype) + pixel_values = torch.cat([item.feature for item in items], dim=0).to( + device=device, dtype=target_dtype + ) + grid_thws = torch.concat([item.grid_thws for item in items], dim=0).to(device) if self.use_data_parallel: image_embeds = run_dp_sharded_mrope_vision_model( @@ -761,15 +760,22 @@ class KimiK25ForConditionalGeneration(nn.Module): @property def start_layer(self) -> int: - return self.language_model.start_layer + return self.language_model.start_layer if self.language_model is not None else 0 @property def end_layer(self) -> int: - return self.language_model.end_layer + if self.language_model is not None: + return self.language_model.end_layer + text_config = getattr(self.config, "text_config", None) + return int(getattr(text_config, "num_hidden_layers", 0)) @property def routed_experts_weights_of_layer(self): - return self.language_model._routed_experts_weights_of_layer.value + return ( + self.language_model._routed_experts_weights_of_layer.value + if self.language_model is not None + else {} + ) def forward( self, @@ -814,19 +820,20 @@ class KimiK25ForConditionalGeneration(nn.Module): # All other weights go to language model language_weights.append((name, loaded_weight)) - # Load vision tower weights - vision_state_dict = dict(vision_weights) - params_dict = dict(self.named_parameters(remove_duplicate=False)) - for name, loaded_weight in vision_state_dict.items(): - if name not in params_dict: - raise ValueError(f"Weight {name} not found in params_dict") - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - # loaded_weight = self._pad_vit_attn_dummy_heads(name, loaded_weight) - weight_loader(param, loaded_weight) + if not self.config.language_only: + # Load vision tower weights + vision_state_dict = dict(vision_weights) + params_dict = dict(self.named_parameters(remove_duplicate=False)) + for name, loaded_weight in vision_state_dict.items(): + if name not in params_dict: + raise ValueError(f"Weight {name} not found in params_dict") + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + # loaded_weight = self._pad_vit_attn_dummy_heads(name, loaded_weight) + weight_loader(param, loaded_weight) # Load language model weights - if language_weights: + if not self.config.encoder_only and language_weights: self.language_model.load_weights(language_weights) @classmethod @@ -842,7 +849,9 @@ class KimiK25ForConditionalGeneration(nn.Module): self, layer_ids: Optional[List[int]] = None ) -> None: """Set the layers to capture for EAGLE3 speculative decoding.""" - if not hasattr(self.language_model, "set_eagle3_layers_to_capture"): + if self.language_model is None or not hasattr( + self.language_model, "set_eagle3_layers_to_capture" + ): raise AttributeError( "language_model does not support EAGLE3 speculative decoding." ) @@ -875,7 +884,9 @@ class KimiK25ForConditionalGeneration(nn.Module): def get_embed_and_head(self) -> Tuple[torch.Tensor, torch.Tensor]: """Get embedding and LM head weights for speculative decoding.""" - if not hasattr(self.language_model, "get_embed_and_head"): + if self.language_model is None or not hasattr( + self.language_model, "get_embed_and_head" + ): raise AttributeError( "language_model does not support get_embed_and_head()." ) @@ -884,7 +895,9 @@ class KimiK25ForConditionalGeneration(nn.Module): def set_embed_and_head(self, embed: torch.Tensor, head: torch.Tensor) -> None: """Set embedding and LM head weights for speculative decoding.""" - if not hasattr(self.language_model, "set_embed_and_head"): + if self.language_model is None or not hasattr( + self.language_model, "set_embed_and_head" + ): raise AttributeError( "language_model does not support set_embed_and_head()." ) diff --git a/python/sglang/srt/multimodal/processors/kimi_k25.py b/python/sglang/srt/multimodal/processors/kimi_k25.py index ece7b02c3..96f8f504f 100644 --- a/python/sglang/srt/multimodal/processors/kimi_k25.py +++ b/python/sglang/srt/multimodal/processors/kimi_k25.py @@ -4,6 +4,7 @@ from typing import Dict, List, Tuple, Union import torch from sglang.srt.managers.schedule_batch import ( + Modality, MultimodalDataItem, MultimodalProcessorOutput, ) @@ -55,6 +56,70 @@ class KimiK2_5VLImageProcessor(SGLangBaseProcessor): im_token_id=self.mm_tokens.image_token_id, ) + def _num_image_tokens_from_grid(self, grid_thw: torch.Tensor) -> int: + # Kimi-K2.5 applies temporal pooling and spatial 2D merge in vision tower. + # The output sequence length per image is h*w/(merge_h*merge_w). + merge_h, merge_w = self.hf_config.vision_config.merge_kernel_size + _t, h, w = grid_thw.tolist() + return (h * w) // (merge_h * merge_w) + + def get_mm_data(self, prompt, embeddings, **kwargs): + img_grid_thw = kwargs.get("img_grid_thw", None) + + if not isinstance(prompt, list): + prompt = self._tokenizer.encode(prompt) + + image_token_id = self.mm_tokens.image_token_id + image_token_counts = [ + self._num_image_tokens_from_grid(grid) for grid in img_grid_thw + ] + + input_ids = [] + offsets = [] + img_idx = 0 + + for token in prompt: + if token != image_token_id: + input_ids.append(token) + continue + + if img_idx >= len(image_token_counts): + raise ValueError( + "The number of image placeholders exceeds img_grid_thw entries." + ) + + num_tokens = image_token_counts[img_idx] + start = len(input_ids) + input_ids.extend([image_token_id] * num_tokens) + offsets.append((start, len(input_ids) - 1)) + img_idx += 1 + + if img_idx != len(image_token_counts): + raise ValueError( + "The number of image placeholders does not match img_grid_thw entries." + ) + + image_embeddings = embeddings[Modality.IMAGE] + mm_items = [] + consumed = 0 + for start, end in offsets: + num_tokens = end - start + 1 + embedding_slice = image_embeddings[consumed : consumed + num_tokens] + consumed += num_tokens + mm_items.append( + MultimodalDataItem( + modality=Modality.IMAGE, + offsets=[(start, end)], + precomputed_embeddings=embedding_slice, + ) + ) + + return MultimodalProcessorOutput( + input_ids=input_ids, + 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]: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index b65922d52..fe00f97df 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -3507,6 +3507,8 @@ class ServerArgs: "Qwen3OmniMoeForConditionalGeneration", "Qwen2AudioForConditionalGeneration", "Qwen2_5OmniForConditionalGeneration", + "KimiVLForConditionalGeneration", + "KimiK25ForConditionalGeneration", ]: raise ValueError( f"Model type {model_arch} is not supported for encoder disaggregation, only Qwen models are supported for now."