PagedAttention: Solving LLM KV Cache Memory Fragmentation

import numpy as np


# Simulating traditional KV cache allocation
def traditional_allocation_simulation(
    max_seq_len: int = 4096, num_requests: int = 8, actual_lengths: list = None
):
    """
    Simulate memory allocation with contiguous pre-allocation.
    Returns wasted memory statistics.
    """
    if actual_lengths is None:
        # Realistic distribution: most sequences are shorter than max
        np.random.seed(42)
        actual_lengths = np.random.exponential(scale=500, size=num_requests)
        actual_lengths = np.clip(actual_lengths, 50, max_seq_len).astype(int)

    # Each request reserves max_seq_len slots
    total_allocated = num_requests * max_seq_len
    total_used = sum(actual_lengths)

    return {
        "max_seq_len": max_seq_len,
        "num_requests": num_requests,
        "allocated": total_allocated,
        "used": total_used,
        "wasted": total_allocated - total_used,
        "utilization": total_used / total_allocated,
        "actual_lengths": actual_lengths,
    }


# Run simulation
stats = traditional_allocation_simulation()

def visualize_fragmentation():
    """
    Visualize how external fragmentation develops over time.
    """
    # Memory represented as slots (simplified)
    memory_size = 40  # Total slots available
    max_seq_alloc = 10  # Each request needs 10 contiguous slots

    # Timeline of events
    events = [
        ("alloc", "A"),  # Request A starts
        ("alloc", "B"),  # Request B starts
        ("alloc", "C"),  # Request C starts
        ("alloc", "D"),  # Request D starts
        ("free", "A"),  # Request A completes
        ("free", "C"),  # Request C completes
        ("alloc", "E"),  # Request E arrives - needs 10 contiguous slots
    ]

    # Track allocations: dict mapping request_id -> (start, end)
    allocations = {}
    next_free = 0
    snapshots = []

    for event_type, request_id in events:
        if event_type == "alloc":
            if (
                request_id not in allocations
                and next_free + max_seq_alloc <= memory_size
            ):
                allocations[request_id] = (next_free, next_free + max_seq_alloc)
                next_free += max_seq_alloc
        else:  # free
            if request_id in allocations:
                del allocations[request_id]

        # Create memory snapshot
        snapshot = ["." for _ in range(memory_size)]
        for req_id, (start, end) in allocations.items():
            for i in range(start, end):
                snapshot[i] = req_id

        snapshots.append(
            (f"{event_type} {request_id}", "".join(snapshot), dict(allocations))
        )

    return snapshots, max_seq_alloc


snapshots, req_size = visualize_fragmentation()

# Pre-calculate display data
display_data = []
for event, memory, _ in snapshots:
    display_data.append((event, memory, memory.count(".")))

total_free = snapshots[-1][1].count(".")

import matplotlib.pyplot as plt

plt.rcParams.update(
    {
        "figure.figsize": (3.0, 2.5),
        "figure.dpi": 300,
        "figure.constrained_layout.use": True,
        "font.family": "sans-serif",
        "font.sans-serif": [
            "Noto Sans CJK SC",
            "Apple SD Gothic Neo",
            "DejaVu Sans",
            "Arial",
        ],
        "font.size": 10,
        "axes.titlesize": 11,
        "axes.titleweight": "bold",
        "axes.titlepad": 8,
        "axes.labelsize": 10,
        "axes.labelpad": 4,
        "xtick.labelsize": 9,
        "ytick.labelsize": 9,
        "legend.fontsize": 9,
        "legend.title_fontsize": 10,
        "legend.frameon": True,
        "legend.loc": "best",
        "lines.linewidth": 1.5,
        "lines.markersize": 5,
        "axes.grid": True,
        "grid.alpha": 0.3,
        "grid.linestyle": "--",
        "axes.spines.top": False,
        "axes.spines.right": False,
        "axes.prop_cycle": plt.cycler(
            color=["#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#7f7f7f"]
        ),
    }
)

# Plot 1: Internal Fragmentation
plt.figure()
allocations = [
    ("Request A", 0, 1000, 300),  # allocated, actual used
    ("Request B", 1000, 2000, 750),
    ("Request C", 2000, 3000, 200),
    ("Request D", 3000, 4000, 600),
]
for label, start, end, used in allocations:
    # Used portion
    plt.barh(
        label,
        used,
        left=start,
        color="#2ecc71",
        edgecolor="black",
        linewidth=0.5,
    )
    # Wasted portion (internal fragmentation)
    plt.barh(
        label,
        (end - start) - used,
        left=start + used,
        color="#e74c3c",
        alpha=0.5,
        edgecolor="black",
        linewidth=0.5,
    )
plt.xlabel("Memory Slots")
plt.title("Internal Fragmentation\n(Unused space within allocations)")
plt.legend(["Used", "Wasted"], loc="upper right")
plt.xlim(0, 4500)
plt.show()

# Plot 2: External Fragmentation
plt.figure()
memory_blocks = [
    ("B: 1000", "#2ecc71"),  # In use
    ("Free: 1000", "#e74c3c"),  # External fragment
    ("D: 1000", "#2ecc71"),  # In use
    ("Free: 1000", "#e74c3c"),  # External fragment
]
bottom = 0
for label, color in memory_blocks:
    size = 1000
    plt.bar(
        ["Memory"],
        size,
        bottom=bottom,
        color=color,
        edgecolor="black",
        linewidth=0.5,
        alpha=0.7 if "Free" in label else 1.0,
    )
    plt.text(0, bottom + size / 2, label, ha="center", va="center")
    bottom += size
plt.ylabel("Memory Slots")
plt.title("External Fragmentation\n(Scattered free blocks)")
plt.ylim(0, 4500)
plt.show()

class PagedKVCacheSimulator:
    """
    Simulates paged KV cache allocation to demonstrate the concept.
    """

    def __init__(self, total_physical_blocks: int, block_size: int = 16):
        self.total_blocks = total_physical_blocks
        self.block_size = block_size  # Tokens per block

        # Physical memory: list of blocks, None if free
        self.physical_memory = [None] * total_physical_blocks

        # Block tables: sequence_id -> list of physical block indices
        self.block_tables = {}

        # Free block list
        self.free_blocks = list(range(total_physical_blocks))

    def allocate_block(self, sequence_id: int) -> int:
        """Allocate a single physical block for a sequence."""
        if not self.free_blocks:
            raise MemoryError("No free blocks available")

        # Get a free physical block
        physical_idx = self.free_blocks.pop(0)
        self.physical_memory[physical_idx] = sequence_id

        # Add to sequence's block table
        if sequence_id not in self.block_tables:
            self.block_tables[sequence_id] = []
        self.block_tables[sequence_id].append(physical_idx)

        return physical_idx

    def free_sequence(self, sequence_id: int):
        """Free all blocks belonging to a sequence."""
        if sequence_id in self.block_tables:
            for physical_idx in self.block_tables[sequence_id]:
                self.physical_memory[physical_idx] = None
                self.free_blocks.append(physical_idx)
            del self.block_tables[sequence_id]

    def add_tokens(self, sequence_id: int, num_tokens: int):
        """Add tokens to a sequence, allocating blocks as needed."""
        if sequence_id not in self.block_tables:
            self.block_tables[sequence_id] = []

        current_tokens = len(self.block_tables[sequence_id]) * self.block_size
        tokens_needed = num_tokens

        while tokens_needed > 0:
            # Check if current block has space
            current_blocks = len(self.block_tables[sequence_id])
            if current_blocks == 0 or current_tokens % self.block_size == 0:
                # Need a new block
                self.allocate_block(sequence_id)

            # Fill current block
            space_in_block = self.block_size - (
                current_tokens % self.block_size
            )
            tokens_to_add = min(space_in_block, tokens_needed)
            tokens_needed -= tokens_to_add
            current_tokens += tokens_to_add

        return current_tokens

    def get_utilization(self) -> dict:
        """Calculate memory utilization statistics."""
        used_blocks = self.total_blocks - len(self.free_blocks)
        return {
            "total_blocks": self.total_blocks,
            "used_blocks": used_blocks,
            "free_blocks": len(self.free_blocks),
            "utilization": used_blocks / self.total_blocks,
            "sequences": len(self.block_tables),
        }


# Demonstrate paged allocation
paged_cache = PagedKVCacheSimulator(total_physical_blocks=200, block_size=16)

# Create a fresh paged cache with enough capacity for this demonstration
# Total tokens needed: A(300) + B(750) + C(200) + D(600) + E(500) = 2350 tokens
# At 16 tokens per block, we need at least 147 blocks, use 200 for headroom
paged_cache = PagedKVCacheSimulator(total_physical_blocks=200, block_size=16)

# Define request sizes
req_a = 300
req_b = 750
req_c = 200
req_d = 600
req_e = 500

# Simulate the same scenario that caused fragmentation earlier
# Request A
paged_cache.add_tokens(sequence_id=0, num_tokens=req_a)
# Request B
paged_cache.add_tokens(sequence_id=1, num_tokens=req_b)
# Request C
paged_cache.add_tokens(sequence_id=2, num_tokens=req_c)
# Request D
paged_cache.add_tokens(sequence_id=3, num_tokens=req_d)

before_free = paged_cache.get_utilization()

# Free requests A and C (simulating completion)
paged_cache.free_sequence(0)
paged_cache.free_sequence(2)

after_free = paged_cache.get_utilization()

# Now allocate a new request E
paged_cache.add_tokens(sequence_id=4, num_tokens=req_e)

after_new = paged_cache.get_utilization()

def visualize_block_tables(cache: PagedKVCacheSimulator):
    """Show the block table mappings for active sequences."""
    print("Block Tables (Logical → Physical mapping)")
    print("-" * 45)

    for seq_id, blocks in sorted(cache.block_tables.items()):
        logical_indices = list(range(len(blocks)))
        print(f"\nSequence {seq_id}:")
        print(f"  Logical blocks:  {logical_indices}")
        print(f"  Physical blocks: {blocks}")

        # Show the mapping
        mapping = [f"L{l}→P{p}" for l, p in enumerate(blocks)]
        print(
            f"  Mapping: {', '.join(mapping[:5])}{'...' if len(mapping) > 5 else ''}"
        )

import torch
import torch.nn.functional as F


def paged_attention_reference(
    query: torch.Tensor,  # Shape: (batch, num_heads, 1, head_dim) for single token
    key_cache: torch.Tensor,  # Shape: (num_blocks, block_size, num_heads, head_dim)
    value_cache: torch.Tensor,  # Shape: (num_blocks, block_size, num_heads, head_dim)
    block_tables: torch.Tensor,  # Shape: (batch, max_blocks) - physical block indices
    context_lens: torch.Tensor,  # Shape: (batch,) - actual sequence lengths
    block_size: int,
    scale: float,
) -> torch.Tensor:
    """
    Reference implementation of PagedAttention for understanding.
    Not optimized - real implementations use custom CUDA kernels.
    """
    batch_size, num_heads, _, head_dim = query.shape
    output = torch.zeros_like(query)

    for b in range(batch_size):
        seq_len = context_lens[b].item()
        num_blocks_used = (seq_len + block_size - 1) // block_size

        # Gather keys and values from physical blocks
        keys_list = []
        values_list = []

        for block_idx in range(num_blocks_used):
            physical_block = block_tables[b, block_idx].item()

            # Calculate how many tokens to use from this block
            start_pos = block_idx * block_size
            end_pos = min(start_pos + block_size, seq_len)
            tokens_in_block = end_pos - start_pos

            # Fetch from physical memory location
            keys_list.append(key_cache[physical_block, :tokens_in_block])
            values_list.append(value_cache[physical_block, :tokens_in_block])

        # Concatenate gathered keys and values
        # Shape: (seq_len, num_heads, head_dim)
        keys = torch.cat(keys_list, dim=0)
        values = torch.cat(values_list, dim=0)

        # Reshape for attention computation
        # keys: (num_heads, seq_len, head_dim)
        keys = keys.permute(1, 0, 2)
        values = values.permute(1, 0, 2)

        # Query shape: (num_heads, 1, head_dim)
        q = query[b]  # (num_heads, 1, head_dim)

        # Compute attention scores
        # (num_heads, 1, head_dim) @ (num_heads, head_dim, seq_len) -> (num_heads, 1, seq_len)
        attn_scores = torch.matmul(q, keys.transpose(-2, -1)) * scale
        attn_probs = F.softmax(attn_scores, dim=-1)

        # Apply attention to values
        # (num_heads, 1, seq_len) @ (num_heads, seq_len, head_dim) -> (num_heads, 1, head_dim)
        attn_output = torch.matmul(attn_probs, values)
        output[b] = attn_output

    return output

# Set up a small example
torch.manual_seed(42)

num_blocks = 8
block_size = 4
num_heads = 2
head_dim = 8
batch_size = 2

# Physical KV cache - blocks can be used by any sequence
key_cache = torch.randn(num_blocks, block_size, num_heads, head_dim)
value_cache = torch.randn(num_blocks, block_size, num_heads, head_dim)

# Block tables - note non-contiguous physical blocks
# Sequence 0 uses physical blocks [0, 3, 5] (scattered!)
# Sequence 1 uses physical blocks [1, 7]
block_tables = torch.tensor(
    [
        [
            0,
            3,
            5,
            0,
            0,
            0,
            0,
            0,
        ],  # Sequence 0: 3 blocks (12 tokens max, using 10)
        [
            1,
            7,
            0,
            0,
            0,
            0,
            0,
            0,
        ],  # Sequence 1: 2 blocks (8 tokens max, using 7)
    ]
)

# Actual sequence lengths
context_lens = torch.tensor([10, 7])

# Query for current token being generated
query = torch.randn(batch_size, num_heads, 1, head_dim)
scale = 1.0 / (head_dim**0.5)

# Run paged attention
output = paged_attention_reference(
    query, key_cache, value_cache, block_tables, context_lens, block_size, scale
)

# Extract block mappings for visualization
sequence_mappings = []
for i, seq_len in enumerate(context_lens):
    # Calculate number of blocks used by this sequence
    n_blocks = (seq_len.item() + block_size - 1) // block_size
    # Slice the block table to get only the used physical blocks
    blocks = block_tables[i, :n_blocks].tolist()
    sequence_mappings.append((i, seq_len.item(), blocks))

# Pseudocode for optimized PagedAttention kernel
# This shows the conceptual approach - actual CUDA code is more complex


def paged_attention_kernel_pseudocode():
    """
    Conceptual flow of the optimized PagedAttention CUDA kernel.

    Key optimizations:
    1. No explicit gather step - compute attention directly from blocks
    2. Each thread block processes a portion of the attention
    3. Leverage shared memory for frequently accessed block table entries
    """
    kernel_description = """
    CUDA Kernel: paged_attention_v1
    
    Grid: (num_heads, batch_size, 1)
    Block: (NUM_THREADS,)
    
    For each (head, batch) pair:
        1. Load query vector for this head into registers
        2. Initialize running softmax numerator and denominator
        
        For each logical block in sequence:
            a. Look up physical block index from block_table
            b. Load key block from global memory (coalesced access)
            c. Compute partial attention scores for this block
            d. Update running softmax using online softmax trick
            
        For each logical block in sequence:
            a. Load value block from global memory
            b. Weight by final attention probabilities
            c. Accumulate to output
        
        Write output for this (batch, head, token)
    """
    return kernel_description


kernel_pseudocode = paged_attention_kernel_pseudocode()

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


@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

# This code shows vLLM usage patterns (requires GPU and vllm installed)
from vllm import LLM, SamplingParams

# vLLM automatically uses PagedAttention
llm = LLM(
    model="meta-llama/Llama-2-7b-hf",
    # Block size can be configured (default is typically 16)
    block_size=16,
    # GPU memory utilization for KV cache
    gpu_memory_utilization=0.9,
    # Maximum number of sequences to batch together
    max_num_seqs=256,
    # CPU swap space in GB (for swapping out preempted requests)
    swap_space=4,
)

# Multiple requests can be batched efficiently
prompts = [
    "Explain quantum computing in simple terms:",
    "Write a haiku about programming:",
    "What are the benefits of exercise?",
    # ... many more prompts
]

sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=256)

# vLLM handles dynamic batching and memory management automatically
outputs = llm.generate(prompts, sampling_params)

import numpy as np


def compare_allocation_strategies(
    num_sequences: int = 50,
    max_seq_len: int = 2048,
    block_size: int = 16,
    seed: int = 42,
):
    """
    Compare traditional contiguous vs paged allocation.
    """
    np.random.seed(seed)

    # Generate realistic sequence lengths (biased toward shorter)
    actual_lengths = np.random.exponential(scale=300, size=num_sequences)
    actual_lengths = np.clip(actual_lengths, 50, max_seq_len).astype(int)

    # Simulate completion order (random)
    completion_order = np.random.permutation(num_sequences)

    results = {
        "traditional": {
            "peak_memory": 0,
            "avg_utilization": [],
            "fragmentation_events": 0,
        },
        "paged": {
            "peak_memory": 0,
            "avg_utilization": [],
            "fragmentation_events": 0,  # Should be 0
        },
    }

    # Traditional: pre-allocate max_seq_len for each
    trad_total_allocated = num_sequences * max_seq_len
    trad_total_used = sum(actual_lengths)
    results["traditional"]["max_seq_len"] = max_seq_len
    results["traditional"]["peak_memory"] = trad_total_allocated
    results["traditional"]["avg_utilization"] = (
        trad_total_used / trad_total_allocated
    )

    # Paged: allocate only what's needed (+ partial block overhead)
    paged_blocks_used = sum(
        (l + block_size - 1) // block_size for l in actual_lengths
    )
    paged_tokens_capacity = paged_blocks_used * block_size
    results["paged"]["peak_memory"] = paged_tokens_capacity
    results["paged"]["avg_utilization"] = (
        trad_total_used / paged_tokens_capacity
    )

    return results, actual_lengths


results, lengths = compare_allocation_strategies()

# Calculate improvement metrics
trad_peak = results["traditional"]["peak_memory"]
paged_peak = results["paged"]["peak_memory"]
memory_reduction = 1 - (paged_peak / trad_peak)
capacity_multiplier = 1 / (1 - memory_reduction)

import matplotlib.pyplot as plt
import numpy as np

plt.rcParams.update(
    {
        "figure.figsize": (6.0, 4.0),
        "figure.dpi": 300,
        "figure.constrained_layout.use": True,
        "font.family": "sans-serif",
        "font.sans-serif": [
            "Noto Sans CJK SC",
            "Apple SD Gothic Neo",
            "DejaVu Sans",
            "Arial",
        ],
        "font.size": 10,
        "axes.titlesize": 11,
        "axes.titleweight": "bold",
        "axes.titlepad": 8,
        "axes.labelsize": 10,
        "axes.labelpad": 4,
        "xtick.labelsize": 9,
        "ytick.labelsize": 9,
        "legend.fontsize": 9,
        "legend.title_fontsize": 10,
        "legend.frameon": True,
        "legend.loc": "best",
        "lines.linewidth": 1.5,
        "lines.markersize": 5,
        "axes.grid": True,
        "grid.alpha": 0.3,
        "grid.linestyle": "--",
        "axes.spines.top": False,
        "axes.spines.right": False,
        "axes.prop_cycle": plt.cycler(
            color=["#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#7f7f7f"]
        ),
    }
)

plt.figure()
plt.hist(lengths, bins=30, edgecolor="black", alpha=0.7, color="#3498db")
plt.axvline(
    x=2048,
    color="#e74c3c",
    linestyle="--",
    linewidth=2,
    label="Max length (allocation)",
)
plt.axvline(
    x=np.mean(lengths),
    color="#2ecc71",
    linestyle="-",
    linewidth=2,
    label=f"Mean length ({np.mean(lengths):.0f})",
)
plt.xlabel("Sequence Length (tokens)")
plt.ylabel("Count")
plt.title("Actual Sequence Length Distribution")
plt.legend()
plt.show()

import matplotlib.pyplot as plt

plt.rcParams.update(
    {
        "figure.figsize": (6.0, 4.0),
        "figure.dpi": 300,
        "figure.constrained_layout.use": True,
        "font.family": "sans-serif",
        "font.sans-serif": [
            "Noto Sans CJK SC",
            "Apple SD Gothic Neo",
            "DejaVu Sans",
            "Arial",
        ],
        "font.size": 10,
        "axes.titlesize": 11,
        "axes.titleweight": "bold",
        "axes.titlepad": 8,
        "axes.labelsize": 10,
        "axes.labelpad": 4,
        "xtick.labelsize": 9,
        "ytick.labelsize": 9,
        "legend.fontsize": 9,
        "legend.title_fontsize": 10,
        "legend.frameon": True,
        "legend.loc": "best",
        "lines.linewidth": 1.5,
        "lines.markersize": 5,
        "axes.grid": True,
        "grid.alpha": 0.3,
        "grid.linestyle": "--",
        "axes.spines.top": False,
        "axes.spines.right": False,
        "axes.prop_cycle": plt.cycler(
            color=["#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#7f7f7f"]
        ),
    }
)

plt.figure()
strategies = ["Traditional\n(Contiguous)", "Paged\n(vLLM)"]
utilizations = [
    results["traditional"]["avg_utilization"] * 100,
    results["paged"]["avg_utilization"] * 100,
]
colors = ["#e74c3c", "#2ecc71"]

bars = plt.bar(strategies, utilizations, color=colors, edgecolor="black")
plt.ylabel("Memory Utilization (%)")
plt.title("KV Cache Memory Utilization")
plt.ylim(0, 100)

## Add percentage labels
for bar, util in zip(bars, utilizations):
    plt.annotate(
        f"{util:.1f}%",
        xy=(bar.get_x() + bar.get_width() / 2, bar.get_height() + 2),
        ha="center",
        fontsize=9,
        fontweight="bold",
    )
plt.show()