[EPD][VLM] Support Kimi K25 EPD (#22269)

Signed-off-by: LHXuuu <xulianhao.xlh@antgroup.com>

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:

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()."
             )

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]:

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."