Paged Attention › vLLM Implementation

from collections import deque
from dataclasses import dataclass
from typing import Dict, List


@dataclass
class PhysicalBlock:
    """Represents a physical block of GPU memory."""

    block_number: int
    ref_count: int = 0  # For copy-on-write sharing


@dataclass
class LogicalBlock:
    """Represents a logical block in a sequence."""

    block_number: int
    num_tokens: int = 0


class BlockSpaceManager:
    """
    Manages allocation of physical blocks to sequences.
    Simplified version of vLLM's BlockSpaceManager.
    """

    def __init__(
        self,
        block_size: int,
        num_gpu_blocks: int,
        num_cpu_blocks: int = 0,  # For CPU offloading
    ):
        self.block_size = block_size
        self.num_gpu_blocks = num_gpu_blocks
        self.num_cpu_blocks = num_cpu_blocks

        # Free block pools
        self.gpu_free_blocks: deque = deque(range(num_gpu_blocks))
        self.cpu_free_blocks: deque = deque(range(num_cpu_blocks))

        # Block tables for each sequence
        self.block_tables: Dict[int, List[int]] = {}

        # Track block reference counts for sharing
        self.block_ref_counts: Dict[int, int] = {}

    def can_allocate(self, num_blocks: int) -> bool:
        """Check if we can allocate the requested number of blocks."""
        return len(self.gpu_free_blocks) >= num_blocks

    def allocate(self, seq_id: int, num_blocks: int) -> List[int]:
        """Allocate blocks for a new sequence."""
        if not self.can_allocate(num_blocks):
            raise MemoryError(f"Cannot allocate {num_blocks} blocks")

        allocated = []
        for _ in range(num_blocks):
            block = self.gpu_free_blocks.popleft()
            allocated.append(block)
            self.block_ref_counts[block] = 1

        self.block_tables[seq_id] = allocated
        return allocated

    def append_block(self, seq_id: int) -> int:
        """Allocate one more block for an existing sequence."""
        if not self.gpu_free_blocks:
            raise MemoryError("No free blocks available")

        block = self.gpu_free_blocks.popleft()
        self.block_tables[seq_id].append(block)
        self.block_ref_counts[block] = 1
        return block

    def free(self, seq_id: int) -> List[int]:
        """Free all blocks belonging to a sequence."""
        if seq_id not in self.block_tables:
            return []

        freed = []
        for block in self.block_tables[seq_id]:
            self.block_ref_counts[block] -= 1
            if self.block_ref_counts[block] == 0:
                self.gpu_free_blocks.append(block)
                del self.block_ref_counts[block]
                freed.append(block)

        del self.block_tables[seq_id]
        return freed

    def get_block_table(self, seq_id: int) -> List[int]:
        """Get the block table for a sequence."""
        return self.block_tables.get(seq_id, [])

    def get_num_free_blocks(self) -> int:
        """Get the number of available blocks."""
        return len(self.gpu_free_blocks)

# Simulate vLLM-style block management
manager = BlockSpaceManager(block_size=16, num_gpu_blocks=100)

# Sequence arrives with prompt
seq_0_prompt_len = 45  # tokens
initial_blocks_needed = (seq_0_prompt_len + 16 - 1) // 16  # Ceiling division
manager.allocate(seq_id=0, num_blocks=initial_blocks_needed)

# Another sequence arrives
seq_1_prompt_len = 128
blocks_needed = (seq_1_prompt_len + 16 - 1) // 16
manager.allocate(seq_id=1, num_blocks=blocks_needed)

status_after_prompts = {
    "free_blocks": manager.get_num_free_blocks(),
    "seq_0_blocks": manager.get_block_table(0),
    "seq_1_blocks": manager.get_block_table(1),
}

# Generation proceeds, each sequence needs more blocks
for _ in range(3):  # Generate 3 more blocks worth of tokens for seq 0
    manager.append_block(seq_id=0)

# Seq 1 completes
manager.free(seq_id=1)

status_final = {
    "free_blocks": manager.get_num_free_blocks(),
    "seq_0_blocks": manager.get_block_table(0),
}

def demonstrate_copy_on_write():
    """
    Show how PagedAttention enables memory-efficient beam search
    through block sharing.
    """
    # Parameters
    prompt_tokens = 48
    block_size = 16
    num_beams = 4

    # Initial sequence (3 blocks at 16 tokens each)
    num_prompt_blocks = prompt_tokens // block_size
    prompt_blocks = [
        f"P{i}" for i in range(num_prompt_blocks)
    ]  # Physical blocks for prompt

    # Beam search with 4 beams
    # All beams share the prompt blocks initially
    beams = {}
    for i in range(num_beams):
        beams[f"beam_{i}"] = {
            "blocks": list(prompt_blocks),
            "ref_count_contribution": 1,
        }

    # Reference counts: each prompt block has 4 references
    block_refs = {block: 4 for block in prompt_blocks}

    # After generating 16 tokens, each beam diverges
    # Each beam gets its own new block
    for i, beam in enumerate(beams.keys()):
        new_block = f"G{i}"  # Generation block
        beams[beam]["blocks"].append(new_block)
        block_refs[new_block] = 1

    return beams, block_refs, prompt_tokens, block_size, num_beams


beams, refs, prompt_tokens, block_size, beam_width = demonstrate_copy_on_write()

# Calculate memory savings statistics
num_beams = len(beams)
blocks_per_beam = len(next(iter(beams.values()))["blocks"])
naive_blocks = num_beams * blocks_per_beam

# Count unique physical blocks used (all active blocks are in refs)
cow_blocks = len(refs)
savings_pct = (1 - cow_blocks / naive_blocks) * 100