lightseekorg--tokenspeed/python/tokenspeed/runtime/pd/epd/prefill_receiver.py

# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""EPD prefill-side receive glue: fill each multimodal item's ``encoded``
tensor from the Mooncake transfer instead of running the vision tower.

Inverse of ``encode_executor.assign_encoded_embeddings``: the encode worker
row-splits the tower output per item and column-splits the deepstack half; here
the prefill allocates a receive buffer sized to each item's post-merge token
count (and deepstack width), registers it with that item's encode worker via
:class:`MooncakeEmbeddingReceiver`, waits for the push, and assigns the result
onto ``item.encoded`` / ``item.encoded_deepstack`` -- the ``skip-ViT`` form the
prefill ``VisionEmbedder`` consumes (it never re-encodes an item whose
``encoded`` is already set). A wrong token count or deepstack width mis-sizes the
RDMA target (silent corruption), so buffer-sizing + assignment is isolated and
CPU-unit-testable with a fake receiver.

Poll-driven state machine (:class:`EmbeddingReceiveJob`): ``start()`` does only
the non-blocking Phase-1 allocation + receiver construction (it must NEVER block
the prefill event loop); ``poll()`` advances every receiver a little each cycle
and publishes ``item.encoded`` only once EVERY handshaked item has landed. The
scheduler holds the job and does not admit the request to a prefill forward until
``poll()`` returns ``DONE``.

``receive_encoded_embeddings`` is a thin BLOCKING wrapper (start + spin poll) for
synchronous callers and the CPU tests.
"""

from __future__ import annotations

import logging
import time
from collections import deque
from collections.abc import Callable, Mapping, Sequence
from typing import Any

import torch
import torch.distributed as dist

logger = logging.getLogger(__name__)

from tokenspeed.runtime.multimodal.embedder import _item_token_count
from tokenspeed.runtime.multimodal.inputs import MultimodalDataItem
from tokenspeed.runtime.pd.base.status import TransferPoll
from tokenspeed.runtime.pd.epd.embedding_transfer import (
    MooncakeEmbeddingReceiver,
)
from tokenspeed.runtime.utils.env import envs

# (manager, bootstrap_addr, bootstrap_room) -> receiver. Defaults to the real
# MooncakeEmbeddingReceiver; overridden in tests with a fake.
ReceiverFactory = Callable[[Any, str, int], Any]

# poll() return values (the per-job lifecycle status, distinct from the
# per-receiver Poll status it aggregates).
PENDING = "pending"
DONE = "done"
FAILED = "failed"

# Deregistering a receive MR the instant its Success notif arrives races the
# NIC's DMA placement tail (Success is the transport ACK, not proof the last
# bytes cleared HCA->PCIe), tripping `local access violation work queue error`.
# So deregistration is DEFERRED: entries hold the tensor ref (the allocator
# cannot reuse a still-registered address -- the no-double-register invariant)
# and are swept after a grace period, on the scheduler loop (no thread, no lock).
_DEREG_DELAY_S = 0.5
_RECV_POOL_QUARANTINE_S = 10.0
# (due_monotonic, engine, [(tensor, ptr), ...]) in due order (delay is constant).
_pending_dereg: deque = deque()


def _record_current_stream_event(tensor: torch.Tensor) -> torch.cuda.Event | None:
    if not tensor.is_cuda:
        return None
    event = torch.cuda.Event()
    torch.cuda.current_stream(tensor.device).record_event(event)
    return event


def _lazy_deregister(engine: Any, tensors: list[tuple[torch.Tensor, int]]) -> None:
    _pending_dereg.append((time.monotonic() + _DEREG_DELAY_S, engine, tensors))
    _sweep_deregister()


def _sweep_deregister() -> None:
    now = time.monotonic()
    while _pending_dereg and _pending_dereg[0][0] <= now:
        _, engine, tensors = _pending_dereg.popleft()
        for _tensor, ptr in tensors:
            try:
                engine.deregister(ptr)
            except Exception:  # noqa: BLE001 -- best-effort; worst case the MR leaks
                pass


class _RecvBufferPool:
    """Pre-registered, lifetime-stable receive slots for E->P embedding lands.

    With per-request buffers, deregistering then re-registering a recycled
    allocator address mints a NEW rkey for the same range while the encode side's
    Mooncake segment cache still resolves it to the OLD rkey -- every reuse risks
    a `local access violation work queue error` that kills the QP. The pool
    registers ONE region for the engine's lifetime, so the sender's cached
    mapping can never go stale; requests lease slots, the publish path clones the
    landed rows out, and CUDA-backed slots return only after that clone completes.

    Failure path: a FAILED job may still have an in-flight remote write targeting
    its slot, which under a single lifetime MR would land SILENTLY in the next
    tenant's data, so failed slots sit in quarantine until the transfer layer's
    timeouts have expired.

    Single-threaded by design: all callers run on the scheduler loop (no locks).
    """

    def __init__(self, engine: Any, device: Any, slot_bytes: int, n_slots: int):
        self.engine = engine
        self.slot_bytes = slot_bytes
        self.buf = torch.empty(n_slots * slot_bytes, dtype=torch.uint8, device=device)
        engine.register(self.buf.data_ptr(), self.buf.numel())
        self._free = list(range(n_slots))
        self._quarantine: deque = deque()  # (release_due_monotonic, slot)
        self._pending_release: deque = deque()  # (cuda_event, slot)

    def _sweep_pending_release(self) -> None:
        kept = deque()
        while self._pending_release:
            event, slot = self._pending_release.popleft()
            try:
                ready = event.query()
            except Exception:  # noqa: BLE001 -- avoid permanently losing a slot
                logger.warning(
                    "EPD recv pool: CUDA event query failed; releasing slot anyway",
                    exc_info=True,
                )
                ready = True
            if ready:
                self._free.append(slot)
            else:
                kept.append((event, slot))
        self._pending_release = kept

    def _sweep_quarantine(self) -> None:
        now = time.monotonic()
        while self._quarantine and self._quarantine[0][0] <= now:
            self._free.append(self._quarantine.popleft()[1])

    def sweep(self) -> None:
        self._sweep_pending_release()
        self._sweep_quarantine()

    def lease(self, nbytes: int) -> int | None:
        self.sweep()
        if nbytes > self.slot_bytes or not self._free:
            return None
        return self._free.pop()

    def view(self, slot: int, nbytes: int, dtype: torch.dtype, shape) -> torch.Tensor:
        off = slot * self.slot_bytes
        return self.buf[off : off + nbytes].view(dtype).reshape(shape)

    def release(self, slot: int) -> None:
        self._free.append(slot)

    def release_after_copy(self, slot: int, copied_tensor: torch.Tensor) -> None:
        event = _record_current_stream_event(copied_tensor)
        if event is None:
            self.release(slot)
        else:
            self._pending_release.append((event, slot))

    def quarantine(self, slot: int, delay_s: float) -> None:
        self._quarantine.append((time.monotonic() + delay_s, slot))


# (id(engine), str(device)) -> _RecvBufferPool | False (False = disabled).
_POOLS: dict = {}


def _get_pool(engine: Any, device: Any) -> _RecvBufferPool | None:
    key = (id(engine), str(device))
    pool = _POOLS.get(key)
    if pool is None:
        n_slots = envs.TOKENSPEED_EPD_RECV_POOL_SLOTS.get()
        slot_mb = envs.TOKENSPEED_EPD_RECV_POOL_SLOT_MB.get()
        if n_slots <= 0 or slot_mb <= 0:
            pool = False
        else:
            pool = _RecvBufferPool(engine, device, slot_mb << 20, n_slots)
            logger.info(
                "EPD recv pool up: %d slots x %d MB (lifetime MR)", n_slots, slot_mb
            )
        _POOLS[key] = pool
    return pool or None


def shard_rows(span: int, shard_rank: int, shard_size: int) -> tuple[int, int]:
    """Balanced contiguous row shard of ``span`` rows across ``shard_size``
    ranks: returns this rank's ``(row_start, row_count)``. The first
    ``span % shard_size`` ranks get one extra row; shards tile ``[0, span)``
    disjointly and in rank order, which BOTH sides of the transfer (the
    receiver's pre_alloc and the post-receive reassembly) must derive from this
    one function so their geometry can never diverge. ``row_count`` may be 0
    when ``span < shard_size`` (tiny images)."""
    base, rem = divmod(span, shard_size)
    start = shard_rank * base + min(shard_rank, rem)
    count = base + (1 if shard_rank < rem else 0)
    return start, count


class _ItemReceive:
    """Per-item receive bookkeeping for one :class:`EmbeddingReceiveJob`.

    Holds the single ``[n_tokens, hidden]`` receive buffer (+ optional deepstack
    column buffer) and one receiver per owned image, each targeting a contiguous
    row sub-range of that buffer. The buffers are filled in place by the encode
    side's RDMA writes; this object publishes them onto the item only once every
    one of its receivers reaches ``Success`` (handled by the owning job).
    """

    __slots__ = (
        "item",
        "recv_main",
        "recv_deepstack",
        "receivers",
        "pre_alloced",
        "main_slice_ptrs",
        "deepstack_slice_ptrs",
        "spans",
        "row_starts",
        "row_counts",
        "sharded",
        "n_tokens",
        "hidden",
        "pool",
        "pool_slot",
    )

    def __init__(
        self,
        item: MultimodalDataItem,
        recv_main: torch.Tensor,
        recv_deepstack: torch.Tensor | None,
        receivers: list[Any],
        main_slice_ptrs: list[int],
        deepstack_slice_ptrs: list[int],
        spans: list[int],
        row_starts: list[int],
        row_counts: list[int],
        sharded: list[bool],
        n_tokens: int,
        hidden: int,
        pool: _RecvBufferPool | None = None,
        pool_slot: int | None = None,
    ):
        self.item = item
        self.recv_main = recv_main
        self.recv_deepstack = recv_deepstack
        self.receivers = receivers
        self.pool = pool
        self.pool_slot = pool_slot
        # Lazy-pre_alloc latch, one per receiver (index-aligned): each receiver's
        # pre_alloc must be issued exactly once, AFTER it reports Bootstrapped, and
        # never again (a second pre_alloc would double-register the transfer).
        self.pre_alloced = [False] * len(receivers)
        self.main_slice_ptrs = main_slice_ptrs
        self.deepstack_slice_ptrs = deepstack_slice_ptrs
        self.spans = spans
        # Per-image shard geometry (index-aligned with receivers/spans): this
        # rank's row sub-range WITHIN the image, and whether the image was
        # sharded at all (identity images need no reassembly broadcast).
        self.row_starts = row_starts
        self.row_counts = row_counts
        self.sharded = sharded
        self.n_tokens = n_tokens
        self.hidden = hidden


class EmbeddingReceiveJob:
    """Poll-driven, non-blocking EPD embedding receive for ONE request.

    Usage (driven by the scheduler/event loop):

        job = start_embedding_receive(items, manager, ...)
        # ... each cycle, until DONE/FAILED:
        status = job.poll()        # PENDING | DONE | FAILED

    ``start()`` (the constructor) runs Phase-1 only: it sizes + allocates each
    handshaked item's receive buffer, registers it with the Mooncake engine, and
    constructs one :class:`MooncakeEmbeddingReceiver` on the item's room (the
    receiver's own ``__init__`` performs the bootstrap handshake). It does NOT
    wait for any transfer and does NOT call ``pre_alloc`` -- that is deferred to
    ``poll()``, which issues each ``pre_alloc`` lazily once its receiver reports
    Bootstrapped and then waits (across cycles, never blocking) for Success.

    HARD CONSTRAINT: one room/receiver PER ITEM. A single request's
    items may be served by DIFFERENT encode workers, so collapsing to one
    receiver per request is not possible. ``poll()`` returns DONE only when every
    item's receiver is Success.

    Buffer lifetime: by default the receive target is a leased slot from the
    lifetime-registered :class:`_RecvBufferPool`; on DONE the landed rows are
    cloned onto ``item.encoded`` and the slot is reused after that copy completes
    (no MR churn, see the pool docstring). Deepstack models, oversized items, pool
    exhaustion, or ``TOKENSPEED_EPD_RECV_POOL_SLOTS=0`` fall back to a per-request
    buffer registered on start and lazily deregistered after publish. Pooled
    slots stay leased until the CUDA copy into ``item.encoded`` has completed,
    so a following request cannot overwrite rows still being cloned. The GPU cost
    per request is roughly ``n_tokens * hidden * dtype.itemsize`` (plus ``* (1 +
    num_deepstack)`` with deepstack); the caller should cap the number of
    in-flight jobs accordingly.

    Idempotent re-`start`/poll: items whose ``item.encoded`` is already set
    (chunked prefill re-runs the receive per Path-4 forward on the same item) are
    SKIPPED at start time -- no receiver is constructed for them.
    """

    def __init__(
        self,
        items: Sequence[MultimodalDataItem],
        manager: Any,
        *,
        hidden: int,
        num_deepstack: int,
        dtype: torch.dtype,
        device: torch.device | str,
        receiver_factory: ReceiverFactory | None = None,
        shard_rank: int = 0,
        shard_size: int = 1,
    ):
        self.manager = manager
        self.hidden = hidden
        self.num_deepstack = num_deepstack
        # Row sharding across the attn-TP group: with shard_size > 1 this rank
        # registers only its shard_rows() sub-range of each image; after the
        # rank-agreed DONE the caller MUST run reassemble() (rank-lockstep) to
        # rebuild the full rows before any forward consumes item.encoded.
        # shard_size <= 1 is a plain full copy (no reassemble).
        self._shard_rank = shard_rank
        self._shard_size = shard_size
        # torch.dtype is used for buffer allocation (no dtype lost in a string
        # round-trip); the string is only the encode-side wire contract in pre_alloc.
        self.dtype = dtype
        self.dtype_str = str(dtype)
        self._factory: ReceiverFactory = receiver_factory or MooncakeEmbeddingReceiver
        # Terminal latch: once DONE/FAILED, poll() is a cheap no-op returning it
        # (the receivers/buffers have been torn down).
        self._status: str = PENDING

        # Phase 1 (non-blocking): for each item carrying an encode handshake
        # (``item.encode_handshake``), allocate ONE [n_tokens, hidden] buffer
        # (+ deepstack columns), register it with the Mooncake engine, and build a
        # single receiver on the item's room. The handshake lives ON the item: the
        # gateway mints one room per item and the encode worker row-splits the
        # concatenated-subgrid embedding per item -- so one item == one room == one
        # embedding.
        self._items: list[_ItemReceive] = []
        for item in items:
            handshake = getattr(item, "encode_handshake", None)
            if handshake is None:
                # No EPD-routed embedding on this item; leave it for the tower.
                continue

            # Chunked prefill re-runs the receive per forward on the SAME item;
            # ``encoded`` is set only after Success, so ``encoded is not None``
            # means fully received -- skip it (re-bootstrapping a Success room
            # would never re-report Bootstrapped and would stick). Mirrors
            # ``VisionEmbedder``, which never re-encodes an encoded item.
            if item.encoded is not None:
                continue

            self._items.append(self._start_item(item, handshake, device))

        # Nothing to receive (text-only / all-already-encoded / no EPD item):
        # the job is immediately DONE so the scheduler admits the request at once.
        if not self._items:
            self._status = DONE

    def _start_item(
        self,
        item: MultimodalDataItem,
        handshake: Mapping[str, Any],
        device: torch.device | str,
    ) -> _ItemReceive:
        """Allocate + register one item's receive buffer and construct its single
        receiver on the item's room, recording the (possibly sharded) destination
        row range so ``poll()`` can lazily issue ``pre_alloc`` once the receiver
        bootstraps. Does NOT block on any transfer (no ``_wait``, no ``pre_alloc``
        here).

        One room per item: the encode worker concatenates the item's subgrid
        tokens into a single ``[n_tokens, hidden]`` embedding and row-splits it per
        prefill TP rank, so the receive geometry spans the item's FULL token count
        (sum of its offset spans), not per offset.
        """
        elt = torch.empty(0, dtype=self.dtype).element_size()

        # The item's image spans all its concatenated subgrid tokens.
        span = _item_token_count(item)
        addr = f"{handshake['bootstrap_host']}:{handshake['bootstrap_port']}"
        # The receiver __init__ performs the bootstrap handshake (HTTP /route fetch
        # + endpoint registration); on return its poll() is already Bootstrapped or
        # Failed. We do NOT pre_alloc here -- poll() does it once Bootstrapped.
        receiver = self._factory(self.manager, addr, int(handshake["bootstrap_room"]))

        # Row sharding across the attn-TP group: this rank registers only its
        # shard_rows() sub-range, or the full image when sharding is off.
        is_sharded = self._shard_size > 1
        if is_sharded:
            row_start, row_count = shard_rows(span, self._shard_rank, self._shard_size)
        else:
            row_start, row_count = 0, span
        # Length-1 per-image lists so poll()/_packed_to_full()/reassemble() iterate
        # unchanged (one item == one image under per-item rooms).
        receivers: list[Any] = [receiver]
        spans: list[int] = [span]
        row_starts: list[int] = [row_start]
        row_counts: list[int] = [row_count]
        sharded: list[bool] = [is_sharded]

        # The RECEIVE buffer holds only THIS rank's shard rows, packed contiguously
        # (image i at rows [packed_cursor, packed_cursor + row_count)); the FULL
        # embedding is rebuilt later in item.encoded (publish scatter + reassemble),
        # which needs no MR. Registering just the shard rows -- not the full image
        # -- keeps a big multi-image item from overflowing a recv pool slot onto
        # the GIL-held per-request register fallback. With nothing sharded,
        # packed_tokens == full_tokens.
        full_tokens = _item_token_count(item)
        packed_tokens = sum(row_counts)
        nbytes = packed_tokens * self.hidden * elt
        pool = None
        pool_slot = None
        recv_main = None
        if self.num_deepstack == 0:
            # Preferred path: lease a lifetime-registered pool slot (see
            # _RecvBufferPool for why per-request register/deregister kills QPs).
            # Deepstack models keep the legacy path (a second buffer per item,
            # out of the pool's scope).
            pool = _get_pool(self.manager.engine, device)
            if pool is not None:
                pool_slot = pool.lease(nbytes)
                if pool_slot is not None:
                    recv_main = pool.view(
                        pool_slot, nbytes, self.dtype, (packed_tokens, self.hidden)
                    )
                else:
                    logger.info(
                        "EPD recv pool: no slot for %d B (free=%d); falling back "
                        "to per-request registration",
                        nbytes,
                        len(pool._free),
                    )
                    pool = None
        if recv_main is None:
            recv_main = torch.empty(
                (packed_tokens, self.hidden), dtype=self.dtype, device=device
            )
            # Legacy fallback (pool disabled/exhausted, oversized item, or
            # deepstack present): register per request; the deregister is DEFERRED
            # on the publish/fail paths (lazy queue) to soften -- not eliminate --
            # the stale-rkey window. A 0-row packed buffer has nothing to register.
            if packed_tokens > 0:
                self.manager.engine.register(
                    recv_main.data_ptr(),
                    recv_main.numel() * recv_main.element_size(),
                )
        recv_deepstack: torch.Tensor | None = None
        if self.num_deepstack > 0:
            recv_deepstack = torch.empty(
                (packed_tokens, self.hidden * self.num_deepstack),
                dtype=self.dtype,
                device=device,
            )
            if packed_tokens > 0:
                self.manager.engine.register(
                    recv_deepstack.data_ptr(),
                    recv_deepstack.numel() * recv_deepstack.element_size(),
                )

        # Destination pointers into the PACKED buffer: the encode side writes its
        # n_tokens(=row_count) source rows CONTIGUOUSLY at this pointer (it does
        # not re-apply row_start -- that only selects which SOURCE rows to read).
        # A 0-row shard still records its (empty) pointer; its pre_alloc is the
        # registration heartbeat the encode fanout gate counts.
        main_slice_ptrs: list[int] = []
        deepstack_slice_ptrs: list[int] = []
        packed_cursor = 0
        for row_count in row_counts:
            main_slice_ptrs.append(
                recv_main[packed_cursor : packed_cursor + row_count].data_ptr()
            )
            if recv_deepstack is not None:
                deepstack_slice_ptrs.append(
                    recv_deepstack[packed_cursor : packed_cursor + row_count].data_ptr()
                )
            else:
                deepstack_slice_ptrs.append(0)
            packed_cursor += row_count

        return _ItemReceive(
            item=item,
            recv_main=recv_main,
            recv_deepstack=recv_deepstack,
            receivers=receivers,
            main_slice_ptrs=main_slice_ptrs,
            deepstack_slice_ptrs=deepstack_slice_ptrs,
            spans=spans,
            row_starts=row_starts,
            row_counts=row_counts,
            sharded=sharded,
            n_tokens=full_tokens,
            hidden=self.hidden,
            pool=pool,
            pool_slot=pool_slot,
        )

    def poll(self) -> str:
        """Advance the receive state machine one cheap step.

        For every still-pending receiver: if it just reached Bootstrapped, issue
        its (single) ``pre_alloc`` so the encode side learns where to write; if it
        reached Failed -> the whole job is FAILED; otherwise leave it. When EVERY
        receiver of EVERY item is Success, deregister the buffers, publish
        ``item.encoded`` / ``item.encoded_deepstack``, and return DONE.

        Returns ``PENDING`` | ``DONE`` | ``FAILED``. Idempotent after a terminal
        result (the buffers/receivers are gone, so it just returns the latch).
        """
        _sweep_deregister()
        if self._status is not PENDING:
            return self._status

        all_success = True
        for it in self._items:
            for idx, receiver in enumerate(it.receivers):
                status = receiver.poll()
                if status == TransferPoll.Failed:
                    self._fail()
                    return FAILED
                if not it.pre_alloced[idx] and status in (
                    TransferPoll.Bootstrapped,
                    TransferPoll.Transferring,
                    TransferPoll.Success,
                ):
                    # Bootstrapped (or already further along on a fast/in-process
                    # transport): issue the one-shot pre_alloc and latch it -- a
                    # second would double-register the transfer on the encode side.
                    # In shard mode n_tokens carries this rank's row COUNT and the
                    # dst pointers are already offset to the shard's first row; a
                    # row_count of 0 is still sent (it doubles as the encode-side
                    # fanout-gate heartbeat).
                    receiver.pre_alloc(
                        dst_embedding_ptr=it.main_slice_ptrs[idx],
                        n_tokens=it.row_counts[idx],
                        hidden=it.hidden,
                        dtype=self.dtype_str,
                        dst_deepstack_ptr=it.deepstack_slice_ptrs[idx],
                        has_deepstack=self.num_deepstack > 0,
                        row_start=it.row_starts[idx],
                        span=it.spans[idx],
                    )
                    it.pre_alloced[idx] = True
                    # Re-poll once after pre_alloc: an in-process/fake transport
                    # may flip straight to Success on pre_alloc.
                    status = receiver.poll()
                    if status == TransferPoll.Failed:
                        self._fail()
                        return FAILED
                if status != TransferPoll.Success:
                    all_success = False

        if not all_success:
            return PENDING

        # Every image of every item has landed; publish and reclaim. Pooled
        # path: clone the landed rows OUT of the slot (item.encoded must outlive
        # the lease) and release the slot only after that copy completes. Legacy
        # path: hand the buffer itself to item.encoded and queue the registration
        # for DEFERRED drop (the lazy entry holds the tensor ref, so the allocator
        # cannot recycle a still-registered address).
        for it in self._items:
            any_sharded = any(it.sharded)
            if it.pool is not None:
                # Pooled path is deepstack-free. Copy the rows OUT of the slot
                # before releasing it. When sharded the buffer is packed (this
                # rank's rows only), so scatter into a full-layout tensor at each
                # image's absolute offset (reassemble fills the rest); when not
                # sharded the buffer is already full -> clone it.
                it.item.encoded = (
                    self._packed_to_full(it, it.recv_main, self.hidden)
                    if any_sharded
                    else it.recv_main.clone()
                )
                it.item.encoded_deepstack = None
                it.pool.release_after_copy(it.pool_slot, it.item.encoded)
            else:
                # Legacy path. When sharded, the packed buffers are smaller than
                # the image and cannot alias item.encoded, so build separate full
                # tensors; when not sharded, hand the (full) buffers over directly.
                if any_sharded:
                    it.item.encoded = self._packed_to_full(
                        it, it.recv_main, self.hidden
                    )
                    it.item.encoded_deepstack = (
                        self._packed_to_full(
                            it, it.recv_deepstack, self.hidden * self.num_deepstack
                        )
                        if it.recv_deepstack is not None
                        else None
                    )
                else:
                    it.item.encoded = it.recv_main
                    it.item.encoded_deepstack = it.recv_deepstack
                pairs = [(it.recv_main, it.recv_main.data_ptr())]
                if it.recv_deepstack is not None:
                    pairs.append((it.recv_deepstack, it.recv_deepstack.data_ptr()))
                _lazy_deregister(self.manager.engine, pairs)
        # Receive concluded: release each room's bookkeeping from the prefill manager.
        self._clear_receivers()
        self._status = DONE
        return DONE

    def _packed_to_full(
        self, it: "_ItemReceive", packed: torch.Tensor, width: int
    ) -> torch.Tensor:
        """Scatter a PACKED shard buffer (this rank's rows, image-contiguous) into
        a full-layout ``[n_tokens, width]`` tensor, placing each image's shard at
        its absolute row offset and leaving the non-owned rows for ``reassemble``
        to fill. The packed buffer holds ``sum(row_counts)`` rows; the full tensor
        holds ``sum(spans)`` rows (== n_tokens). Identity images (row_count==span,
        row_start==0) copy whole, so a buffer with no real sharding round-trips
        unchanged."""
        full = torch.empty(
            (it.n_tokens, width), dtype=packed.dtype, device=packed.device
        )
        packed_cursor = 0
        full_cursor = 0
        for span, row_start, row_count in zip(it.spans, it.row_starts, it.row_counts):
            if row_count > 0:
                full[full_cursor + row_start : full_cursor + row_start + row_count] = (
                    packed[packed_cursor : packed_cursor + row_count]
                )
            packed_cursor += row_count
            full_cursor += span
        return full

    def reassemble(self, nccl_group: Any, group_ranks: Sequence[int]) -> None:
        """Rebuild full embeddings from row shards via per-image broadcasts.

        Must be called RANK-LOCKSTEP on every attn-TP rank, only after the
        rank-agreed DONE (the drain's MIN all-reduce) and BEFORE any forward
        consumes ``item.encoded``: until then each rank's buffer holds only its
        own shard rows, the rest is uninitialized memory. All ranks iterate the
        identical items/images in identical order issuing identical collectives,
        which also requires the caller to run on the NON-overlap event loop (a
        second thread launching forward collectives concurrently would break the
        cross-rank launch-order guarantee NCCL needs across communicators).

        ``group_ranks`` is the attn-TP group as GLOBAL ranks
        (``mapping.attn.tp_group``): ``dist.broadcast`` takes a global src rank,
        and group rank p == global rank p only in the no-DP single-engine case.
        Per-image sub-range broadcasts (2 x shard_size x n_images launches per
        request).

        No-op for identity images and when sharding is off; safe after _fail
        (items were dropped). Idempotence is NOT required: called exactly once
        per admitted request by the drain.

        Broadcasts target ``item.encoded`` (the PUBLISHED tensor), never
        ``recv_main``: on the pooled path the publish step cloned the rows out
        and queued the slot for reuse after the clone completes, so ``recv_main``
        may later belong to the next tenant; on the legacy path ``item.encoded``
        IS ``recv_main``, so the two are equivalent there.
        """
        if self._shard_size <= 1 or self._status is not DONE:
            return
        for it in self._items:
            main = it.item.encoded
            deep = it.item.encoded_deepstack
            cursor = 0
            for idx, span in enumerate(it.spans):
                if not it.sharded[idx]:
                    cursor += span
                    continue
                for p in range(self._shard_size):
                    start, count = shard_rows(span, p, self._shard_size)
                    if count == 0:
                        continue
                    src = group_ranks[p]
                    dist.broadcast(
                        main[cursor + start : cursor + start + count],
                        src=src,
                        group=nccl_group,
                    )
                    if deep is not None:
                        dist.broadcast(
                            deep[cursor + start : cursor + start + count],
                            src=src,
                            group=nccl_group,
                        )
                cursor += span

    def _clear_receivers(self) -> None:
        """Release each receiver's per-room manager bookkeeping (terminal paths only;
        no-op for fake receivers without clear())."""
        for it in self._items:
            for receiver in it.receivers:
                clear = getattr(receiver, "clear", None)
                if clear is not None:
                    clear()

    def _fail(self) -> None:
        """Tear down on failure and drop our references. A SIBLING image's
        write may still be in flight into the item buffer, so reclamation is
        deferred on both paths: pooled slots go to quarantine (under the
        lifetime MR a late write would land SILENTLY in the next tenant's
        data), legacy buffers go to the lazy-deregistration queue.
        ``item.encoded`` is left unset -- the request is being failed, not
        served.
        """
        for it in self._items:
            if it.pool is not None:
                it.pool.quarantine(it.pool_slot, _RECV_POOL_QUARANTINE_S)
            else:
                pairs = [(it.recv_main, it.recv_main.data_ptr())]
                if it.recv_deepstack is not None:
                    pairs.append((it.recv_deepstack, it.recv_deepstack.data_ptr()))
                _lazy_deregister(self.manager.engine, pairs)
        self._clear_receivers()
        self._items = []
        self._status = FAILED

    @property
    def status(self) -> str:
        return self._status

    def release(self) -> None:
        """Best-effort teardown for an externally-driven abort.

        Used when a rank-agreed FAILED is reached but THIS rank had not yet seen a
        local Failed poll (so its buffers are still registered): free/deregister
        them so a reused allocator address is never left double-registered.
        Idempotent and a no-op once terminal (poll() already tore a DONE/FAILED job
        down). After release() the job is FAILED.
        """
        if self._status is not PENDING:
            return
        self._fail()


def start_embedding_receive(
    items: Sequence[MultimodalDataItem],
    manager: Any,
    *,
    hidden: int,
    num_deepstack: int,
    dtype: torch.dtype,
    device: torch.device | str,
    receiver_factory: ReceiverFactory | None = None,
    shard_rank: int = 0,
    shard_size: int = 1,
) -> EmbeddingReceiveJob:
    """Begin (non-blocking) the per-item EPD embedding receive for one request.

    See :class:`EmbeddingReceiveJob`. The handshake for each EPD-routed item
    rides on ``item.encode_handshake`` (a dict ``{bootstrap_room, bootstrap_host,
    bootstrap_port}``); items without one are left to the vision tower. Returns a
    job whose ``poll()`` the caller drives every cycle until DONE/FAILED; the
    request must not be admitted to a prefill forward before then. With
    ``shard_size > 1`` the caller must also run ``job.reassemble()`` rank-
    lockstep after the rank-agreed DONE, before the request's first forward.
    """
    return EmbeddingReceiveJob(
        items,
        manager,
        hidden=hidden,
        num_deepstack=num_deepstack,
        dtype=dtype,
        device=device,
        receiver_factory=receiver_factory,
        shard_rank=shard_rank,
        shard_size=shard_size,
    )


def receive_encoded_embeddings(
    items: Sequence[MultimodalDataItem],
    manager: Any,
    *,
    hidden: int,
    num_deepstack: int,
    dtype: torch.dtype,
    device: torch.device | str,
    timeout: float = 60.0,
    receiver_factory: ReceiverFactory | None = None,
) -> None:
    """BLOCKING wrapper: fill ``item.encoded`` from the per-item transfers.

    For synchronous callers and the CPU unit tests: the poll-driven
    :class:`EmbeddingReceiveJob` spun to completion in place (start, then
    busy-poll until DONE, raising on FAILED or timeout). Event-loop code should
    drive ``poll()`` once per cycle instead.

    The handshake for each EPD-routed item rides on ``item.encode_handshake``
    (``{bootstrap_room, bootstrap_host, bootstrap_port}``). ``hidden`` is the main
    embedding width, ``num_deepstack`` the deepstack level count (0 if absent),
    and ``dtype`` must match the encode worker's tower output dtype (asserted on
    the encode side before the unchecked RDMA write).
    """
    job = start_embedding_receive(
        items,
        manager,
        hidden=hidden,
        num_deepstack=num_deepstack,
        dtype=dtype,
        device=device,
        receiver_factory=receiver_factory,
    )
    deadline = time.monotonic() + timeout
    while True:
        status = job.poll()
        if status == DONE:
            return
        if status == FAILED:
            raise RuntimeError("encode->prefill embedding transfer failed")
        if time.monotonic() > deadline:
            raise TimeoutError(
                f"embedding receive timed out after {timeout:.1f}s "
                f"(job still {status})"
            )
        time.sleep(0.0005)


def build_prefill_embedding_manager(server_args, global_rank, is_multimodal_active):
    """Stand up the EPD encode->prefill embedding sink, if applicable.

    Only a multimodal *prefill* node receives embeddings from encode workers. Built
    as one TP group (the embedding transport allows prefill_tp = any multiple of
    encode_tp) with no fixed base buffer (receive buffers are allocated per image
    at receive time). Construction spawns a daemon status thread, so build it
    exactly once per rank. Returns None for decode/encode/text-only nodes.
    """
    if server_args.disaggregation_mode != "prefill" or not is_multimodal_active:
        return None

    from tokenspeed.runtime.pd.epd.embedding_transfer import (
        MooncakeEmbeddingManagerPrefill,
    )
    from tokenspeed.runtime.pd.epd.entities import EmbeddingArgs, EmbeddingManagerArgs

    emb_mgr_args = EmbeddingManagerArgs(
        bootstrap_port=server_args.disaggregation_bootstrap_port,
        tp_size=server_args.mapping.attn.tp_size,
    )
    emb_args = EmbeddingArgs(
        engine_rank=global_rank,
        gpu_id=global_rank,
        ib_device=server_args.disaggregation_ib_device,
        embedding_data_ptr=0,
        embedding_data_len=0,
    )
    return MooncakeEmbeddingManagerPrefill(emb_mgr_args, emb_args)


class EpdPrefillAdmission:
    """EPD prefill-side embedding receive + rank-synced admission.

    Owns the encode->prefill embedding sink (MooncakeEmbeddingManagerPrefill), the
    set of requests whose per-image embeddings are still arriving (``_pending``),
    and the optional NCCL row-shard reassembly. Exists ONLY on a multimodal prefill
    node; the EventLoop holds one (or None) and drives it each non-overlap cycle.

    DECIDE/ACT split: ``drain()`` polls the staged receives, applies the rank-
    lockstep MIN all-reduce + reassembly, and RETURNS ``(admitted, failed)``
    decisions. The EventLoop performs the acts those imply (P->D sender
    register/abort, scheduler submit, output-processor finish) -- they touch
    EventLoop collaborators, so they stay there.
    """

    def __init__(
        self,
        *,
        manager,
        device,
        hidden,
        num_deepstack,
        dtype,
        attn_tp_rank,
        attn_tp_size,
        attn_tp_cpu_group,
        attn_tp_group,
        pg_manager,
    ):
        self._manager = manager
        self._device = device
        self._hidden = hidden
        self._num_deepstack = num_deepstack
        self._dtype = dtype
        self._attn_tp_rank = attn_tp_rank
        self._attn_tp_size = attn_tp_size
        self._attn_tp_cpu_group = attn_tp_cpu_group
        # Requests whose per-image encode->prefill embeddings are still arriving:
        # registered but NOT yet submitted to the scheduler; drain() polls them each
        # cycle and admits (rank-synced) only once ready. Rank-identical in
        # length+order across attn-TP ranks (recv_reqs broadcasts the new-request
        # set), which drain()'s MIN all-reduce relies on.
        self._pending: list = []
        # Deadline for an EPD request's per-image embeddings to all arrive; past it
        # the request is aborted (not waited on forever). Reuse the PD KV-receive
        # wait knob (default 300s): the prefill waiting on the encode->prefill
        # embedding transfer is the direct analog of the decode waiting on the
        # prefill->decode KV transfer, so one operator knob covers both.
        self._embed_timeout: float = float(
            envs.TOKENSPEED_DISAGGREGATION_WAITING_TIMEOUT.get()
        )

        # EPD embedding row-sharding: each attn-TP rank receives only 1/N of every
        # image's rows over the wire; the full embedding is rebuilt by an NCCL
        # all-gather in drain() (job.reassemble), which runs on the prefill's
        # NON-overlap loop so the drain and the forward launch from one thread in
        # the same order on every rank -- the cross-rank launch-order consistency
        # NCCL requires across communicators.
        self._shard_embeddings = False
        self._nccl_group = None
        self._group_ranks = tuple(attn_tp_group)
        shard_flag = False
        if attn_tp_size > 1:
            shard_flag = bool(envs.TOKENSPEED_EPD_EMBEDDING_SHARD.get())
            # The flag is a per-process env read but gates a GROUP collective:
            # torn across ranks (e.g. set on one node of a multi-node prefill),
            # flag-on ranks would join the warmup broadcast below while flag-off
            # ranks never do -- a silent boot hang. Agree first, loud.
            flag_t = torch.tensor(
                [int(shard_flag), -int(shard_flag)], dtype=torch.int32
            )
            dist.all_reduce(flag_t, op=dist.ReduceOp.MIN, group=attn_tp_cpu_group)
            if flag_t[0].item() != -flag_t[1].item():
                raise RuntimeError(
                    "TOKENSPEED_EPD_EMBEDDING_SHARD differs across attn-TP ranks; "
                    "set it identically on every node of the prefill engine"
                )
        if shard_flag:
            self._nccl_group = pg_manager.get_process_group("nccl", attn_tp_group)
            self._shard_embeddings = True
            # Warm the communicator at startup: NCCL initializes lazily on the
            # first collective, and nothing else issues torch.distributed NCCL ops
            # on this group in the scheduler process -- without this the FIRST
            # admitted EPD request pays communicator init (hundreds of ms, all
            # ranks) inside the drain, and a misconfigured group would surface on a
            # customer request instead of at boot.
            warmup = torch.zeros(1, device=device)
            dist.broadcast(warmup, src=self._group_ranks[0], group=self._nccl_group)
            torch.cuda.current_stream().synchronize()
            logger.info(
                "EPD embedding row-sharding enabled (attn_tp=%d, NCCL group warm)",
                attn_tp_size,
            )

    def stage(self, request_id, mm_items) -> None:
        """Begin the non-blocking per-image embedding receive and stage it by
        request_id until its embeddings land (polled in drain()). The request payload
        (spec/state/bootstrap) stays with the caller, keyed by request_id; this
        controller tracks only the receive job (mirrors kv_transfer)."""
        job = start_embedding_receive(
            items=mm_items,
            manager=self._manager,
            hidden=self._hidden,
            num_deepstack=self._num_deepstack,
            dtype=self._dtype,
            device=self._device,
            shard_rank=self._attn_tp_rank,
            shard_size=self._attn_tp_size if self._shard_embeddings else 1,
        )
        self._pending.append((request_id, job, time.time()))

    def has_pending(self) -> bool:
        return bool(self._pending)

    def drain(self):
        """Poll staged EPD embedding receives; return ``(admitted_ids, failed_ids)``.

        poll/timeout/MIN-all-reduce/reassemble/release + ``_pending`` bookkeeping
        happen here, rank-lockstep, keyed by request_id. The caller maps the ids back
        to its staged request payloads and performs the acts (kv_transfer
        register/abort, scheduler submit, output-processor finish).

        - admitted_ids: DONE on every rank, reassembled.
        - failed_ids:   FAILED/timed-out, job released.
        """
        if not self._pending:
            return [], []  # rank-identical emptiness -> all ranks skip the collective

        code_of = {FAILED: 0, PENDING: 1, DONE: 2}
        codes = [code_of[job.poll()] for (_rid, job, _ts) in self._pending]

        # Timeout: a still-PENDING request whose per-image embeddings have not all
        # arrived within the deadline is marked FAILED (-> rank-agreed abort below).
        # Without this the prefill waits FOREVER if an embedding is ever lost (a
        # degraded/dead encode worker, a network drop). Folded into the SAME MIN
        # all-reduce: a timed-out job -> code 0 -> all ranks abort it together;
        # union-of-timeout is rank-safe even if ranks cross the deadline a cycle
        # apart (any rank's 0 propagates via MIN).
        _now = time.time()
        for _i in range(len(self._pending)):
            if codes[_i] == 1 and (_now - self._pending[_i][2]) > self._embed_timeout:
                codes[_i] = 0
                logger.warning(
                    "EPD embedding receive timed out after %.0fs for rid=%s; aborting",
                    self._embed_timeout,
                    self._pending[_i][0],
                )

        if self._attn_tp_size > 1:
            t = torch.tensor(codes, dtype=torch.uint8, device="cpu")
            dist.all_reduce(t, op=dist.ReduceOp.MIN, group=self._attn_tp_cpu_group)
            codes = t.tolist()

        admitted = []
        failed = []
        leftover = []
        for (request_id, job, start_ts), code in zip(self._pending, codes):
            if code == 2:  # DONE on every rank
                # Sharded receive: rebuild the full rows from the per-rank shards
                # FIRST -- item.encoded is shard-only until this runs. Rank-lockstep-
                # safe: codes are identical post-MIN and _pending is rank-identical
                # in length/order, so every rank issues identical collectives in
                # identical order this cycle.
                if self._shard_embeddings:
                    job.reassemble(self._nccl_group, self._group_ranks)
                admitted.append(request_id)
            elif code == 0:  # FAILED/timed-out on some rank -> abort everywhere
                job.release()
                failed.append(request_id)
            else:  # still pending on some rank
                leftover.append((request_id, job, start_ts))
        self._pending = leftover
        return admitted, failed


def make_epd_prefill_admission(
    server_args,
    global_rank,
    *,
    model_config,
    model_executor,
    mapping,
    attn_tp_rank,
    attn_tp_size,
    attn_tp_cpu_group,
    pg_manager,
):
    """Build the EPD-prefill admission controller, or None for non-EPD nodes.

    Returns None unless this is a multimodal *prefill* node (the only node that
    receives encode->prefill embeddings)."""
    manager = build_prefill_embedding_manager(
        server_args, global_rank, model_config.is_multimodal_active
    )
    if manager is None:
        return None
    # Extract the narrow model facts the controller needs (vision dtype, hidden
    # width, deepstack width, device) here -- the controller holds these, not the
    # whole model_executor (mirrors how create_kv_transfer takes a kv_args struct).
    model = model_executor.model_runner.model
    return EpdPrefillAdmission(
        manager=manager,
        device=model_executor.device,
        hidden=model.config.hidden_size,
        num_deepstack=getattr(model, "num_deepstack_embeddings", 0),
        dtype=(getattr(model, "visual", None) or model.vision_tower).dtype,
        attn_tp_rank=attn_tp_rank,
        attn_tp_size=attn_tp_size,
        attn_tp_cpu_group=attn_tp_cpu_group,
        attn_tp_group=mapping.attn.tp_group,
        pg_manager=pg_manager,
    )