Efficient Memory Management for Large Language Model Serving with PagedAttention

def __init__(self, num_blocks, block_size, key_dim, value_dim,
             allocation_order=None):
    self.num_blocks = positive_integer(num_blocks, "num_blocks")
    self.block_size = positive_integer(block_size, "block_size")
    self.key_dim = positive_integer(key_dim, "key_dim")
    self.value_dim = positive_integer(value_dim, "value_dim")
    try:
        order = (list(range(num_blocks)) if allocation_order is None
                 else list(allocation_order))
    except TypeError as error:
        raise ValueError(
            "allocation_order must be a permutation") from error
    if (any(type(x) is not int for x in order)
            or sorted(order) != list(range(num_blocks))):
        raise ValueError("allocation_order must be a full permutation")
    self.order = tuple(order)
    self.keys = np.zeros((num_blocks, block_size, key_dim))
    self.values = np.zeros((num_blocks, block_size, value_dim))
    self.refs = np.zeros(num_blocks, dtype=np.int64)
    self.free = set(range(num_blocks))
    self.tables, self.lengths, self.contexts = {}, {}, {}

def create(self, sequence, context):
    """Create an empty table; do not reserve a per-request block."""
    self._new_identity(sequence)
    context = cache_context(context)
    self.tables[sequence] = []
    self.lengths[sequence] = 0
    self.contexts[sequence] = context

def address(self, sequence, token):
    """Translate only live, zero-based token indices."""
    self._existing(sequence)
    if (type(token) is not int
            or not 0 <= token < self.lengths[sequence]):
        raise ValueError("token index must address a live token")
    logical, offset = divmod(token, self.block_size)
    return self.tables[sequence][logical], offset

def gather(self, sequence):
    """Gather live K/V copies in logical order, excluding unused tails."""
    self._existing(sequence)
    addresses = [self.address(sequence, i)
                 for i in range(self.lengths[sequence])]
    keys = np.empty((len(addresses), self.key_dim))
    values = np.empty((len(addresses), self.value_dim))
    for i, (physical, offset) in enumerate(addresses):
        keys[i] = self.keys[physical, offset]
        values[i] = self.values[physical, offset]
    return keys, values

def append(self, sequence, key, value):
    """Validate first; allocate on a boundary or copy a shared partial tail.

    A single-token append needs at most one free block. Exhaustion is
    checked before changing tables, refcounts, data or the free set.
    """
    self._existing(sequence)
    key = vector(key, self.key_dim)
    value = vector(value, self.value_dim)
    length = self.lengths[sequence]
    offset = length % self.block_size
    if offset == 0:
        physical = self._allocate()
        self.tables[sequence].append(physical)
    else:
        physical = self._private_tail(sequence)
    self.keys[physical, offset] = key
    self.values[physical, offset] = value
    self.lengths[sequence] = length + 1

def fork(self, parent, child, context=None):
    """Fork actual KV data in the same model/cache/position context."""
    self._existing(parent)
    self._new_identity(child)
    inherited = self.contexts[parent]
    if context is not None and cache_context(context) != inherited:
        raise ValueError("fork requires identical cache context")
    table = self.tables[parent].copy()
    self.tables[child] = table
    self.lengths[child] = self.lengths[parent]
    self.contexts[child] = inherited
    for physical in table:
        self.refs[physical] += 1

def _private_tail(self, sequence):
    """Copy a shared partial tail; leave completed prefix blocks shared."""
    old = self.tables[sequence][-1]
    if self.refs[old] == 1:
        return old
    # Allocate first: MemoryError leaves every ownership field unchanged.
    new = self._allocate()
    self.keys[new] = self.keys[old]
    self.values[new] = self.values[old]
    self.tables[sequence][-1] = new
    self.refs[old] -= 1
    return new

def attention(self, sequence, query):
    """Gather live rows and compute Eq. 3, not the blockwise CUDA kernel.

    The sole query represents the next decoding position and can attend
    to the entire existing cache. The authored empty-cache result is zero.
    """
    self._existing(sequence)
    query = vector(query, self.key_dim)
    keys, values = self.gather(sequence)
    if not len(keys):
        return np.zeros(self.value_dim)
    with np.errstate(over="ignore", invalid="ignore"):
        scores = keys @ query / np.sqrt(self.key_dim)
    if not np.isfinite(scores).all():
        raise ValueError("attention scores overflow float64")
    # Finite opposite-sign scores may subtract to -inf: exp(-inf)=0.
    with np.errstate(over="ignore"):
        shifted = scores - scores.max()
    weights = np.exp(shifted)
    weights /= weights.sum()
    output = weights @ values
    if not np.isfinite(output).all():
        raise ValueError("attention output overflow float64")
    return output

def release(self, sequence):
    """Free only blocks with no remaining owners; reject double release."""
    self._existing(sequence)
    for physical in self.tables[sequence]:
        self.refs[physical] -= 1
        if self.refs[physical] == 0:
            self.free.add(physical)
            self.keys[physical].fill(0)
            self.values[physical].fill(0)
    del self.tables[sequence]
    del self.lengths[sequence]
    del self.contexts[sequence]

    def snapshot(self, event):
        """Measure occupied slots once per physical block, even with sharing."""
        occupied = {}
        live = {}
        for sequence, length in self.lengths.items():
            keys, values = self.gather(sequence)
            live[sequence] = {"keys": keys.tolist(), "values": values.tolist()}
            for logical, physical in enumerate(self.tables[sequence]):
                filled = min(self.block_size,
                             length - logical * self.block_size)
                occupied[physical] = max(occupied.get(physical, 0), filled)
        unique = len(occupied)
        live_slots = sum(occupied.values())
        return {"event": event,
                "tables": {s: t.copy() for s, t in self.tables.items()},
                "lengths": self.lengths.copy(), "refs": self.refs.tolist(),
                "free": sorted(self.free), "live_data": live,
                "logical_live_slots": sum(self.lengths.values()),
                "unique_live_slots": live_slots, "unique_blocks": unique,
                "tail_waste_slots": unique * self.block_size - live_slots}


def ownership_trace(seed):
    """Execute a same-context fork, independent appends and both releases."""
    rng = np.random.default_rng(seed)
    keys = rng.normal(size=(10, 3))
    values = rng.normal(size=(10, 2))
    pool = KVPool(8, 4, 3, 2, [5, 2, 7, 0, 3, 6, 1, 4])
    pool.create("A", ("synthetic-model", "single-head-float64", "origin-0"))
    trace = [pool.snapshot("empty")]
    for key, value in zip(keys[:7], values[:7]):
        pool.append("A", key, value)
    trace.append(pool.snapshot("prompt-7"))
    pool.fork("A", "B")
    trace.append(pool.snapshot("fork"))
    pool.append("B", keys[7], values[7])
    trace.append(pool.snapshot("B-copy-on-write-8"))
    pool.append("A", keys[8], values[8])
    trace.append(pool.snapshot("A-private-8"))
    pool.append("B", keys[9], values[9])
    trace.append(pool.snapshot("B-boundary-9"))
    query = rng.normal(size=3)
    outputs = {s: pool.attention(s, query).tolist() for s in ("A", "B")}
    pool.release("A")
    trace.append(pool.snapshot("release-A"))
    pool.release("B")
    trace.append(pool.snapshot("release-B"))
    return trace, outputs