vllm.utils.extensible_tensor

Growable device byte buffers backed by GPU virtual memory management.

Classes:

• ExtensibleTensor –

  A device byte buffer that grows without moving its base pointer.

Functions:

• granule_size –

  Commit granule for device_index: driver granularity, floored at 2 MiB.

• uint8_tensor_from_ptr –

  Wrap device memory at ptr as a uint8 tensor; the caller keeps it mapped.

ExtensibleTensor

A device byte buffer that grows without moving its base pointer.

The reservation is divided into contiguous segments whose committed bytes form a prefix of each; num_segments equal segments by default, or the explicit segment_capacities (in bytes, summing to max_num_bytes).

Methods:

• full_view –

  Uint8 view spanning the whole reservation.

• release_physical –

  Drop all physical pages but keep the reservation and its views.

• resize_per_segment_ –

  Grow every segment's committed prefix to the same size.

• resize_segments_ –

  Grow each segment's committed prefix; optionally zero the new bytes.

• segment_view –

  Uint8 tensor over one segment's committed prefix, sized exactly to it.

• segments_backed_by_one_chunk –

  Whether each segment's first sizes[i] bytes lie within a single

Attributes:

• bytes_per_segment (int) –

  Current committed prefix size of each segment in bytes (equal segments).

• num_bytes (int) –

  Current committed size in bytes, summed over all segments.

• num_physical_chunks (int) –

  Number of driver allocations currently mapped into the reservation.

• physical_bytes (int) –

  Physically mapped bytes, including granule rounding.

• segment_capacity_bytes (int) –

  Maximum size of each segment (equal segments).

• segment_offsets (list[int]) –

  Byte offset of each segment within the reservation.

Source code in vllm/utils/extensible_tensor.py

class ExtensibleTensor:
    """A device byte buffer that grows without moving its base pointer.

    The reservation is divided into contiguous segments whose committed bytes
    form a prefix of each; num_segments equal segments by default, or the
    explicit segment_capacities (in bytes, summing to max_num_bytes).
    """

    def __init__(
        self,
        max_num_bytes: int,
        device: torch.device | str | int | None = None,
        num_segments: int = 1,
        segment_capacities: Sequence[int] | None = None,
    ) -> None:
        if max_num_bytes < 0:
            raise ValueError("max_num_bytes must be non-negative.")
        if segment_capacities is None:
            if num_segments < 1:
                raise ValueError(f"num_segments must be positive, got {num_segments}.")
            if max_num_bytes % num_segments != 0:
                raise ValueError(
                    f"max_num_bytes ({max_num_bytes}) must be divisible by "
                    f"num_segments ({num_segments})."
                )
            segment_capacities = [max_num_bytes // num_segments] * num_segments
        elif not segment_capacities or any(c <= 0 for c in segment_capacities):
            raise ValueError("segment_capacities must be non-empty and positive.")
        elif sum(segment_capacities) != max_num_bytes:
            raise ValueError(
                f"segment_capacities sum to {sum(segment_capacities)}, not "
                f"max_num_bytes ({max_num_bytes})."
            )

        if device is None:
            device = torch.accelerator.current_device_index()
        dev = device if isinstance(device, torch.device) else torch.device(device)
        if dev.type != "cuda":
            raise ValueError(f"ExtensibleTensor requires a cuda device, got {dev}.")
        self._device_index: int = (
            dev.index
            if dev.index is not None
            else torch.accelerator.current_device_index()
        )

        self._max_num_bytes: int = max_num_bytes
        self._segment_capacities: list[int] = list(segment_capacities)
        self._segment_offsets: list[int] = list(
            itertools.accumulate([0, *self._segment_capacities[:-1]])
        )
        self._committed: list[int] = [0] * len(self._segment_capacities)
        self._buffer: _VirtualBuffer = _VirtualBuffer(max_num_bytes, self._device_index)

    def full_view(self) -> torch.Tensor:
        """Uint8 view spanning the whole reservation."""
        return uint8_tensor_from_ptr(
            self._buffer.base_ptr, self._max_num_bytes, self._device_index
        )

    def segment_view(self, index: int) -> torch.Tensor:
        """Uint8 tensor over one segment's committed prefix, sized exactly to it."""
        if not 0 <= index < self.num_segments:
            raise IndexError(f"Segment {index} out of range ({self.num_segments}).")
        return uint8_tensor_from_ptr(
            self._buffer.base_ptr + self._segment_offsets[index],
            self._committed[index],
            self._device_index,
        )

    def resize_per_segment_(
        self, bytes_per_segment: int, zero_new: bool = False
    ) -> None:
        """Grow every segment's committed prefix to the same size."""
        self.resize_segments_([bytes_per_segment] * self.num_segments, zero_new)

    def resize_segments_(self, sizes: Sequence[int], zero_new: bool = False) -> None:
        """Grow each segment's committed prefix; optionally zero the new bytes."""
        if len(sizes) != self.num_segments:
            raise ValueError(f"Expected {self.num_segments} sizes, got {len(sizes)}.")
        for old, new, capacity in zip(self._committed, sizes, self._segment_capacities):
            if new < old:
                raise ValueError(
                    f"ExtensibleTensor is grow-only: cannot resize from {old} "
                    f"to {new} bytes per segment."
                )
            if new > capacity:
                raise ValueError(
                    f"Requested {new} bytes per segment exceeds the "
                    f"segment capacity {capacity}."
                )
        grown = [
            (start + old, start + new)
            for start, old, new in zip(self._segment_offsets, self._committed, sizes)
            if new > old
        ]
        for start, end in grown:
            self._buffer.ensure_committed_range(start, end)
        self._committed = list(sizes)
        if zero_new and grown:
            full = self.full_view()
            for start, end in grown:
                full[start:end].zero_()

    def release_physical(self) -> None:
        """Drop all physical pages but keep the reservation and its views."""
        self._buffer.release_physical()
        self._committed = [0] * self.num_segments

    @property
    def num_bytes(self) -> int:
        """Current committed size in bytes, summed over all segments."""
        return sum(self._committed)

    @property
    def bytes_per_segment(self) -> int:
        """Current committed prefix size of each segment in bytes (equal segments)."""
        return self._uniform(self._committed)

    @property
    def num_segments(self) -> int:
        return len(self._segment_capacities)

    @property
    def segment_capacity_bytes(self) -> int:
        """Maximum size of each segment (equal segments)."""
        return self._uniform(self._segment_capacities)

    @property
    def segment_capacities(self) -> list[int]:
        return list(self._segment_capacities)

    @property
    def segment_offsets(self) -> list[int]:
        """Byte offset of each segment within the reservation."""
        return list(self._segment_offsets)

    @staticmethod
    def _uniform(values: list[int]) -> int:
        if any(v != values[0] for v in values):
            raise ValueError(f"Segments differ: {values}.")
        return values[0]

    @property
    def capacity_bytes(self) -> int:
        return self._buffer.reserved_size

    @property
    def physical_bytes(self) -> int:
        """Physically mapped bytes, including granule rounding."""
        return self._buffer.committed_bytes

    @property
    def num_physical_chunks(self) -> int:
        """Number of driver allocations currently mapped into the reservation."""
        return len(self._buffer._handles)

    def segments_backed_by_one_chunk(self, sizes: Sequence[int]) -> list[bool]:
        """Whether each segment's first `sizes[i]` bytes lie within a single
        driver allocation, as RDMA registration of the segment requires."""
        chunks = [(offset, offset + size) for _, offset, size in self._buffer._handles]
        return [
            any(start <= offset and offset + size <= end for start, end in chunks)
            for offset, size in zip(self._segment_offsets, sizes)
        ]

    @property
    def granularity(self) -> int:
        return self._buffer.granularity

    @property
    def base_ptr(self) -> int:
        return self._buffer.base_ptr

    @property
    def device(self) -> torch.device:
        return torch.device("cuda", self._device_index)

    def free(self) -> None:
        self._buffer.free()
        self._committed = [0] * self.num_segments

bytes_per_segment property

Current committed prefix size of each segment in bytes (equal segments).

num_bytes property

Current committed size in bytes, summed over all segments.

num_physical_chunks property

Number of driver allocations currently mapped into the reservation.

physical_bytes property

Physically mapped bytes, including granule rounding.

segment_capacity_bytes property

Maximum size of each segment (equal segments).

segment_offsets property

Byte offset of each segment within the reservation.

full_view()

Uint8 view spanning the whole reservation.

Source code in vllm/utils/extensible_tensor.py

def full_view(self) -> torch.Tensor:
    """Uint8 view spanning the whole reservation."""
    return uint8_tensor_from_ptr(
        self._buffer.base_ptr, self._max_num_bytes, self._device_index
    )

release_physical()

Drop all physical pages but keep the reservation and its views.

Source code in vllm/utils/extensible_tensor.py

def release_physical(self) -> None:
    """Drop all physical pages but keep the reservation and its views."""
    self._buffer.release_physical()
    self._committed = [0] * self.num_segments

resize_per_segment_(bytes_per_segment, zero_new=False)

Grow every segment's committed prefix to the same size.

Source code in vllm/utils/extensible_tensor.py

def resize_per_segment_(
    self, bytes_per_segment: int, zero_new: bool = False
) -> None:
    """Grow every segment's committed prefix to the same size."""
    self.resize_segments_([bytes_per_segment] * self.num_segments, zero_new)

resize_segments_(sizes, zero_new=False)

Grow each segment's committed prefix; optionally zero the new bytes.

Source code in vllm/utils/extensible_tensor.py

def resize_segments_(self, sizes: Sequence[int], zero_new: bool = False) -> None:
    """Grow each segment's committed prefix; optionally zero the new bytes."""
    if len(sizes) != self.num_segments:
        raise ValueError(f"Expected {self.num_segments} sizes, got {len(sizes)}.")
    for old, new, capacity in zip(self._committed, sizes, self._segment_capacities):
        if new < old:
            raise ValueError(
                f"ExtensibleTensor is grow-only: cannot resize from {old} "
                f"to {new} bytes per segment."
            )
        if new > capacity:
            raise ValueError(
                f"Requested {new} bytes per segment exceeds the "
                f"segment capacity {capacity}."
            )
    grown = [
        (start + old, start + new)
        for start, old, new in zip(self._segment_offsets, self._committed, sizes)
        if new > old
    ]
    for start, end in grown:
        self._buffer.ensure_committed_range(start, end)
    self._committed = list(sizes)
    if zero_new and grown:
        full = self.full_view()
        for start, end in grown:
            full[start:end].zero_()

segment_view(index)

Uint8 tensor over one segment's committed prefix, sized exactly to it.

Source code in vllm/utils/extensible_tensor.py

def segment_view(self, index: int) -> torch.Tensor:
    """Uint8 tensor over one segment's committed prefix, sized exactly to it."""
    if not 0 <= index < self.num_segments:
        raise IndexError(f"Segment {index} out of range ({self.num_segments}).")
    return uint8_tensor_from_ptr(
        self._buffer.base_ptr + self._segment_offsets[index],
        self._committed[index],
        self._device_index,
    )

segments_backed_by_one_chunk(sizes)

Whether each segment's first sizes[i] bytes lie within a single driver allocation, as RDMA registration of the segment requires.

Source code in vllm/utils/extensible_tensor.py

def segments_backed_by_one_chunk(self, sizes: Sequence[int]) -> list[bool]:
    """Whether each segment's first `sizes[i]` bytes lie within a single
    driver allocation, as RDMA registration of the segment requires."""
    chunks = [(offset, offset + size) for _, offset, size in self._buffer._handles]
    return [
        any(start <= offset and offset + size <= end for start, end in chunks)
        for offset, size in zip(self._segment_offsets, sizes)
    ]

_VirtualBuffer

One device VA reservation and the physical chunks mapped into it.

Memory is committed in granules via ensure_committed_range; already-mapped granules are skipped, so ranges may abut or overlap.

Methods:

• ensure_committed_range –

  Map physical pages so that the byte range [start, end) is backed.

• release_physical –

  Unmap and release all physical memory, keeping the VA reservation.

Attributes:

• committed_bytes (int) –

  Total physically mapped bytes (a multiple of the granularity).

Source code in vllm/utils/extensible_tensor.py

class _VirtualBuffer:
    """One device VA reservation and the physical chunks mapped into it.

    Memory is committed in granules via `ensure_committed_range`; already-mapped
    granules are skipped, so ranges may abut or overlap.
    """

    def __init__(self, max_bytes: int, device_index: int) -> None:
        self._driver = get_vmm_driver()
        self._driver.ensure_context(device_index)
        self.device_index = device_index

        self.granularity: int = granule_size(device_index)
        self.reserved_size: int = _round_up(max(max_bytes, 1), self.granularity)
        self.base_ptr: int = self._driver.reserve(self.reserved_size)

        # Granule indices (VA offset // granularity) that have physical
        # memory mapped.
        self._mapped_granules: set[int] = set()
        # Each entry is (handle, va_offset, size) for one mapped physical chunk.
        self._handles: list[tuple[int, int, int]] = []
        self._freed: bool = False

    @property
    def committed_bytes(self) -> int:
        """Total physically mapped bytes (a multiple of the granularity)."""
        return len(self._mapped_granules) * self.granularity

    def ensure_committed_range(self, start: int, end: int) -> None:
        """Map physical pages so that the byte range `[start, end)` is backed."""
        if not 0 <= start <= end:
            raise ValueError(f"Invalid range [{start}, {end}).")
        if end > self.reserved_size:
            raise ValueError(
                f"Requested range end {end} exceeds reserved capacity "
                f"{self.reserved_size}."
            )
        if start == end:
            return
        first = start // self.granularity
        last = (end + self.granularity - 1) // self.granularity  # exclusive
        run_start: int | None = None
        for g in range(first, last + 1):
            unmapped = g < last and g not in self._mapped_granules
            if unmapped and run_start is None:
                run_start = g
            elif not unmapped and run_start is not None:
                self._map_chunk_at(
                    run_start * self.granularity, (g - run_start) * self.granularity
                )
                # Grant before recording the run as mapped: a failure in a
                # later run must not leave these granules mapped but
                # inaccessible, since later commits would skip them.
                self._grant_access(*self._run_bounds(run_start, g))
                self._mapped_granules.update(range(run_start, g))
                run_start = None

    def _run_bounds(self, first: int, last: int) -> tuple[int, int]:
        """Bounds of the maximal contiguous mapped run covering `[first, last)`."""
        while first - 1 in self._mapped_granules:
            first -= 1
        while last in self._mapped_granules:
            last += 1
        return first, last

    def _grant_access(self, first: int, last: int) -> None:
        """Grant device access over the granule run `[first, last)`.

        Per run rather than per chunk: ROCm rejects a set-access range that starts
        inside an already-mapped region. Re-granting is a cheap no-op.
        """
        self._driver.set_access(
            self.base_ptr + first * self.granularity,
            (last - first) * self.granularity,
            self.device_index,
        )

    def _map_chunk_at(self, offset: int, size: int) -> None:
        """Create one physical chunk of `size` bytes and map it at `offset`."""
        driver = self._driver
        driver.ensure_context(self.device_index)
        try:
            handle = driver.create(size, self.device_index)
        except RuntimeError:
            # The VMM allocator cannot reuse memory idling in torch's
            # caching allocator; return it to the driver and retry once.
            torch.accelerator.empty_cache()
            handle = driver.create(size, self.device_index)

        addr = self.base_ptr + offset
        try:
            driver.map(addr, size, handle)
        except RuntimeError:
            driver.release(handle)
            raise
        # Access is granted by the caller, per contiguous run.
        self._handles.append((handle, offset, size))

    def release_physical(self) -> None:
        """Unmap and release all physical memory, keeping the VA reservation."""
        driver = self._driver
        driver.ensure_context(self.device_index)
        if self._handles:
            torch.accelerator.synchronize(self.device_index)
        for handle, offset, size in self._handles:
            driver.unmap(self.base_ptr + offset, size)
            driver.release(handle)
        self._handles = []
        self._mapped_granules = set()

    def free(self) -> None:
        if self._freed:
            return
        self._freed = True
        self.release_physical()
        if self.base_ptr:
            self._driver.free_reserved(self.base_ptr, self.reserved_size)
        self.base_ptr = 0

    def __del__(
        self, _suppress: Any = suppress, _is_finalizing: Any = sys.is_finalizing
    ) -> None:
        # Bound as defaults: module globals may be gone at interpreter exit.
        # The process frees its mappings on exit, and the HIP runtime may
        # already be torn down, where a driver call aborts the process.
        if _suppress is None or _is_finalizing():
            return
        with _suppress(Exception):
            self.free()

committed_bytes property

Total physically mapped bytes (a multiple of the granularity).

_grant_access(first, last)

Grant device access over the granule run [first, last).

Per run rather than per chunk: ROCm rejects a set-access range that starts inside an already-mapped region. Re-granting is a cheap no-op.

Source code in vllm/utils/extensible_tensor.py

def _grant_access(self, first: int, last: int) -> None:
    """Grant device access over the granule run `[first, last)`.

    Per run rather than per chunk: ROCm rejects a set-access range that starts
    inside an already-mapped region. Re-granting is a cheap no-op.
    """
    self._driver.set_access(
        self.base_ptr + first * self.granularity,
        (last - first) * self.granularity,
        self.device_index,
    )

_map_chunk_at(offset, size)

Create one physical chunk of size bytes and map it at offset.

Source code in vllm/utils/extensible_tensor.py

def _map_chunk_at(self, offset: int, size: int) -> None:
    """Create one physical chunk of `size` bytes and map it at `offset`."""
    driver = self._driver
    driver.ensure_context(self.device_index)
    try:
        handle = driver.create(size, self.device_index)
    except RuntimeError:
        # The VMM allocator cannot reuse memory idling in torch's
        # caching allocator; return it to the driver and retry once.
        torch.accelerator.empty_cache()
        handle = driver.create(size, self.device_index)

    addr = self.base_ptr + offset
    try:
        driver.map(addr, size, handle)
    except RuntimeError:
        driver.release(handle)
        raise
    # Access is granted by the caller, per contiguous run.
    self._handles.append((handle, offset, size))

_run_bounds(first, last)

Bounds of the maximal contiguous mapped run covering [first, last).

Source code in vllm/utils/extensible_tensor.py

def _run_bounds(self, first: int, last: int) -> tuple[int, int]:
    """Bounds of the maximal contiguous mapped run covering `[first, last)`."""
    while first - 1 in self._mapped_granules:
        first -= 1
    while last in self._mapped_granules:
        last += 1
    return first, last

ensure_committed_range(start, end)

Map physical pages so that the byte range [start, end) is backed.

Source code in vllm/utils/extensible_tensor.py

def ensure_committed_range(self, start: int, end: int) -> None:
    """Map physical pages so that the byte range `[start, end)` is backed."""
    if not 0 <= start <= end:
        raise ValueError(f"Invalid range [{start}, {end}).")
    if end > self.reserved_size:
        raise ValueError(
            f"Requested range end {end} exceeds reserved capacity "
            f"{self.reserved_size}."
        )
    if start == end:
        return
    first = start // self.granularity
    last = (end + self.granularity - 1) // self.granularity  # exclusive
    run_start: int | None = None
    for g in range(first, last + 1):
        unmapped = g < last and g not in self._mapped_granules
        if unmapped and run_start is None:
            run_start = g
        elif not unmapped and run_start is not None:
            self._map_chunk_at(
                run_start * self.granularity, (g - run_start) * self.granularity
            )
            # Grant before recording the run as mapped: a failure in a
            # later run must not leave these granules mapped but
            # inaccessible, since later commits would skip them.
            self._grant_access(*self._run_bounds(run_start, g))
            self._mapped_granules.update(range(run_start, g))
            run_start = None

release_physical()

Unmap and release all physical memory, keeping the VA reservation.

Source code in vllm/utils/extensible_tensor.py

def release_physical(self) -> None:
    """Unmap and release all physical memory, keeping the VA reservation."""
    driver = self._driver
    driver.ensure_context(self.device_index)
    if self._handles:
        torch.accelerator.synchronize(self.device_index)
    for handle, offset, size in self._handles:
        driver.unmap(self.base_ptr + offset, size)
        driver.release(handle)
    self._handles = []
    self._mapped_granules = set()

granule_size(device_index)

Commit granule for device_index: driver granularity, floored at 2 MiB.

Source code in vllm/utils/extensible_tensor.py

def granule_size(device_index: int) -> int:
    """Commit granule for device_index: driver granularity, floored at 2 MiB."""
    driver_granularity = get_vmm_driver().granularity(device_index)
    return _round_up(max(driver_granularity, _MIN_GRANULARITY), driver_granularity)

uint8_tensor_from_ptr(ptr, num_bytes, device_index)

Wrap device memory at ptr as a uint8 tensor; the caller keeps it mapped.

The DLManagedTensor is malloc'd and never freed: torch reads it when the view is destroyed, which can be after this module's globals are cleared at interpreter exit. Its deleter is NULL, so nothing calls back into Python.

Source code in vllm/utils/extensible_tensor.py

def uint8_tensor_from_ptr(ptr: int, num_bytes: int, device_index: int) -> torch.Tensor:
    """Wrap device memory at `ptr` as a uint8 tensor; the caller keeps it mapped.

    The DLManagedTensor is malloc'd and never freed: torch reads it when the
    view is destroyed, which can be after this module's globals are cleared at
    interpreter exit. Its deleter is NULL, so nothing calls back into Python.
    """
    managed_size = ctypes.sizeof(_DLManagedTensor)
    addr = _malloc(managed_size + ctypes.sizeof(ctypes.c_int64))
    if not addr:
        raise MemoryError("Failed to allocate a DLManagedTensor.")
    ctypes.memset(addr, 0, managed_size + ctypes.sizeof(ctypes.c_int64))
    shape_addr = addr + managed_size
    ctypes.c_int64.from_address(shape_addr).value = num_bytes

    managed = _DLManagedTensor.from_address(addr)
    managed.dl_tensor.data = ctypes.c_void_p(ptr)
    device_type = get_vmm_driver().dlpack_device_type
    managed.dl_tensor.device = _DLDevice(device_type, device_index)
    managed.dl_tensor.ndim = 1
    managed.dl_tensor.dtype = _DLDataType(_K_DL_UINT, _UINT8_BITS, 1)
    managed.dl_tensor.shape = ctypes.cast(shape_addr, ctypes.POINTER(ctypes.c_int64))

    capsule = _PyCapsule_New(addr, b"dltensor", None)
    return torch.from_dlpack(capsule)