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: