# How the AI Engineering Curriculum Covers KV Cache and Flash Attention for Inference Optimization

> Learn how the AI Engineering curriculum optimizes LLM inference with KV cache and Flash Attention, reducing complexity from O(N²) to O(N) and overcoming memory bottlenecks.

- Repository: [Rohit Ghumare/ai-engineering-from-scratch](https://github.com/rohitg00/ai-engineering-from-scratch)
- Tags: deep-dive
- Published: 2026-08-29

---

**The AI Engineering from Scratch curriculum teaches KV cache and Flash Attention as complementary inference optimization techniques that reduce LLM generation complexity from O(N²) to O(N) and eliminate memory bandwidth bottlenecks through SRAM-resident tiling.**

The `rohitg00/ai-engineering-from-scratch` repository provides a hands-on deep dive into transformer inference efficiency. In the dedicated lesson "KV Cache, Flash Attention & Inference Optimization," located at `phases/07-transformers-deep-dive/12-kv-cache-flash-attention/`, the curriculum bridges theoretical complexity analysis with minimal working implementations to demonstrate how modern serving stacks achieve sub-quadratic inference scaling.

## Understanding KV Cache Mechanics

The KV cache eliminates redundant computation in autoregressive generation by storing key (K) and value (V) tensors after their initial calculation.

### Memory Footprint and Byte Cost

The curriculum derives the exact memory cost for caching K and V vectors across transformer layers. For each token, layer, and attention head, the cache requires:

```text
bytes_per_token_per_layer = 2 × d_head × dtype_size

```

The lesson provides concrete examples in [`phases/07-transformers-deep-dive/12-kv-cache-flash-attention/docs/en.md`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/07-transformers-deep-dive/12-kv-cache-flash-attention/docs/en.md) (lines 27-50). A 7B parameter model with 32 layers, 32 heads, `d_head=128`, and fp16 precision consumes 512 bytes per token, totaling 16 KB for a 32K context window. For Llama-3 70B using Grouped Query Attention (80 layers, 8 KV heads), the cost rises to 4 KB per token, requiring 10.4 GB for 32K context.

### Implementation in the Curriculum

The repository provides a minimal pure-standard-library `KVCache` class in [`phases/07-transformers-deep-dive/12-kv-cache-flash-attention/code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/07-transformers-deep-dive/12-kv-cache-flash-attention/code/main.py) (lines 59-67):

```python
class KVCache:
    def __init__(self):
        self.K = []          # list of key vectors

        self.V = []          # list of value vectors

    def append(self, k, v):
        self.K.append(k)
        self.V.append(v)

```

This class appends K/V vectors per generation step and serves them to subsequent attention computations without recalculation.

### Computational Complexity Impact

Without caching, naive autoregressive decoding recomputes attention over the entire prefix at each step, resulting in O(N²) total operations. The curriculum demonstrates in [`main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/main.py) (lines 72-97) that the KV cache reduces per-token attention cost to linear O(N), as only a single query vector must be scored against the growing cache rather than reconstructing all past keys.

## Flash Attention Tiling Implementation

Flash Attention addresses the memory bandwidth bottleneck that dominates transformer inference for long contexts.

### The Tiling Trick for Memory Bandwidth

Standard attention materializes the full N×N score matrix in High Bandwidth Memory (HBM), causing three round-trips per layer. Flash Attention processes the matrix in tiles (e.g., 128×128), loading each block into fast on-chip SRAM to compute the softmax and matrix multiplication while maintaining running statistics (maximum values and scaling factors). This approach reduces memory traffic from O(N²) to O(N), as detailed in [`docs/en.md`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/docs/en.md) (lines 62-84).

### Reference Implementation

The `tiled_softmax_dot` function in [`code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/code/main.py) (lines 35-56) implements this algorithm without allocating the full attention matrix:

```python
def tiled_softmax_dot(q, Ks, Vs, tile=4):
    """Flash-Attention-style incremental softmax(qK^T)V."""
    d_head = len(Vs[0])
    scale = 1.0 / math.sqrt(len(q))
    m = float("-inf")
    s = 0.0
    out = [0.0] * d_head
    for start in range(0, len(Ks), tile):
        k_block = Ks[start:start + tile]
        v_block = Vs[start:start + tile]
        scores = [sum(qi * ki for qi, ki in zip(q, k)) * scale for k in k_block]
        new_m = max(m, *scores)
        exp_old = math.exp(m - new_m) if m != float("-inf") else 0.0
        exp_new = [math.exp(sc - new_m) for sc in scores]
        s = s * exp_old + sum(exp_new)
        for j in range(d_head):
            out[j] = out[j] * exp_old + sum(e * v[j] for e, v in zip(exp_new, v_block))
        m = new_m
    return [o / s for o in out]

```

### Numerical Verification

The curriculum verifies mathematical equivalence between the tiled and full-matrix implementations. Code in [`main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/main.py) (lines 28-34) confirms that Flash Attention produces identical outputs to standard softmax attention within floating-point reassociation error, demonstrating that the memory savings introduce no approximation.

## Combining Optimizations for Production Inference

The curriculum frames KV cache and Flash Attention as complementary solutions to distinct bottlenecks. The cache eliminates redundant FLOPs, while Flash Attention minimizes the memory traffic required for the remaining necessary computations.

Together, these optimizations yield **2–4× speed-up on A100** GPUs and **5–10× on H100** (FP8) hardware, according to the lesson documentation. The curriculum emphasizes that by 2026, both techniques are considered universal requirements in production serving stacks such as vLLM, TensorRT-LLM, SGLang, and llama.cpp.

### Advanced: PagedAttention

The lesson extends KV cache concepts to **PagedAttention**, vLLM's virtual-memory allocator that stores cache in 16-token pages rather than contiguous blocks. This enables prefix sharing across beams and hot-swapping, providing approximately **4× throughput gains** over naive contiguous allocation.

## Summary

- The **KV cache** reduces autoregressive generation complexity from O(N²) to O(N) by storing key and value tensors, costing `2 × d_head × dtype_size` bytes per token per layer.
- **Flash Attention** eliminates HBM traffic through SRAM-resident tiling, reducing memory operations from O(N²) to O(N) while maintaining numerical fidelity via online softmax statistics.
- The `KVCache` class and `tiled_softmax_dot` function in [`phases/07-transformers-deep-dive/12-kv-cache-flash-attention/code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/07-transformers-deep-dive/12-kv-cache-flash-attention/code/main.py) provide minimal, educational implementations of both techniques.
- Combined, these optimizations enable 2–4× inference speedup on A100 and 5–10× on H100 hardware for long-context generation.
- The curriculum extends these foundations to **PagedAttention** and **speculative decoding**, positioning KV cache and Flash Attention as essential knowledge for modern LLM deployment.

## Frequently Asked Questions

### What is the exact memory cost of the KV cache in a Llama-3 70B model?

For Llama-3 70B with 80 layers and Grouped Query Attention (8 KV heads), the cache requires approximately 4 KB per token at fp16 precision. A 32,000-token context window therefore consumes roughly 10.4 GB of GPU memory, as calculated in [`phases/07-transformers-deep-dive/12-kv-cache-flash-attention/docs/en.md`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/07-transformers-deep-dive/12-kv-cache-flash-attention/docs/en.md) (lines 27-50).

### How does Flash Attention maintain accuracy while reducing memory usage?

Flash Attention uses **online softmax statistics**—maintaining running maximum values and scaling factors across tiles—to compute mathematically identical results to standard attention. The `tiled_softmax_dot` function in [`code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/code/main.py) demonstrates this approach, with numerical verification confirming equivalence to within floating-point reassociation error (lines 28-34).

### Why are KV cache and Flash Attention considered complementary optimizations?

**KV cache** eliminates redundant computations by reusing previously calculated key-value pairs from prior steps, while **Flash Attention** solves the memory bandwidth bottleneck caused by accessing the full attention matrix for each new query. The cache reduces FLOPs, whereas Flash Attention optimizes how the remaining FLOPs access memory, together enabling practical sub-quadratic inference scaling.

### Where can I find the implementation code for these techniques in the repository?

The reference implementations reside in [`phases/07-transformers-deep-dive/12-kv-cache-flash-attention/code/main.py`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/phases/07-transformers-deep-dive/12-kv-cache-flash-attention/code/main.py), containing the `KVCache` class (lines 59-67) and `tiled_softmax_dot` function (lines 35-56). The full mathematical derivations, memory tables, and algorithmic explanations appear in the corresponding [`docs/en.md`](https://github.com/rohitg00/ai-engineering-from-scratch/blob/main/docs/en.md) file within the same directory.