perf(multimodal): optimize SHM input transport

python/tokenspeed/runtime/engine/data_parallel_controller.py

@@ -311,6 +311,9 @@ def launch_tensor_parallel_group(
         self.max_num_seqs = scheduler_info[0]["max_num_seqs"]
         self.chunked_prefill_size = scheduler_info[0]["chunked_prefill_size"]
         self.max_model_len = scheduler_info[0]["max_model_len"]
+        self.multimodal_encoder_dtype = scheduler_info[0].get(
+            "multimodal_encoder_dtype"
+        )
 
     def round_robin_scheduler(self, req: Req):
         if self.server_args.disaggregation_mode == "null":
@@ -364,6 +367,7 @@ def run_data_parallel_controller_process(
                 "max_num_seqs": controller.max_num_seqs,
                 "chunked_prefill_size": controller.chunked_prefill_size,
                 "max_model_len": controller.max_model_len,
+                "multimodal_encoder_dtype": controller.multimodal_encoder_dtype,
             }
         )
         if server_args.node_rank == 0:

python/tokenspeed/runtime/engine/event_loop.py

@@ -181,6 +181,7 @@ def __init__(
         target, draft = create_model_runner(
             server_args, self.model_config, draft_model_config, gpu_id, global_rank
         )
+        self.multimodal_encoder_dtype = target.multimodal_encoder_dtype
         self.use_overlap_schedule = should_use_overlap_schedule(
             disable_overlap_schedule=server_args.disable_overlap_schedule,
             disaggregation_mode=server_args.disaggregation_mode,
@@ -1893,6 +1894,7 @@ def run_event_loop(
                 "max_num_seqs": server_args.max_num_seqs,
                 "chunked_prefill_size": server_args.chunked_prefill_size,
                 "max_model_len": event_loop.model_config.context_len,
+                "multimodal_encoder_dtype": event_loop.multimodal_encoder_dtype,
             }
         )
 

python/tokenspeed/runtime/execution/model_runner.py

@@ -41,6 +41,26 @@
 logger = get_colorful_logger(__name__)
 
 
+def infer_multimodal_encoder_dtype(model: torch.nn.Module) -> str | None:
+    """Return the dtype of the loaded vision tower when it is discoverable."""
+    for path in (
+        ("visual",),
+        ("vision_tower",),
+        ("model", "visual"),
+        ("model", "vision_tower"),
+    ):
+        component = model
+        for name in path:
+            component = getattr(component, name, None)
+            if component is None:
+                break
+        else:
+            dtype = getattr(component, "dtype", None)
+            if isinstance(dtype, torch.dtype):
+                return str(dtype).removeprefix("torch.")
+    return None
+
+
 class ModelRunner:
     def __init__(
         self,
@@ -115,6 +135,10 @@ def load_model(self):
             self.model, "spec_step_idx"
         )
 
+    @property
+    def multimodal_encoder_dtype(self) -> str | None:
+        return infer_multimodal_encoder_dtype(self.model)
+
     @staticmethod
     def _forward_accepts_kwarg(model, name: str) -> bool:
         try:

python/tokenspeed/runtime/models/kimi_k25.py

@@ -789,7 +789,7 @@ def __init__(
                 self.mm_projector = self.mm_projector.to(dtype=target_dtype)
 
             # image_encoder may be swapped to a cudagraph wrapper by ModelExecutor.
-            self.vision_embedder = VisionEmbedder()
+            self.vision_embedder = VisionEmbedder(mapping.vision.tp_group)
             self.image_encoder = self.get_image_feature
         else:
             self.vision_embedder = None

python/tokenspeed/runtime/models/qwen3_5.py

@@ -1193,7 +1193,7 @@ def __init__(
             self.num_deepstack_embeddings = len(self.deepstack_visual_indexes)
             # Encoder callables may be swapped to cudagraph wrappers by
             # ModelExecutor.
-            self.vision_embedder = VisionEmbedder()
+            self.vision_embedder = VisionEmbedder(mapping.vision.tp_group)
             self.image_encoder = self.get_image_feature
             self.video_encoder = self.get_video_feature
 

python/tokenspeed/runtime/multimodal/embedder.py

@@ -53,6 +53,7 @@
 from __future__ import annotations
 
 import logging
+import math
 import time
 from collections import defaultdict
 from collections.abc import Callable
@@ -75,6 +76,14 @@
 
 logger = logging.getLogger(__name__)
 LOG_MM_TIMING = envs.TOKENSPEED_LOG_MM_TIMING.get()
+# Small transfers are faster after staging the whole batch; larger transfers
+# benefit from overlapping each H2D enqueue with the next SHM-to-pinned copy.
+_INTERLEAVED_H2D_MIN_AVERAGE_BYTES = 1024 * 1024
+# One rank can stage the input and distribute it faster than every TP rank
+# repeating the host copy once the payload reaches this TP-scaled threshold.
+# Local B200 measurements put the break-even points near 128 MiB for TP2 and
+# 64 MiB for TP4.
+_TP_BROADCAST_BASE_MIN_BYTES = 256 * 1024 * 1024
 
 
 @dataclass
@@ -182,8 +191,20 @@ def _item_token_count(item: MultimodalDataItem) -> int:
 class MultimodalEmbedder:
     """Multimodal input embedding pipeline for one model executor."""
 
-    def __init__(self) -> None:
+    def __init__(self, vision_tp_group: tuple[int, ...] | None = None) -> None:
         self._h2d_stream: torch.cuda.Stream | None = None
+        self._vision_tp_group = vision_tp_group
+        self._vision_tp_process_group = None
+        self._vision_tp_src_rank: int | None = None
+        if vision_tp_group is not None and len(vision_tp_group) > 1:
+            from tokenspeed.runtime.distributed.process_group_manager import (
+                process_group_manager as pg_manager,
+            )
+
+            self._vision_tp_process_group = pg_manager.get_process_group(
+                "nccl", vision_tp_group
+            )
+            self._vision_tp_src_rank = vision_tp_group[0]
 
     # --- public entry point ------------------------------------------------
 
@@ -504,24 +525,109 @@ def _move_features_to_device(
         if not pending:
             return
 
-        for it in pending:
-            if isinstance(it.feature, ShmTensorHandle):
-                it.feature = it.feature.consume()
-
         if device.type != "cuda":
             for it in pending:
+                if isinstance(it.feature, ShmTensorHandle):
+                    it.feature = it.feature.consume()
                 if isinstance(it.feature, torch.Tensor):
                     it.feature = it.feature.to(device, non_blocking=True)
             return
 
+        shm_count = 0
+        shm_nbytes = 0
+        for item in pending:
+            if isinstance(item.feature, ShmTensorHandle):
+                shm_count += 1
+                shm_nbytes += item.feature.nbytes
+        use_tp_broadcast = self._should_move_shm_via_tp_broadcast(pending)
+        interleave_h2d = shm_nbytes > shm_count * _INTERLEAVED_H2D_MIN_AVERAGE_BYTES
+        defer_shm_cleanup = shm_count == 1 and interleave_h2d
+        if not use_tp_broadcast and not interleave_h2d:
+            for item in pending:
+                if isinstance(item.feature, ShmTensorHandle):
+                    item.feature = item.feature.consume()
+
+        # Keep collectives on the model stream so their ordering matches other
+        # TP collectives queued by the forward pass on every rank.
+        if use_tp_broadcast:
+            self._move_shm_via_tp_broadcast(pending, device)
+            return
+
         h2d = self._h2d_stream_on(device)
         current = torch.cuda.current_stream(device)
         with torch.cuda.stream(h2d):
             for it in pending:
+                if isinstance(it.feature, ShmTensorHandle):
+                    if defer_shm_cleanup:
+                        handle = it.feature
+                        try:
+                            it.feature = handle.copy_to_pinned()
+                            it.feature = it.feature.to(device, non_blocking=True)
+                        finally:
+                            handle.release()
+                        continue
+                    it.feature = it.feature.consume()
                 if isinstance(it.feature, torch.Tensor):
                     it.feature = it.feature.to(device, non_blocking=True)
         current.wait_stream(h2d)
 
+    def _should_move_shm_via_tp_broadcast(
+        self, items: list[MultimodalDataItem]
+    ) -> bool:
+        tp_group = self._vision_tp_group
+        if (
+            tp_group is None
+            or self._vision_tp_process_group is None
+            or self._vision_tp_src_rank is None
+            or not items
+            or not all(isinstance(item.feature, ShmTensorHandle) for item in items)
+        ):
+            return False
+
+        handles = [item.feature for item in items]
+        dtype = handles[0].dtype
+        if any(handle.dtype != dtype for handle in handles):
+            return False
+        total_nbytes = sum(handle.nbytes for handle in handles)
+        return total_nbytes >= _TP_BROADCAST_BASE_MIN_BYTES // len(tp_group)
+
+    def _move_shm_via_tp_broadcast(
+        self,
+        items: list[MultimodalDataItem],
+        device: torch.device,
+    ) -> None:
+        tp_group = self._vision_tp_group
+        process_group = self._vision_tp_process_group
+        src_rank = self._vision_tp_src_rank
+        assert tp_group is not None
+        assert process_group is not None
+        assert src_rank is not None
+
+        handles = [item.feature for item in items]
+        dtype = handles[0].dtype
+
+        element_lengths = [math.prod(handle.shape) for handle in handles]
+        base = torch.empty(sum(element_lengths), dtype=dtype, device=device)
+        is_source = torch.distributed.get_rank() == src_rank
+        offset = 0
+        if is_source:
+            for handle, length in zip(handles, element_lengths, strict=True):
+                if len(handles) == 1:
+                    handle.copy_into(base.narrow(0, offset, length))
+                else:
+                    source = handle.consume().reshape(-1)
+                    base.narrow(0, offset, length).copy_(source, non_blocking=True)
+                offset += length
+        else:
+            for handle in handles:
+                handle.release()
+
+        torch.distributed.broadcast(base, src=src_rank, group=process_group)
+        offset = 0
+        for item, handle, length in zip(items, handles, element_lengths, strict=True):
+            item.feature = base.narrow(0, offset, length).view(handle.shape)
+            offset += length
+
     @staticmethod
     def _drop_raw_feature(item: MultimodalDataItem) -> bool:
         if item.feature is None:

python/tokenspeed/runtime/multimodal/shm_transport.py

@@ -31,6 +31,7 @@
 from __future__ import annotations
 
 import logging
+import math
 import time
 from dataclasses import dataclass, field
 from multiprocessing import shared_memory
@@ -54,6 +55,12 @@ class ShmTensorHandle:
         default=None, init=False, repr=False, compare=False
     )
 
+    @property
+    def nbytes(self) -> int:
+        if self._segment is not None:
+            return self._segment.size
+        return math.prod(self.shape) * self.dtype.itemsize
+
     @classmethod
     def publish(cls, tensor: torch.Tensor) -> ShmTensorHandle:
         nbytes = tensor.numel() * tensor.element_size()
@@ -80,38 +87,88 @@ def consume(self) -> torch.Tensor:
         """Copy into a pinned tensor (so downstream non_blocking H2D is real),
         close this rank's FD, and unlink. ``attach()`` must have run.
         """
+        started = time.perf_counter() if LOG_MM_TIMING else None
+        try:
+            dst = self._copy_to_pinned()
+        finally:
+            self._close_and_unlink()
+        if LOG_MM_TIMING and started is not None:
+            logger.info(
+                "mm_timing shm_consume_ms name=%s elapsed=%.3f shape=%s dtype=%s",
+                self.shm_name,
+                (time.perf_counter() - started) * 1000,
+                list(self.shape),
+                self.dtype,
+            )
+        return dst
+
+    def copy_to_pinned(self) -> torch.Tensor:
+        """Copy into pinned memory while retaining this rank's SHM ownership.
+
+        The caller must subsequently call :meth:`release`. This allows an
+        asynchronous H2D copy to be enqueued before close/unlink cleanup.
+        """
+        started = time.perf_counter() if LOG_MM_TIMING else None
+        dst = self._copy_to_pinned()
+        if LOG_MM_TIMING and started is not None:
+            logger.info(
+                "mm_timing shm_copy_to_pinned_ms name=%s elapsed=%.3f shape=%s dtype=%s",
+                self.shm_name,
+                (time.perf_counter() - started) * 1000,
+                list(self.shape),
+                self.dtype,
+            )
+        return dst
+
+    def copy_into(self, destination: torch.Tensor) -> None:
+        """Synchronously copy into an existing tensor and release the SHM segment."""
         if self._segment is None:
             raise RuntimeError(
                 f"ShmTensorHandle({self.shm_name!r}) must be attach()'d "
-                "before consume() (or has already been consumed on this rank)"
+                "before copying (or has already been released on this rank)"
             )
-        segment = self._segment
         started = time.perf_counter() if LOG_MM_TIMING else None
+        source = torch.frombuffer(self._segment.buf, dtype=self.dtype).reshape(
+            self.shape
+        )
         try:
-            dst = torch.empty(self.shape, dtype=self.dtype, pin_memory=True)
-            src = torch.frombuffer(segment.buf, dtype=self.dtype).reshape(self.shape)
-            dst.copy_(src)
+            if source.dtype != destination.dtype:
+                raise ValueError(
+                    "SHM source and destination dtypes differ: "
+                    f"{source.dtype} != {destination.dtype}"
+                )
+            if source.shape != destination.shape:
+                if source.numel() != destination.numel():
+                    raise ValueError(
+                        "SHM source and destination element counts differ: "
+                        f"{source.numel()} != {destination.numel()}"
+                    )
+                source = source.reshape(destination.shape)
+            destination.copy_(source)
         finally:
-            self._segment = None
-            segment.close()
-            try:
-                segment.unlink()
-            except FileNotFoundError:
-                # Another rank already won the unlink race; benign.
-                pass
+            del source
+            self._close_and_unlink()
         if LOG_MM_TIMING and started is not None:
             logger.info(
-                "mm_timing shm_consume_ms name=%s elapsed=%.3f shape=%s dtype=%s",
+                "mm_timing shm_copy_into_ms name=%s elapsed=%.3f shape=%s dtype=%s",
                 self.shm_name,
                 (time.perf_counter() - started) * 1000,
                 list(self.shape),
                 self.dtype,
             )
+
+    def _copy_to_pinned(self) -> torch.Tensor:
+        if self._segment is None:
+            raise RuntimeError(
+                f"ShmTensorHandle({self.shm_name!r}) must be attach()'d "
+                "before copying (or has already been released on this rank)"
+            )
+        dst = torch.empty(self.shape, dtype=self.dtype, pin_memory=True)
+        src = torch.frombuffer(self._segment.buf, dtype=self.dtype).reshape(self.shape)
+        dst.copy_(src)
         return dst
 
-    def release(self) -> None:
-        """Close and unlink a SHM segment without materializing the tensor."""
-        started = time.perf_counter() if LOG_MM_TIMING else None
+    def _close_and_unlink(self) -> None:
         segment = self._segment
         self._segment = None
         try:
@@ -121,9 +178,15 @@ def release(self) -> None:
             try:
                 segment.unlink()
             except FileNotFoundError:
+                # Another rank already won the unlink race; benign.
                 pass
         except FileNotFoundError:
             pass
+
+    def release(self) -> None:
+        """Close and unlink a SHM segment without materializing the tensor."""
+        started = time.perf_counter() if LOG_MM_TIMING else None
+        self._close_and_unlink()
         if LOG_MM_TIMING and started is not None:
             logger.info(
                 "mm_timing shm_release_ms name=%s elapsed=%.3f shape=%s dtype=%s",