vllm.v1.token_buffer

Classes:

• TokenBuffer –

  Mutable int32 token ids with no per-token Python objects.

• TokenIdsView –

  tokens followed by suffix, shortened from the end by pop().

Functions:

• extend_from_slice –

  dst.extend(src[start:stop]), copied without Python ints when both are

• rfind_token_sequences –

  Return (index, k) for the newest occurrence in token_ids[start:stop]

• rfind_tokens –

  Return the index of the last token in token_ids[start:stop] that is

• token_slice –

  tokens[start:stop]; for a TokenBuffer an int32 array, so that no

TokenBuffer

Bases: array

Mutable int32 token ids with no per-token Python objects.

Source code in vllm/v1/token_buffer.py

class TokenBuffer(array):
    """Mutable int32 token ids with no per-token Python objects."""

    __slots__ = ()

    def __new__(cls, tokens: Iterable[int] = ()):
        if isinstance(tokens, np.ndarray):
            # A 32-bit int array, as the frontends send, copied as one block.
            assert tokens.dtype.kind in "iu" and tokens.dtype.itemsize == 4
            buf = super().__new__(cls, "i")
            buf.frombytes(memoryview(np.ascontiguousarray(tokens)).cast("B"))
            return buf
        # From another array of the same type this is a memcpy.
        return super().__new__(cls, "i", tokens)

    def __getitem__(self, index):
        item = array.__getitem__(self, index)
        return item.tolist() if isinstance(index, slice) else item

    def __eq__(self, other):
        if isinstance(other, (list, tuple)):
            return len(self) == len(other) and self.tolist() == as_list(other)
        return array.__eq__(self, other)

    def __ne__(self, other):
        return not self == other

    def __add__(self, other):
        return self.tolist() + as_list(other)

    def __radd__(self, other):
        return as_list(other) + self.tolist()

    def __iadd__(self, other):
        self.extend(other)
        return self

    def copy(self) -> list[int]:
        return self.tolist()

    def __reduce_ex__(self, protocol):
        # A PickleBuffer lets MessageQueue send large ids out of band. It wraps
        # a copy since zmq may hold it past enqueue, blocking resizes of self.
        data = self.tobytes()
        return _rebuild_token_buffer, (
            PickleBuffer(data) if protocol >= 5 else data,
        )

TokenIdsView

Bases: Sequence[int]

tokens followed by suffix, shortened from the end by pop().

Copied into a list only when read element-wise, so a caller that ignores it, or searches it with rfind_tokens, doesn't copy tokens.

Source code in vllm/v1/token_buffer.py

class TokenIdsView(Sequence[int]):
    """``tokens`` followed by ``suffix``, shortened from the end by pop().

    Copied into a list only when read element-wise, so a caller that ignores
    it, or searches it with rfind_tokens, doesn't copy ``tokens``.
    """

    def __init__(self, tokens: Sequence[int], suffix: Iterable[int] = ()) -> None:
        self._tokens = tokens
        self._suffix = list(suffix)
        self._len = len(tokens) + len(self._suffix)
        self._list: list[int] | None = None

    def __len__(self) -> int:
        return self._len

    def _as_list(self) -> list[int]:
        if self._list is None:
            num_tokens = min(len(self._tokens), self._len)
            self._list = list(self._tokens[:num_tokens])
            self._list.extend(self._suffix[: self._len - num_tokens])
        return self._list

    @overload
    def __getitem__(self, index: int) -> int: ...

    @overload
    def __getitem__(self, index: slice) -> list[int]: ...

    def __getitem__(self, index: int | slice) -> int | list[int]:
        return self._as_list()[index]

    def __iter__(self) -> Iterator[int]:
        return iter(self._as_list())

    def __contains__(self, value: object) -> bool:
        return value in self._as_list()

    def pop(self) -> None:
        if self._list is not None:
            self._list.pop()
        self._len -= 1

    def rfind(self, targets: Collection[int | None], start: int, stop: int) -> int:
        if self._list is not None:
            return rfind_tokens(self._list, targets, start, stop)
        num_tokens = len(self._tokens)
        stop = min(stop, self._len)
        for i in range(stop - 1, max(start, num_tokens) - 1, -1):
            if self._suffix[i - num_tokens] in targets:
                return i
        return rfind_tokens(self._tokens, targets, start, min(stop, num_tokens))

extend_from_slice(dst, src, start, stop)

dst.extend(src[start:stop]), copied without Python ints when both are TokenBuffers.

Source code in vllm/v1/token_buffer.py

def extend_from_slice(
    dst: list[int], src: Sequence[int], start: int, stop: int
) -> None:
    """``dst.extend(src[start:stop])``, copied without Python ints when both are
    TokenBuffers."""
    if isinstance(dst, array) and isinstance(src, array):
        dst.extend(array.__getitem__(src, slice(start, stop)))
    else:
        dst.extend(src[start:stop])

rfind_token_sequences(token_ids, markers, start=0, stop=None)

Return (index, k) for the newest occurrence in token_ids[start:stop] of any of the multi-token markers, k being its position in markers (the first listed wins at the same index), or (-1, -1).

Candidates are found with rfind_tokens on the markers' first tokens, so a TokenBuffer is searched with numpy and only candidate positions are read.

Source code in vllm/v1/token_buffer.py

def rfind_token_sequences(
    token_ids: Sequence[int],
    markers: Sequence[Sequence[int | None]],
    start: int = 0,
    stop: int | None = None,
) -> tuple[int, int]:
    """Return ``(index, k)`` for the newest occurrence in ``token_ids[start:stop]``
    of any of the multi-token ``markers``, ``k`` being its position in
    ``markers`` (the first listed wins at the same index), or ``(-1, -1)``.

    Candidates are found with rfind_tokens on the markers' first tokens, so a
    TokenBuffer is searched with numpy and only candidate positions are read.
    """
    if stop is None:
        stop = len(token_ids)
    marker_lists = [list(marker) for marker in markers]
    first_ids = {marker[0] for marker in marker_lists if marker}
    end = stop
    while (i := rfind_tokens(token_ids, first_ids, start, end)) >= 0:
        head = token_ids[i]
        for k, marker in enumerate(marker_lists):
            if (marker and marker[0] == head and i + len(marker) <= stop) and (
                len(marker) == 1
                or as_list(token_ids[i + 1 : i + len(marker)]) == marker[1:]
            ):
                return i, k
        end = i
    return -1, -1

rfind_tokens(token_ids, targets, start=0, stop=None)

Return the index of the last token in token_ids[start:stop] that is one of targets (None entries, for absent tokens, are ignored), or -1.

An int32 buffer (a TokenBuffer) is searched with numpy, in chunks growing back from stop so a recent match is found quickly, without creating a Python int per token. Other sequences are scanned backwards in Python.

Source code in vllm/v1/token_buffer.py

def rfind_tokens(
    token_ids: Sequence[int],
    targets: Collection[int | None],
    start: int = 0,
    stop: int | None = None,
) -> int:
    """Return the index of the last token in ``token_ids[start:stop]`` that is
    one of ``targets`` (``None`` entries, for absent tokens, are ignored), or -1.

    An int32 buffer (a TokenBuffer) is searched with numpy, in chunks growing
    back from ``stop`` so a recent match is found quickly, without creating a
    Python int per token. Other sequences are scanned backwards in Python.
    """
    if stop is None:
        stop = len(token_ids)
    if isinstance(token_ids, TokenIdsView):
        return token_ids.rfind(targets, start, stop)
    if isinstance(token_ids, array) and token_ids.typecode == "i":
        ids = np.frombuffer(token_ids, dtype=np.int32)
        wanted = np.array([t for t in targets if t is not None], dtype=np.int32)
        chunk = 256
        while stop > start:
            lo = max(start, stop - chunk)
            hits = np.flatnonzero(np.isin(ids[lo:stop], wanted))
            if hits.size:
                return lo + int(hits[-1])
            stop, chunk = lo, chunk * 8
        return -1
    for i in range(stop - 1, start - 1, -1):
        if token_ids[i] in targets:
            return i
    return -1

token_slice(tokens, start, stop)

tokens[start:stop]; for a TokenBuffer an int32 array, so that no Python ints are made until it is read.

Source code in vllm/v1/token_buffer.py

def token_slice(tokens: Sequence[int], start: int, stop: int) -> Sequence[int]:
    """``tokens[start:stop]``; for a TokenBuffer an int32 array, so that no
    Python ints are made until it is read."""
    if isinstance(tokens, array):
        return array.__getitem__(tokens, slice(start, stop))
    return tokens[start:stop]