fix: fix Kimi-VL encoder parallelism (#30869)

python/sglang/srt/model_executor/model_runner.py

@@ -2725,8 +2725,11 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             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"
             )
@@ -2739,9 +2742,12 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             )
             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,6 +43,7 @@
 
 import copy
 import logging
+import math
 from dataclasses import dataclass
 from typing import Iterable, List, Optional, Tuple
 
@@ -50,7 +51,6 @@ 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.model_loader.weight_utils import (
 )
 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 @@ class KimiVLForConditionalGeneration(nn.Module):
         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 @@ class KimiVLForConditionalGeneration(nn.Module):
         ):
             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 @@ class KimiVLForConditionalGeneration(nn.Module):
             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:

python/sglang/srt/models/kimi_vl_moonvit.py

@@ -61,11 +61,19 @@ except ImportError:
 
 from sglang.srt.configs import MoonViTConfig
 from sglang.srt.layers.conv import Conv2dLayer
-from sglang.srt.layers.linear import ReplicatedLinear
+from sglang.srt.layers.linear import (
+    ColumnParallelLinear,
+    QKVParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
 from sglang.srt.layers.quantization import QuantizationConfig
 from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig
+from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import add_prefix, get_device
 
+_MAX_INFERENCE_POS_EMB_CACHE_ENTRIES = 256
+
 
 @debug_kernel_api
 def multihead_attention(
@@ -74,6 +82,7 @@ def multihead_attention(
     v: torch.Tensor,
     q_cu_seqlens: Optional[torch.Tensor] = None,
     k_cu_seqlens: Optional[torch.Tensor] = None,
+    max_seqlen: Optional[int] = None,
 ):
     """Multi-head attention using flash attention 2.
     This function is used to handle the case where the query, key, and value are packed.
@@ -94,25 +103,28 @@ def multihead_attention(
         )
     # Unified format legal check
     assert q.dim() == k.dim() == v.dim() == 3, "q, k, v must have 3 dims"
-    assert q_cu_seqlens[-1] == q.shape[0], "q_cu_seqlens must sum to q.shape[0]"
-    assert (
-        k_cu_seqlens[-1] == k.shape[0] == v.shape[0]
-    ), "k_cu_seqlens must sum to k.shape[0]"
+    # Keep validation on CPU for debugging, but avoid synchronizing the GPU
+    # once per MoonViT layer in the normal packed CUDA path.
+    if not q_cu_seqlens.is_cuda:
+        assert q_cu_seqlens[-1] == q.shape[0], "q_cu_seqlens must sum to q.shape[0]"
+        assert (
+            k_cu_seqlens[-1] == k.shape[0] == v.shape[0]
+        ), "k_cu_seqlens must sum to k.shape[0]"
     assert q.dtype in [
         torch.bfloat16,
         torch.float16,
     ], f"unsupported dtype {q.dtype} for multihead attn"
 
-    max_seqlen_q = (q_cu_seqlens[1:] - q_cu_seqlens[:-1]).max().item()
-    max_seqlen_k = (k_cu_seqlens[1:] - k_cu_seqlens[:-1]).max().item()
+    if max_seqlen is None:
+        max_seqlen = (q_cu_seqlens[1:] - q_cu_seqlens[:-1]).max().item()
     attn_out = flash_attn_varlen_func(
         q,
         k,
         v,
         q_cu_seqlens,
         k_cu_seqlens,
-        max_seqlen_q,
-        max_seqlen_k,
+        max_seqlen,
+        max_seqlen,
         causal=False,
     )
     attn_out = attn_out.flatten(start_dim=-2)
@@ -126,6 +138,7 @@ def sdpa_attention(
     v: torch.Tensor,
     q_cu_seqlens: Optional[torch.Tensor] = None,
     k_cu_seqlens: Optional[torch.Tensor] = None,
+    max_seqlen: Optional[int] = None,
 ) -> torch.Tensor:
     """Multi-head attention using torch scaled dot product attention.
     This function is used to handle the case where the query, key, and value are packed.
@@ -208,6 +221,13 @@ class Learnable2DInterpPosEmb(nn.Module):
         self.width = width
         self.interpolation_mode = interpolation_mode
         self.weight = nn.Parameter(torch.empty(height, width, dim))
+        # In serving, MoonViT weights are immutable and image grids commonly
+        # repeat. Avoid launching bicubic interpolation for every request.
+        # Keep this as a plain cache (rather than a buffer) so it is neither
+        # serialized nor used during training.
+        self._interpolated_pos_emb_cache: dict[
+            tuple[tuple[int, int], torch.dtype, torch.device], torch.Tensor
+        ] = {}
         self.reset_parameters()
 
     def reset_parameters(self):
@@ -216,19 +236,33 @@ class Learnable2DInterpPosEmb(nn.Module):
     def forward(self, x: torch.Tensor, grid_hws: torch.Tensor) -> torch.Tensor:
         pos_embs = []
         for shape in grid_hws.tolist():
+            shape = tuple(shape)
             if shape == self.weight.shape[:-1]:
                 pos_embs.append(self.weight.flatten(end_dim=1))
             else:
-                pos_embs.append(
-                    F.interpolate(
-                        self.weight.permute((2, 0, 1)).unsqueeze(0),
-                        size=shape,
-                        mode=self.interpolation_mode,
+                cache_key = (shape, self.weight.dtype, self.weight.device)
+                pos_emb = self._interpolated_pos_emb_cache.get(cache_key)
+                if pos_emb is None:
+                    pos_emb = (
+                        F.interpolate(
+                            self.weight.permute((2, 0, 1)).unsqueeze(0),
+                            size=shape,
+                            mode=self.interpolation_mode,
+                        )
+                        .squeeze(0)
+                        .permute((1, 2, 0))
+                        .flatten(end_dim=1)
                     )
-                    .squeeze(0)
-                    .permute((1, 2, 0))
-                    .flatten(end_dim=1)
-                )
+                    if not self.training:
+                        if (
+                            len(self._interpolated_pos_emb_cache)
+                            >= _MAX_INFERENCE_POS_EMB_CACHE_ENTRIES
+                        ):
+                            self._interpolated_pos_emb_cache.pop(
+                                next(iter(self._interpolated_pos_emb_cache))
+                            )
+                        self._interpolated_pos_emb_cache[cache_key] = pos_emb
+                pos_embs.append(pos_emb)
         out = x + torch.cat(pos_embs)
         return out
 
@@ -408,11 +442,16 @@ class MLP2(nn.Module):
         bias: bool = True,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
+        use_data_parallel: bool = False,
+        use_tensor_parallel: bool = False,
     ):
         super().__init__()
         assert len(dims) == 3
 
         self.quant_config = quant_config
+        use_tensor_parallel = use_tensor_parallel and not use_data_parallel
+        tp_size = get_parallel().attn_tp_size if use_tensor_parallel else 1
+        tp_rank = get_parallel().attn_tp_rank if use_tensor_parallel else 0
         if isinstance(self.quant_config, ModelSlimConfig):
             self.fc0 = ReplicatedLinear(
                 dims[0],
@@ -428,6 +467,23 @@ class MLP2(nn.Module):
                 quant_config=quant_config,
                 prefix=add_prefix("fc1", prefix),
             )
+        elif use_tensor_parallel:
+            self.fc0 = ColumnParallelLinear(
+                dims[0],
+                dims[1],
+                bias=bias,
+                prefix=add_prefix("fc0", prefix),
+                tp_rank=tp_rank,
+                tp_size=tp_size,
+            )
+            self.fc1 = RowParallelLinear(
+                dims[1],
+                dims[2],
+                bias=bias,
+                prefix=add_prefix("fc1", prefix),
+                tp_rank=tp_rank,
+                tp_size=tp_size,
+            )
         else:
             self.fc0 = nn.Linear(dims[0], dims[1], bias=bias)
             self.fc1 = nn.Linear(dims[1], dims[2], bias=bias)
@@ -443,6 +499,10 @@ class MLP2(nn.Module):
             x, _ = self.fc0(x)
             x = self.activation(x)
             x, _ = self.fc1(x)
+        elif isinstance(self.fc0, ColumnParallelLinear):
+            x, _ = self.fc0(x)
+            x = self.activation(x)
+            x, _ = self.fc1(x)
         else:
             x = self.fc0(x)
             x = self.activation(x)
@@ -461,6 +521,9 @@ class MoonVitEncoderLayer(nn.Module):
         attn_implementation: str = "flash_attention_2",  # use fa2 in sglang by default
         activation=F.gelu,
         attn_bias: bool = False,
+        prefix: str = "",
+        use_data_parallel: bool = False,
+        use_tensor_parallel: bool = False,
     ):
         super().__init__()
         self.num_heads = num_heads
@@ -468,28 +531,63 @@ class MoonVitEncoderLayer(nn.Module):
         self.hidden_size_per_attention_head = self.hidden_dim // self.num_heads
         self.attn_implementation = attn_implementation
 
+        self.use_tensor_parallel = use_tensor_parallel and not use_data_parallel
+        tp_size = get_parallel().attn_tp_size if self.use_tensor_parallel else 1
+        tp_rank = get_parallel().attn_tp_rank if self.use_tensor_parallel else 0
+        self.num_attention_heads_per_partition = self.num_heads // tp_size
+
         self.norm0 = nn.LayerNorm(hidden_dim)
         self.norm1 = nn.LayerNorm(hidden_dim)
-        self.mlp = MLP2([hidden_dim, mlp_dim, hidden_dim], activation)
-        self.wqkv = nn.Linear(hidden_dim, hidden_dim * 3, bias=attn_bias)
-        self.wo = nn.Linear(hidden_dim, hidden_dim, bias=attn_bias)
+        self.mlp = MLP2(
+            [hidden_dim, mlp_dim, hidden_dim],
+            activation,
+            prefix=add_prefix("mlp", prefix),
+            use_data_parallel=use_data_parallel,
+            use_tensor_parallel=self.use_tensor_parallel,
+        )
+        if self.use_tensor_parallel:
+            self.wqkv = QKVParallelLinear(
+                hidden_size=hidden_dim,
+                head_size=self.hidden_size_per_attention_head,
+                total_num_heads=num_heads,
+                total_num_kv_heads=num_heads,
+                bias=attn_bias,
+                prefix=add_prefix("wqkv", prefix),
+                tp_rank=tp_rank,
+                tp_size=tp_size,
+            )
+            self.wo = RowParallelLinear(
+                hidden_dim,
+                hidden_dim,
+                bias=attn_bias,
+                prefix=add_prefix("wo", prefix),
+                tp_rank=tp_rank,
+                tp_size=tp_size,
+            )
+        else:
+            self.wqkv = nn.Linear(hidden_dim, hidden_dim * 3, bias=attn_bias)
+            self.wo = nn.Linear(hidden_dim, hidden_dim, bias=attn_bias)
 
     def attention_qkvpacked(
         self,
         x: torch.Tensor,
         cu_seqlens: torch.Tensor,
         rope_freqs_cis: Optional[torch.Tensor] = None,
+        max_seqlen: Optional[int] = None,
     ):
         """
         Args:
             x (torch.Tensor): (batch_size, seqlen, hidden_dim)
             cu_seqlens (torch.Tensor):
         """
-        xqkv = self.wqkv(x)
+        if self.use_tensor_parallel:
+            xqkv, _ = self.wqkv(x)
+        else:
+            xqkv = self.wqkv(x)
 
         qkv_shape = xqkv.size()[:-1] + (
             3,
-            self.num_heads,
+            self.num_attention_heads_per_partition,
             self.hidden_size_per_attention_head,
         )
         # xqkv: (batch_size, seqlen, 3, nheads, headdim)
@@ -500,10 +598,18 @@ class MoonVitEncoderLayer(nn.Module):
 
         attn_func = VL_VISION_ATTENTION_FUNCTIONS[self.attn_implementation]
         attn_out = attn_func(
-            xq, xk, xv, q_cu_seqlens=cu_seqlens, k_cu_seqlens=cu_seqlens
+            xq,
+            xk,
+            xv,
+            q_cu_seqlens=cu_seqlens,
+            k_cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
         )
 
-        attn_out = self.wo(attn_out)
+        if self.use_tensor_parallel:
+            attn_out, _ = self.wo(attn_out)
+        else:
+            attn_out = self.wo(attn_out)
         return attn_out
 
     def forward(
@@ -511,6 +617,7 @@ class MoonVitEncoderLayer(nn.Module):
         hidden_states: torch.Tensor,
         cu_seqlens: torch.Tensor,
         rope_freqs_cis: Union[torch.Tensor, None] = None,
+        max_seqlen: Optional[int] = None,
     ) -> torch.Tensor:
         """
         Args:
@@ -522,7 +629,10 @@ class MoonVitEncoderLayer(nn.Module):
         residual = hidden_states
         hidden_states = self.norm0(hidden_states)
         attn_out = self.attention_qkvpacked(
-            hidden_states, cu_seqlens, rope_freqs_cis=rope_freqs_cis
+            hidden_states,
+            cu_seqlens,
+            rope_freqs_cis=rope_freqs_cis,
+            max_seqlen=max_seqlen,
         )
         hidden_states = residual + attn_out
 
@@ -539,6 +649,9 @@ class MoonVitEncoder(nn.Module):
         hidden_dim: int,
         num_layers: int,
         block_cfg: dict,
+        prefix: str = "",
+        use_data_parallel: bool = False,
+        use_tensor_parallel: bool = False,
     ) -> None:
         super().__init__()
 
@@ -546,12 +659,23 @@ class MoonVitEncoder(nn.Module):
             block_cfg["hidden_dim"] // block_cfg["num_heads"], 512, 512
         )
         self.blocks = nn.ModuleList(
-            [MoonVitEncoderLayer(**block_cfg) for _ in range(num_layers)]
+            [
+                MoonVitEncoderLayer(
+                    prefix=add_prefix(f"blocks.{layer_idx}", prefix),
+                    use_data_parallel=use_data_parallel,
+                    use_tensor_parallel=use_tensor_parallel,
+                    **block_cfg,
+                )
+                for layer_idx in range(num_layers)
+            ]
         )
         self.final_layernorm = nn.LayerNorm(hidden_dim)
 
     def forward(
-        self, hidden_states: torch.Tensor, grid_hw: torch.Tensor
+        self,
+        hidden_states: torch.Tensor,
+        grid_hw: torch.Tensor,
+        max_seqlen: Optional[int] = None,
     ) -> torch.Tensor:
         rope_freqs_cis = self.rope_2d.get_freqs_cis_by_seqlens(grid_hws=grid_hw)
 
@@ -562,10 +686,15 @@ class MoonVitEncoder(nn.Module):
             )
         )
         cu_seqlens = lengths.cumsum(dim=0, dtype=torch.int32)
+        if max_seqlen is None:
+            max_seqlen = (grid_hw[:, 0] * grid_hw[:, 1]).max().item()
 
         for _, block in enumerate(self.blocks):
             hidden_states = block(
-                hidden_states, cu_seqlens, rope_freqs_cis=rope_freqs_cis
+                hidden_states,
+                cu_seqlens,
+                rope_freqs_cis=rope_freqs_cis,
+                max_seqlen=max_seqlen,
             )
 
         hidden_states = self.final_layernorm(hidden_states)
@@ -635,7 +764,15 @@ class MoonVitPretrainedModel(PreTrainedModel):
     _supports_flash_attn_2 = True
     _supports_sdpa = True
 
-    def __init__(self, config: MoonViTConfig, *inputs, **kwargs):
+    def __init__(
+        self,
+        config: MoonViTConfig,
+        prefix: str = "",
+        use_data_parallel: bool = False,
+        use_tensor_parallel: bool = False,
+        *inputs,
+        **kwargs,
+    ):
         from transformers.activations import GELUTanh
 
         super().__init__(config, *inputs, **kwargs)
@@ -660,10 +797,16 @@ class MoonVitPretrainedModel(PreTrainedModel):
                 "attn_bias": True,
                 "attn_implementation": config._attn_implementation,
             },
+            prefix=add_prefix("encoder", prefix),
+            use_data_parallel=use_data_parallel,
+            use_tensor_parallel=use_tensor_parallel,
         )
 
     def forward(
-        self, pixel_values: torch.Tensor, grid_hw: torch.Tensor
+        self,
+        pixel_values: torch.Tensor,
+        grid_hw: torch.Tensor,
+        max_seqlen: Optional[int] = None,
     ) -> torch.Tensor:
         """
         Args:
@@ -674,7 +817,7 @@ class MoonVitPretrainedModel(PreTrainedModel):
             torch.Tensor: The output tokens.
         """
         hidden_states = self.patch_embed(pixel_values, grid_hw)
-        hidden_states = self.encoder(hidden_states, grid_hw)
+        hidden_states = self.encoder(hidden_states, grid_hw, max_seqlen=max_seqlen)
         hidden_states = patch_merger(
             hidden_states, grid_hw, merge_kernel_size=self.merge_kernel_size
         )

python/sglang/srt/multimodal/mm_utils.py

@@ -503,7 +503,26 @@ def run_dp_sharded_mrope_vision_model(
     """
     tp_size = get_parallel().attn_tp_size
     if tp_size == 1:
-        return vision_model(pixel_values, grid_thw=torch.tensor(grid_thw_list))
+        grid_thw = torch.tensor(
+            grid_thw_list,
+            # MoonViT's 2D RoPE implementation combines the grid metadata
+            # with CUDA activations. Keep the metadata colocated in that
+            # path; other encoders retain their existing CPU contract.
+            device=pixel_values.device if rope_type == "rope_2d" else None,
+        )
+        if rope_type == "rope_2d":
+            image_embeds = vision_model(
+                pixel_values,
+                grid_hw=grid_thw,
+                max_seqlen=max(math.prod(grid) for grid in grid_thw_list),
+            )
+            # MoonViT returns one tensor per image. The multi-GPU path below
+            # already concatenates these tensors before returning, so keep the
+            # TP=1 DP-encoder path on the same projector-facing contract.
+            if isinstance(image_embeds, list):
+                return torch.cat(image_embeds, dim=0)
+            return image_embeds
+        return vision_model(pixel_values, grid_thw=grid_thw)
 
     # GPU_0 tp_rank_local = 0
     # GPU_1 tp_rank_local = 1
@@ -567,8 +586,13 @@ def run_dp_sharded_mrope_vision_model(
     # Run the vision model on the local pixel_values_local
     if rope_type == "rope_2d":
         if pixel_values_local.shape[0] > 0:
+            local_grid_thw = torch.tensor(
+                local_grid_thw_list, device=pixel_values_local.device
+            )
             image_embeds_local = vision_model(
-                pixel_values_local, torch.tensor(local_grid_thw_list)
+                pixel_values_local,
+                grid_hw=local_grid_thw,
+                max_seqlen=max(math.prod(grid) for grid in local_grid_thw_list),
             )
             if isinstance(image_embeds_local, list):
                 image_embeds_local = torch.cat(image_embeds_local, dim=0)

python/sglang/srt/multimodal/vit_cuda_graph_runner.py

@@ -116,6 +116,18 @@ class ViTCudaGraphRunner:
         # x_3d: [S, B, H], B=1, S as graph_key
         return x_3d.shape[0]
 
+    def _capture_context(self):
+        # A DP-sharded encoder intentionally lets each rank capture only the
+        # images it owns (and some ranks can own none). Entering the TP
+        # communication capture in that case requires every TP peer to enter
+        # the same collective capture sequence, which deadlocks on an uneven
+        # image assignment. The encoder's output all-gather is outside this
+        # graph, and all layers are local in DP mode, so capture locally.
+        if getattr(self.vit, "use_data_parallel", False):
+            return nullcontext()
+        ca_comm = get_tp_group().ca_comm
+        return ca_comm.capture() if ca_comm is not None else nullcontext()
+
     def _create_graph(
         self,
         graph_key: int,
@@ -141,11 +153,7 @@ class ViTCudaGraphRunner:
 
         override_backend = get_server_args().mm_attention_backend
 
-        tp_group = get_tp_group()
-        ca_comm = tp_group.ca_comm
-        capture_ctx = ca_comm.capture() if ca_comm is not None else nullcontext()
-
-        with capture_ctx, torch.cuda.graph(graph):
+        with self._capture_context(), torch.cuda.graph(graph):
             y = None
             deepstack_outs: List[torch.Tensor] = []
             deepstack_capture_idx = 0

python/sglang/srt/utils/hf_transformers_patches.py

@@ -248,8 +248,10 @@ def _patch_image_processor_kwargs():
     (e.g. KimiVL) that defines ``preprocess()`` without ``**kwargs`` will
     crash with ``TypeError``.
 
-    Fix: wrap ``__call__`` to catch ``TypeError`` and retry with only the
-    kwargs that ``preprocess()`` actually accepts.
+    Fix: wrap ``__call__`` and filter unsupported kwargs before invoking
+    ``preprocess()``.  The accepted-kwargs set is cached per processor class:
+    apart from avoiding the exception/logging slow path, this matters for VLM
+    requests that preprocess many images on the request critical path.
 
     TODO(upstream): KimiVL image_processing_kimi_vl.py needs ``**kwargs``.
     """
@@ -257,30 +259,40 @@ def _patch_image_processor_kwargs():
         from transformers.image_processing_utils import BaseImageProcessor
 
         original = BaseImageProcessor.__call__
+        accepted_kwargs_cache = {}
+        warned_unsupported_kwargs = set()
 
         def safe_call(self, images, *args, **kwargs):
-            try:
-                return original(self, images, *args, **kwargs)
-            except TypeError as e:
-                if "unexpected keyword argument" not in str(e):
-                    raise
+            processor_type = type(self)
+            accepted_kwargs = accepted_kwargs_cache.get(processor_type)
+            if accepted_kwargs is None and processor_type not in accepted_kwargs_cache:
                 sig = inspect.signature(self.preprocess)
                 params = sig.parameters
                 if any(
                     p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()
                 ):
-                    raise
-                dropped = {k for k in kwargs if k not in params}
-                if dropped:
+                    accepted_kwargs = None
+                else:
+                    accepted_kwargs = frozenset(params)
+                accepted_kwargs_cache[processor_type] = accepted_kwargs
+
+            if accepted_kwargs is None:
+                return original(self, images, *args, **kwargs)
+
+            dropped = frozenset(kwargs) - accepted_kwargs
+            if dropped:
+                warning_key = (processor_type, dropped)
+                if warning_key not in warned_unsupported_kwargs:
                     logger.warning(
                         "Image processor %s.preprocess() does not accept %s; "
-                        "retrying without them. Update the model's image processor "
-                        "to accept **kwargs.",
-                        type(self).__name__,
-                        dropped,
+                        "filtering them before preprocessing. Update the model's image "
+                        "processor to accept **kwargs.",
+                        processor_type.__name__,
+                        sorted(dropped),
                     )
-                valid = {k: v for k, v in kwargs.items() if k in params}
-                return original(self, images, *args, **valid)
+                    warned_unsupported_kwargs.add(warning_key)
+                kwargs = {k: v for k, v in kwargs.items() if k in accepted_kwargs}
+            return original(self, images, *args, **kwargs)
 
         BaseImageProcessor.__call__ = safe_call
     except ImportError:

test/registered/unit/models/test_kimi_vl.py

@@ -0,0 +1,159 @@
+"""CPU-only coverage for Kimi-VL encoder parallelism wiring."""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.nn as nn
+
+from sglang.srt.layers.linear import (
+    ColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
+from sglang.srt.models.kimi_vl import KimiVLForConditionalGeneration
+from sglang.srt.models.kimi_vl_moonvit import MoonVitEncoderLayer, multihead_attention
+from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model
+from sglang.srt.runtime_context import get_parallel
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class _VisionTower:
+    dtype = torch.float32
+    device = torch.device("cpu")
+
+    def __init__(self):
+        self.calls = []
+
+    def __call__(
+        self, pixel_values, image_grid_hws=None, max_seqlen=None, grid_hw=None
+    ):
+        image_grid_hws = grid_hw if image_grid_hws is None else image_grid_hws
+        self.calls.append((pixel_values, image_grid_hws))
+        return [
+            torch.full((1, 4, 2), index + 1.0)
+            for index in range(image_grid_hws.shape[0])
+        ]
+
+
+class _Projector:
+    def __init__(self):
+        self.input = None
+
+    def __call__(self, image_features):
+        self.input = image_features
+        return image_features
+
+
+class _GridRecordingVisionTower:
+    def __call__(self, pixel_values, grid_hw, max_seqlen=None):
+        self.grid_thw = grid_hw
+        self.max_seqlen = max_seqlen
+        return pixel_values
+
+
+def _bare_model(*, use_data_parallel: bool):
+    model = KimiVLForConditionalGeneration.__new__(KimiVLForConditionalGeneration)
+    nn.Module.__init__(model)
+    model.config = SimpleNamespace(text_config=SimpleNamespace(hidden_size=16))
+    model.use_data_parallel = use_data_parallel
+    model.vision_tower = _VisionTower()
+    model.multi_modal_projector = _Projector()
+    return model
+
+
+def _image_item(feature, grid_hws):
+    return MultimodalDataItem(
+        modality=Modality.IMAGE,
+        offsets=[(0, 1)],
+        feature=feature,
+        model_specific_data={"image_grid_hws": torch.tensor(grid_hws)},
+    )
+
+
+class TestKimiVLEncoderParallelism(CustomTestCase):
+    def test_moonvit_uses_tensor_parallel_layers(self):
+        with get_parallel().override(
+            tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0
+        ):
+            layer = MoonVitEncoderLayer(
+                num_heads=2,
+                hidden_dim=8,
+                mlp_dim=16,
+                prefix="vision_tower.encoder.blocks.0",
+                use_tensor_parallel=True,
+            )
+
+        self.assertIsInstance(layer.wqkv, QKVParallelLinear)
+        self.assertIsInstance(layer.wo, RowParallelLinear)
+        self.assertIsInstance(layer.mlp.fc0, ColumnParallelLinear)
+        self.assertIsInstance(layer.mlp.fc1, RowParallelLinear)
+
+    def test_encoder_dp_uses_existing_mrope_sharding_helper(self):
+        model = _bare_model(use_data_parallel=True)
+        items = [
+            _image_item(torch.randn(4, 2), [[2, 2]]),
+            _image_item(torch.randn(8, 2), [[4, 2]]),
+        ]
+        sharded_features = torch.randn(3, 4, 2)
+
+        with patch(
+            "sglang.srt.models.kimi_vl.run_dp_sharded_mrope_vision_model",
+            return_value=sharded_features,
+        ) as run_dp:
+            output = model.get_image_feature(items)
+
+        run_dp.assert_called_once()
+        _, pixel_values, grid_hws = run_dp.call_args.args
+        self.assertEqual(pixel_values.shape, (12, 2))
+        self.assertEqual(grid_hws, [[2, 2], [4, 2]])
+        self.assertEqual(run_dp.call_args.kwargs, {"rope_type": "rope_2d"})
+        self.assertIs(output, sharded_features)
+
+    def test_encoder_dp_keeps_moonvit_grid_metadata_on_vision_device(self):
+        vision_tower = _GridRecordingVisionTower()
+        pixel_values = torch.randn(4, 2)
+
+        with get_parallel().override(
+            tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0
+        ):
+            output = run_dp_sharded_mrope_vision_model(
+                vision_tower, pixel_values, [[2, 2]], rope_type="rope_2d"
+            )
+
+        self.assertIs(output, pixel_values)
+        self.assertEqual(vision_tower.grid_thw.device, pixel_values.device)
+        self.assertEqual(vision_tower.max_seqlen, 4)
+
+    def test_encoder_dp_tp1_concatenates_moonvit_image_outputs(self):
+        vision_tower = _VisionTower()
+        pixel_values = torch.randn(4, 2)
+        with get_parallel().override(
+            tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0
+        ):
+            output = run_dp_sharded_mrope_vision_model(
+                vision_tower, pixel_values, [[2, 2]], rope_type="rope_2d"
+            )
+        self.assertIsInstance(output, torch.Tensor)
+        self.assertEqual(output.shape, (1, 4, 2))
+
+    def test_moonvit_attention_accepts_precomputed_max_seqlen(self):
+        q = torch.randn(4, 2, 4, dtype=torch.bfloat16)
+        cu_seqlens = torch.tensor([0, 4], dtype=torch.int32)
+        fake_output = torch.randn_like(q)
+        with patch(
+            "sglang.srt.models.kimi_vl_moonvit.flash_attn_varlen_func",
+            return_value=fake_output,
+        ) as flash_attn:
+            output = multihead_attention(q, q, q, cu_seqlens, cu_seqlens, max_seqlen=4)
+        self.assertTrue(torch.equal(output, fake_output.flatten(start_dim=-2)))
+        self.assertEqual(flash_attn.call_args.args[5:7], (4, 4))
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-v"]))

test/registered/unit/models/test_kimi_vl_moonvit.py

@@ -0,0 +1,64 @@
+import pytest
+import torch
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=3, suite="base-a-test-cpu")
+
+from sglang.srt.models import kimi_vl_moonvit
+from sglang.srt.models.kimi_vl_moonvit import Learnable2DInterpPosEmb
+
+
+def test_learnable_2d_pos_emb_caches_inference_interpolation(monkeypatch):
+    module = Learnable2DInterpPosEmb(height=2, width=2, dim=4).eval()
+    inputs = torch.zeros(6, 4)
+    grid_hw = torch.tensor([[2, 3]])
+    calls = 0
+    original_interpolate = torch.nn.functional.interpolate
+
+    def counting_interpolate(*args, **kwargs):
+        nonlocal calls
+        calls += 1
+        return original_interpolate(*args, **kwargs)
+
+    monkeypatch.setattr(torch.nn.functional, "interpolate", counting_interpolate)
+    first = module(inputs, grid_hw)
+    second = module(inputs, grid_hw)
+    torch.testing.assert_close(first, second)
+    assert calls == 1
+
+
+def test_learnable_2d_pos_emb_does_not_cache_training_interpolation(monkeypatch):
+    module = Learnable2DInterpPosEmb(height=2, width=2, dim=4).train()
+    inputs = torch.zeros(6, 4)
+    grid_hw = torch.tensor([[2, 3]])
+    calls = 0
+    original_interpolate = torch.nn.functional.interpolate
+
+    def counting_interpolate(*args, **kwargs):
+        nonlocal calls
+        calls += 1
+        return original_interpolate(*args, **kwargs)
+
+    monkeypatch.setattr(torch.nn.functional, "interpolate", counting_interpolate)
+    module(inputs, grid_hw)
+    module(inputs, grid_hw)
+    assert calls == 2
+
+
+def test_learnable_2d_pos_emb_evicts_oldest_inference_cache_entry(monkeypatch):
+    monkeypatch.setattr(kimi_vl_moonvit, "_MAX_INFERENCE_POS_EMB_CACHE_ENTRIES", 1)
+    module = Learnable2DInterpPosEmb(height=2, width=2, dim=4).eval()
+    inputs = torch.zeros(6, 4)
+
+    module(inputs, torch.tensor([[2, 3]]))
+    module(inputs, torch.tensor([[3, 2]]))
+
+    assert len(module._interpolated_pos_emb_cache) == 1
+    assert ((3, 2), module.weight.dtype, module.weight.device) in (
+        module._interpolated_pos_emb_cache
+    )
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-v"]))

test/registered/unit/multimodal/test_vit_cuda_graph_runner.py

@@ -0,0 +1,59 @@
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=3, suite="base-a-test-cpu")
+
+from sglang.srt.multimodal.vit_cuda_graph_runner import ViTCudaGraphRunner
+
+
+class _Block:
+    def forward(self, x):
+        return x
+
+
+def _runner(*, use_data_parallel: bool) -> ViTCudaGraphRunner:
+    vit = SimpleNamespace(
+        blocks=[_Block()],
+        deepstack_visual_indexes=[],
+        deepstack_merger_list=None,
+        use_data_parallel=use_data_parallel,
+    )
+    return ViTCudaGraphRunner(vit)
+
+
+def test_dp_vit_graph_capture_does_not_enter_tp_communication_capture():
+    runner = _runner(use_data_parallel=True)
+    with patch(
+        "sglang.srt.multimodal.vit_cuda_graph_runner.get_tp_group",
+        side_effect=AssertionError("DP capture must be rank-local"),
+    ):
+        with runner._capture_context():
+            pass
+
+
+def test_non_dp_vit_graph_capture_uses_tp_communication_capture():
+    entered = []
+
+    class Capture:
+        def __enter__(self):
+            entered.append(True)
+
+        def __exit__(self, *args):
+            return False
+
+    group = SimpleNamespace(ca_comm=SimpleNamespace(capture=lambda: Capture()))
+    runner = _runner(use_data_parallel=False)
+    with patch(
+        "sglang.srt.multimodal.vit_cuda_graph_runner.get_tp_group", return_value=group
+    ):
+        with runner._capture_context():
+            pass
+    assert entered == [True]
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-v"]))

test/registered/unit/utils/test_hf_transformers.py

@@ -4,12 +4,16 @@ Tests cover the pure utility functions (compat patches, config helpers,
 context length, GGUF detection, etc.) that don't require actual model files.
 """
 
+import inspect
 import tempfile
 import unittest
 from types import SimpleNamespace
+from unittest.mock import patch
 
 from transformers import PretrainedConfig
+from transformers.image_processing_utils import BaseImageProcessor
 
+from sglang.srt.utils import hf_transformers_patches
 from sglang.srt.utils.hf_transformers.common import (
     _is_deepseek_ocr2_model,
     _is_deepseek_ocr_model,
@@ -27,6 +31,33 @@ from sglang.test.ci.ci_register import register_cpu_ci
 register_cpu_ci(est_time=6, suite="base-a-test-cpu")
 
 
+# ---------------------------------------------------------------------------
+# _patch_image_processor_kwargs
+# ---------------------------------------------------------------------------
+
+
+class TestImageProcessorKwargsPatch(unittest.TestCase):
+    def test_filters_unsupported_kwargs_and_caches_signature(self):
+        class StrictImageProcessor(BaseImageProcessor):
+            model_input_names = ["pixel_values"]
+
+            def preprocess(self, images, accepted=None):
+                return {"images": images, "accepted": accepted}
+
+        processor = StrictImageProcessor()
+        with patch.object(
+            hf_transformers_patches.inspect,
+            "signature",
+            wraps=inspect.signature,
+        ) as signature:
+            first = processor("first", accepted=True, device="cuda")
+            second = processor("second", accepted=False, device="cuda")
+
+        self.assertEqual(first, {"images": "first", "accepted": True})
+        self.assertEqual(second, {"images": "second", "accepted": False})
+        self.assertEqual(signature.call_count, 1)
+
+
 # ---------------------------------------------------------------------------
 # normalize_rope_scaling_compat
 # ---------------------------------------------------------------------------