Skip to content

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

Reusable pool of (src_ptrs, dst_ptrs, sizes) int64 tensor triples.

Used by ECCPUWorker to batch swap_blocks_batch descriptors without per-step allocation overhead.

Classes:

DescriptorBufferPool

Pool of descriptor buffer triples for swap_blocks_batch.

Each buffer is a DescriptorBuffers namedtuple of three 1-D int64 tensors of equal length. Buffers are recycled across steps; if a returned buffer is too small it is discarded and a fresh one allocated.

Methods:

  • acquire –

    Get a buffer triple with capacity >= n.

  • release –

    Return a buffer triple to the pool for reuse.

Source code in vllm/distributed/ec_transfer/ec_connector/cpu/worker/descriptor_buffers.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
class DescriptorBufferPool:
    """Pool of descriptor buffer triples for swap_blocks_batch.

    Each buffer is a `DescriptorBuffers` namedtuple of three 1-D int64
    tensors of equal length. 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
        return DescriptorBuffers(
            torch.empty(n, dtype=torch.int64),
            torch.empty(n, dtype=torch.int64),
            torch.empty(n, dtype=torch.int64),
        )

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

acquire(n)

Get a buffer triple with capacity >= n.

Source code in vllm/distributed/ec_transfer/ec_connector/cpu/worker/descriptor_buffers.py 
32
33
34
35
36
37
38
39
40
41
42
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
    return DescriptorBuffers(
        torch.empty(n, dtype=torch.int64),
        torch.empty(n, dtype=torch.int64),
        torch.empty(n, dtype=torch.int64),
    )

release(bufs)

Return a buffer triple to the pool for reuse.

Source code in vllm/distributed/ec_transfer/ec_connector/cpu/worker/descriptor_buffers.py 
44
45
46
def release(self, bufs: DescriptorBuffers) -> None:
    """Return a buffer triple to the pool for reuse."""
    self._pool.append(bufs)