diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index ce769549f..5b0953887 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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 diff --git a/python/sglang/srt/models/kimi_vl.py b/python/sglang/srt/models/kimi_vl.py index d9a929e56..27d1e1845 100644 --- a/python/sglang/srt/models/kimi_vl.py +++ b/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: diff --git a/python/sglang/srt/models/kimi_vl_moonvit.py b/python/sglang/srt/models/kimi_vl_moonvit.py index 7b431987d..48cb383ce 100644 --- a/python/sglang/srt/models/kimi_vl_moonvit.py +++ b/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 ) diff --git a/python/sglang/srt/multimodal/mm_utils.py b/python/sglang/srt/multimodal/mm_utils.py index 7334a772e..1a4b9db09 100644 --- a/python/sglang/srt/multimodal/mm_utils.py +++ b/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) diff --git a/python/sglang/srt/multimodal/vit_cuda_graph_runner.py b/python/sglang/srt/multimodal/vit_cuda_graph_runner.py index 28d50c7df..2401084a1 100644 --- a/python/sglang/srt/multimodal/vit_cuda_graph_runner.py +++ b/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 diff --git a/python/sglang/srt/utils/hf_transformers_patches.py b/python/sglang/srt/utils/hf_transformers_patches.py index 17a394190..614fc95ba 100644 --- a/python/sglang/srt/utils/hf_transformers_patches.py +++ b/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: diff --git a/test/registered/unit/models/test_kimi_vl.py b/test/registered/unit/models/test_kimi_vl.py new file mode 100644 index 000000000..828192f23 --- /dev/null +++ b/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"])) diff --git a/test/registered/unit/models/test_kimi_vl_moonvit.py b/test/registered/unit/models/test_kimi_vl_moonvit.py new file mode 100644 index 000000000..af2f34c17 --- /dev/null +++ b/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"])) diff --git a/test/registered/unit/multimodal/test_vit_cuda_graph_runner.py b/test/registered/unit/multimodal/test_vit_cuda_graph_runner.py new file mode 100644 index 000000000..7a7df286c --- /dev/null +++ b/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"])) diff --git a/test/registered/unit/utils/test_hf_transformers.py b/test/registered/unit/utils/test_hf_transformers.py index d15d8630d..61cea60b2 100644 --- a/test/registered/unit/utils/test_hf_transformers.py +++ b/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 # ---------------------------------------------------------------------------