fix: fix Kimi-VL encoder parallelism

python/sglang/srt/model_executor/model_runner.py

@@ -2730,8 +2730,11 @@ def init_prefill_cuda_graph(self, force_for_draft_worker: bool = False):
             self.prefill_cuda_graph_runner = self.eager_runner
             return
 
-        # Disable prefill CUDA graph for non-language models
-        if not hasattr(self.model, "model"):
+        # Resolve the decoder once. Some VLM wrappers (for example Kimi-VL)
+        # expose it as ``language_model`` rather than ``model``.
+        try:
+            language_model = resolve_language_model(self.model)
+        except AttributeError:
             logger.warning(
                 "Disable prefill CUDA graph because the model is not a language model"
             )
@@ -2744,9 +2747,12 @@ def init_prefill_cuda_graph(self, force_for_draft_worker: bool = False):
             )
             return
 
-        # Collect attention layers and moe layers from the model
-        self.model.model = resolve_language_model(self.model)
-        language_model = getattr(self.model, "language_model", self.model)
+        # Collect attention layers and moe layers from the model. Keep a VLM
+        # wrapper that exposes ``language_model`` unchanged: assigning it to
+        # ``model`` would register a duplicate module alias and duplicate the
+        # model's state-dict namespace.
+        if hasattr(self.model, "model"):
+            self.model.model = language_model
 
         # Find the module that owns the decoder `layers`. Models wrap it at
         # varying depths: a direct text model exposes `.layers`, a CausalLM

python/sglang/srt/models/kimi_vl.py

@@ -43,14 +43,14 @@
 
 import copy
 import logging
+import math
 from dataclasses import dataclass
 from typing import Iterable, List, Optional, Tuple
 
 import torch
 from torch import nn
 from transformers.activations import GELUActivation
 
-from sglang.srt.configs import KimiVLConfig
 from sglang.srt.configs.deepseekvl2 import DeepseekV2Config
 from sglang.srt.configs.kimi_vl import KimiVLConfig
 from sglang.srt.configs.kimi_vl_moonvit import MoonViTConfig
@@ -73,6 +73,8 @@
 )
 from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM
 from sglang.srt.models.kimi_vl_moonvit import MoonVitPretrainedModel
+from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model
+from sglang.srt.runtime_context import get_server_args
 from sglang.srt.utils import add_prefix
 
 logger = logging.getLogger(__name__)
@@ -124,7 +126,13 @@ def __init__(
         self.config = config
         assert isinstance(config.vision_config, MoonViTConfig)
 
-        self.vision_tower = MoonVitPretrainedModel(config.vision_config)
+        self.use_data_parallel = get_server_args().mm_enable_dp_encoder
+        self.vision_tower = MoonVitPretrainedModel(
+            config.vision_config,
+            prefix=add_prefix("vision_tower", prefix),
+            use_data_parallel=self.use_data_parallel,
+            use_tensor_parallel=not self.use_data_parallel,
+        )
 
         self.multi_modal_projector = KimiVLMultiModalProjector(config=config)
         self.quant_config = quant_config
@@ -152,13 +160,26 @@ def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
         ):
             return pixel_values
 
-        image_grid_hws = torch.cat([item.image_grid_hws for item in items], dim=0).to(
-            self.vision_tower.device
-        )
-        image_features = self.vision_tower(pixel_values, image_grid_hws)
-        assert isinstance(image_features, list)
-        # lengths = [x.shape[0] for x in image_features]
-        res = self.multi_modal_projector(torch.cat(image_features))  # .split(lengths)
+        image_grid_hws = torch.cat([item.image_grid_hws for item in items], dim=0)
+        image_grid_hws_list = image_grid_hws.tolist()
+        if self.use_data_parallel:
+            image_features = run_dp_sharded_mrope_vision_model(
+                self.vision_tower,
+                pixel_values,
+                image_grid_hws_list,
+                rope_type="rope_2d",
+            )
+        else:
+            image_grid_hws = image_grid_hws.to(self.vision_tower.device)
+            image_features = self.vision_tower(
+                pixel_values,
+                image_grid_hws,
+                max_seqlen=max(math.prod(grid) for grid in image_grid_hws_list),
+            )
+            assert isinstance(image_features, list)
+            image_features = torch.cat(image_features)
+
+        res = self.multi_modal_projector(image_features)
         return res
 
     def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs):
@@ -242,8 +263,6 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
             use_default_weight_loading = False
             if "vision" in name:
                 if self.vision_tower is not None:
-                    # We only do sharding for language model and
-                    # not vision model for now.
                     use_default_weight_loading = True
             else:
                 for param_name, weight_name, shard_id in stacked_params_mapping: