vllm.distributed.ec_transfer.ec_connector.cpu.worker.descriptor_buffers ¶

class DescriptorBufferPool:
    """Pool of descriptor buffer triples for swap_blocks_batch.

    Each buffer is a `DescriptorBuffers` namedtuple of three 1-D tensors
    (dtype `_PTR_DTYPE`, platform-dependent) of equal length, paired with
    numpy aliases used to fill them. Buffers are recycled across steps; if a
    returned buffer is too small it is discarded and a fresh one allocated.
    """

    def __init__(self) -> None:
        # LIFO stack of idle buffer triples.
        self._pool: list[DescriptorBuffers] = []

    def acquire(self, n: int) -> DescriptorBuffers:
        """Get a buffer triple with capacity >= *n*."""
        if self._pool:
            bufs = self._pool.pop()
            if bufs.src_ptrs.numel() >= n:
                return bufs
        src, dst, sizes = (torch.empty(n, dtype=_PTR_DTYPE) for _ in range(3))
        return DescriptorBuffers(src, dst, sizes, src.numpy(), dst.numpy())

    def release(self, bufs: DescriptorBuffers) -> None:
        """Return a buffer triple to the pool for reuse."""
        self._pool.append(bufs)

class DescriptorBuffers(NamedTuple):
    src_ptrs: torch.Tensor
    dst_ptrs: torch.Tensor
    sizes: torch.Tensor
    # Numpy aliases of src_ptrs/dst_ptrs, written via set_ptrs().
    src_np: np.ndarray
    dst_np: np.ndarray

    def set_ptrs(self, idx: int, src: int, dst: int) -> None:
        """Record the source and destination address of descriptor *idx*.

        TODO(torch>=2.14): drop this indirection and assign the tensors
        directly once the minimum supported torch is 2.14. The numpy detour
        exists only because torch's setitem unpacks the value as a signed
        long long before 2.14 (pytorch#191458), rejecting XPU USM pointers
        >= 2**63, while the two's-complement rewrite those versions accept is
        in turn rejected by uint64 from 2.14 on. Numpy casts against the
        array dtype and so works on either. `sizes` holds byte counts and
        needs no such care.
        """
        self.src_np[idx] = src
        self.dst_np[idx] = dst