Inference Systems: KV Caching & Context Management

Explain self-attention in your own words to a reader who knows linear algebra but not ML. Must answer: what Q, K, V represent; why multiple heads; and concretely why cost is quadratic in sequence length and what that implies for long inputs.

Implement scaled dot-product attention, then multi-head attention, from scratch in NumPy or raw PyTorch ops — no nn.MultiheadAttention. Include a causal mask. Verify against PyTorch's built-in to ~1e-5.

Skim nanoGPT (and Karpathy's "Let's build GPT from scratch" video). Train the char-level model for a few minutes so the full stack is concrete.

Derive and explain why decode is bottlenecked by memory bandwidth rather than FLOPs, while prefill is the opposite. Tie it to arithmetic intensity. State the implication: where does optimization effort pay off in each phase?

Implement greedy generation twice: recomputing attention over the full sequence each step (naive) vs. appending to a KV cache. Plot latency-per-token vs. sequence length for both. The naive curve climbs; the cached one stays roughly flat per step.

Write a CLI taking a model config (n_layers, n_kv_heads, head_dim, dtype) plus batch_size and seq_len, printing KV-cache bytes, weight bytes, and the crossover point where cache exceeds weights. Run it for an 8B-class config. This is your first Rust/Go in the repo and primes you for paged allocation.

Explain PagedAttention by explicit analogy to OS paging. Define block, block table, and the fragmentation it eliminates. Then explain how copy-on-write lets two sequences share blocks — and predict how you'll exploit that for shared prompt prefixes.

Implement a fixed pool of N blocks (each = capacity for K tokens), per-sequence block tables, and ops: allocate_sequence, append_token (grab a new block when the current fills), free_sequence, and fork(seq) with copy-on-write. Model KV data as opaque byte buffers — no tensors. Add a stats method for pool utilization and wasted bytes.

Place MHA, MQA, GQA, and MLA on a 2-axis chart (KV-cache size vs. quality) and justify each placement. Explain mechanically how GQA differs from MLA — sharing heads vs. compressing into a latent.

Extend your Module 0.1 attention to take a configurable n_kv_heads, yielding MHA (= n_heads), MQA (= 1), and GQA (between) via repeat_interleave of the KV heads. Tabulate KV-cache bytes per token as you sweep n_kv_heads.

Write one paragraph on why MLA's low-rank latent is clever for long-context serving specifically, and what it costs at compute time.

Eviction vs. quantization vs. low-rank vs. hybrid-memory. Give one representative method per category and the workload where you'd reach for each. This doubles as your 2-minute whiteboard answer to "how would you reduce KV-cache pressure?"

Implement a StreamingLLM-style cache on a small model: keep the first k "attention sink" tokens plus a rolling window of the last w, evicting the middle. Show that keeping the sinks preserves coherent generation while a naive "keep only the last w" degrades. Report a quality proxy (perplexity or eyeball coherence).

Explain (1) how the radix tree enables automatic prefix sharing across requests, (2) why evicting leaves first preserves shared prefixes, and (3) what reference counting protects against. Connect it back to the copy-on-write you built in Module 1.1.

No model attached: insert(token_ids), match_prefix(token_ids) returning the longest cached prefix and matched node, LRU-leaf eviction under a capacity bound, and refcount locking so in-use nodes can't be evicted. Unit-test prefix matching and the leaf-eviction ordering.

Port the radix cache into core/ and wire it to the Module 1.1 allocator: each tree node points at the KV blocks for its token span, and two requests sharing a prefix share those blocks (copy-on-write) until they diverge. API: given a request's token IDs, return (reused_blocks, suffix_to_compute).

Design the L1/L2/L3 extension for local/edge hardware (your differentiator: replace "distributed cluster storage" with "local NVMe"). Specify per-node metadata per tier; the promotion/demotion policy; what you deliberately won't track for L3 and why; and how async prefetch hides load latency. Note where your single-machine design diverges from Mooncake/HiCache. This is literally your capstone's design doc.

Add L2 (in-memory host pool) and L3 (on-disk — files or an embedded store like sled/RocksDB) to your prefix cache. On L1 eviction, demote blocks to L2 then L3; on a prefix match whose blocks live in L3, load them back up; and add an async prefetch that begins pulling needed blocks before the scheduler asks. Instrument everything with timing.

Contrast exact prefix (radix), modular reuse (Prompt Cache), and selective-recompute (CacheBlend), plus CacheGen's compress-and-stream. For each, name the accuracy risk it accepts and explain why breaking strict prefix order hurts attention (hint: attention sinks and inter-chunk dependencies).

Add CacheGen-style quantize-and-serialize for KV blocks in your L3 tier: compress before writing to NVMe, decompress on load. Measure size reduction vs. the recompute you avoid.

A blog-style post: the architecture (L1/L2/L3 radix cache for local parallel agents), the design decisions and where they diverge from Mooncake/HiCache, and your benchmark numbers (cache-hit speedups, NVMe restore vs. recompute, fragmentation under paging). This writeup plus the repo is the artifact you put in front of interviewers — and the launchpad for an OSS contribution.