Skip to content

vllm.models.deepseek_v41.amd.mono.common.peer_memory

Symmetric peer memory for the in-kernel TP all-reduces, and the agreement every rank must reach before anything collective (a rank that turned the mono layer off alone would leave the others spinning on the GPU).

Classes:

  • MonoUnsupported –

    The loaded model or runtime configuration is outside what mono serves.

  • PeerBuffer –

    One uncached buffer per rank, every rank holding every peer's address

Functions:

  • tp_agree –

    True when every rank of group passed ok true. Every rank must call

MonoUnsupported

Bases: Exception

The loaded model or runtime configuration is outside what mono serves.

Source code in vllm/models/deepseek_v41/amd/mono/common/peer_memory.py
class MonoUnsupported(Exception):
    """The loaded model or runtime configuration is outside what mono serves."""

PeerBuffer

One uncached buffer per rank, every rank holding every peer's address (addresses, int64, indexed by TP rank). What lives where, and why no rank overwrites a peer's unread data, is the model's layout's to say.

The constructor is collective over group (the handle exchange, then an agreement that every rank mapped every peer): every rank must reach it, so a runner reaches it only after tp_agree. A failed mapping on any rank raises MonoUnsupported on all of them.

Methods:

  • close –

    Close the peer mappings and free the local buffer; idempotent. Only once no

Source code in vllm/models/deepseek_v41/amd/mono/common/peer_memory.py
class PeerBuffer:
    """One uncached buffer per rank, every rank holding every peer's address
    (``addresses``, int64, indexed by TP rank). What lives where, and why no rank
    overwrites a peer's unread data, is the model's layout's to say.

    The constructor is collective over ``group`` (the handle exchange, then an
    agreement that every rank mapped every peer): every rank must reach it, so a
    runner reaches it only after ``tp_agree``. A failed mapping on any rank raises
    ``MonoUnsupported`` on all of them.
    """

    def __init__(
        self,
        nbytes: int,
        group,
        rank: int,
        npes: int,
        device: torch.device,
    ):
        self.local = UncachedIpcHeap.alloc_uncached(nbytes)
        # a view, not an owner: close() frees the memory under it
        self.bytes = torch.as_tensor(_DeviceBytes(self.local, nbytes), device=device)
        self.rank = rank
        self._opened: list[int] = []
        addresses = [self.local]
        if npes > 1:
            handles = UncachedIpcHeap.gather_object_list_via_broadcast(
                group, UncachedIpcHeap.get_mem_handle_bytes(self.local)
            )
            addresses, failure = [], None
            try:
                for peer, handle in enumerate(handles):
                    if peer == rank:
                        addresses.append(self.local)
                    else:
                        base = UncachedIpcHeap.open_mem_handle(handle)
                        self._opened.append(base)
                        addresses.append(base)
            except RuntimeError as err:  # every rank must still reach the agreement
                failure = err
            if not tp_agree(failure is None, group):
                self.close()
                raise MonoUnsupported(
                    "peer memory handshake failed on "
                    f"{'this' if failure else 'another'}"
                    " TP rank"
                ) from failure
        self.addresses = torch.tensor(addresses, dtype=torch.int64, device=device)

    def close(self) -> None:
        """Close the peer mappings and free the local buffer; idempotent. Only once no
        kernel of any rank can still touch it (the runner is being dropped)."""
        opened, self._opened = self._opened, []
        for base in opened:
            UncachedIpcHeap.close_mem_handle(base)
        if self.local:
            self.bytes = None
            UncachedIpcHeap.free_device_mem(self.local)
            self.local = 0

close()

Close the peer mappings and free the local buffer; idempotent. Only once no kernel of any rank can still touch it (the runner is being dropped).

Source code in vllm/models/deepseek_v41/amd/mono/common/peer_memory.py
def close(self) -> None:
    """Close the peer mappings and free the local buffer; idempotent. Only once no
    kernel of any rank can still touch it (the runner is being dropped)."""
    opened, self._opened = self._opened, []
    for base in opened:
        UncachedIpcHeap.close_mem_handle(base)
    if self.local:
        self.bytes = None
        UncachedIpcHeap.free_device_mem(self.local)
        self.local = 0

_DeviceBytes

__cuda_array_interface__ of nbytes bytes at a raw device address, for a non-owning torch view.

Source code in vllm/models/deepseek_v41/amd/mono/common/peer_memory.py
class _DeviceBytes:
    """``__cuda_array_interface__`` of ``nbytes`` bytes at a raw device address, for a
    non-owning torch view."""

    def __init__(self, ptr: int, nbytes: int):
        self.__cuda_array_interface__ = {
            "shape": (nbytes,),
            "typestr": "|u1",
            "data": (ptr, False),
            "version": 3,
        }

tp_agree(ok, group)

True when every rank of group passed ok true. Every rank must call it at the same point; group is a CPU (gloo) process group, or None for a single rank.

Source code in vllm/models/deepseek_v41/amd/mono/common/peer_memory.py
def tp_agree(ok: bool, group) -> bool:
    """True when every rank of ``group`` passed ``ok`` true. Every rank must call
    it at the same point; ``group`` is a CPU (gloo) process group, or None for a
    single rank."""
    if group is None:
        return ok
    flag = torch.tensor([int(ok)], dtype=torch.int32)
    dist.all_reduce(flag, op=dist.ReduceOp.MIN, group=group)
    return bool(flag.item())