[Feature] add support kimi vl model (#5383)

Co-authored-by: wenju.li <[email protected]>

Author: liwenju0

docs/supported_models/vision_language_models.md

@@ -28,4 +28,5 @@ python3 -m sglang.launch_server \
 | **LLaVA** (v1.5 & v1.6)        | *e.g.* `liuhaotian/llava-v1.5-13b`               | `vicuna_v1.1`        | Open vision-chat models that add an image encoder to LLaMA/Vicuna (e.g. LLaMA2 13B) for following multimodal instruction prompts. |
 | **LLaVA-NeXT** (8B, 72B)       | `lmms-lab/llava-next-72b`                        | `chatml-llava`       | Improved LLaVA models (with an 8B Llama3 version and a 72B version) offering enhanced visual instruction-following and accuracy on multimodal benchmarks. |
 | **LLaVA-OneVision**            | `lmms-lab/llava-onevision-qwen2-7b-ov`           | `chatml-llava`       | Enhanced LLaVA variant integrating Qwen as the backbone; supports multiple images (and even video frames) as inputs via an OpenAI Vision API-compatible format. |
-| **Gemma 3 (Multimodal)**       | `google/gemma-3-4b-it`                           | `gemma-it`           | Gemma 3’s larger models (4B, 12B, 27B) accept images (each image encoded as 256 tokens) alongside text in a combined 128K-token context. |
+| **Gemma 3 (Multimodal)**       | `google/gemma-3-4b-it`                           | `gemma-it`           | Gemma 3’s larger models (4B, 12B, 27B) accept images (each image encoded as 256 tokens) alongside text in a combined 128K-token context. |
+| **Kimi-VL** (A3B)              | `moonshotai/Kimi-VL-A3B-Instruct`                | `kimi-vl`            | Kimi-VL is a multimodal model that can understand and generate text from images. |

python/pyproject.toml

@@ -42,6 +42,7 @@ runtime_common = [
     "uvicorn",
     "uvloop",
     "xgrammar==0.1.17",
+    "blobfile==3.0.0"
 ]
 
 srt = [

python/sglang/srt/configs/__init__.py

@@ -3,11 +3,15 @@
 from sglang.srt.configs.deepseekvl2 import DeepseekVL2Config
 from sglang.srt.configs.exaone import ExaoneConfig
 from sglang.srt.configs.janus_pro import MultiModalityConfig
+from sglang.srt.configs.kimi_vl import KimiVLConfig
+from sglang.srt.configs.kimi_vl_moonvit import MoonViTConfig
 
 __all__ = [
     "ExaoneConfig",
     "ChatGLMConfig",
     "DbrxConfig",
     "DeepseekVL2Config",
     "MultiModalityConfig",
+    "KimiVLConfig",
+    "MoonViTConfig",
 ]

python/sglang/srt/configs/kimi_vl.py

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from https://huggingface.co/moonshotai/Kimi-VL-A3B-Instruct/blob/main/configuration_kimi_vl.py
+from typing import Optional, Union
+
+from transformers.configuration_utils import PretrainedConfig
+
+from sglang.srt.configs.deepseekvl2 import DeepseekV2Config
+from sglang.srt.configs.kimi_vl_moonvit import MoonViTConfig
+
+
+class KimiVLConfig(PretrainedConfig):
+    model_type = "kimi_vl"
+
+    def __init__(
+        self,
+        vision_config: Optional[Union[dict, MoonViTConfig]] = None,
+        text_config: Optional[Union[dict, DeepseekV2Config]] = None,
+        ignore_index: int = -100,
+        media_placeholder_token_id: int = 163605,
+        pad_token_id: int = 0,
+        **kwargs
+    ):
+        if vision_config is None:
+            vision_config = MoonViTConfig()
+        elif isinstance(vision_config, dict):
+            vision_config = MoonViTConfig(**vision_config)
+        self.vision_config = vision_config
+
+        if text_config is None:
+            text_config = DeepseekV2Config()
+        elif isinstance(text_config, dict):
+            text_config = DeepseekV2Config(**text_config)
+        self.text_config = text_config
+
+        self.ignore_index = ignore_index
+        self.media_placeholder_token_id = media_placeholder_token_id
+
+        super().__init__(pad_token_id=pad_token_id, **kwargs)

python/sglang/srt/configs/kimi_vl_moonvit.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from https://huggingface.co/moonshotai/Kimi-VL-A3B-Instruct/blob/main/configuration_kimi_vl.py
+from transformers.configuration_utils import PretrainedConfig
+
+
+class MoonViTConfig(PretrainedConfig):
+    model_type = "moonvit"
+
+    def __init__(
+        self,
+        patch_size: int = 14,
+        init_pos_emb_height: int = 64,
+        init_pos_emb_width: int = 64,
+        num_attention_heads: int = 16,
+        num_hidden_layers: int = 27,
+        hidden_size: int = 1152,
+        intermediate_size: int = 4304,
+        merge_kernel_size: tuple[int, int] = (2, 2),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.patch_size = patch_size
+        # Positional embedding config
+        self.init_pos_emb_height = init_pos_emb_height
+        self.init_pos_emb_width = init_pos_emb_width
+        # Transformer config
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        # Patch merger config
+        self.merge_kernel_size = merge_kernel_size

python/sglang/srt/configs/model_config.py

@@ -176,6 +176,13 @@ def __init__(
             self.attention_arch = AttentionArch.MLA
             self.kv_lora_rank = self.hf_text_config.kv_lora_rank
             self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim
+        elif "KimiVLForConditionalGeneration" in self.hf_config.architectures:
+            self.head_dim = 256
+            self.attention_arch = AttentionArch.MLA
+            self.kv_lora_rank = self.hf_text_config.kv_lora_rank
+            self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim
+            self.v_head_dim = self.hf_text_config.v_head_dim
+            self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim
         else:
             self.attention_arch = AttentionArch.MHA
 
@@ -530,6 +537,7 @@ def is_generation_model(model_architectures: List[str], is_embedding: bool = Fal
     "Qwen2VLForConditionalGeneration",
     "Qwen2_5_VLForConditionalGeneration",
     "CLIPModel",
+    "KimiVLForConditionalGeneration",
 ]
 
 

python/sglang/srt/conversation.py

@@ -806,6 +806,24 @@ def generate_chat_conv(
     )
 )
 
+# Reference: https://huggingface.co/moonshotai/Kimi-VL-A3B-Instruct/blob/main/chat_template.jinja
+register_conv_template(
+    Conversation(
+        name="kimi-vl",
+        system_message="You are a helpful assistant",
+        system_template="<|im_system|>system<|im_middle|>{system_message}",
+        roles=(
+            "<|im_user|>user<|im_middle|>",
+            "<|im_assistant|>assistant<|im_middle|>",
+        ),
+        messages=[],
+        sep="<|im_end|>",
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        stop_str="<|im_end|>",
+        image_token="<|media_start|>image<|media_content|><|media_pad|><|media_end|>",
+    )
+)
+
 
 @register_conv_template_matching_function
 def match_deepseek_janus_pro(model_path: str):
@@ -888,3 +906,10 @@ def match_openbmb_minicpm(model_path: str):
         return "minicpmv"
     elif "minicpm-o" in model_path:
         return "minicpmo"
+
+
+@register_conv_template_matching_function
+def match_moonshot_kimivl(model_path: str):
+    model_path = model_path.lower()
+    if "kimi" in model_path and "vl" in model_path:
+        return "kimi-vl"

python/sglang/srt/hf_transformers_utils.py

@@ -35,6 +35,7 @@
     DbrxConfig,
     DeepseekVL2Config,
     ExaoneConfig,
+    KimiVLConfig,
     MultiModalityConfig,
 )
 from sglang.srt.connector import create_remote_connector
@@ -46,6 +47,7 @@
     ExaoneConfig.model_type: ExaoneConfig,
     DeepseekVL2Config.model_type: DeepseekVL2Config,
     MultiModalityConfig.model_type: MultiModalityConfig,
+    KimiVLConfig.model_type: KimiVLConfig,
 }
 
 for name, cls in _CONFIG_REGISTRY.items():

python/sglang/srt/managers/multimodal_processors/kimi_vl.py

@@ -0,0 +1,73 @@
+import asyncio
+import math
+from typing import List, Union
+
+import torch
+from PIL import Image
+
+from sglang.srt.managers.multimodal_processors.base_processor import (
+    BaseMultimodalProcessor as SGLangBaseProcessor,
+)
+from sglang.srt.managers.multimodal_processors.base_processor import (
+    MultimodalSpecialTokens,
+)
+from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
+from sglang.srt.models.kimi_vl import KimiVLForConditionalGeneration
+
+
+# Compatible with KimiVLForConditionalGeneration
+class KimiVLImageProcessor(SGLangBaseProcessor):
+    models = [KimiVLForConditionalGeneration]
+
+    def __init__(self, hf_config, server_args, _processor):
+        super().__init__(hf_config, server_args, _processor)
+        self.IMAGE_TOKEN = "<|media_pad|>"
+        self.im_token_id = _processor.tokenizer.convert_tokens_to_ids(self.IMAGE_TOKEN)
+
+        self.im_start = "<|media_start|>"
+        self.im_start_id = _processor.tokenizer.convert_tokens_to_ids(self.im_start)
+
+        self.im_end = "<|media_end|>"
+        self.im_end_id = _processor.tokenizer.convert_tokens_to_ids(self.im_end)
+
+        self.im_content = "<|media_content|>"
+        self.im_content_id = _processor.tokenizer.convert_tokens_to_ids(self.im_content)
+
+    async def process_mm_data_async(
+        self,
+        image_data: List[Union[str, bytes]],
+        input_text,
+        request_obj,
+        max_req_input_len,
+        *args,
+        **kwargs,
+    ):
+        if not image_data:
+            return None
+        if isinstance(image_data, str):
+            image_data = [image_data]
+
+        base_output = self.load_mm_data(
+            prompt=input_text,
+            image_data=image_data,
+            multimodal_tokens=MultimodalSpecialTokens(image_token=self.IMAGE_TOKEN),
+            max_req_input_len=max_req_input_len,
+        )
+        ret = self.process_mm_data(
+            input_text=base_output.input_text,
+            images=base_output.images,
+        )
+        return {
+            "input_ids": ret["input_ids"].flatten().tolist(),
+            "mm_items": [
+                MultimodalDataItem(
+                    pixel_values=ret["pixel_values"],
+                    image_grid_thws=ret["image_grid_hws"],
+                    modality=Modality.IMAGE,
+                )
+            ],
+            "im_token_id": self.im_token_id,
+            "im_start_id": self.im_start_id,
+            "im_end_id": self.im_end_id,
+            "im_content_id": self.im_content_id,
+        }

python/sglang/srt/models/deepseek_v2.py

@@ -752,7 +752,7 @@ def forward_absorb(
         q_nope_out = q_nope_out.transpose(0, 1)
 
         k_nope = latent_cache[..., : self.kv_lora_rank]
-        k_nope = self.kv_a_layernorm(k_nope).unsqueeze(1)
+        k_nope = self.kv_a_layernorm(k_nope.contiguous()).unsqueeze(1)
         k_pe = latent_cache[..., self.kv_lora_rank :].unsqueeze(1)
 
         q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
@@ -1391,6 +1391,9 @@ def __init__(
 
         self.dp_size = get_attention_dp_size()
 
+    def get_input_embeddings(self) -> torch.Tensor:
+        return self.embed_tokens
+
     def forward(
         self,
         input_ids: torch.Tensor,