model: TP-wide single-owner image encoding for DeepSeek V4.1

ViT and Aligner are replicated per TP rank, encoding each image eight times with TP8 on both CP1 and CP8. Elect one owner per image and use ordered full-span broadcasts with a six-phase agreement protocol. Both CP1 and CP8 benefit while local cache hits preserve collective order.
2026-09-23 14:36:01 +08:00
parent bfeb7cd9b2
commit b48e2cb1eb
6 changed files with 1663 additions and 54 deletions
@@ -0,0 +1,486 @@
+"""One owner rank encodes each image span and broadcasts it to the ranks that
+run the same prefill chunk; every agreement precedes the payload it guards."""
+
+from __future__ import annotations
+
+import logging
+from contextlib import contextmanager
+from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple
+
+import msgspec
+import torch
+
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    disable_symmetric_memory_context,
+    restore_symmetric_memory_context,
+)
+from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache
+
+logger = logging.getLogger(__name__)
+
+SpanKey = Tuple[Optional[int], int]
+SpanEncoder = Callable[[List[Any]], torch.Tensor | List[torch.Tensor]]
+SpanSignature = Callable[[Any, int], Tuple[Any, ...]]
+
+LOCAL_HIT = 0
+OWNER_CACHE_BROADCAST = 1
+OWNER_ENCODE_BROADCAST = 2
+
+PHASE_PREPARE = "prepare"
+PHASE_FEATURES = "features"
+PHASE_FINALIZE = "finalize"
+
+
+class MmOwnerProtocolError(RuntimeError):
+    """Raised with identical text on every group member after a group-agreed failure."""
+
+
+class ImageSpanRequest(msgspec.Struct, frozen=True):
+    hash: Optional[int]
+    span_len: int
+    item: Any
+    inside_chunk: bool
+    duplicates: List[Any] = []
+
+
+class ImageSpanKey(msgspec.Struct, frozen=True):
+    hash: Optional[int]
+    span_len: int
+    geometry: Optional[Tuple[Any, ...]]
+
+
+class RankManifest(msgspec.Struct, frozen=True):
+    rank: int
+    keys: List[ImageSpanKey]
+    cached: List[bool]
+    dtype: str
+    width: int
+    rids: List[str]
+    error: Optional[str] = None
+
+
+class OwnerPlan(msgspec.Struct, frozen=True):
+    actions: List[int]
+    owners: List[int]
+    error: Optional[str] = None
+
+
+class RankStatus(msgspec.Struct, frozen=True):
+    rank: int
+    error: Optional[str] = None
+
+
+def select_owner_group(parallel) -> Optional[Any]:
+    """The group whose members all execute the same requests, or None when a
+    single rank already encodes every image it sees."""
+    replication = parallel.tp_size // parallel.attn_dp_size
+    if replication <= 1:
+        return None
+    if parallel.attn_cp_size == 1:
+        group = parallel.attn_tp_group
+    elif parallel.attn_dp_size == 1 and parallel.attn_cp_size == parallel.tp_size:
+        group = parallel.attn_cp_group
+    else:
+        return None
+    return group if group.world_size == replication else None
+
+
+def has_owner_span_work(
+    mm_inputs: Sequence[Any],
+    extend_prefix_lens: Sequence[int],
+    extend_seq_lens: Sequence[int],
+) -> bool:
+    """Host-side mirror of the per-image scheduling path: does any raw
+    single-span image overlap the chunk on every rank of the group."""
+    for mm_input, prefix_len, extend_len in zip(
+        mm_inputs, extend_prefix_lens, extend_seq_lens
+    ):
+        if mm_input is None or extend_len <= 0:
+            continue
+        items = [item for item in mm_input.mm_items if item is not None]
+        if not items or any(
+            item.precomputed_embeddings is not None or len(item.offsets) != 1
+            for item in items
+        ):
+            continue
+        for item in items:
+            start, end = item.offsets[0]
+            if end >= prefix_len and start < prefix_len + extend_len:
+                return True
+    return False
+
+
+class MmOwnerSession(msgspec.Struct):
+    group: Any
+    device: Any
+    dtype: Any
+    width: int
+    rids: List[str]
+    signature: Any
+    engaged: bool
+    phase: str = PHASE_PREPARE
+    in_collective: bool = False
+
+    def resolve(
+        self,
+        requests: Sequence[ImageSpanRequest],
+        cache: MultiModalStaticCache,
+        encode: SpanEncoder,
+    ) -> Dict[SpanKey, torch.Tensor]:
+        if not self.engaged:
+            raise RuntimeError(
+                "owner protocol reached for a chunk whose host metadata has no image span"
+            )
+        # Owners allocate different amounts than receivers, so none of these
+        # buffers may come out of a symmetric pool.
+        saved_context = disable_symmetric_memory_context()
+        try:
+            return _resolve_owner_features(self, requests, cache, encode)
+        finally:
+            restore_symmetric_memory_context(saved_context)
+
+    def features_ready(self) -> None:
+        self._complete()
+        self.phase = PHASE_FINALIZE
+
+    @contextmanager
+    def uncaptured(self) -> Iterator[None]:
+        # A failure inside a collective leaves the group in an unknown state;
+        # no later exchange may try to agree on it.
+        self.in_collective = True
+        yield
+        self.in_collective = False
+
+    @contextmanager
+    def fence(self) -> Iterator[None]:
+        try:
+            yield
+        except Exception as exc:
+            self._fail(exc)
+            raise
+        self._complete()
+
+    def _fail(self, exc: BaseException) -> None:
+        if (
+            not self.engaged
+            or self.in_collective
+            or isinstance(exc, MmOwnerProtocolError)
+        ):
+            raise exc
+        text = _describe(self, self.phase, exc)
+        if self.phase == PHASE_PREPARE:
+            try:
+                _exchange_manifest(self, _manifest(self, [], [], error=text))
+            except MmOwnerProtocolError as agreed:
+                raise agreed from exc
+        _exchange_status(self, text, exc)
+
+    def _complete(self) -> None:
+        if not self.engaged:
+            return
+        if self.phase == PHASE_PREPARE:
+            raise RuntimeError(
+                f"owner protocol {self.phase} completed without a manifest exchange"
+            )
+        error = None
+        cause = None
+        try:
+            _synchronize(self.device)
+        except Exception as exc:
+            cause = exc
+            error = _describe(self, self.phase, exc)
+        _exchange_status(self, error, cause)
+
+
+def _manifest(
+    session: MmOwnerSession,
+    keys: List[ImageSpanKey],
+    cached: List[bool],
+    error: Optional[str] = None,
+) -> RankManifest:
+    return RankManifest(
+        rank=session.group.rank_in_group,
+        keys=keys,
+        cached=cached,
+        dtype=str(session.dtype),
+        width=session.width,
+        rids=list(session.rids),
+        error=error,
+    )
+
+
+def _exchange_manifest(session: MmOwnerSession, manifest: RankManifest) -> OwnerPlan:
+    group = session.group
+    with session.uncaptured():
+        manifests = group.all_gather_object(manifest)
+        plan = _plan_or_error(session, manifests) if group.rank_in_group == 0 else None
+        plan = group.broadcast_object(plan, src=0)
+    session.phase = PHASE_FEATURES
+    if plan.error is not None:
+        raise MmOwnerProtocolError(plan.error)
+    return plan
+
+
+def _plan_or_error(session: MmOwnerSession, manifests: List[RankManifest]) -> OwnerPlan:
+    try:
+        return _make_plan(manifests)
+    except Exception as exc:
+        return OwnerPlan(actions=[], owners=[], error=_describe(session, "plan", exc))
+
+
+def _exchange_status(
+    session: MmOwnerSession, error: Optional[str], cause: Optional[BaseException]
+) -> None:
+    with session.uncaptured():
+        statuses = session.group.all_gather_object(
+            RankStatus(rank=session.group.rank_in_group, error=error)
+        )
+    _raise_first_error(statuses, cause)
+
+
+def _resolve_owner_features(
+    session: MmOwnerSession,
+    requests: Sequence[ImageSpanRequest],
+    cache: MultiModalStaticCache,
+    encode: SpanEncoder,
+) -> Dict[SpanKey, torch.Tensor]:
+    group = session.group
+    features: Dict[SpanKey, torch.Tensor] = {}
+    keys: List[ImageSpanKey] = []
+    cached: List[bool] = []
+    error = None
+    try:
+        keys, cached = _pin_local_cache(session, requests, cache, features)
+    except Exception as exc:
+        error = _describe(session, "manifest", exc)
+    plan = _exchange_manifest(session, _manifest(session, keys, cached, error))
+
+    if all(action == LOCAL_HIT for action in plan.actions):
+        return features
+
+    buffers: Dict[int, torch.Tensor] = {}
+    error = None
+    try:
+        buffers = _prepare_transfers(session, requests, keys, plan, features, encode)
+        _synchronize(session.device)
+    except Exception as exc:
+        error = _describe(session, "encode", exc)
+    _exchange_status(session, error, None)
+
+    with session.uncaptured():
+        for index, (action, owner) in enumerate(zip(plan.actions, plan.owners)):
+            if action != LOCAL_HIT:
+                group.broadcast(buffers[index], src=owner)
+
+    for index, key in enumerate(keys):
+        if plan.actions[index] == LOCAL_HIT:
+            continue
+        span = buffers[index]
+        features[(key.hash, key.span_len)] = span
+        cache.set(key.hash, EmbeddingResult(embedding=span))
+    return features
+
+
+def _pin_local_cache(
+    session: MmOwnerSession,
+    requests: Sequence[ImageSpanRequest],
+    cache: MultiModalStaticCache,
+    features: Dict[SpanKey, torch.Tensor],
+) -> Tuple[List[ImageSpanKey], List[bool]]:
+    keys: List[ImageSpanKey] = []
+    cached: List[bool] = []
+    for request in requests:
+        if request.hash is None:
+            raise ValueError(
+                f"image span of {request.span_len} tokens has no content hash"
+            )
+        geometry = session.signature(request.item, request.span_len)
+        for duplicate in request.duplicates:
+            other = session.signature(duplicate, request.span_len)
+            if other != geometry:
+                raise ValueError(
+                    f"image hash {request.hash} ({request.span_len} tokens) occurs "
+                    f"with different geometry: {geometry} vs {other}"
+                )
+        keys.append(
+            ImageSpanKey(
+                hash=request.hash, span_len=request.span_len, geometry=geometry
+            )
+        )
+        span = _valid_cached_span(session, cache, request)
+        if span is not None:
+            features[(request.hash, request.span_len)] = span
+        cached.append(span is not None)
+    return keys, cached
+
+
+def _valid_cached_span(
+    session: MmOwnerSession,
+    cache: MultiModalStaticCache,
+    request: ImageSpanRequest,
+) -> Optional[torch.Tensor]:
+    entry = cache.get_single(request.hash)
+    if entry is None:
+        return None
+    span = entry.embedding
+    if (
+        span.dim() == 2
+        and span.shape[0] == request.span_len
+        and span.shape[1] == session.width
+        and span.dtype == session.dtype
+        and span.device == session.device
+    ):
+        return span
+    logger.warning(
+        "Discarding cached multimodal embedding that cannot serve the current "
+        "image span: cache_key=%s expected=(%d, %d, %s) cached=(%s, %s).",
+        request.hash,
+        request.span_len,
+        session.width,
+        session.dtype,
+        tuple(span.shape),
+        span.dtype,
+    )
+    cache.free(request.hash, None)
+    return None
+
+
+def _make_plan(manifests: List[RankManifest]) -> OwnerPlan:
+    for manifest in manifests:
+        if manifest.error is not None:
+            return OwnerPlan(actions=[], owners=[], error=manifest.error)
+    lead = manifests[0]
+    for manifest in manifests[1:]:
+        if (manifest.keys, manifest.dtype, manifest.width, manifest.rids) != (
+            lead.keys,
+            lead.dtype,
+            lead.width,
+            lead.rids,
+        ):
+            return OwnerPlan(
+                actions=[],
+                owners=[],
+                error=(
+                    "image manifest mismatch between group ranks 0 and "
+                    f"{manifest.rank}: rids={lead.rids} vs {manifest.rids}, "
+                    f"keys={lead.keys} vs {manifest.keys}, "
+                    f"dtype={lead.dtype} vs {manifest.dtype}, "
+                    f"width={lead.width} vs {manifest.width}"
+                ),
+            )
+    replication = len(manifests)
+    actions: List[int] = []
+    owners: List[int] = []
+    for index, key in enumerate(lead.keys):
+        owner = key.hash % replication
+        if all(manifest.cached[index] for manifest in manifests):
+            action = LOCAL_HIT
+        elif manifests[owner].cached[index]:
+            action = OWNER_CACHE_BROADCAST
+        else:
+            action = OWNER_ENCODE_BROADCAST
+        actions.append(action)
+        owners.append(owner)
+    return OwnerPlan(actions=actions, owners=owners)
+
+
+def _prepare_transfers(
+    session: MmOwnerSession,
+    requests: Sequence[ImageSpanRequest],
+    keys: List[ImageSpanKey],
+    plan: OwnerPlan,
+    features: Dict[SpanKey, torch.Tensor],
+    encode: SpanEncoder,
+) -> Dict[int, torch.Tensor]:
+    rank = session.group.rank_in_group
+    buffers: Dict[int, torch.Tensor] = {}
+    owned: List[int] = []
+    for index, (action, owner) in enumerate(zip(plan.actions, plan.owners)):
+        if action == LOCAL_HIT:
+            continue
+        if owner != rank:
+            try:
+                buffers[index] = _new_span_buffer(session, keys[index])
+            except Exception as exc:
+                raise RuntimeError(
+                    f"receive buffer for image hash {keys[index].hash} shape "
+                    f"{(keys[index].span_len, session.width)} {session.dtype} "
+                    f"failed: {type(exc).__name__}: {exc}"
+                ) from exc
+        elif action == OWNER_CACHE_BROADCAST:
+            key = (keys[index].hash, keys[index].span_len)
+            buffers[index] = features[key].contiguous()
+        else:
+            owned.append(index)
+    if owned:
+        owned_hashes = [keys[index].hash for index in owned]
+        try:
+            encoded = encode([requests[index].item for index in owned])
+        except Exception as exc:
+            raise RuntimeError(
+                f"owner encode of image hashes {owned_hashes} failed: "
+                f"{type(exc).__name__}: {exc}"
+            ) from exc
+        spans = _split_spans(encoded, [keys[index].span_len for index in owned])
+        for index, span in zip(owned, spans):
+            buffers[index] = _validated_span(session, keys[index], span)
+    return buffers
+
+
+def _new_span_buffer(session: MmOwnerSession, key: ImageSpanKey) -> torch.Tensor:
+    return torch.empty(
+        (key.span_len, session.width), device=session.device, dtype=session.dtype
+    )
+
+
+def _split_spans(
+    encoded: torch.Tensor | List[torch.Tensor], span_lens: List[int]
+) -> List[torch.Tensor]:
+    if isinstance(encoded, list):
+        if len(encoded) != len(span_lens):
+            raise ValueError(
+                f"encoder returned {len(encoded)} spans for {len(span_lens)} images"
+            )
+        return [span.reshape(-1, span.shape[-1]) for span in encoded]
+    encoded = encoded.reshape(-1, encoded.shape[-1])
+    if encoded.shape[0] != sum(span_lens):
+        raise ValueError(
+            f"encoder returned {encoded.shape[0]} rows for spans of {span_lens}"
+        )
+    return list(torch.split(encoded, span_lens, dim=0))
+
+
+def _validated_span(
+    session: MmOwnerSession, key: ImageSpanKey, span: torch.Tensor
+) -> torch.Tensor:
+    expected = (key.span_len, session.width)
+    if tuple(span.shape) != expected or span.dtype != session.dtype:
+        raise ValueError(
+            f"encoded span for hash={key.hash} has shape {tuple(span.shape)} "
+            f"dtype {span.dtype}; expected {expected} {session.dtype}"
+        )
+    if span.device != session.device:
+        span = span.to(session.device)
+    return span.contiguous()
+
+
+def _synchronize(device) -> None:
+    if device.type == "cuda":
+        torch.cuda.current_stream(device).synchronize()
+
+
+def _describe(session: MmOwnerSession, stage: str, exc: BaseException) -> str:
+    return (
+        f"multimodal owner protocol failed during {stage} on group rank "
+        f"{session.group.rank_in_group} (global rank "
+        f"{session.group.ranks[session.group.rank_in_group]}, rids={list(session.rids)}): "
+        f"{type(exc).__name__}: {exc}"
+    )
+
+
+def _raise_first_error(
+    statuses: List[RankStatus], cause: Optional[BaseException]
+) -> None:
+    for status in statuses:
+        if status.error is not None:
+            raise MmOwnerProtocolError(status.error) from cause
@@ -5,6 +5,7 @@ from typing import Callable, Dict, List, Optional, Tuple

 import torch

+from sglang.srt.managers.mm_owner_embedding import ImageSpanRequest, MmOwnerSession
 from sglang.srt.managers.schedule_batch import MultimodalDataItem
 from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache
 from sglang.srt.multimodal.evs import EVSEmbeddingResult
@@ -339,43 +340,24 @@ def _batch_encode_per_image_misses(
    unique_misses: Dict[Tuple[Optional[int], int], Tuple[MultimodalDataItem, int]] = {}
    hash_to_embedding: Dict[Tuple[Optional[int], int], torch.Tensor] = {}

-    # Phase 1a: find overlapping items per request and collect cache misses
-    for req_info in per_image_requests:
-        chunk_start = req_info.extend_prefix_len
-        chunk_end = chunk_start + req_info.extend_seq_len  # exclusive
-        overlapping = []
-        if req_info.extend_seq_len > 0:
-            for idx, (item, (start, end)) in enumerate(
-                zip(req_info.items, req_info.items_offset)
-            ):
-                if end >= chunk_start and start < chunk_end:
-                    overlapping.append((idx, item, start, end))
-        req_info.overlapping = overlapping
-
-        for _idx, item, start, end in overlapping:
-            expected_token_count = end - start + 1
-            cache_key = (item.hash, expected_token_count)
-            if cache_key in hash_to_embedding:
+    # Phase 1a: collect cache misses over the unique overlapping spans
+    for span in _collect_image_span_requests(per_image_requests):
+        cache_key = (span.hash, span.span_len)
+        cached = embedding_cache.get_single(span.hash)
+        if cached is not None:
+            cached_embedding = cached.embedding
+            cached_token_count = _embedding_token_count(cached_embedding)
+            if cached_token_count == span.span_len:
+                hash_to_embedding[cache_key] = cached_embedding
                continue
-            cached = embedding_cache.get_single(item.hash)
-            if cached is not None:
-                cached_embedding = cached.embedding
-                cached_token_count = _embedding_token_count(cached_embedding)
-                if cached_token_count == expected_token_count:
-                    hash_to_embedding[cache_key] = cached_embedding
-                else:
-                    _discard_mismatched_cached_embedding(
-                        item.hash, expected_token_count, cached_token_count
-                    )
-                    unique_misses[cache_key] = (item, expected_token_count)
-            elif cache_key not in unique_misses:
-                if (
-                    start >= chunk_start
-                    and end < chunk_end
-                    and item.can_defer_cuda_ipc_feature_reconstruction()
-                ):
-                    item.model_specific_data[BORROW_CUDA_IPC_FEATURE_KEY] = True
-                unique_misses[cache_key] = (item, expected_token_count)
+            _discard_mismatched_cached_embedding(
+                span.hash, span.span_len, cached_token_count
+            )
+        elif (
+            span.inside_chunk and span.item.can_defer_cuda_ipc_feature_reconstruction()
+        ):
+            span.item.model_specific_data[BORROW_CUDA_IPC_FEATURE_KEY] = True
+        unique_misses[cache_key] = (span.item, span.span_len)

    # Phase 1b: single ViT call for all unique cache misses
    if unique_misses:
@@ -412,6 +394,52 @@ def _batch_encode_per_image_misses(
    return hash_to_embedding


+def _collect_image_span_requests(
+    per_image_requests: List[PerImageRequestInfo],
+) -> List[ImageSpanRequest]:
+    spans: Dict[
+        Tuple[Optional[int], int],
+        Tuple[MultimodalDataItem, bool, List[MultimodalDataItem]],
+    ] = {}
+    for req_info in per_image_requests:
+        chunk_start = req_info.extend_prefix_len
+        chunk_end = chunk_start + req_info.extend_seq_len  # exclusive
+        overlapping = []
+        if req_info.extend_seq_len > 0:
+            for idx, (item, (start, end)) in enumerate(
+                zip(req_info.items, req_info.items_offset)
+            ):
+                if end >= chunk_start and start < chunk_end:
+                    overlapping.append((idx, item, start, end))
+        req_info.overlapping = overlapping
+
+        for _idx, item, start, end in overlapping:
+            cache_key = (item.hash, end - start + 1)
+            if cache_key in spans:
+                spans[cache_key][2].append(item)
+                continue
+            spans[cache_key] = (item, start >= chunk_start and end < chunk_end, [])
+    return [
+        ImageSpanRequest(
+            hash=item_hash,
+            span_len=span_len,
+            item=item,
+            inside_chunk=inside_chunk,
+            duplicates=duplicates,
+        )
+        for (item_hash, span_len), (item, inside_chunk, duplicates) in spans.items()
+    ]
+
+
+def _owner_span_encoder(data_embedding_func: DataEmbeddingFunc, device: torch.device):
+    def encode(items: List[MultimodalDataItem]):
+        if not _can_skip_pre_embed_feature_move(data_embedding_func):
+            _move_items_to_device(items, device)
+        return data_embedding_func(items)
+
+    return encode
+
+
 def _get_chunked_embedding_by_item(
    data_embedding_func: DataEmbeddingFunc,
    embedding_items_per_req: List[MultimodalDataItem],
@@ -537,6 +565,7 @@ def _get_chunked_prefill_embedding(
    extend_length: List[int],
    items_offset_list: List[List[Tuple[int, int]]],
    input_ids: torch.Tensor,
+    mm_owner: Optional[MmOwnerSession] = None,
 ) -> tuple[torch.Tensor | None, torch.Tensor]:
    """
    Chunked prefill embedding: encode items across all requests and extract
@@ -598,7 +627,22 @@ def _get_chunked_prefill_embedding(

    # Phase 1: batch encode all per-image cache misses in ONE ViT call
    hash_to_embedding: Dict[Tuple[Optional[int], int], torch.Tensor] = {}
-    if per_image_requests:
+    if per_image_requests and mm_owner is not None:
+        # The owner protocol must see every overlapping span before any local
+        # cache filtering: a rank-local hit can never skip a group collective.
+        span_requests = _collect_image_span_requests(per_image_requests)
+        if mm_owner.engaged:
+            hash_to_embedding = mm_owner.resolve(
+                span_requests,
+                cache=embedding_cache,
+                encode=_owner_span_encoder(data_embedding_func, device),
+            )
+        elif span_requests:
+            raise RuntimeError(
+                "owner eligibility saw no image span in this chunk, but "
+                f"scheduling found {len(span_requests)}"
+            )
+    elif per_image_requests:
        hash_to_embedding = _batch_encode_per_image_misses(
            data_embedding_func, per_image_requests, device
        )
@@ -701,6 +745,7 @@ def get_embedding_and_mask(
    prefix_length: List[int],
    extend_length: List[int],
    items_offset_list: List[List[Tuple[int, int]]],
+    mm_owner: Optional[MmOwnerSession] = None,
 ) -> Tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor]:
    """
    Generate multimodal embeddings and create a mask for identifying their positions in the input sequence.
@@ -741,6 +786,7 @@ def get_embedding_and_mask(
            extend_length,
            items_offset_list,
            input_ids,
+            mm_owner=mm_owner,
        )
        if embedding is None:
            return None, None, input_ids
@@ -10,6 +10,7 @@ import pickle
 import sys
 from abc import abstractmethod
 from collections import defaultdict
+from contextlib import nullcontext
 from multiprocessing import shared_memory
 from typing import Any, Dict, List, Optional, Tuple

@@ -24,6 +25,7 @@ from sglang.srt.managers.io_struct import (
    TokenizedEmbeddingReqInput,
    TokenizedGenerateReqInput,
 )
+from sglang.srt.managers.mm_owner_embedding import MmOwnerSession

 # Preserve the existing initialization import for downstream callers.
 from sglang.srt.managers.mm_schedule import (
@@ -397,6 +399,7 @@ def embed_mm_inputs(
    data_embedding_func_mapping: Dict[Modality, DataEmbeddingFunc] = None,
    placeholder_tokens: dict[Modality, List[int]] = None,
    use_deepstack: Dict[Modality, bool] = {},
+    mm_owner: Optional[MmOwnerSession] = None,
 ) -> Optional[torch.Tensor]:
    """
    Embed multimodal inputs and integrate them with text token embeddings.
@@ -478,6 +481,7 @@ def embed_mm_inputs(
                prefix_length=extend_prefix_lens,
                extend_length=extend_seq_lens,
                items_offset_list=items_offsets,
+                mm_owner=mm_owner,
            )

            if use_deepstack.get(modality, None) and embedding is not None:
@@ -498,7 +502,12 @@ def embed_mm_inputs(
    # filled with the hash values of the multimodal for the prefix matching in the radix attention.
    # There values are useless because their embeddings will be replaced by vision embeddings anyway.
    input_ids.clamp_(min=0, max=vocab_size - 1)
-    input_embeds = input_embedding(input_ids)
+    if mm_owner is not None:
+        # The text embedding may all-reduce across TP; a rank-local failure in
+        # feature preparation has to be agreed on before any rank enters it.
+        mm_owner.features_ready()
+    with mm_owner.uncaptured() if mm_owner is not None else nullcontext():
+        input_embeds = input_embedding(input_ids)

    # deepstack embedding
    if use_deepstack:
@@ -525,7 +534,9 @@ def embed_mm_inputs(
        _scatter_mm_embedding(dest=input_embeds, mask=mask, src=embedding)
        if use_deepstack.get(modality, None):
            _scatter_mm_embedding(
-                dest=input_deepstack_embeds, mask=mask, src=deepstack_embeddings[i]
+                dest=input_deepstack_embeds,
+                mask=mask,
+                src=deepstack_embeddings[i],
            )

    return input_embeds, other_info
@@ -121,6 +121,11 @@ from sglang.srt.layers.quantization.mxfp8_input import Mxfp8SwizzledInput
 from sglang.srt.layers.rotary_embedding import get_rope_wrapper
 from sglang.srt.layers.utils import PPMissingLayer, get_layer_id
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from sglang.srt.managers.mm_owner_embedding import (
+    MmOwnerSession,
+    has_owner_span_work,
+    select_owner_group,
+)
 from sglang.srt.managers.mm_utils import (
    MultiModalityDataPaddingPatternMultimodalTokens,
    embed_mm_inputs,
@@ -4917,6 +4922,13 @@ class DeepseekV4ForCausalLM(nn.Module):
            self.image_start = nn.Parameter(torch.empty(config.hidden_size))
            self.image_end = nn.Parameter(torch.empty(config.hidden_size))
            self.image_newline = nn.Parameter(torch.empty(config.hidden_size))
+        # Ranks of this group run identical image chunks; one owner encodes each
+        # span and broadcasts it. None keeps the replicated encoder.
+        self.mm_owner_group = (
+            select_owner_group(get_parallel())
+            if self.vision is not None and _is_cuda
+            else None
+        )
        self.model = DeepseekV4Model(
            config, quant_config, prefix=add_prefix("model", prefix)
        )
@@ -5051,7 +5063,42 @@ class DeepseekV4ForCausalLM(nn.Module):
            spans.append(span)
        return spans

-    def _prepare_mm_embeddings(self, input_ids, forward_batch):
+    def _image_span_signature(self, item, span_len: int):
+        h, w = int(item.n_vit_h), int(item.n_vit_w)
+        r = self.config.vision_downsample_ratio
+        expected = len(image_token_types((h + r - 1) // r, (w + r - 1) // r))
+        if expected != span_len:
+            raise ValueError(
+                f"image grid {(h, w)} yields {expected} span tokens, "
+                f"placeholder has {span_len}"
+            )
+        plan = item.model_specific_data.get(GPU_PLAN_KEY)
+        feature = item.feature
+        return (
+            h,
+            w,
+            tuple(feature.shape) if isinstance(feature, torch.Tensor) else None,
+            None if plan is None else tuple(sorted(plan.items())),
+        )
+
+    def _mm_owner_session(self, forward_batch) -> Optional[MmOwnerSession]:
+        if self.mm_owner_group is None:
+            return None
+        return MmOwnerSession(
+            group=self.mm_owner_group,
+            device=self.image_start.device,
+            dtype=self.image_start.dtype,
+            width=self.config.hidden_size,
+            rids=list(forward_batch.rids or ()),
+            signature=self._image_span_signature,
+            engaged=has_owner_span_work(
+                forward_batch.mm_inputs,
+                forward_batch.extend_prefix_lens_cpu,
+                forward_batch.extend_seq_lens_cpu,
+            ),
+        )
+
+    def _prepare_mm_embeddings(self, input_ids, forward_batch, mm_owner):
        # Keep scheduler hash IDs intact: the shared embedder clamps its input in place.
        input_embeds, _ = embed_mm_inputs(
            mm_inputs_list=[
@@ -5063,6 +5110,7 @@ class DeepseekV4ForCausalLM(nn.Module):
            input_ids=input_ids.clone(),
            input_embedding=self.get_input_embeddings(),
            multimodal_model=self,
+            mm_owner=mm_owner,
        )
        forward_batch.mm_input_embeds = input_embeds
        return input_embeds
@@ -5078,24 +5126,31 @@ class DeepseekV4ForCausalLM(nn.Module):
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        if self.vision is None:
            return input_ids, input_embeds
-        if (
+        has_images = (
            not forward_batch.forward_mode.is_decode()
            and not forward_batch.forward_mode.is_target_verify()
            and forward_batch.mm_inputs is not None
            and any(x is not None for x in forward_batch.mm_inputs)
-        ):
-            if input_embeds is not None:
-                raise ValueError("Cannot combine input_embeds and image inputs")
-            input_embeds = self._prepare_mm_embeddings(input_ids, forward_batch)
-        if not (
-            forward_batch.forward_mode.is_decode_or_idle()
-            or forward_batch.forward_mode.is_target_verify()
-        ):
-            # Decode/verify IDs are already vocabulary IDs; remap prompt image
-            # hashes for Engram and routing.
-            input_ids = input_ids.masked_fill(
-                input_ids >= MM_PAD_SHIFT_VALUE, self.config.image_token_id
-            )
+        )
+        if has_images and input_embeds is not None:
+            raise ValueError("Cannot combine input_embeds and image inputs")
+        mm_owner = self._mm_owner_session(forward_batch) if has_images else None
+        # Peers may only enter the body or the CP shard once every rank has
+        # finished all of its fallible input preparation, the remap included.
+        with mm_owner.fence() if mm_owner is not None else nullcontext():
+            if has_images:
+                input_embeds = self._prepare_mm_embeddings(
+                    input_ids, forward_batch, mm_owner
+                )
+            if not (
+                forward_batch.forward_mode.is_decode_or_idle()
+                or forward_batch.forward_mode.is_target_verify()
+            ):
+                # Decode/verify IDs are already vocabulary IDs; remap prompt image
+                # hashes for Engram and routing.
+                input_ids = input_ids.masked_fill(
+                    input_ids >= MM_PAD_SHIFT_VALUE, self.config.image_token_id
+                )
        return input_ids, input_embeds

    def set_dspark_layers_to_capture(self, layer_ids: List[int]) -> None: