Implement KV Cache for Efficient Decoding

Implementationmedium~35 min

Objective

Complete the KVCacheDecoder class so it produces identical outputs to the naive decoder but avoids redundant KV computation during autoregressive generation.

Background

Autoregressive transformer decoding generates one token at a time. Without a KV cache, every new token requires recomputing K and V for ALL previous tokens, resulting in O(n^2) total compute for n tokens. A KV cache stores previously computed K and V values, so each decode step only computes Q/K/V for the single new token and reuses everything else. This is the standard optimization used in every production LLM serving system (vLLM, TGI, etc.). Your task is to implement this from scratch using NumPy.

Requirements

1.Initialize empty KV cache arrays
2.Implement prefill to compute and cache K, V for all prompt tokens
3.Implement decode_step to compute Q/K/V for one new token and append K/V to cache
4.Compute attention using new Q against all cached K, V
5.Implement the generate loop using prefill + decode_step

Evaluation (100 points)

KV cache initialized

KVCacheDecoder must initialize k_cache and v_cache storage

20pt

Prefill stores KV pairs

prefill() must compute K, V for all tokens and store in cache

20pt

Decode step appends to cache

decode_step() must grow the cache by appending new K, V

25pt

Single-token Q computation

decode_step computes Q only for the new token

20pt

Generate loop uses prefill + decode

generate() must call prefill then decode_step in a loop

15pt

Hints

Select a file to start editing

Terminal

AI Assistant50K tokens left

Ask me about the code, bugs, or concepts.
I'll guide you in plain English, no code output.
Budget: 50K tokens per lab