optimization: shard kimi dp image feature transport and misc optimizations (#31227)

python/sglang/srt/layers/attention/vision.py

@@ -174,6 +174,23 @@ def resolve_max_seqlen(source, cu_seqlens: torch.Tensor) -> int:
     return int(seq_lens.max().item())
 
 
+def resolve_precomputed_max_seqlen(
+    cu_seqlens: torch.Tensor, max_seqlen: int | torch.Tensor | None
+) -> int:
+    """Use an encoder-provided max sequence length when one is available.
+
+    Packed vision encoders execute many attention blocks for one image batch.
+    Deriving the max from GPU ``cu_seqlens`` in every block synchronizes the
+    launch stream, whereas the encoder can materialize this host scalar once.
+    """
+    if max_seqlen is None:
+        seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+        return int(seq_lens.max().item())
+    if isinstance(max_seqlen, torch.Tensor):
+        return int(max_seqlen.item())
+    return int(max_seqlen)
+
+
 class VisionSdpaAttention(nn.Module):
     r"""
     Scaled Dot Product Attention inner product
@@ -392,15 +409,21 @@ class VisionTritonAttention(nn.Module):
             # [b * s, head, head_size]
             output = torch.empty_like(q)
 
-            seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
-            max_seqlen = seq_lens.max().item()
+            seq_lens = kwargs.get("sequence_lengths")
+            if seq_lens is None:
+                seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
+            else:
+                seq_lens = seq_lens.to(device=q.device, dtype=torch.int32)
+            max_seqlen = resolve_precomputed_max_seqlen(
+                cu_seqlens, kwargs.get("max_seqlen")
+            )
             context_attention_fwd(
                 q,
                 k,
                 v,
                 output,
                 cu_seqlens.to(q.device),
-                seq_lens.to(q.device),
+                seq_lens,
                 max_seqlen,
                 is_causal=False,
                 sm_scale=softmax_scale,
@@ -458,19 +481,9 @@ class VisionFlash3Attention(nn.Module):
         else:
             cu_seqlens = resolve_seqlens(cu_seqlens, bsz, seq_len, device=q.device)
             cu_seqlens = cu_seqlens.to(dtype=torch.int32).to(q.device)
-            # Some vision encoders precompute this scalar once per encoder
-            # forward and share it across all of their attention blocks.  Use
-            # that value when available: deriving it here requires a
-            # GPU-to-host sync, so repeating it per block serializes the ViT
-            # launch stream for variable-size images.
-            max_seqlen = kwargs.get("max_seqlen")
-            if max_seqlen is None:
-                seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
-                max_seqlen = int(seq_lens.max().item())
-            elif isinstance(max_seqlen, torch.Tensor):
-                max_seqlen = int(max_seqlen.item())
-            else:
-                max_seqlen = int(max_seqlen)
+            max_seqlen = resolve_precomputed_max_seqlen(
+                cu_seqlens, kwargs.get("max_seqlen")
+            )
 
             fa_kwargs = dict(
                 cu_seqlens_q=cu_seqlens,
@@ -523,8 +536,9 @@ class VisionFlash4Attention(nn.Module):
             cu_seqlens = cu_seqlens.get_data()
 
         cu_seqlens = cu_seqlens.to(dtype=torch.int32).to(q.device)
-        seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
-        max_seqlen = seq_lens.max().item()
+        max_seqlen = resolve_precomputed_max_seqlen(
+            cu_seqlens, kwargs.get("max_seqlen")
+        )
 
         output = flash_attn_varlen_func(
             q,

python/sglang/srt/managers/mm_utils.py

@@ -32,7 +32,7 @@ from sglang.srt.managers.schedule_batch import (
 from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.multimodal.evs import EVSEmbeddingResult
-from sglang.srt.runtime_context import get_server_args
+from sglang.srt.runtime_context import get_parallel, get_server_args
 from sglang.srt.utils import flatten_nested_list, is_npu, print_warning_once
 from sglang.srt.utils.stale_shm_cleanup import make_shm_name
 from sglang.utils import logger
@@ -478,7 +478,7 @@ DataEmbeddingFunc = Callable[
 
 
 def _can_skip_pre_embed_feature_move(data_embedding_func: DataEmbeddingFunc) -> bool:
-    """qwen-vl visual forward already moves batched features to the target device.
+    """Models that materialize and batch visual features inside their encoder.
 
     instead of performing multiple H2D for each mm feature from all mm_items (followed by concatenation on device),
     for some models which internally performs H2D on concated mm feature, these small H2D calls could be replaced with a single big H2D
@@ -496,6 +496,7 @@ def _can_skip_pre_embed_feature_move(data_embedding_func: DataEmbeddingFunc) ->
         "Qwen3VLMoeForConditionalGeneration",
         "Qwen3_5ForConditionalGeneration",
         "Qwen3_5MoeForConditionalGeneration",
+        "KimiK25ForConditionalGeneration",
     }
 
 
@@ -508,6 +509,27 @@ def _move_items_to_device(
             item.feature = item.feature.to(device, non_blocking=True)
 
 
+def _acknowledge_deferred_cuda_ipc_cache_hits(
+    items: List[MultimodalDataItem],
+) -> None:
+    """Release lazy Kimi IPC slices when a cached embedding skips ViT.
+
+    On an encoder-DP miss, exactly one rank copies an image and acknowledges
+    the full TP group.  On a cache hit no rank copies it, so rank zero performs
+    the equivalent single acknowledgement.  This preserves the fixed-pool
+    lifecycle without reintroducing an unnecessary GPU-to-GPU copy.
+    """
+    parallel = get_parallel()
+    if parallel.attn_tp_rank != 0:
+        return
+    server_args = get_server_args()
+    # The pool's recycler uses ServerArgs.tp_size, so its acknowledgement must
+    # match that count even when an attention subgroup is smaller.
+    consumer_count = max(getattr(server_args, "tp_size", parallel.attn_tp_size), 1)
+    for item in items:
+        item.acknowledge_deferred_cuda_ipc_feature(consumer_count)
+
+
 def _get_chunked_embedding_full(
     data_embedding_func: DataEmbeddingFunc,
     embedding_items_per_req: List[MultimodalDataItem],
@@ -535,6 +557,8 @@ def _get_chunked_embedding_full(
             else embedding
         )
         embedding_cache.set(embedding_items_hash, embedding_per_req)
+    else:
+        _acknowledge_deferred_cuda_ipc_cache_hits(embedding_items_per_req)
 
     if isinstance(embedding_per_req, EVSEmbeddingResult):
         item = embedding_items_per_req[0]
@@ -594,13 +618,15 @@ def _get_chunked_embedding_by_item(
         cached = embedding_cache.get_single(item.hash)
         if cached is not None:
             cached_embeddings[idx] = cached.embedding
+            _acknowledge_deferred_cuda_ipc_cache_hits([item])
         else:
             miss_items.append((idx, item, start, end))
 
     # 3. Batch encode all cache-miss items in one ViT call
     if miss_items:
         miss_item_list = [item for _, item, _, _ in miss_items]
-        _move_items_to_device(miss_item_list, device)
+        if not _can_skip_pre_embed_feature_move(data_embedding_func):
+            _move_items_to_device(miss_item_list, device)
         all_miss_embedding = data_embedding_func(miss_item_list)
         all_miss_embedding = all_miss_embedding.reshape(
             -1, all_miss_embedding.shape[-1]

python/sglang/srt/managers/schedule_batch.py

@@ -111,7 +111,10 @@ from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
 from sglang.srt.sampling.sampling_params import SamplingParams
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import flatten_nested_list
-from sglang.srt.utils.cuda_ipc_transport_utils import CudaIpcTensorTransportProxy
+from sglang.srt.utils.cuda_ipc_transport_utils import (
+    DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY,
+    CudaIpcTensorTransportProxy,
+)
 
 if TYPE_CHECKING:
     from typing import Any, Dict
@@ -359,10 +362,15 @@ class MultimodalDataItem:
             )
         )
 
-    def reconstruct(self, target_device: int):
+    def reconstruct(self, target_device: int, ipc_consumer_count: int = 1):
         """materialize cuda ipc proxy tensors in-place on target_device"""
         if isinstance(self.feature, CudaIpcTensorTransportProxy):
-            self.feature = self.feature.reconstruct_on_target_device(target_device)
+            if ipc_consumer_count == 1:
+                self.feature = self.feature.reconstruct_on_target_device(target_device)
+            else:
+                self.feature = self.feature.reconstruct_on_target_device(
+                    target_device, consumer_count=ipc_consumer_count
+                )
         if isinstance(self.precomputed_embeddings, CudaIpcTensorTransportProxy):
             self.precomputed_embeddings = (
                 self.precomputed_embeddings.reconstruct_on_target_device(target_device)
@@ -376,6 +384,32 @@ class MultimodalDataItem:
                 ].reconstruct_on_target_device(target_device)
                 self.model_specific_data[extra_key] = extra_data
 
+    def can_defer_cuda_ipc_feature_reconstruction(self) -> bool:
+        """Whether a DP-aware model will materialize this feature lazily.
+
+        Hashing and pad-value generation must already have completed on the
+        tokenizer worker.  Any additional IPC proxy would still need eager
+        reconstruction, so keep the narrow fast path feature-only.
+        """
+        return (
+            self.model_specific_data.get(
+                DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY, False
+            )
+            and self.hash is not None
+            and self.pad_value is not None
+            and isinstance(self.feature, CudaIpcTensorTransportProxy)
+            and not isinstance(self.precomputed_embeddings, CudaIpcTensorTransportProxy)
+            and not any(
+                isinstance(value, CudaIpcTensorTransportProxy)
+                for value in self.model_specific_data.values()
+            )
+        )
+
+    def acknowledge_deferred_cuda_ipc_feature(self, consumer_count: int = 1):
+        """Release a lazy IPC feature when an embedding-cache hit skips ViT."""
+        if isinstance(self.feature, CudaIpcTensorTransportProxy):
+            self.feature.acknowledge_consumption(consumer_count)
+
 
 @dataclasses.dataclass
 class MultimodalProcessorOutput:
@@ -510,7 +544,10 @@ class MultimodalInputs:
         # try reconstructing from cuda-ipc
         reconstruct_device = None
         for mm_item in mm_items:
-            if mm_item.has_cuda_ipc_proxy():
+            if (
+                mm_item.has_cuda_ipc_proxy()
+                and not mm_item.can_defer_cuda_ipc_feature_reconstruction()
+            ):
                 if reconstruct_device is None:
                     reconstruct_device = torch.cuda.current_device()
                 mm_item.reconstruct(reconstruct_device)

python/sglang/srt/models/kimi_k25.py

@@ -30,8 +30,11 @@ from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.deepseek_v2 import DeepseekV3ForCausalLM
 from sglang.srt.models.kimi_vl_moonvit import MLP2
 from sglang.srt.models.utils import WeightsMapper
-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.multimodal.mm_utils import (
+    materialize_multimodal_features,
+    run_dp_sharded_mrope_vision_model,
+)
+from sglang.srt.runtime_context import get_parallel, get_server_args
 from sglang.srt.utils import add_prefix, is_npu
 
 logger = logging.getLogger(__name__)
@@ -142,6 +145,7 @@ class MoonViTEncoderLayer(nn.Module):
         cu_seqlens: torch.Tensor,
         max_seqlen: int,
         rope_freqs_cis: torch.Tensor | None = None,
+        sequence_lengths: torch.Tensor | None = None,
     ):
         residual = hidden_states
         hidden_states = self.norm0(hidden_states)
@@ -151,6 +155,7 @@ class MoonViTEncoderLayer(nn.Module):
             cu_seqlens=cu_seqlens,
             position_embeddings=rope_freqs_cis,
             max_seqlen=max_seqlen,
+            sequence_lengths=sequence_lengths,
         )
 
         hidden_states = residual + hidden_states
@@ -463,10 +468,13 @@ class MoonViT3dEncoder(nn.Module):
             grid_thws=grid_thws, device=hidden_states.device
         )
 
+        sequence_lengths = (grid_thws[:, 0] * grid_thws[:, 1] * grid_thws[:, 2]).to(
+            device=hidden_states.device, dtype=torch.int32
+        )
         lengths = torch.cat(
             (
-                torch.zeros(1, dtype=grid_thws.dtype, device=grid_thws.device),
-                grid_thws[:, 0] * grid_thws[:, 1] * grid_thws[:, 2],
+                torch.zeros(1, dtype=torch.int32, device=hidden_states.device),
+                sequence_lengths,
             )
         )
 
@@ -478,7 +486,11 @@ class MoonViT3dEncoder(nn.Module):
 
         for block in self.blocks:
             hidden_states = block(
-                hidden_states, cu_seqlens, max_seqlen, rope_freqs_cis=rope_freqs_cis
+                hidden_states,
+                cu_seqlens,
+                max_seqlen,
+                rope_freqs_cis=rope_freqs_cis,
+                sequence_lengths=sequence_lengths,
             )
 
         hidden_states = self.final_layernorm(hidden_states)
@@ -685,28 +697,69 @@ class KimiK25ForConditionalGeneration(nn.Module):
     def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
         device = self.vision_tower.device
         target_dtype = self.vision_tower.patch_embed.proj.weight.dtype
-        pixel_values = torch.cat([item.feature for item in items], dim=0).to(
-            device=device, dtype=target_dtype
-        )
         image_grid_thws = []
         for item in items:
             grid_thw = item.model_specific_data.get("image_grid_thw")
             if grid_thw is None:
                 grid_thw = item.model_specific_data["grid_thws"]
             image_grid_thws.append(grid_thw)
-        grid_thws = torch.concat(image_grid_thws, dim=0).to(device)
+        grid_thws = torch.concat(image_grid_thws, dim=0)
+
+        def materialize_item_features(image_indices: List[int]) -> torch.Tensor:
+            """Move only this encoder-DP rank's images to its local GPU.
+
+            CUDA IPC features are intentionally reconstructed after the image
+            assignment.  Each image therefore crosses the tokenizer/scheduler
+            boundary once instead of once per TP rank.  The selected consumer
+            acknowledges the entire TP group so the bounded IPC pool remains
+            recyclable.
+            """
+            parallel = get_parallel()
+            server_args = get_server_args()
+            # Match MmItemMemoryPool.try_to_recycle(), which waits for the
+            # server TP size rather than the attention subgroup size.
+            ipc_consumer_count = max(
+                getattr(server_args, "tp_size", parallel.attn_tp_size), 1
+            )
+            device_index = device.index
+            if device.type == "cuda" and device_index is None:
+                device_index = torch.cuda.current_device()
+
+            features = []
+            for image_index in image_indices:
+                item = items[image_index]
+                if device.type == "cuda":
+                    item.reconstruct(
+                        device_index, ipc_consumer_count=ipc_consumer_count
+                    )
+                feature = item.feature
+                if not isinstance(feature, torch.Tensor):
+                    raise TypeError(
+                        "Kimi-K2.5/K2.7 image feature must be a torch.Tensor, "
+                        f"got {type(feature)}"
+                    )
+                features.append(feature)
+            return materialize_multimodal_features(
+                features, device=device, dtype=target_dtype
+            )
 
         if self.use_data_parallel:
             image_embeds = run_dp_sharded_mrope_vision_model(
                 self.vision_tower,
-                pixel_values,
+                None,
                 grid_thws.tolist(),
-                rope_type="rope_2d",
+                # MoonViT3d uses 2D RoPE and returns packed patch embeddings.
+                # Its grid metadata is a positional argument, unlike Kimi-VL.
+                rope_type="rope_2d_packed",
+                load_local_pixel_values=materialize_item_features,
+                pixel_values_device=device,
+                pixel_values_dtype=target_dtype,
             )
             image_features = self.mm_projector(image_embeds)
             return image_features
 
-        image_embeds = self.vision_tower(pixel_values, grid_thws)
+        pixel_values = materialize_item_features(list(range(len(items))))
+        image_embeds = self.vision_tower(pixel_values, grid_thws.to(device))
         proj_out = mm_projection_auto(self.mm_projector, image_embeds)
         return torch.cat(proj_out, dim=0)
 

python/sglang/srt/multimodal/mm_utils.py

@@ -33,7 +33,7 @@ import itertools
 import math
 import re
 from io import BytesIO
-from typing import Literal
+from typing import Callable, Literal, Optional, Sequence
 
 import numpy as np
 import pybase64
@@ -58,6 +58,59 @@ def has_valid_data(data) -> bool:
     return True
 
 
+def materialize_multimodal_features(
+    features: Sequence[torch.Tensor],
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Concatenate variable-length feature tensors into one destination buffer.
+
+    A multimodal item can arrive as a CPU tensor, a CUDA-IPC reconstruction,
+    or an already resident tensor with a different dtype.  Calling ``to`` on
+    every item and then ``torch.cat`` creates one temporary tensor per item
+    before allocating the final packed input.  Allocate the final buffer once
+    and copy each item directly into its slice instead; ``copy_`` performs the
+    required device and dtype conversion in the destination copy.
+
+    All tensors must agree on dimensions after the leading token dimension.
+    The leading dimension may differ because images commonly have different
+    numbers of vision patches.
+    """
+
+    if not features:
+        raise ValueError("features must contain at least one tensor")
+
+    first = features[0]
+    if not isinstance(first, torch.Tensor):
+        raise TypeError(f"expected torch.Tensor, got {type(first)}")
+    if first.ndim == 0:
+        raise ValueError("multimodal feature tensors must have a leading dimension")
+    trailing_shape = first.shape[1:]
+    total_tokens = 0
+    for feature in features:
+        if not isinstance(feature, torch.Tensor):
+            raise TypeError(f"expected torch.Tensor, got {type(feature)}")
+        if feature.ndim == 0 or feature.shape[1:] != trailing_shape:
+            raise ValueError(
+                "multimodal feature tensors must have matching trailing shapes: "
+                f"expected {trailing_shape}, got {feature.shape}"
+            )
+        total_tokens += feature.shape[0]
+
+    output = torch.empty(
+        (total_tokens, *trailing_shape),
+        device=device,
+        dtype=dtype,
+    )
+    offset = 0
+    for feature in features:
+        length = feature.shape[0]
+        output[offset : offset + length].copy_(feature, non_blocking=True)
+        offset += length
+    return output
+
+
 def select_best_resolution(original_size, possible_resolutions):
     """
     Selects the best resolution from a list of possible resolutions based on the original size.
@@ -470,10 +523,13 @@ def run_dp_sharded_vision_model(
 # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/vision.py
 def run_dp_sharded_mrope_vision_model(
     vision_model: torch.nn.Module,
-    pixel_values: torch.Tensor,
+    pixel_values: Optional[torch.Tensor],
     grid_thw_list: list,
     *,
-    rope_type: Literal["rope_3d", "rope_2d"],
+    rope_type: Literal["rope_3d", "rope_2d", "rope_2d_packed"],
+    load_local_pixel_values: Optional[Callable[[list[int]], torch.Tensor]] = None,
+    pixel_values_device: Optional[torch.device] = None,
+    pixel_values_dtype: Optional[torch.dtype] = None,
 ):
     """Run a vision model with data parallelism (DP) sharding.
     The function will shard the input image tensor on the
@@ -487,7 +543,9 @@ def run_dp_sharded_mrope_vision_model(
         rope_type: Type of rope used in the vision model.
                    Different rope types have different dimension to do ViT.
                    "rope_3d" for 3D rope (e.g., Qwen2.5-VL)
-                   "rope_2d" for 2D rope (e.g., Kimi-VL)
+                   "rope_2d" for packed 2D rope outputs (e.g., Kimi-VL)
+                   "rope_2d_packed" for packed 2D rope outputs that accept
+                   ``grid_thws`` positionally (e.g., Kimi-K2.5/K2.7)
     Returns:
         torch.Tensor: Output image embeddings
 
@@ -495,14 +553,29 @@ def run_dp_sharded_mrope_vision_model(
         ```
         vision_model.out_hidden_size = 64
         vision_model.spatial_merge_size = 2
-        pixel_values.shape = (1350, channel)
+        pixel_values.shape = (1350, channel), or a local loader supplies
+        per-image features after the data-parallel assignment is known.
         grid_thw_list = [[1, 10, 100], [1, 10, 10], [1, 10, 20], [1, 50]]
         tp_size = 2
         ```
 
     """
+    if pixel_values is None and load_local_pixel_values is None:
+        raise ValueError("pixel_values or load_local_pixel_values must be provided")
+
+    input_device = (
+        pixel_values.device if pixel_values is not None else pixel_values_device
+    )
+    input_dtype = pixel_values.dtype if pixel_values is not None else pixel_values_dtype
+    if input_device is None or input_dtype is None:
+        raise ValueError(
+            "pixel_values_device and pixel_values_dtype are required with a local loader"
+        )
+
     tp_size = get_parallel().attn_tp_size
     if tp_size == 1:
+        if pixel_values is None:
+            pixel_values = load_local_pixel_values(list(range(len(grid_thw_list))))
         grid_thw = torch.tensor(
             grid_thw_list,
             # MoonViT's 2D RoPE implementation combines the grid metadata
@@ -522,6 +595,11 @@ def run_dp_sharded_mrope_vision_model(
             if isinstance(image_embeds, list):
                 return torch.cat(image_embeds, dim=0)
             return image_embeds
+        if rope_type == "rope_2d_packed":
+            image_embeds = vision_model(pixel_values, grid_thw)
+            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
@@ -553,21 +631,23 @@ def run_dp_sharded_mrope_vision_model(
 
     # Get the pixel values for the local images based on the image_idxs_local
     if len(image_idxs_local) > 0:
-        pixel_values_local = torch.cat(
-            [
-                pixel_values[cum_patches_per_image[i] : cum_patches_per_image[i + 1]]
-                for i in image_idxs_local
-            ]
-        )
+        if load_local_pixel_values is not None:
+            pixel_values_local = load_local_pixel_values(image_idxs_local)
+        else:
+            assert pixel_values is not None
+            pixel_values_local = torch.cat(
+                [
+                    pixel_values[
+                        cum_patches_per_image[i] : cum_patches_per_image[i + 1]
+                    ]
+                    for i in image_idxs_local
+                ]
+            )
     else:
-        # Handle case where this rank has no images
-        pixel_values_local = torch.empty(
-            (0, pixel_values.shape[1]),
-            device=pixel_values.device,
-            dtype=pixel_values.dtype,
-        )
+        pixel_values_local = None
     # embed_dim_reduction_factor = 2 * 2
-    if rope_type == "rope_2d":
+    packed_2d_rope = rope_type in ("rope_2d", "rope_2d_packed")
+    if packed_2d_rope:
         embed_dim_reduction_factor = (
             vision_model.merge_kernel_size[0] * vision_model.merge_kernel_size[1]
         )
@@ -584,37 +664,45 @@ def run_dp_sharded_mrope_vision_model(
     local_grid_thw_list = [grid_thw_list[i] for i in image_idxs_local]
 
     # Run the vision model on the local pixel_values_local
-    if rope_type == "rope_2d":
-        if pixel_values_local.shape[0] > 0:
+    if packed_2d_rope:
+        if pixel_values_local is not None and 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,
-                grid_hw=local_grid_thw,
-                max_seqlen=max(math.prod(grid) for grid in local_grid_thw_list),
-            )
+            if rope_type == "rope_2d":
+                image_embeds_local = vision_model(
+                    pixel_values_local,
+                    grid_hw=local_grid_thw,
+                    max_seqlen=max(math.prod(grid) for grid in local_grid_thw_list),
+                )
+            else:
+                image_embeds_local = vision_model(pixel_values_local, local_grid_thw)
             if isinstance(image_embeds_local, list):
                 image_embeds_local = torch.cat(image_embeds_local, dim=0)
         else:
             out_dim = getattr(vision_model.config, "hidden_size", None)
             image_embeds_local = torch.empty(
                 (0, embed_dim_reduction_factor, out_dim),
-                device=pixel_values.device,
-                dtype=pixel_values.dtype,
+                device=input_device,
+                dtype=input_dtype,
             )
     else:
-        if pixel_values_local.shape[0] > 0:
+        if pixel_values_local is not None and pixel_values_local.shape[0] > 0:
             # print(f"{local_grid_thw_list = }", flush=True)
             image_embeds_local = vision_model(
                 pixel_values_local, torch.tensor(local_grid_thw_list)
             )
+            if isinstance(image_embeds_local, list):
+                image_embeds_local = torch.cat(image_embeds_local, dim=0)
         else:
             # Handle empty case
+            out_dim = getattr(vision_model, "out_hidden_size", None)
+            if out_dim is None:
+                out_dim = vision_model.config.hidden_size
             image_embeds_local = torch.empty(
-                (0, vision_model.out_hidden_size),
-                device=pixel_values.device,
-                dtype=pixel_values.dtype,
+                (0, out_dim),
+                device=input_device,
+                dtype=input_dtype,
             )
 
     # Pad the output based on max_len_per_rank
@@ -622,7 +710,7 @@ def run_dp_sharded_mrope_vision_model(
     current_len = image_embeds_local.shape[0]
     if current_len < max_len_per_rank:
         padding_size = max_len_per_rank - current_len
-        if rope_type == "rope_2d":
+        if packed_2d_rope:
             padding = torch.empty(
                 (
                     padding_size,

python/sglang/srt/multimodal/processors/base_processor.py

@@ -35,6 +35,7 @@ from sglang.srt.utils.cuda_ipc_transport_utils import (
     MM_ITEM_MEMORY_POOL_RECYCLE_INTERVAL,
     CudaIpcTensorTransportProxy,
     MmItemMemoryPool,
+    get_mm_feature_pool_size_per_worker,
 )
 
 _is_cpu = is_cpu()
@@ -280,16 +281,19 @@ class BaseMultimodalProcessor(ABC):
             # tokenizer workers. Each worker gets an equal share so that adding
             # workers doesn't multiply the GPU-side footprint.
             worker_num = self.server_args.tokenizer_worker_num
-            per_worker_pool_size = max(
-                MM_FEATURE_CACHE_SIZE // worker_num,
-                128 * 1024 * 1024,
+            per_worker_pool_size = get_mm_feature_pool_size_per_worker(
+                MM_FEATURE_CACHE_SIZE, worker_num
             )
+            total_pool_size = per_worker_pool_size * worker_num
             logger.info(
-                "MmItemMemoryPool size per tokenizer worker: %.0f MiB "
-                "(budget %.0f MiB / %d worker(s))",
+                "CUDA IPC multimodal feature pools reserve %.0f MiB total on "
+                "GPU %d (%.0f MiB per tokenizer worker × %d; configured "
+                "budget %.0f MiB).",
+                total_pool_size / (1024 * 1024),
+                self.server_args.base_gpu_id,
                 per_worker_pool_size / (1024 * 1024),
-                MM_FEATURE_CACHE_SIZE / (1024 * 1024),
                 worker_num,
+                MM_FEATURE_CACHE_SIZE / (1024 * 1024),
             )
             self.cudaipc_mmfeature_pool = MmItemMemoryPool(
                 per_worker_pool_size,

python/sglang/srt/multimodal/processors/kimi_k25.py

@@ -19,6 +19,9 @@ from sglang.srt.multimodal.processors.base_processor import (
     MultimodalSpecialTokens,
 )
 from sglang.srt.multimodal.processors.kimi_common import KimiGridMMDataMixin
+from sglang.srt.utils.cuda_ipc_transport_utils import (
+    DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY,
+)
 
 # ---------------------------------------------------------------------------
 # GPU image preprocessing utilities (resize, pad, normalize, patchify on CUDA)
@@ -143,6 +146,49 @@ def _process_single_image(
     return x, grid_thw
 
 
+def _resize_images_by_source_shape(
+    indexed_images: list[tuple[int, torch.Tensor]],
+    target_height: int,
+    target_width: int,
+) -> list[torch.Tensor]:
+    """Resize images while batching only inputs with an identical source layout.
+
+    A NaViT target-size group can still contain images with different source
+    dimensions.  Interpolation requires a rectangular batch, so preserve the
+    individual path for those images and batch only equal ``(shape, dtype)``
+    inputs.  The returned tensors retain the caller's original image order.
+    """
+    by_source_shape = defaultdict(list)
+    for index, image in indexed_images:
+        by_source_shape[(tuple(image.shape), image.dtype)].append((index, image))
+
+    resized_by_index = {}
+    for images in by_source_shape.values():
+        if len(images) == 1:
+            index, image = images[0]
+            resized_by_index[index] = F.interpolate(
+                image.unsqueeze(0).float(),
+                size=(target_height, target_width),
+                mode="bicubic",
+                align_corners=False,
+            )
+            continue
+
+        source_batch = torch.cat(
+            [image.unsqueeze(0) for _, image in images], dim=0
+        ).float()
+        resized_batch = F.interpolate(
+            source_batch,
+            size=(target_height, target_width),
+            mode="bicubic",
+            align_corners=False,
+        )
+        for local_index, (index, _) in enumerate(images):
+            resized_by_index[index] = resized_batch[local_index : local_index + 1]
+
+    return [resized_by_index[index] for index, _ in indexed_images]
+
+
 def _gpu_preprocess_images(
     images: list[Union[torch.Tensor, Image.Image]],
     resize_configs: list[dict],
@@ -182,21 +228,22 @@ def _gpu_preprocess_images(
             all_patches[idx] = patches
             all_grids[idx] = grid
         else:
-            tensors = []
-            for _, image, _ in group:
+            indexed_images = []
+            for idx, image, _ in group:
                 if isinstance(image, Image.Image):
                     image = _pil_to_cuda_chw(image)
                 else:
                     image = _ensure_chw_rgb(image)
-                tensors.append(image.unsqueeze(0).float())
+                indexed_images.append((idx, image))
 
-            resized = []
-            for t in tensors:
-                r = F.interpolate(
-                    t, size=(target_h, target_w), mode="bicubic", align_corners=False
-                )
-                resized.append(r)
-            batch = torch.cat(resized, dim=0)
+            # One NaViT target group can include several original resolutions.
+            # Batch only source-compatible images, which removes redundant
+            # bicubic launches for common multi-image requests without padding
+            # random-size inputs to a larger source resolution.
+            batch = torch.cat(
+                _resize_images_by_source_shape(indexed_images, target_h, target_w),
+                dim=0,
+            )
 
             pad_h = padded_h - target_h
             pad_w = padded_w - target_w
@@ -413,6 +460,16 @@ class KimiK2_5VLImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor):
             base_output, self.mm_tokens
         )
 
+        # K2.5/K2.7 encoder-DP assigns an image to exactly one TP rank. Keep
+        # its IPC proxy lazy until that assignment is known, avoiding a full
+        # image copy to every rank. The scheduler only honors this marker once
+        # the processor has already set the item's hash and pad value.
+        if self.use_cuda_ipc and self.server_args.mm_enable_dp_encoder:
+            for item in mm_items:
+                item.model_specific_data[DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY] = (
+                    True
+                )
+
         return MultimodalProcessorOutput(
             input_ids=input_ids.tolist(),
             mm_items=mm_items,

python/sglang/srt/utils/cuda_ipc_transport_utils.py

@@ -22,6 +22,32 @@ MM_ITEM_MEMORY_POOL_RECYCLE_INTERVAL = (
 
 SHM_LOCK_FILE = "/tmp/shm_wr_lock.lock"
 
+# Processors set this marker only when their encoder consumes each IPC feature
+# on a single TP rank.  The scheduler then keeps the feature lazy until the
+# model has computed the data-parallel assignment.
+DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY = (
+    "_sglang_defer_cuda_ipc_feature_reconstruction"
+)
+
+
+def get_mm_feature_pool_size_per_worker(
+    total_pool_size: int, tokenizer_worker_num: int
+) -> int:
+    """Split the CUDA IPC feature-pool budget without exceeding it.
+
+    Each tokenizer worker owns a distinct CUDA allocation, even though all pools
+    are created on ``base_gpu_id``.  Therefore a minimum per-worker allocation
+    would make the aggregate HBM reservation larger than the configured budget.
+    Keep the configured value as a hard per-node cap and leave at most
+    ``tokenizer_worker_num - 1`` bytes unused when it is not evenly divisible.
+    """
+    if total_pool_size <= 0:
+        raise ValueError("total_pool_size must be positive")
+    if tokenizer_worker_num <= 0:
+        raise ValueError("tokenizer_worker_num must be positive")
+
+    return total_pool_size // tokenizer_worker_num
+
 
 # Cache for pool-level IPC handles on the consumer side.
 # Key: the pool CUDA IPC handle tuple. Value: opened UntypedStorage.
@@ -330,6 +356,7 @@ class CudaIpcTensorTransportProxy:
         self.reconstruct_tensor = None
         self.sync_data_meta = sync_buffer_meta
         self.sync_buffer = None
+        self._consumer_acknowledged = False
 
     @property
     def get_sync_flag(self):
@@ -404,32 +431,53 @@ class CudaIpcTensorTransportProxy:
 
         return slice_tensor, target_device, cache_key, storage_to_cache
 
+    def _acknowledge_consumption(self, consumer_count: int = 1):
+        """Mark this IPC feature as consumed without necessarily copying it.
+
+        A normal TP execution reconstructs a feature once per rank, so each
+        consumer contributes one acknowledgement.  Encoder-DP can instead
+        route a feature to exactly one rank; that rank acknowledges all TP
+        consumers after its copy completes.  Keeping this acknowledgement
+        idempotent is important for chunked-prefill cache hits, where the same
+        proxy may be visited more than once.
+        """
+        if getattr(self, "_consumer_acknowledged", False):
+            return
+        if consumer_count <= 0:
+            raise ValueError("consumer_count must be positive")
+        if self.sync_data_meta is not None:
+            open(SHM_LOCK_FILE, "a").close()
+            # Keep the counter update atomic across scheduler processes.
+            with open(SHM_LOCK_FILE, "w+") as f:
+                fcntl.flock(f, fcntl.LOCK_EX)
+                sync_flag = self.get_sync_flag
+                sync_flag += consumer_count
+                fcntl.flock(f, fcntl.LOCK_UN)
+            self.close_shm()
+        self._consumer_acknowledged = True
+
+    def acknowledge_consumption(self, consumer_count: int = 1):
+        """Release an IPC-pool slice when a cache hit needs no tensor copy."""
+        self._acknowledge_consumption(consumer_count)
+
     def _copy_slice_tensor_to_target(
         self,
         slice_tensor: torch.Tensor,
         rebuild_device: torch.device,
         recons_shape,
         recons_dtype,
+        consumer_count: int,
     ):
         with torch.cuda.device(rebuild_device):
             reconstructed_tensor = torch.empty(
                 recons_shape, dtype=recons_dtype, device=rebuild_device
             ).contiguous()
             reconstructed_tensor.view(torch.int8).view(-1).copy_(slice_tensor)
-
-            open(SHM_LOCK_FILE, "a").close()
-            # write the shm_sync_buffer with a file lock
-            with open(SHM_LOCK_FILE, "w+") as f:
-                fcntl.flock(f, fcntl.LOCK_EX)
-                sync_flag = self.get_sync_flag
-                sync_flag += 1
-                fcntl.flock(f, fcntl.LOCK_UN)
-
-            self.close_shm()
+            self._acknowledge_consumption(consumer_count)
 
         return reconstructed_tensor
 
-    def reconstruct_on_target_device(self, rebuild_device_idx):
+    def reconstruct_on_target_device(self, rebuild_device_idx, consumer_count: int = 1):
         rebuild_device = torch.device(f"cuda:{rebuild_device_idx}")
         if (
             isinstance(self.reconstruct_tensor, torch.Tensor)
@@ -501,7 +549,11 @@ class CudaIpcTensorTransportProxy:
                     raise
 
             reconstructed_tensor = self._copy_slice_tensor_to_target(
-                slice_tensor, rebuild_device, recons_shape, recons_dtype
+                slice_tensor,
+                rebuild_device,
+                recons_shape,
+                recons_dtype,
+                consumer_count,
             )
         elif isinstance(self.proxy_state["tensor_data"], torch.Tensor):
             reconstructed_tensor = self.proxy_state["tensor_data"].to(