Module 10. KV Caches and Efficient Inference
Read the explanations and code in the existing notebooks, or download the files. Code does not run on this site.
Download course files (.zip) 2588.1 KBThe archive preserves the source folder structure, both languages, and supporting files.
Notebook 10. KV cache and serving: bytes, pages, and ceilings
Main module notebook
Download .ipynb module_10_notebook_kvcache.ipynb · 70.1 KB
Show notebook text
Notebook 10. KV cache and serving: bytes, pages, and ceilings
Module 10 — hands-on notebook
This is engineering-disciplined educational code, not a vLLM implementation or a CUDA kernel. The mandatory path takes one KV geometry through four levels:
-
exact payload per token and context;
-
an arithmetic decode ceiling and batch knee;
-
a page table with prefix sharing, copy-on-write, and block release;
-
prefix caching, chunked prefill, compression, and eviction.
The unit contract is explicit:
KiB/MiB/GiB/TiBuse powers of two;GB/TBand bandwidth use powers of ten;- payload, upper bound, and measured device occupancy are different quantities.
Colab contract: Runtime → Run all completes on CPU in a clean session. The notebook
uses the standard library for arithmetic and NumPy/PyTorch for B17. It reads no repository files and needs no upload, Google Drive, API key, or paid service. Network access is needed only to install missing dependencies. The frozen M10 numerical contract is embedded
and checked by SHA-256.
0. Embedded numerical contract
The reference inputs belong to M10 rather than mutable lecture prose. Compact JSON is
embedded in the notebook; SHA-256 prevents accidental drift among the builder, answer
sheet, and published .ipynb.
Python
from __future__ import annotations
from dataclasses import dataclass
import hashlib
import json
import math
from typing import Hashable, Sequence
REFERENCE_JSON = '{"capacity_upper_bound":{"device_binary_gib":80,"max_requests":{"1000000":0,"131072":3,"32768":14,"8192":57},"runtime_reserve_binary_gib":8},"compression":{"base_bits_per_coordinate":16,"base_kib_per_token":128.0,"ceiling_128k_tokens_per_second":174.286494109339,"compressed_bits_per_coordinate":3,"compressed_kib_per_token":24.0},"context_memory_70b_gqa":{"1000000":{"binary_gib":305.17578125,"bytes":327680000000,"decimal_gb":327.68},"10000000":{"binary_gib":3051.7578125,"bytes":3276800000000,"decimal_gb":3276.8},"131072":{"binary_gib":40.0,"bytes":42949672960,"decimal_gb":42.94967296},"32768":{"binary_gib":10.0,"bytes":10737418240,"decimal_gb":10.73741824},"8192":{"binary_gib":2.5,"bytes":2684354560,"decimal_gb":2.68435456}},"eviction":{"attention_mass":[0.625,0.897,0.776,0.225,0.3,0.874,0.005,0.821,0.797,0.468,0.303,0.278,0.255,0.445,0.505,0.553],"heavy_hitters":[0,1,2,5,7,8,14,15],"lost_heavy_hitters":[5,7,8],"streaming":[0,1,2,3,12,13,14,15]},"hardware":{"bandwidth_decimal_tb_per_second":3.35,"device_binary_gib":80.0,"peak_decimal_tflops":990.0,"runtime_reserve_binary_gib":8.0,"weight_binary_gib":14.901161193848,"weight_bytes":16000000000,"weight_decimal_gb":16.0},"kv_ladder":{"70B MHA-equivalent (H=64)":{"bytes_per_token":2621440,"kib_per_token":2560.0},"DeepSeek-V3 MLA (latent+RoPE)":{"bytes_per_token":70272,"kib_per_token":68.625},"LLaMA-3.3-70B GQA":{"bytes_per_token":327680,"kib_per_token":320.0},"Llama-3-8B GQA":{"bytes_per_token":131072,"kib_per_token":128.0}},"paged":{"exact_2048_waste_fraction":0.0,"expected_paged_waste_fraction":0.003648747263,"naive_waste_fraction":0.9375,"preallocation_capacity_factor":16.0,"trace":{"after_a":{"allocated_blocks":2,"free_blocks":6,"internal_waste_tokens":2,"logical_tokens":6,"physical_capacity_tokens":8,"physical_payload_tokens":6,"sequences":{"A":{"blocks":[0,1],"length":6}},"sharing_savings_tokens":0},"after_a_release":{"allocated_blocks":3,"free_blocks":5,"internal_waste_tokens":3,"logical_tokens":9,"physical_capacity_tokens":12,"physical_payload_tokens":9,"sequences":{"B":{"blocks":[0,2,3],"length":9}},"sharing_savings_tokens":0},"after_b_append":{"allocated_blocks":4,"free_blocks":4,"internal_waste_tokens":5,"logical_tokens":15,"physical_capacity_tokens":16,"physical_payload_tokens":11,"sequences":{"A":{"blocks":[0,1],"length":6},"B":{"blocks":[0,2,3],"length":9}},"sharing_savings_tokens":4},"after_share":{"allocated_blocks":2,"free_blocks":6,"internal_waste_tokens":2,"logical_tokens":12,"physical_capacity_tokens":8,"physical_payload_tokens":6,"sequences":{"A":{"blocks":[0,1],"length":6},"B":{"blocks":[0,1],"length":6}},"sharing_savings_tokens":6}}},"prefix_and_interference":{"arithmetic_speedup":16.0,"cached_causal_suffix_ms":10.426705047014,"cached_suffix_prefill_ms":8.274747474747,"causal_arithmetic_ratio":14.402250086446,"chunked_itl_ms":13.050866877733,"decode_weight_floor_ms":4.776119402985,"full_causal_prefill_ms":150.168013664711,"full_prefill_ms":132.39595959596,"illustrative_input_cost":{"amortized_cost_factor":3.786982248521,"baseline":81920,"cached":21632.0,"warm_request_cost_ratio":0.15625},"monolithic_spike_ms":137.172078998945},"roofline":{"batch_32_8k":{"aggregate_tokens_per_second":2128.684609452179,"ideal_without_kv_tokens_per_second":6700.0,"kv_tax_factor":3.147483648},"batch_knees":{"131072":0.931322574615,"32768":3.725290298462,"8192":14.901161193848},"ceilings":{"0":{"aggregate_tokens_per_second":209.375,"kv_bytes":0},"1000000":{"aggregate_tokens_per_second":22.777959094865,"kv_bytes":131072000000},"131072":{"aggregate_tokens_per_second":100.96483447305,"kv_bytes":17179869184},"8192":{"aggregate_tokens_per_second":196.207722626508,"kv_bytes":1073741824}},"intensity_flops_per_byte_128k":0.482220104946,"intensity_including_attention_128k":2.553339685162,"ridge_flops_per_byte":295.522388059701,"ridge_to_full_intensity":115.739550744859,"ridge_to_intensity":612.837136047761},"schema_version":3,"unit_policy":{"binary":"KiB/MiB/GiB/TiB use powers of two","decimal":"GB/TB and bandwidth use powers of ten"}}'
REFERENCE_SHA256 = '1a283dbcb4061ab38b42bdf2d48f2ad1b1c844ffe88faab9137703985ee9ca5b'
actual_sha256 = hashlib.sha256(REFERENCE_JSON.encode("utf-8")).hexdigest()
assert actual_sha256 == REFERENCE_SHA256
REF = json.loads(REFERENCE_JSON)
assert REF["schema_version"] == 3
KIB, MIB, GIB, TIB = 2**10, 2**20, 2**30, 2**40
DECIMAL_GB, DECIMAL_TB = 10**9, 10**12
print('The embedded M10 contract is verified', actual_sha256[:12])1. Typed KV geometry and units
For ordinary MHA/GQA/MQA:
The factor two is K plus V. For MLA we count the stored latent and separate positional component rather than imaginary materialized K/V. A hybrid model must pass the number of attention layers, not total layers.
Python
def positive_int(name: str, value: int) -> int:
"""Validate a strictly positive integer without accepting bools."""
if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
raise ValueError(f"{name} must be a positive integer")
return value
def non_negative_int(name: str, value: int) -> int:
"""Validate a non-negative integer without accepting bools."""
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
raise ValueError(f"{name} must be a non-negative integer")
return value
def positive_finite(name: str, value: float) -> float:
"""Validate a positive finite numeric scalar."""
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise TypeError(f"{name} must be numeric")
number = float(value)
if not math.isfinite(number) or number <= 0.0:
raise ValueError(f"{name} must be positive and finite")
return number
@dataclass(frozen=True)
class AttentionKVGeometry:
"""KV geometry for layers that materialize ordinary K and V tensors."""
name: str
attention_layers: int
kv_heads: int
head_dim: int
bytes_per_element: int = 2
def __post_init__(self) -> None:
"""Validate the declared ordinary-attention geometry."""
if not isinstance(self.name, str) or not self.name.strip():
raise ValueError("name must be a non-empty string")
positive_int("attention_layers", self.attention_layers)
positive_int("kv_heads", self.kv_heads)
positive_int("head_dim", self.head_dim)
positive_int("bytes_per_element", self.bytes_per_element)
@property
def bytes_per_token(self) -> int:
"""Return K+V bytes for one token across all attention layers."""
return (2 * self.attention_layers * self.kv_heads
* self.head_dim * self.bytes_per_element)
@dataclass(frozen=True)
class MLAKVGeometry:
"""KV geometry for a stored latent and separate positional component."""
name: str
attention_layers: int
latent_dim: int
rope_dim: int
bytes_per_element: int = 2
def __post_init__(self) -> None:
"""Validate the declared latent-attention geometry."""
if not isinstance(self.name, str) or not self.name.strip():
raise ValueError("name must be a non-empty string")
positive_int("attention_layers", self.attention_layers)
positive_int("latent_dim", self.latent_dim)
non_negative_int("rope_dim", self.rope_dim)
positive_int("bytes_per_element", self.bytes_per_element)
@property
def bytes_per_token(self) -> int:
"""Return latent-plus-position bytes for one token across all layers."""
return (self.attention_layers * (self.latent_dim + self.rope_dim)
* self.bytes_per_element)
def kv_payload_bytes(geometry, context_tokens: int, batch_size: int = 1,
kv_scale: float = 1.0) -> int:
"""Return an exact integer KV payload for equal-length requests."""
tokens = non_negative_int("context_tokens", context_tokens)
batch = positive_int("batch_size", batch_size)
scale = positive_finite("kv_scale", kv_scale)
raw = geometry.bytes_per_token * tokens * batch * scale
rounded = round(raw)
if not math.isclose(raw, rounded, rel_tol=0.0, abs_tol=1e-9):
raise ValueError("scaled KV payload is not an integer byte count")
return int(rounded)
GEOMETRIES = {
"70B MHA-equivalent (H=64)": AttentionKVGeometry(
"70B MHA-equivalent (H=64)", 80, 64, 128, 2),
"LLaMA-3.3-70B GQA": AttentionKVGeometry(
"LLaMA-3.3-70B GQA", 80, 8, 128, 2),
"Llama-3-8B GQA": AttentionKVGeometry(
"Llama-3-8B GQA", 32, 8, 128, 2),
"DeepSeek-V3 MLA (latent+RoPE)": MLAKVGeometry(
"DeepSeek-V3 MLA (latent+RoPE)", 61, 512, 64, 2),
}
print('KV geometries and unit helpers are defined')2. The KV/token ladder and context growth
We display the same payload in decimal GB and binary GiB. For example, 70B GQA at
128K is exactly 40 GiB, or 42.95 GB; both numbers are valid, but they must not carry
the same unit label.
Python
for name, geometry in GEOMETRIES.items():
expected = REF["kv_ladder"][name]
assert geometry.bytes_per_token == expected["bytes_per_token"]
assert math.isclose(geometry.bytes_per_token / KIB,
expected["kib_per_token"], abs_tol=1e-12)
print(f"{name:36s} {geometry.bytes_per_token / KIB:8.3f} KiB/token")
g70 = GEOMETRIES["LLaMA-3.3-70B GQA"]
for token_text, expected in REF["context_memory_70b_gqa"].items():
tokens = int(token_text)
payload = kv_payload_bytes(g70, tokens)
assert payload == expected["bytes"]
assert math.isclose(payload / DECIMAL_GB, expected["decimal_gb"], abs_tol=1e-12)
assert math.isclose(payload / GIB, expected["binary_gib"], abs_tol=1e-12)
print(f"N={tokens:>10,}: {payload / DECIMAL_GB:9.3f} GB"
f" = {payload / GIB:9.3f} GiB")
print('The ladder and context payload are reproduced')Python
"""Small, dependency-free FLOP and billing models; no hardware timing claims."""
import math
def attention_flops(layers, query_heads, head_dim, queries, cached_prefix=0):
"""QK and AV for causal suffix queries, two FLOPs per multiply-add.
Query j sees cached_prefix + j + 1 keys. Softmax, norms, projections,
cache writes, and kernel padding are outside this count.
"""
for name, value in (("layers", layers), ("query_heads", query_heads),
("head_dim", head_dim), ("queries", queries),
("cached_prefix", cached_prefix)):
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
raise ValueError(f"{name} must be a non-negative integer")
if min(layers, query_heads, head_dim) == 0:
raise ValueError("attention geometry must be positive")
pairs = queries * cached_prefix + queries * (queries + 1) // 2
return 4 * layers * query_heads * head_dim * pairs
def decode_flops(parameters, layers, query_heads, head_dim, keys, batch=1):
"""Projection estimate plus QK/AV for one query attending to `keys`."""
for value in (parameters, batch, keys):
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
raise ValueError("counts must be non-negative integers")
if parameters == 0 or batch == 0:
raise ValueError("parameters and batch must be positive")
# Zero keys is a synthetic weights-only baseline, not a real decode query.
attn = attention_flops(layers, query_heads, head_dim, 1, max(keys - 1, 0))
return batch * (2 * parameters + (attn if keys else 0))
def prefill_flops(parameters, layers, query_heads, head_dim, suffix, prefix=0):
"""Only suffix projections are new; suffix attention still reads the prefix."""
if isinstance(parameters, bool) or not isinstance(parameters, int) or parameters <= 0:
raise ValueError("parameters must be a positive integer")
return (attention_flops(layers, query_heads, head_dim, suffix, prefix)
+ 2 * parameters * suffix)
def cached_input_cost(prefix, suffix, requests, read_rate=0.1, write_rate=1.25):
"""Input-token equivalents: one prefix write, then complete prefix hits.
Each suffix is uncached. No output cost, expiry, or extra cache writes.
Multipliers are illustrative inputs, not an API tariff lookup.
"""
for value in (prefix, suffix, requests):
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
raise ValueError("counts must be non-negative integers")
if requests == 0 or prefix + suffix == 0:
raise ValueError("need a request with non-empty input")
for rate in (read_rate, write_rate):
if isinstance(rate, bool) or not isinstance(rate, (int, float)) or not math.isfinite(rate) or rate <= 0:
raise ValueError("rates must be positive and finite")
baseline = requests * (prefix + suffix)
cached = prefix * (write_rate + (requests - 1) * read_rate) + requests * suffix
warm = prefix * read_rate + suffix
return {"baseline": baseline, "cached": cached,
"amortized_cost_factor": baseline / cached,
"warm_request_cost_ratio": warm / (prefix + suffix)}3. Decode roofline, intensity, and batch knees
The scenario reads the weights once per aggregate step and each request's KV:
This is a bandwidth ceiling, not a wall-clock benchmark. Activations, temporary buffers, communication, allocator overhead, and measured hardware utilization are outside it.
First count projections, then add attention using query heads: GQA storage depends on Hkv, whereas QK/AV work depends on Hq.
Python
@dataclass(frozen=True)
class ServingHardware:
"""Inputs for arithmetic serving ceilings and capacity upper bounds."""
weight_bytes: int
bandwidth_bytes_per_second: float
peak_flops_per_second: float
parameter_count: int
device_capacity_bytes: int
runtime_reserve_bytes: int = 0
def __post_init__(self) -> None:
"""Validate arithmetic hardware inputs and the reserve budget."""
positive_int("weight_bytes", self.weight_bytes)
positive_finite("bandwidth_bytes_per_second", self.bandwidth_bytes_per_second)
positive_finite("peak_flops_per_second", self.peak_flops_per_second)
positive_int("parameter_count", self.parameter_count)
positive_int("device_capacity_bytes", self.device_capacity_bytes)
non_negative_int("runtime_reserve_bytes", self.runtime_reserve_bytes)
if self.weight_bytes + self.runtime_reserve_bytes > self.device_capacity_bytes:
raise ValueError("weights plus reserve exceed device capacity")
@dataclass(frozen=True)
class DecodeCeiling:
"""Traffic and token-rate fields for one declared decode scenario."""
batch_size: int
context_tokens: int
kv_bytes: int
total_bytes_per_step: int
aggregate_tokens_per_second: float
per_request_tokens_per_second: float
def decode_ceiling(hardware: ServingHardware, geometry, context_tokens: int,
batch_size: int = 1, kv_scale: float = 1.0) -> DecodeCeiling:
"""Compute a weight-plus-KV bandwidth ceiling for one decode step."""
batch = positive_int("batch_size", batch_size)
tokens = non_negative_int("context_tokens", context_tokens)
kv_bytes = kv_payload_bytes(geometry, tokens, batch, kv_scale)
total = hardware.weight_bytes + kv_bytes
aggregate = batch * hardware.bandwidth_bytes_per_second / total
return DecodeCeiling(batch, tokens, kv_bytes, total, aggregate,
aggregate / batch)
def arithmetic_intensity(hardware: ServingHardware, geometry,
context_tokens: int, batch_size: int = 1,
kv_scale: float = 1.0) -> float:
"""Projection-only intensity: context attention is added in the next cell."""
scenario = decode_ceiling(hardware, geometry, context_tokens,
batch_size, kv_scale)
return (2 * hardware.parameter_count * batch_size
/ scenario.total_bytes_per_step)
def batch_knee(hardware: ServingHardware, geometry,
context_tokens: int, kv_scale: float = 1.0) -> float:
"""Return B* where one request's KV traffic equals weight traffic."""
tokens = positive_int("context_tokens", context_tokens)
return hardware.weight_bytes / kv_payload_bytes(
geometry, tokens, 1, kv_scale)
HARDWARE = ServingHardware(
weight_bytes=16_000_000_000,
bandwidth_bytes_per_second=3.35e12,
peak_flops_per_second=990e12,
parameter_count=8_000_000_000,
device_capacity_bytes=80 * GIB,
runtime_reserve_bytes=8 * GIB,
)
g8 = GEOMETRIES["Llama-3-8B GQA"]
for token_text, expected in REF["roofline"]["ceilings"].items():
scenario = decode_ceiling(HARDWARE, g8, int(token_text))
assert scenario.kv_bytes == expected["kv_bytes"]
assert math.isclose(scenario.aggregate_tokens_per_second,
expected["aggregate_tokens_per_second"], abs_tol=1e-9)
print(f"N={int(token_text):>9,}:"
f" {scenario.aggregate_tokens_per_second:8.3f} tok/s")
intensity = arithmetic_intensity(HARDWARE, g8, 131_072)
ridge = HARDWARE.peak_flops_per_second / HARDWARE.bandwidth_bytes_per_second
assert math.isclose(intensity,
REF["roofline"]["intensity_flops_per_byte_128k"], abs_tol=1e-12)
assert math.isclose(ridge, REF["roofline"]["ridge_flops_per_byte"], abs_tol=1e-12)
for token_text, expected in REF["roofline"]["batch_knees"].items():
assert math.isclose(batch_knee(HARDWARE, g8, int(token_text)),
expected, abs_tol=1e-12)
b32 = decode_ceiling(HARDWARE, g8, 8_192, batch_size=32)
assert math.isclose(b32.aggregate_tokens_per_second,
REF["roofline"]["batch_32_8k"]["aggregate_tokens_per_second"],
abs_tol=1e-9)
print(f"128K intensity={intensity:.6f} FLOP/B; ridge={ridge:.3f} FLOP/B")
full_intensity = decode_flops(8_000_000_000, 32, 32, 128, 131072) / decode_ceiling(HARDWARE, g8, 131072).total_bytes_per_step
assert math.isclose(full_intensity, REF["roofline"]["intensity_including_attention_128k"], abs_tol=1e-12)
print(f"Including QK/AV: {full_intensity:.6f} FLOP/B")
print(f"Batch 32: aggregate={b32.aggregate_tokens_per_second:.3f} tok/s; per request={b32.per_request_tokens_per_second:.3f} tok/s")
print('The roofline, intensity, and batch knees are verified')4. An upper bound on resident requests
We subtract the weight payload and an explicit 8 GiB runtime reserve from 80 GiB,
then divide the remainder by one request's KV. The result is a KV-only upper bound:
real capacity and goodput are lower because of workspaces, pages, fragmentation, SLOs,
and scheduler policy.
Python
def max_resident_requests(hardware: ServingHardware, geometry,
context_tokens: int, kv_scale: float = 1.0) -> int:
"""Return a KV-only capacity upper bound after weights and reserve."""
tokens = positive_int("context_tokens", context_tokens)
free = (hardware.device_capacity_bytes - hardware.weight_bytes
- hardware.runtime_reserve_bytes)
return free // kv_payload_bytes(geometry, tokens, 1, kv_scale)
for token_text, expected in REF["capacity_upper_bound"]["max_requests"].items():
actual = max_resident_requests(HARDWARE, g8, int(token_text))
assert actual == expected
print(f"N={int(token_text):>9,}: at most {actual} resident requests"
" (KV-only upper bound)")
assert math.isclose(HARDWARE.weight_bytes / GIB,
REF["hardware"]["weight_binary_gib"], abs_tol=1e-12)
print('KV-only capacity upper bounds are verified')5. Page pool: prefix sharing and copy-on-write
The simulator stores logical lengths, page tables, fill, and reference counts, not real tensors. If two sequences share a partial final block, a writer first makes a private copy. Allocation is preflighted, so OOM cannot leave a partially mutated pool.
Python
_positive_int = positive_int
_non_negative_int = non_negative_int
@dataclass
class _PhysicalBlock:
"""Internal physical block state for the teaching page-pool simulator."""
filled_slots: int = 0
ref_count: int = 0
@dataclass
class _SequenceState:
"""Internal logical sequence state."""
length: int
blocks: list[int]
class PagedKVPool:
"""A validated page-table simulator with prefix sharing and copy-on-write.
It models block ownership and logical lengths; it does not store real tensors. A
shared partial final block is copied before either sequence appends into it. All
allocations are preflighted so an out-of-memory error leaves the pool unchanged.
"""
def __init__(self, total_blocks: int, block_size: int = 16) -> None:
self.total_blocks = _positive_int("total_blocks", total_blocks)
self.block_size = _positive_int("block_size", block_size)
self._blocks = [_PhysicalBlock() for _ in range(self.total_blocks)]
self._free = list(range(self.total_blocks - 1, -1, -1))
self._sequences: dict[Hashable, _SequenceState] = {}
def _require_sequence(self, sequence_id: Hashable) -> _SequenceState:
try:
return self._sequences[sequence_id]
except KeyError as exc:
raise KeyError(f"unknown sequence: {sequence_id!r}") from exc
def create(self, sequence_id: Hashable) -> None:
"""Create an empty sequence with a unique hashable identifier."""
try:
hash(sequence_id)
except TypeError as exc:
raise TypeError("sequence_id must be hashable") from exc
if sequence_id in self._sequences:
raise ValueError(f"sequence already exists: {sequence_id!r}")
self._sequences[sequence_id] = _SequenceState(length=0, blocks=[])
def _blocks_needed_for_append(self, state: _SequenceState, token_count: int) -> int:
if token_count == 0:
return 0
needed = 0
remainder = state.length % self.block_size
available_in_last = 0
if state.blocks and remainder:
last = self._blocks[state.blocks[-1]]
if last.ref_count > 1:
needed += 1 # copy-on-write replacement for the shared partial block
available_in_last = self.block_size - remainder
remaining = max(0, token_count - available_in_last)
needed += math.ceil(remaining / self.block_size)
return needed
def _allocate_block(self, filled_slots: int = 0) -> int:
block_id = self._free.pop()
block = self._blocks[block_id]
block.filled_slots = filled_slots
block.ref_count = 1
return block_id
def append(self, sequence_id: Hashable, token_count: int) -> None:
"""Append logical tokens, allocating arbitrary physical blocks as needed."""
count = _non_negative_int("token_count", token_count)
state = self._require_sequence(sequence_id)
required = self._blocks_needed_for_append(state, count)
if required > len(self._free):
raise MemoryError(
f"need {required} free blocks, but only {len(self._free)} remain"
)
if count == 0:
return
remainder = state.length % self.block_size
if state.blocks and remainder:
last_id = state.blocks[-1]
last = self._blocks[last_id]
if last.ref_count > 1:
last.ref_count -= 1
replacement = self._allocate_block(filled_slots=remainder)
state.blocks[-1] = replacement
last_id = replacement
last = self._blocks[last_id]
take = min(count, self.block_size - remainder)
last.filled_slots = remainder + take
state.length += take
count -= take
while count:
take = min(count, self.block_size)
block_id = self._allocate_block(filled_slots=take)
state.blocks.append(block_id)
state.length += take
count -= take
self._refresh_fills()
self.validate()
def share_prefix(
self,
source_id: Hashable,
target_id: Hashable,
prefix_tokens: int,
) -> None:
"""Create ``target_id`` by sharing source blocks that cover a logical prefix."""
source = self._require_sequence(source_id)
prefix = _non_negative_int("prefix_tokens", prefix_tokens)
if prefix > source.length:
raise ValueError("prefix_tokens exceeds source length")
if target_id in self._sequences:
raise ValueError(f"target sequence already exists: {target_id!r}")
try:
hash(target_id)
except TypeError as exc:
raise TypeError("target_id must be hashable") from exc
block_count = math.ceil(prefix / self.block_size) if prefix else 0
table = source.blocks[:block_count]
for block_id in table:
self._blocks[block_id].ref_count += 1
self._sequences[target_id] = _SequenceState(prefix, list(table))
self.validate()
def release(self, sequence_id: Hashable) -> None:
"""Release a sequence and return unreferenced physical blocks to the free list."""
state = self._require_sequence(sequence_id)
for block_id in state.blocks:
block = self._blocks[block_id]
block.ref_count -= 1
if block.ref_count == 0:
block.filled_slots = 0
self._free.append(block_id)
del self._sequences[sequence_id]
self._refresh_fills()
self.validate()
def _refresh_fills(self) -> None:
"""Count live slots, including unequal prefixes sharing a physical block."""
fills = [0] * self.total_blocks
for live in self._sequences.values():
for j, block_id in enumerate(live.blocks):
fills[block_id] = max(fills[block_id], min(self.block_size, live.length - j * self.block_size))
for block_id, block in enumerate(self._blocks):
block.filled_slots = fills[block_id]
def sequence_length(self, sequence_id: Hashable) -> int:
"""Return logical length in tokens."""
return self._require_sequence(sequence_id).length
def page_table(self, sequence_id: Hashable) -> tuple[int, ...]:
"""Return an immutable page table for inspection."""
return tuple(self._require_sequence(sequence_id).blocks)
@property
def allocated_blocks(self) -> int:
"""Return the number of physical blocks with at least one reference."""
return self.total_blocks - len(self._free)
@property
def free_blocks(self) -> int:
"""Return the number of currently free physical blocks."""
return len(self._free)
@property
def physical_capacity_tokens(self) -> int:
"""Return token slots reserved by allocated physical blocks."""
return self.allocated_blocks * self.block_size
@property
def physical_payload_tokens(self) -> int:
"""Return filled token slots stored once per physical block."""
return sum(block.filled_slots for block in self._blocks if block.ref_count)
@property
def logical_tokens(self) -> int:
"""Return total logical tokens across all sequences, counting shared prefixes."""
return sum(state.length for state in self._sequences.values())
@property
def internal_waste_tokens(self) -> int:
"""Return unused slots inside allocated physical blocks."""
return self.physical_capacity_tokens - self.physical_payload_tokens
@property
def sharing_savings_tokens(self) -> int:
"""Return logical tokens not duplicated physically because of prefix sharing."""
return self.logical_tokens - self.physical_payload_tokens
def snapshot(self) -> dict[str, object]:
"""Return a deterministic JSON-friendly state summary."""
sequence_items = sorted(self._sequences.items(), key=lambda item: repr(item[0]))
return {
"sequences": {
str(sequence_id): {
"length": state.length,
"blocks": list(state.blocks),
}
for sequence_id, state in sequence_items
},
"allocated_blocks": self.allocated_blocks,
"free_blocks": self.free_blocks,
"physical_capacity_tokens": self.physical_capacity_tokens,
"physical_payload_tokens": self.physical_payload_tokens,
"logical_tokens": self.logical_tokens,
"internal_waste_tokens": self.internal_waste_tokens,
"sharing_savings_tokens": self.sharing_savings_tokens,
}
def validate(self) -> None:
"""Check page-table, reference-count, fill, and capacity invariants."""
if len(set(self._free)) != len(self._free):
raise AssertionError("free list contains duplicates")
if any(block_id < 0 or block_id >= self.total_blocks for block_id in self._free):
raise AssertionError("free list contains an invalid block id")
expected_refs = [0] * self.total_blocks
for state in self._sequences.values():
if state.length < 0:
raise AssertionError("sequence length is negative")
expected_count = math.ceil(state.length / self.block_size) if state.length else 0
if len(state.blocks) != expected_count:
raise AssertionError("page-table length does not match sequence length")
if len(set(state.blocks)) != len(state.blocks):
raise AssertionError("one sequence references a physical block twice")
for block_id in state.blocks:
if block_id < 0 or block_id >= self.total_blocks:
raise AssertionError("page table contains an invalid block id")
expected_refs[block_id] += 1
free_set = set(self._free)
for block_id, block in enumerate(self._blocks):
if block.ref_count != expected_refs[block_id]:
raise AssertionError("physical ref_count does not match page tables")
if block.ref_count == 0:
if block_id not in free_set or block.filled_slots != 0:
raise AssertionError("unreferenced block is not clean and free")
else:
if block_id in free_set:
raise AssertionError("referenced block appears in free list")
if not 1 <= block.filled_slots <= self.block_size:
raise AssertionError("referenced block has invalid fill")
if self.logical_tokens < self.physical_payload_tokens:
raise AssertionError("physical payload exceeds logical payload")6. Page trace and fragmentation
Sequence A receives 6 tokens at block size 4. B shares the entire prefix, then appends
3 tokens: the shared partial block goes through copy-on-write. Releasing A returns its
unreferenced block to the free list.
The 0.365% figure is the ratio of expected unused to expected allocated slots,
assuming mean length 2048 and uniform residues modulo 16. It is neither the mean
of individual waste fractions nor a measurement of a concrete workload.
A fixed length of 2048 is divisible by 16 and wastes zero slots. The 0.365% expectation assumes varying lengths; its denominator is allocated slots.
Python
def naive_preallocation_waste(max_tokens: int, used_tokens: int) -> float:
"""Return unused fraction of a fixed maximum-length allocation."""
maximum = positive_int("max_tokens", max_tokens)
used = non_negative_int("used_tokens", used_tokens)
if used > maximum:
raise ValueError("used_tokens cannot exceed max_tokens")
return (maximum - used) / maximum
def expected_paged_waste(avg_tokens: int, block_size: int) -> float:
"""Return expected final-block waste under a uniform remainder."""
average = positive_int("avg_tokens", avg_tokens)
block = positive_int("block_size", block_size)
tail = (block - 1) / 2
return tail / (average + tail)
pool = PagedKVPool(total_blocks=8, block_size=4)
pool.create("A")
pool.append("A", 6)
assert pool.snapshot() == REF["paged"]["trace"]["after_a"]
pool.share_prefix("A", "B", 6)
assert pool.snapshot() == REF["paged"]["trace"]["after_share"]
pool.append("B", 3)
assert pool.snapshot() == REF["paged"]["trace"]["after_b_append"]
assert pool.page_table("A")[-1] != pool.page_table("B")[1]
pool.release("A")
assert pool.snapshot() == REF["paged"]["trace"]["after_a_release"]
naive = naive_preallocation_waste(32_768, 2_048)
paged = expected_paged_waste(2_048, 16)
assert math.isclose(naive, REF["paged"]["naive_waste_fraction"], abs_tol=1e-12)
assert math.isclose(paged, REF["paged"]["expected_paged_waste_fraction"], abs_tol=1e-12)
print(f"naive waste={naive:.2%}; expected paged waste={paged:.3%}")
print('Prefix sharing, COW, release, and waste are verified')7. Prefix reuse: work saved and input cost
The 2PN / peak baseline counts projections only: 132.396 ms for 8192 tokens and 8.275 ms for a 512-token suffix, a ratio of 16. The suffix still attends to all 7680 cached positions. Add causal QK and AV operations, then check that the work of all 16 chunks adds up to the full prefill. Later chunks cost more than earlier ones. These idealized times do not predict a scheduler's ITL.
Keep billing separate: write the prefix once at 1.25, read it nine times at 0.1, and charge ten uncached suffixes. Ignore output tokens and additional cache writes.
Python
def prefill_floor_seconds(parameter_count: int, prompt_tokens: int,
peak_flops_per_second: float) -> float:
"""Return the dense 2*P*N/peak_FLOPs arithmetic floor."""
parameters = positive_int("parameter_count", parameter_count)
tokens = non_negative_int("prompt_tokens", prompt_tokens)
peak = positive_finite("peak_flops_per_second", peak_flops_per_second)
return 2 * parameters * tokens / peak
full = prefill_floor_seconds(8_000_000_000, 8_192, 990e12)
suffix = prefill_floor_seconds(8_000_000_000, 512, 990e12)
decode_weight = 16_000_000_000 / 3.35e12
spike = full + decode_weight
chunked = suffix + decode_weight
expected = REF["prefix_and_interference"]
assert math.isclose(full * 1e3, expected["full_prefill_ms"], abs_tol=1e-9)
assert math.isclose(suffix * 1e3, expected["cached_suffix_prefill_ms"], abs_tol=1e-9)
assert math.isclose(full / suffix, expected["arithmetic_speedup"], abs_tol=1e-12)
assert math.isclose(spike * 1e3, expected["monolithic_spike_ms"], abs_tol=1e-9)
assert math.isclose(chunked * 1e3, expected["chunked_itl_ms"], abs_tol=1e-9)
print(f"prefill={full*1e3:.3f} ms → suffix={suffix*1e3:.3f} ms")
print(f"monolithic spike={spike*1e3:.3f} ms; chunked floor={chunked*1e3:.3f} ms")
full_causal = prefill_flops(8_000_000_000, 32, 32, 128, 8192) / 990e12
suffix_causal = prefill_flops(8_000_000_000, 32, 32, 128, 512, 7680) / 990e12
assert math.isclose(full_causal*1e3, expected["full_causal_prefill_ms"], abs_tol=1e-9)
assert math.isclose(suffix_causal*1e3, expected["cached_causal_suffix_ms"], abs_tol=1e-9)
# All chunks together contain exactly the same causal query-key pairs.
chunks = [prefill_flops(8_000_000_000, 32, 32, 128, 512, p) for p in range(0,8192,512)]
assert sum(chunks) == prefill_flops(8_000_000_000, 32, 32, 128, 8192)
assert chunks[-1] > chunks[0]
cost = cached_input_cost(7680, 512, 10)
for key, value in cost.items():
assert math.isclose(value, expected["illustrative_input_cost"][key], abs_tol=1e-9)
print(f"Including attention: {full_causal*1e3:.3f} / {suffix_causal*1e3:.3f} ms; ratio={full_causal/suffix_causal:.3f}")
print("Input-token equivalents:", cost)
print('Prefix/interference arithmetic is reproduced')8. Compression and eviction
Moving from 16 to 3 bits per coordinate reduces 128 KiB/token to 24 KiB/token. This is a storage scenario; the notebook does not model quantization quality.
Eviction is compared on one frozen trace: content-aware top-mass versus a sink-plus-recent-window rule. Tie-breaking and the budget are explicit.
Python
def heavy_hitter_indices(attention_mass: Sequence[float], budget: int) -> tuple[int, ...]:
"""Return top-mass indices with deterministic lower-index tie-breaking."""
values = tuple(attention_mass)
if not values:
raise ValueError("attention_mass must be non-empty")
keep = positive_int("budget", budget)
if keep > len(values):
raise ValueError("budget exceeds token count")
normalized = []
for index, value in enumerate(values):
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise TypeError(f"attention_mass[{index}] must be numeric")
number = float(value)
if not math.isfinite(number) or number < 0:
raise ValueError("attention masses must be finite and non-negative")
normalized.append(number)
ranked = sorted(range(len(values)), key=lambda i: (-normalized[i], i))
return tuple(sorted(ranked[:keep]))
def streaming_keep_indices(length: int, budget: int,
sink_tokens: int) -> tuple[int, ...]:
"""Return sink tokens plus the most recent window under a fixed budget."""
n = non_negative_int("length", length)
keep = positive_int("budget", budget)
sinks = non_negative_int("sink_tokens", sink_tokens)
if sinks > keep:
raise ValueError("sink_tokens cannot exceed budget")
if n <= keep:
return tuple(range(n))
sink_count = min(sinks, n)
recent_count = keep - sink_count
recent_start = max(sink_count, n - recent_count)
result = tuple(range(sink_count)) + tuple(range(recent_start, n))
assert len(result) == keep and len(set(result)) == keep
return result
compression = REF["compression"]
scale = compression["compressed_bits_per_coordinate"] / compression["base_bits_per_coordinate"]
compressed = decode_ceiling(HARDWARE, g8, 131_072, kv_scale=scale)
assert math.isclose(g8.bytes_per_token * scale / KIB,
compression["compressed_kib_per_token"], abs_tol=1e-12)
assert math.isclose(compressed.aggregate_tokens_per_second,
compression["ceiling_128k_tokens_per_second"], abs_tol=1e-9)
mass = REF["eviction"]["attention_mass"]
heavy = heavy_hitter_indices(mass, 8)
streaming = streaming_keep_indices(len(mass), 8, 4)
lost = tuple(index for index in heavy if index not in streaming)
assert list(heavy) == REF["eviction"]["heavy_hitters"]
assert list(streaming) == REF["eviction"]["streaming"]
assert list(lost) == REF["eviction"]["lost_heavy_hitters"]
print(f"KV: {g8.bytes_per_token/KIB:g} → {g8.bytes_per_token*scale/KIB:g} KiB/token")
print(f"heavy={heavy}; streaming={streaming}; lost={lost}")
print('Compression and the eviction trace are verified')9. Negative tests
We reject invalid shapes and boundaries: zero geometry sizes, negative context, a
non-integral byte payload after scaling, an oversubscribed device, unknown sequence,
duplicate ID, prefix longer than its source, OOM with no partial mutation, an invalid
eviction budget, and NaN.
Python
def expect_error(error_type, fn, *args, **kwargs) -> None:
"""Require one declared exception from a negative contract test."""
try:
fn(*args, **kwargs)
except error_type:
return
raise AssertionError(f"expected {error_type.__name__} from {fn.__name__}")
expect_error(ValueError, AttentionKVGeometry, "bad", 0, 8, 128, 2)
expect_error(ValueError, kv_payload_bytes, g8, -1)
expect_error(ValueError, kv_payload_bytes, g8, 1, 1, 1/3)
expect_error(ValueError, ServingHardware, 80*GIB, 1.0, 1.0, 1, 80*GIB, 1)
expect_error(ValueError, heavy_hitter_indices, [0.1, 0.2], 3)
expect_error(ValueError, heavy_hitter_indices, [0.1, float("nan")], 1)
expect_error(ValueError, streaming_keep_indices, 16, 4, 5)
test_pool = PagedKVPool(total_blocks=2, block_size=4)
expect_error(KeyError, test_pool.append, "missing", 1)
test_pool.create("x")
expect_error(ValueError, test_pool.create, "x")
test_pool.append("x", 6)
before_oom = test_pool.snapshot()
expect_error(MemoryError, test_pool.append, "x", 3)
assert test_pool.snapshot() == before_oom # preflight prevents partial mutation
expect_error(ValueError, test_pool.share_prefix, "x", "y", 7)
test_pool.validate()
# A short child survives release or COW of its longer parent.
for action in ("release", "append"):
unequal = PagedKVPool(4, 4)
unequal.create("parent"); unequal.append("parent", 3)
unequal.share_prefix("parent", "child", 1)
if action == "release":
unequal.release("parent")
assert unequal.physical_payload_tokens == 1
else:
unequal.append("parent", 1)
assert unequal.physical_payload_tokens == 5
unequal.append("child", 2)
unequal.validate()
fixed = PagedKVPool(128,16); fixed.create("fixed"); fixed.append("fixed",2048)
assert fixed.internal_waste_tokens == 0
expect_error(ValueError, attention_flops, 32, 32, 128, -1)
expect_error(ValueError, cached_input_cost, 7680, 512, 0)
print('Negative tests pass')Lecture-level core
Python
"""Lecture-level geometry, bandwidth bound, and private page allocation."""
from dataclasses import dataclass, field
import math
def count(value, minimum=0):
if isinstance(value, bool) or not isinstance(value, int) or value < minimum:
raise ValueError(f"expected an integer >= {minimum}")
return value
@dataclass(frozen=True)
class KVGeometry:
attention_layers: int
kv_heads: int
head_dim: int
bytes_per_element: int = 2
def __post_init__(self):
for value in (self.attention_layers, self.kv_heads, self.head_dim, self.bytes_per_element):
count(value, 1)
@property
def bytes_per_token(self):
return 2 * self.attention_layers * self.kv_heads * self.head_dim * self.bytes_per_element
def decode_bandwidth_ceiling(context_tokens, weight_bytes, geometry, bandwidth_Bps,
batch_size=1, kv_scale=1.0):
"""Aggregate tok/s, assuming one weight read and one KV read per step."""
count(context_tokens); count(weight_bytes, 1); count(batch_size, 1)
for value in (bandwidth_Bps, kv_scale):
if isinstance(value, bool) or not math.isfinite(value) or value <= 0:
raise ValueError("bandwidth and scale must be positive and finite")
kv = geometry.bytes_per_token * context_tokens * batch_size * kv_scale
if not math.isclose(kv, round(kv), rel_tol=0, abs_tol=1e-9):
raise ValueError("the scaled payload must contain an integer number of bytes")
return batch_size * bandwidth_Bps / (weight_bytes + round(kv))
@dataclass
class PagedKV:
"""Private logical pages only; the notebook adds sharing and copy-on-write."""
n_blocks: int
block_size: int = 16
free: list = field(init=False)
tables: dict = field(default_factory=dict, init=False)
lengths: dict = field(default_factory=dict, init=False)
def __post_init__(self):
count(self.n_blocks, 1); count(self.block_size, 1)
self.free = list(range(self.n_blocks - 1, -1, -1))
def append_to(self, seq_id, total_tokens):
count(total_tokens)
hash(seq_id)
if total_tokens < self.lengths.get(seq_id, 0):
raise ValueError("append_to cannot shorten a sequence")
table = self.tables.get(seq_id, [])
needed = (total_tokens + self.block_size - 1) // self.block_size - len(table)
if needed > len(self.free):
raise MemoryError("not enough free blocks")
self.tables.setdefault(seq_id, table)
table.extend(self.free.pop() for _ in range(needed))
self.lengths[seq_id] = total_tokens
def release(self, seq_id):
self.free.extend(self.tables.pop(seq_id))
del self.lengths[seq_id]Python
geometry = KVGeometry(32,8,128)
assert geometry.bytes_per_token == 131072
assert math.isclose(decode_bandwidth_ceiling(8192,16_000_000_000,geometry,3.35e12),196.207722626508)
pages = PagedKV(128,16); pages.append_to("a",2048)
assert len(pages.tables["a"]) == 128
before = (list(pages.free), dict(pages.lengths), {k:list(v) for k,v in pages.tables.items()})
expect_error(MemoryError, pages.append_to, "a", 2049)
assert before == (pages.free,pages.lengths,pages.tables)
pages.release("a"); assert len(pages.free) == 128
expect_error(ValueError, pages.append_to, "x", -1)
print("Lecture core: PASS")PyTorch bridge · B17 — KV cache: append, window, and rollback
This cell embeds the canonical NumPy→PyTorch bridge and runs its parity check. It is self-contained and does not import repository files.
Python
import importlib.util
import subprocess
import sys
missing = []
if importlib.util.find_spec("numpy") is None:
missing.append("numpy>=1.26")
if importlib.util.find_spec("torch") is None:
missing.append("torch>=2.5")
if missing:
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *missing])
import numpy as np
import torch
import torch.nn.functional as F
import math
class NumpyKVCacheExplicit:
"""Owned [B,H,T,D] floating state; truncation cannot recover evicted tokens."""
def __init__(self, *, window: int | None = None) -> None:
if window is not None and (isinstance(window, bool) or not isinstance(window, int) or window <= 0):
raise ValueError("window must be a positive integer or None")
self.window = window
self.k = self.v = None
self.start_pos = self.next_pos = 0
def append(self, k_new, v_new):
k_new, v_new = np.asarray(k_new), np.asarray(v_new)
if k_new.ndim != 4 or k_new.shape != v_new.shape or min(k_new.shape) <= 0:
raise ValueError("K and V need the same non-empty [B,H,T,D] shape")
if k_new.dtype != v_new.dtype or not np.issubdtype(k_new.dtype, np.floating):
raise ValueError("K and V need the same floating dtype")
if self.k is not None:
if self.k.shape[:2] + self.k.shape[3:] != k_new.shape[:2] + k_new.shape[3:]:
raise ValueError("cache geometry changed")
if self.k.dtype != k_new.dtype:
raise ValueError("cache dtype changed")
# Validate first; keep old state intact if an input is rejected.
k = k_new.copy() if self.k is None else np.concatenate([self.k, k_new], axis=-2)
v = v_new.copy() if self.v is None else np.concatenate([self.v, v_new], axis=-2)
drop = max(0, k.shape[-2] - self.window) if self.window is not None else 0
# A slice alone retains its backing allocation. Make owned bounded copies.
self.k = k[..., drop:, :].copy() if drop else k
self.v = v[..., drop:, :].copy() if drop else v
self.start_pos += drop
self.next_pos += k_new.shape[-2]
return self.k, self.v
def truncate(self, absolute_end: int) -> None:
if isinstance(absolute_end, bool) or not isinstance(absolute_end, int):
raise ValueError("absolute_end must be an integer")
if not self.start_pos <= absolute_end <= self.next_pos:
raise ValueError("rollback lies outside the retained interval")
keep = absolute_end - self.start_pos
if self.k is not None:
self.k = self.k[..., :keep, :].copy()
self.v = self.v[..., :keep, :].copy()
self.next_pos = absolute_end
@property
def positions(self):
return np.arange(self.start_pos, self.next_pos, dtype=np.int64)
class TorchKVCacheExplicit:
"""Owned [B,H,T,D] floating state; truncation cannot recover evicted tokens."""
def __init__(self, *, window: int | None = None) -> None:
if window is not None and (isinstance(window, bool) or not isinstance(window, int) or window <= 0):
raise ValueError("window must be a positive integer or None")
self.window = window
self.k = self.v = None
self.start_pos = self.next_pos = 0
def append(self, k_new, v_new):
if not isinstance(k_new, torch.Tensor) or not isinstance(v_new, torch.Tensor):
raise TypeError("K and V must be tensors")
if k_new.ndim != 4 or k_new.shape != v_new.shape or min(k_new.shape) <= 0:
raise ValueError("K and V need the same non-empty [B,H,T,D] shape")
if k_new.dtype != v_new.dtype or k_new.device != v_new.device or not k_new.is_floating_point():
raise ValueError("K and V need the same floating dtype and device")
if k_new.requires_grad or v_new.requires_grad:
raise ValueError("this cache is for inference; detach inputs first")
if self.k is not None:
if self.k.shape[:2] + self.k.shape[3:] != k_new.shape[:2] + k_new.shape[3:]:
raise ValueError("cache geometry changed")
if self.k.dtype != k_new.dtype or self.k.device != k_new.device:
raise ValueError("cache dtype or device changed")
# Validate first; keep old state intact if an input is rejected.
k = k_new.clone() if self.k is None else torch.cat([self.k, k_new], dim=-2)
v = v_new.clone() if self.v is None else torch.cat([self.v, v_new], dim=-2)
drop = max(0, k.shape[-2] - self.window) if self.window is not None else 0
# A slice alone retains its backing allocation. Make owned bounded copies.
self.k = k[..., drop:, :].clone() if drop else k
self.v = v[..., drop:, :].clone() if drop else v
self.start_pos += drop
self.next_pos += k_new.shape[-2]
return self.k, self.v
def truncate(self, absolute_end: int) -> None:
if isinstance(absolute_end, bool) or not isinstance(absolute_end, int):
raise ValueError("absolute_end must be an integer")
if not self.start_pos <= absolute_end <= self.next_pos:
raise ValueError("rollback lies outside the retained interval")
keep = absolute_end - self.start_pos
if self.k is not None:
self.k = self.k[..., :keep, :].clone()
self.v = self.v[..., :keep, :].clone()
self.next_pos = absolute_end
@property
def positions(self):
device = self.k.device if self.k is not None else None
return torch.arange(self.start_pos, self.next_pos, dtype=torch.long, device=device)
np_cache = NumpyKVCacheExplicit(window=3); t_cache = TorchKVCacheExplicit(window=3)
for value in range(5):
k_np = np.full((1,1,1,2), value, dtype=np.float64); v_np = k_np + 100
np_cache.append(k_np, v_np); t_cache.append(torch.tensor(k_np), torch.tensor(v_np))
np.testing.assert_allclose(np_cache.k, t_cache.k.numpy())
np.testing.assert_array_equal(np_cache.positions, t_cache.positions.numpy())
assert (np_cache.start_pos, np_cache.next_pos) == (2,5)
np_cache.truncate(4); t_cache.truncate(4)
np.testing.assert_allclose(np_cache.k, t_cache.k.numpy())
rng = np.random.default_rng(17)
keys = rng.standard_normal((1,1,6,4)); values = rng.standard_normal((1,1,6,4)); queries = rng.standard_normal((1,1,6,4))
cache = TorchKVCacheExplicit(window=3)
for step in range(6):
cache.append(torch.tensor(keys[..., step:step+1, :]), torch.tensor(values[..., step:step+1, :]))
cached = F.scaled_dot_product_attention(torch.tensor(queries[..., step:step+1, :]), cache.k, cache.v)
lo = max(0, step - 2)
full = F.scaled_dot_product_attention(torch.tensor(queries[..., step:step+1, :]), torch.tensor(keys[..., lo:step+1, :]), torch.tensor(values[..., lo:step+1, :]))
torch.testing.assert_close(cached, full)
# Rejected input must not mutate state; rollback cannot resurrect evicted tokens.
for cache_type, convert in ((NumpyKVCacheExplicit, np.asarray), (TorchKVCacheExplicit, torch.tensor)):
c = cache_type(window=3)
c.truncate(0)
x = convert(np.ones((1,1,5,2), dtype=np.float64))
c.append(x,x)
assert (c.start_pos,c.next_pos) == (2,5)
for bad in (-1,1,6,2.5,True):
try: c.truncate(bad)
except ValueError: pass
else: raise AssertionError("invalid rollback accepted")
assert (c.start_pos,c.next_pos) == (2,5)
try: c.append(convert(np.zeros((1,2,1,2))),convert(np.zeros((1,2,1,2))))
except ValueError: pass
else: raise AssertionError("geometry change accepted")
c.truncate(2)
assert c.k.shape[-2] == 0
c.append(convert(np.ones((1,1,1,2))),convert(np.ones((1,1,1,2))))
assert (c.start_pos,c.next_pos) == (2,3)
for window in (0,-1,2.5,True):
try: cache_type(window=window)
except ValueError: pass
else: raise AssertionError("invalid window accepted")
a = NumpyKVCacheExplicit(window=3)
a.append(np.ones((1,1,10,2)),np.ones((1,1,10,2)))
assert a.k.flags.owndata and a.k.nbytes == 3*2*8
t = TorchKVCacheExplicit(window=3)
t.append(torch.ones(1,1,10,2),torch.ones(1,1,10,2))
assert t.k.untyped_storage().nbytes() == t.k.numel()*t.k.element_size()
print("B17 window, ownership, rollback, and invalid inputs: PASS")Exercises
-
[auto] Add a hybrid configuration with 12 attention layers out of 36 and compare its slope with a dense-attention model of the same width.
-
[auto] Add
truncatetoPagedKVPool: blocks beyond the new length must return to the pool while a shared partial block keeps correct reference counts. -
[auto] Build a 128K capacity table for bf16, FP8, and 3-bit KV; label payload upper bound and runtime reserve separately.
-
[auto] Simulate a workload with concrete lengths and compare measured internal fragmentation with the uniform-remainder expectation.
-
[optional extension] Repeat the scenarios in a local serving engine and separate arithmetic ceiling, measured throughput, TTFT, ITL, and goodput.
The mandatory path is fully offline. It validates memory accounting and page-table invariants, but it does not replace a GPU benchmark, scheduler, or production runtime.
Module 10 project-exercise demo — Contextual cache keys and attention-output error
Project demonstration
Download .ipynb module_10_project_demo.ipynb · 10.6 KB
Show notebook text
Module 10 project-exercise demo — Contextual cache keys and attention-output error
Implementation demo, not a completed full-scale project. Demo 1 models pages and exact token-prefix matches, distinguishing request hits from reused-token fraction. Demo 2 measures attention-output error after teaching quantization and eviction. These CPU mechanisms use synthetic data; they are not vLLM, KIVI, TurboQuant, or H2O implementations. All inputs are small, deterministic, offline, and CPU-friendly. The goal is to validate interfaces, metrics, invariants, and report structure before spending GPU or API budget.
Python
from __future__ import annotations
import importlib.util, subprocess, sys
for package in ("numpy", "matplotlib"):
if importlib.util.find_spec(package) is None:
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])
import math
import random
import time
from collections import Counter, defaultdict
import matplotlib.pyplot as plt
import numpy as np
SEED = 20260801
random.seed(SEED)
np.random.seed(SEED)
print(f"reproducible smoke tier; seed={SEED}")
DEMO_LANGUAGE = 'en'
plt.rcParams.update({"font.family":"DejaVu Sans", "font.size":12, "axes.titlesize":15, "axes.labelsize":13, "xtick.labelsize":11, "ytick.labelsize":11, "figure.dpi":150})
Demo 1. Unused tails and matching prefixes
Allocate 16-slot pages for 500 random lengths and compare them with reserving 4096 slots per request. Then issue five 48-token requests for each of 100 users. The first 32 tokens come from one of ten shared templates; the remainder distinguishes user and turn. Cache identity includes the entire preceding prefix, access scope, and model revision. Matching an isolated block is insufficient.
Python
import hashlib
rng = np.random.default_rng(SEED)
lengths = rng.integers(64, 4097, size=500)
page = 16
allocated = ((lengths + page - 1) // page) * page
reserved = np.full_like(lengths, 4096)
assert np.all(allocated - lengths < page)
print("waste / allocated:", (allocated-lengths).sum()/allocated.sum(),
"fixed-reservation waste:", (reserved-lengths).sum()/reserved.sum())
assert (2048 + page - 1) // page * page == 2048
# Complete block boundaries. The key includes ALL preceding token IDs.
# No real KV tensors, tokenizer, network, TTL, eviction, or serving engine.
def prefix_key(tokens, end, scope, revision):
return (scope, revision, tuple(tokens[:end]))
def run_prefix_trace(isolated):
cache = set(); hit_requests = cached_tokens = input_tokens = 0
for user in range(100):
for turn in range(5):
tokens = [100 + user % 10] + list(range(31)) + [1000+user, 2000+turn] + list(range(14))
scope = user if isolated else "shared-public-template"
hit = 0
for end in range(page, len(tokens)+1, page):
if prefix_key(tokens,end,scope,"model-r1") in cache:
hit = end
else:
break
for end in range(page,len(tokens)+1,page):
cache.add(prefix_key(tokens,end,scope,"model-r1"))
hit_requests += hit > 0
cached_tokens += hit
input_tokens += len(tokens)
return hit_requests/500, cached_tokens/input_tokens
shared = run_prefix_trace(False); isolated = run_prefix_trace(True)
assert np.isclose(shared[0], .98) and np.isclose(isolated[0], .8)
assert np.isclose(shared[1], 490*32/(500*48))
print("request-hit fraction / reused-input fraction:", shared, isolated)
a = list(range(32)); b = a.copy(); b[0] = 999
assert a[16:32] == b[16:32] # same block content, different preceding context
assert prefix_key(a,32,"scope","r1") != prefix_key(b,32,"scope","r1")
assert prefix_key(a,32,"scope","r1") != prefix_key(a,32,"scope","r2")
print("Context and model-revision identity: PASS")The shared scope yields a 98% request-hit fraction; per-user isolation yields 80%. Reused-input fractions are only about 65.33% and 53.33%, because each request still has sixteen new tokens. None of these rates alone predicts speedup. In the allocation experiment, exactly 2048 tokens fill every page; the average waste across random lengths belongs to a different distribution.
Demo 2. Attention error when the query changes
Generate 128 keys and values of width 32 and twelve queries. The teaching quantizer uses symmetric integer codes with one scale per row. Reconstruct the arrays and compare attention outputs. A second experiment makes an earlier query favor position 0 and the next query require position 90. Keeping only 16 positions exposes the limits of decisions based on past attention or token position alone.
Python
rng = np.random.default_rng(SEED + 10)
n, d = 128, 32
keys = rng.normal(size=(n,d)); values = rng.normal(size=(n,d))
queries = rng.normal(size=(12,d))
def attention(q,k,v):
logits = q @ k.T / np.sqrt(k.shape[1])
logits -= logits.max(axis=-1,keepdims=True)
prob = np.exp(logits); prob /= prob.sum(axis=-1,keepdims=True)
return prob @ v, prob
def fake_quantize_rows(x,bits):
limit = 2**(bits-1)-1
scale = np.max(np.abs(x),axis=-1,keepdims=True)/limit
scale = np.where(scale == 0, 1.0, scale)
codes = np.clip(np.rint(x/scale),-limit,limit).astype(np.int8)
# Float64 dequantization here is a numerical experiment, not a packed kernel.
return codes.astype(float)*scale
reference,_ = attention(queries,keys,values)
for bits in (8,3):
kq,vq = fake_quantize_rows(keys,bits),fake_quantize_rows(values,bits)
out,_ = attention(queries,kq,vq)
error = np.linalg.norm(out-reference)/np.linalg.norm(reference)
ideal_codes = 2*n*d*bits/8
fp32_scales = 2*n*4
print(f"{bits} bits: relative output error={error:.5f}; ideal packed codes+scales={ideal_codes+fp32_scales:.0f} B")
if bits == 8: assert error < .03
np.testing.assert_array_equal(fake_quantize_rows(np.zeros((2,4)),3),0)
# A controlled query shift: past attention favors token 0; the next query needs token 90.
k = np.zeros((n,d)); k[0,0] = 12; k[90,1] = 12
v = np.zeros((n,2)); v[0,0] = 1; v[90,1] = 1
past = np.zeros((1,d)); past[0,0] = 4
future = np.zeros((1,d)); future[0,1] = 4
_, mass = attention(past,k,v)
budget = 16
ranked = np.argsort(-mass[0],kind="stable")[:budget]
streaming = np.r_[np.arange(4),np.arange(n-(budget-4),n)]
exact,_ = attention(future,k,v)
for name,idx in (("past top-mass",ranked),("sink + recent",streaming)):
out,_ = attention(future,k[idx],v[idx])
error = np.linalg.norm(out-exact)
print(name,"retains future target:",90 in idx,"output error:",error)
assert 90 not in idx and error > .9
# Keeping the target has a low error on this example, but requires future knowledge.
oracle = np.r_[90,np.arange(budget-1)]
recovered,_ = attention(future,k[oracle],v[oracle])
assert np.linalg.norm(recovered-exact) < .03
print("Query-shift failure and oracle control: PASS")The byte estimate includes ideal packed codes and float32 scales. It is not the size of the NumPy arrays, whose reconstructed values are float64. After the query shift, both retention rules discard position 90 and incur a large output error. The oracle control keeps that position and recovers the output, but assumes knowledge of the future query. It diagnoses the failure rather than supplying a deployable prediction rule.
Scaling to the full project
For E4, first vary length distribution and template count. For E5, vary bit width, retention budget, and future target position. An optional engine experiment should pin arrivals, warm-up, revisions, and SLOs, then measure TTFT/ITL, goodput, memory, and answer quality. Retained K/V norm is not a substitute for response quality.
Minimum report
-
Data, model, and code revisions, plus the random seed.
-
The primary metric, plus quality, cost, latency, and memory where applicable.
-
A baseline, a matched budget, and at least one negative control.
-
Failed runs and the limits of the supported conclusions.