Neural Mastery
← Practice
PagedAttention Block Allocatormedium

PagedAttention Block Allocator

medium Inside an LLM 20–25 min

NumPy allowed · plain Python earns 10 bonus points

Attention step 1/4

Tokens attend to tokens

Q · Kᵀ / √d → softmax → weighted values, computed in the page.
Open lesson →
01tokens
02Q K V
03scores
04attention
modelslearnrelationships
models32%
learn38%
relationships30%

Q · Kᵀ / √d scores normalized with softmax

Manage non-contiguous physical memory allocation for LLM key-value caches using PagedAttention virtual memory mapping.

Task

Implement allocate_paged_blocks(seq_len, new_tokens, block_table, free_blocks, block_size). Calculate total tokens total_tokens = seq_len + new_tokens. Total physical blocks needed is ceil(total_tokens / block_size). If existing len(block_table) is insufficient, allocate additional blocks from the front of free_blocks. If not enough free blocks remain, raise MemoryError. Calculate internal_fragmentation_tokens (unused token slots in the last block). Return a dict with block_table, remaining_free_blocks, and internal_fragmentation_tokens.

Function Signature

allocate_paged_blocks(seq_len: int, new_tokens: int, block_table: list[int], free_blocks: list[int], block_size: int) -> dict

Examples

Example 1: Allocate additional blocks
Input: seq_len = 14 new_tokens = 5 block_table = [101] free_blocks = [102, 103, 104] block_size = 16
Output: {'block_table': [101, 102], 'remaining_free_blocks': [103, 104], 'internal_fragmentation_tokens': 13}
Example 2: No new blocks required
Input: seq_len = 2 new_tokens = 1 block_table = [42] free_blocks = [43, 44] block_size = 4
Output: {'block_table': [42], 'remaining_free_blocks': [43, 44], 'internal_fragmentation_tokens': 1}
Example 3: Out of memory error
Input: seq_len = 0 new_tokens = 33 block_table = [] free_blocks = [1] block_size = 16
Output: Raises MemoryError

Constraints

  • Insufficient free blocks must raise MemoryError.
  • seq_len < 0, new_tokens < 0, or block_size <= 0 must raise ValueError.
  • Blocks must be allocated sequentially from free_blocks.
  • A plain-Python solution earns 10 bonus points.
Python3Saved ✓
def allocate_paged_blocks(seq_len: int, new_tokens: int, block_table: list[int], free_blocks: list[int], block_size: int) -> dict:
"""Allocate physical blocks for logical token growth in PagedAttention.
Args:
seq_len: Current token length of the sequence.
new_tokens: Number of incoming tokens to be stored.
block_table: List of physical block IDs currently allocated.
free_blocks: Available unallocated physical block IDs.
block_size: Capacity of each physical block in tokens.
Returns:
dict with 'block_table', 'remaining_free_blocks', and 'internal_fragmentation_tokens'.
"""
# Your implementation here
pass

Case 1: Allocate additional blocks
Input: seq_len = 14 new_tokens = 5 block_table = [101] free_blocks = [102, 103, 104] block_size = 16
Expected: {'block_table': [101, 102], 'remaining_free_blocks': [103, 104], 'internal_fragmentation_tokens': 13}
Case 2: No new blocks required
Input: seq_len = 2 new_tokens = 1 block_table = [42] free_blocks = [43, 44] block_size = 4
Expected: {'block_table': [42], 'remaining_free_blocks': [43, 44], 'internal_fragmentation_tokens': 1}
Case 3: Out of memory error
Input: seq_len = 0 new_tokens = 33 block_table = [] free_blocks = [1] block_size = 16
Expected: Raises MemoryError