# KV-Cache System for Autoregressive Transformer Inference

Pure NumPy implementation — no frameworks. Demonstrates the complete KV-cache pipeline from data structures through GPU mapping.

## Architecture

```
┌─────────────────────────────────────────────────────────────────┐
│                    INFERENCE PIPELINE                            │
│                                                                 │
│  Prompt ──→ [Prefill] ──→ KV Cache populated ──→ [Generate]   │
│              O(n²) attn       O(1) per token      O(seq) attn   │
│                                                                 │
│  Per generation step:                                           │
│    1. Embed + positional encoding                               │
│    2. For each layer:                                           │
│       a. LayerNorm → QKV projection                             │
│       b. Store K,V in cache (append at write_pos)               │
│       c. Cached attention: Q @ K_cache^T → softmax → @ V_cache │
│       d. Output projection → MLP → residual                     │
│    3. LM head → logits → sample next token                      │
└─────────────────────────────────────────────────────────────────┘
```

## Files

| File | Purpose |
|------|---------|
| `kv_cache.py` | Core KV-cache data structures (`KVCache`, `BatchedKVCache`) |
| `attention.py` | Attention computation (standard, cached, GQA, masked) |
| `transformer.py` | Full transformer decoder layer + model with KV-cache integration |
| `optimizations.py` | Paged attention, quantization, chunked prefill |
| `memory_analysis.py` | Memory growth formulas, model size comparisons, GPU limits |
| `gpu_mapping.py` | GPU kernel design, Tensor Core analysis, multi-GPU strategies |
| `demo.py` | 10 end-to-end demos exercising every component |

## 1. Data Structure Layout

### Memory Format

```
cache_k[batch, num_heads, max_seq_len, head_dim]    # float16
cache_v[batch, num_heads, max_seq_len, head_dim]    # float16
lengths[batch]                                      # int32 (actual seq len per item)
write_pos                                           # int (global write pointer)
```

**Why this layout:**
- `batch` first → enables batched GEMM on GPU
- `heads` second → parallel head computation
- `seq_len` third → contiguous scan for Q @ K^T
- `head_dim` last → inner product dimension, coalesced access

### Per-Token Memory Cost

For a 7B model (32 layers, 32 heads, head_dim=128, fp16):

```
Per token per layer: 2 × 32 × 128 × 2 bytes = 16 KB
Per token (all layers): 16 KB × 32 = 512 KB
At 32K context: 512 KB × 32,768 = 16 GB
```

## 2. Update Logic Per Step

```python
# Each generation step:
pos = cache.write_pos
cache.cache_k[:, :, pos, :] = new_k[:, :, 0, :]   # O(1) write
cache.cache_v[:, :, pos, :] = new_v[:, :, 0, :]   # O(1) write
cache.write_pos += 1
```

The write is a simple memory copy — no computation needed. The cache grows by exactly `2 × heads × head_dim × elem_bytes` per token per layer.

## 3. Attention Computation Using Cache

```python
# Retrieve all cached K, V
cached_k, cached_v = cache.get_all()  # (batch, heads, seq_so_far, head_dim)

# Q @ K^T: (batch, heads, 1, head_dim) × (batch, heads, head_dim, seq)
scores = einsum("bhqd,bhkd->bhqk", q, cached_k) / sqrt(head_dim)

# Softmax (no mask needed — cache only has past tokens)
attn = softmax(scores, axis=-1)

# Attn @ V: (batch, heads, 1, seq) × (batch, heads, seq, head_dim)
output = einsum("bhqk,bhkd->bhqd", attn, cached_v)
```

**Key insight:** During generation, the cache naturally enforces causality — it only contains past tokens, so no explicit mask is needed.

## 4. Memory Growth Analysis

### Linear Growth Formula

```
KV_cache(bytes) = 2 × batch × layers × heads × seq_len × head_dim × elem_bytes
```

### 7B Model (batch=1, fp16)

| Context | KV Cache | Total (params + KV) | KV Fraction |
|---------|----------|---------------------|-------------|
| 256     | 0.12 GB  | 7.04 GB             | 1.8%        |
| 4,096   | 2.00 GB  | 8.91 GB             | 22.4%       |
| 8,192   | 4.00 GB  | 10.91 GB            | 36.7%       |
| 32,768  | 16.00 GB | 22.91 GB            | 69.8%       |

### Maximum Context by GPU (7B model, batch=1)

| GPU | Max Context |
|-----|-------------|
| RTX 4090 (24 GB) | 6,690 tokens |
| A100-40GB | 39,458 tokens |
| A100-80GB / H100-80GB | 121,378 tokens |

### Batch Size Impact

KV cache scales linearly with batch size. At batch=4, the 7B model on an A100-80GB can only handle ~30K context instead of 121K.

## 5. Optimizations

### Optimization 1: Paged Attention (vLLM-style)

**Problem:** Contiguous allocation wastes memory when sequences have variable lengths. A batch with one 32K sequence and three 100-token sequences still allocates 32K for all.

**Solution:** Divide memory into fixed-size blocks (pages). Each sequence maintains a page table mapping logical blocks to physical pages.

```
Physical page pool: (total_pages, heads, block_size, head_dim)
Page table: (batch, max_blocks) → logical → physical mapping
```

**Benefits:**
- Zero memory fragmentation
- Supports speculative decoding and branching
- Enables prefix caching (share common prefixes)
- No need to pre-allocate max_seq_len

**Trade-off:** Page table indirection adds complexity to the attention kernel (gather from non-contiguous pages).

### Optimization 2: Quantization

**Problem:** fp16 KV cache dominates memory for long contexts.

**Solution:** Store K/V in int8 with per-channel affine dequantization: `x ≈ scale × q + zero`

```
int8 data: 1 byte per element (vs 2 for fp16)
fp16 scales + zeros: shared per channel (not per token)
Net savings: ~50% memory with <1% accuracy loss
```

**Production approach:** Shared per-channel scales (not per-position) stored in fp16. The per-position approach in this codebase is for correctness demonstration but has higher overhead.

### Optimization 3: Chunked Prefill

**Problem:** Processing a 32K prompt requires materializing a 32K × 32K attention matrix (4 GB in fp32).

**Solution:** Process the prompt in chunks of size C. Each chunk attends to all previous tokens + causal within chunk.

```
Peak memory: O(C × seq_len) instead of O(seq_len²)
For C=512, seq=4096: 8 MB vs 64 MB (8× savings)
```

### Combined: Paged + Quantized

Together these give 2-4× memory reduction, enabling 2-4× longer contexts in the same GPU memory.

## 6. GPU Execution Mapping

### Memory Hierarchy

| Level | Size | Latency | Usage |
|-------|------|---------|-------|
| Registers | 64 KB/SM | 1 cycle | Thread-local, warp computation |
| Shared memory | 166 KB/SM (H100) | 1-3 cycles | Tiling, softmax intermediates |
| L2 cache | 50 MB (H100) | ~20 cycles | Automatic global memory caching |
| HBM | 80 GB (H100) | ~300-400 cycles | Model weights, KV cache, activations |

### Cached Attention Kernel Design

```
Grid: (batch_size, num_heads, 1)
Block: (32, 32) = 1024 threads

Shared memory per block (~16-20 KB):
  - Q tile: 1 × head_dim (512 bytes fp16)
  - K tile: 32 × head_dim (8 KB fp16)
  - Score tile: 32 × 32 (4 KB fp16)
```

**Optimization strategies:**
1. Coalesced global memory access (warp-level consecutive addresses)
2. Tiled GEMM with shared memory
3. Persistent kernels (keep blocks alive until all tiles processed)
4. Async copy (H100 `cp.async`) to overlap memory transfer with computation
5. Tensor Cores (`mma.sync`) for matmul operations
6. Fusion: merge softmax with attention score computation

### Arithmetic Intensity

For single-token generation (batch=1, heads=32, seq=4096):
- **FLOPs:** 0.02 GFLOPs
- **Memory traffic:** 16.79 MB
- **Arithmetic intensity:** 1.0 FLOPs/byte
- **→ Memory-bound** (H100 peak: 1,970 TFLOPS, 3.35 TB/s)

The cached attention is fundamentally memory-bound — the bottleneck is reading the KV cache from HBM, not computation. This is why bandwidth-optimized kernels (FlashAttention-style tiling) matter more than compute optimization.

### Multi-GPU Strategies

| Strategy | KV Cache Impact | Best For |
|----------|----------------|----------|
| Tensor parallelism | Split K/V by head_dim across GPUs | Large models |
| Pipeline parallelism | Each GPU holds its layer shard's K/V | Very large models |
| Sequence parallelism | Split K/V by sequence dimension | Long context prefill |
| Expert parallelism | KV cache shared; only MLP experts sharded | MoE models |

## Running

```bash
# Run all 10 demos
python demo.py

# Run memory analysis standalone
python memory_analysis.py

# Run GPU mapping report
python gpu_mapping.py
```

## Key Takeaways

1. **KV cache grows linearly** with sequence length and batch size — it's the dominant memory cost for long-context inference.

2. **Generation is memory-bound** — reading the KV cache from HBM dominates latency, not the attention computation itself.

3. **Paged attention eliminates fragmentation** — critical for serving variable-length sequences efficiently.

4. **Quantization gives ~50% savings** with minimal accuracy loss when using shared per-channel scales.

5. **FlashAttention-style tiling** reduces HBM traffic by processing K/V in tiles that fit in shared memory, cutting memory bandwidth requirements by 2-4×.