vLLM demonstration - managing vram under the hood; Paged KV Cache and PagedAttention mechanism

PagedKVCache.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import List, Optional

class PagedKVCache:
    """
    advanced paged key-value cache manager for large language models.
    simulates the memory management backend of frameworks like vllm.
    """
    def __init__(
        self,
        num_blocks: int,
        block_size: int,
        num_heads: int,
        head_dim: int,
        dtype: torch.dtype = torch.float16,
        device: str = "cuda"
    ):
        self.block_size = block_size
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.device = device
        
        # [num_blocks, num_heads, block_size, head_dim]
        # pre-allocate the entire physical memory pool on the gpu
        self.key_cache = torch.zeros(
            (num_blocks, num_heads, block_size, head_dim), 
            dtype=dtype, device=device
        )
        self.value_cache = torch.zeros(
            (num_blocks, num_heads, block_size, head_dim), 
            dtype=dtype, device=device
        )
        
        # track free blocks using a simple stack
        self.free_blocks: List[int] = list(range(num_blocks)[::-1])
        self.num_free_blocks = num_blocks

    def allocate(self) -> int:
        # pop a physical block index from the free pool
        if not self.free_blocks:
            raise RuntimeError("kv cache memory pool exhausted. out of memory.")
        self.num_free_blocks -= 1
        return self.free_blocks.pop()

    def free(self, physical_block_indices: List[int]):
        # release blocks back to the pool when a sequence completes
        self.free_blocks.extend(physical_block_indices)
        self.num_free_blocks += len(physical_block_indices)

    def write(
        self, 
        keys: torch.Tensor, 
        values: torch.Tensor, 
        block_tables: torch.Tensor, 
        context_lengths: torch.Tensor
    ):
        """
        writes incoming batched keys/values into the scattered physical blocks.
        supports both prefill (multiple tokens) and decode (single token).
        keys/values shape: [batch_size, num_tokens, num_heads, head_dim]
        block_tables shape: [batch_size, max_blocks_per_seq]
        context_lengths shape: [batch_size] (length before these new tokens)
        """
        batch_size = keys.size(0)
        num_tokens = keys.size(1)
        
        for i in range(batch_size):
            # current sequence length before adding new tokens
            start_len = context_lengths[i].item()
            
            for t in range(num_tokens):
                current_pos = start_len + t
                logical_block_idx = current_pos // self.block_size
                block_offset = current_pos % self.block_size
                
                # dynamically allocate a new block if we crossed a boundary
                if block_tables[i, logical_block_idx] == -1:
                    block_tables[i, logical_block_idx] = self.allocate()
                    
                physical_block_idx = block_tables[i, logical_block_idx].item()
                
                # insert the key and value into the specific physical slot
                self.key_cache[physical_block_idx, :, block_offset, :] = keys[i, t]
                self.value_cache[physical_block_idx, :, block_offset, :] = values[i, t]


class PagedAttention(nn.Module):
    """
    attention mechanism that reads directly from non-contiguous physical memory blocks.
    """
    def __init__(self, num_heads: int, head_dim: int, scale: Optional[float] = None):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.scale = scale if scale is not None else 1.0 / math.sqrt(head_dim)

    def forward(
        self,
        query: torch.Tensor,
        kv_cache: PagedKVCache,
        block_tables: torch.Tensor,
        context_lengths: torch.Tensor
    ) -> torch.Tensor:
        """
        query shape: [batch_size, 1, num_heads, head_dim]
        block_tables shape: [batch_size, max_blocks_per_seq]
        """
        batch_size = query.size(0)
        outputs = torch.empty_like(query)
        
        for i in range(batch_size):
            seq_len = context_lengths[i].item()
            
            # figure out how many blocks this specific sequence is currently using
            num_blocks = math.ceil(seq_len / kv_cache.block_size)
            physical_indices = block_tables[i, :num_blocks].long()
            
            # gather the scattered physical blocks for this sequence
            # shape: [num_blocks, num_heads, block_size, head_dim]
            k_blocks = kv_cache.key_cache[physical_indices]
            v_blocks = kv_cache.value_cache[physical_indices]
            
            # flatten out the blocks to reconstruct the contiguous sequence
            # target shape: [num_heads, padded_seq_len, head_dim]
            k_seq = k_blocks.permute(1, 0, 2, 3).reshape(self.num_heads, -1, self.head_dim)
            v_seq = v_blocks.permute(1, 0, 2, 3).reshape(self.num_heads, -1, self.head_dim)
            
            # slice off the padding from the final partially-filled block
            k_seq = k_seq[:, :seq_len, :]
            v_seq = v_seq[:, :seq_len, :]
            
            # isolate the query for this specific sequence in the batch
            # shape: [num_heads, 1, head_dim]
            q_seq = query[i].transpose(0, 1)
            
            # compute standard scaled dot-product attention
            scores = torch.matmul(q_seq, k_seq.transpose(-2, -1)) * self.scale
            attn_weights = F.softmax(scores, dim=-1)
            
            # compute final context vector and store in output
            # shape: [num_heads, 1, head_dim] -> [1, num_heads, head_dim]
            out = torch.matmul(attn_weights, v_seq)
            outputs[i] = out.transpose(0, 1)
            
        return outputs


# --- execution harness ---
if __name__ == "__main__":
    # mock configuration parameters
    batch_size = 2
    num_heads = 8
    head_dim = 64
    block_size = 16
    max_blocks = 100
    
    # initialize the centralized cache pool and attention module
    cache = PagedKVCache(max_blocks, block_size, num_heads, head_dim, device="cpu")
    attention = PagedAttention(num_heads, head_dim)
    
    # simulate the logical to physical mapping table for a batch of requests
    # initialize with -1 to indicate unallocated slots
    block_tables = torch.full((batch_size, 10), -1, dtype=torch.long)
    
    # current lengths of the sequences in the batch (start at 0)
    context_lengths = torch.tensor([0, 0]) 
    
    print(f"initial free blocks: {cache.num_free_blocks}")
    
    # --- phase 1: prefill (processing the initial prompt) ---
    # prompt length of 18 for both sequences (crosses a 16-token block boundary)
    prefill_len = 18
    k_prompt = torch.randn(batch_size, prefill_len, num_heads, head_dim, dtype=torch.float16, device="cpu")
    v_prompt = torch.randn(batch_size, prefill_len, num_heads, head_dim, dtype=torch.float16, device="cpu")
    
    # write the prompt to the cache (this handles dynamic allocation automatically)
    cache.write(k_prompt, v_prompt, block_tables, context_lengths)
    
    # update context lengths after prefill
    context_lengths += prefill_len
    
    print(f"free blocks after prefill: {cache.num_free_blocks} (allocated {max_blocks - cache.num_free_blocks})")
    
    # --- phase 2: decode (generating the next token) ---
    # generate a single new token query, key, and value
    # shape: [batch, 1, num_heads, head_dim]
    q_decode = torch.randn(batch_size, 1, num_heads, head_dim, dtype=torch.float16, device="cpu")
    k_decode = torch.randn(batch_size, 1, num_heads, head_dim, dtype=torch.float16, device="cpu")
    v_decode = torch.randn(batch_size, 1, num_heads, head_dim, dtype=torch.float16, device="cpu")
    
    # write the new token to the cache
    cache.write(k_decode, v_decode, block_tables, context_lengths)
    
    # update context lengths
    context_lengths += 1
    
    # compute attention reading from the scattered memory
    out = attention(q_decode, cache, block_tables, context_lengths)
    
    print("=== paged kv cache simulation complete ===")
    print(f"attention output shape: {out.shape}")
    print(f"block table layout:\n{block_tables}")