Homework

Task 1: KV Cache Memory Calculator

Goal: Build intuition for the memory costs of serving transformer models by calculating KV cache requirements for real model configurations.

Instructions

Write a Python function kv_cache_memory(n_layers, n_kv_heads, d_head, seq_len, batch_size, dtype_bytes=2) that computes the total KV cache memory in bytes. The formula is:

memory = 2 * n_layers * n_kv_heads * d_head * seq_len * batch_size * dtype_bytes
         ^                                                              ^
    (K and V)                                                    (float16 = 2 bytes)

Calculate KV cache memory for each of the following configurations at a sequence length of 4096 tokens and batch size of 1:

Model	Layers	KV Heads	d_head
Llama 2 7B	32	32	128
Llama 2 13B	40	40	128
Llama 2 70B (GQA)	80	8	128
Mistral 7B (GQA)	32	8	128

Then answer these questions:

1. How much KV cache memory does Llama 2 70B save compared to what it would use with full MHA (80 KV heads instead of 8)?
2. If you are serving Llama 2 70B with 80 GB of GPU memory available and the model weights take 35 GB (in 4-bit quantisation), how many concurrent requests at 4096 tokens can you fit in the remaining memory?
3. If you double the context length to 8192, how does the concurrent request capacity change?

Self-check

• ✓ Your Llama 2 7B result is approximately 1 GB at 4096 tokens.
• ✓ Your Llama 2 70B (GQA) result is significantly smaller than you might expect for a 70B model, thanks to GQA.
• ✓ You can explain why KV cache — not model weights — is often the bottleneck for concurrent serving.

Task 2: Attention Complexity Analysis

Goal: Quantify the computational cost of self-attention at different sequence lengths and understand why efficient attention matters.

Instructions

The FLOPs for standard self-attention (single head, single layer) are approximately:

FLOPs = 2 * n^2 * d_k   (for Q @ K^T)
      + 2 * n^2 * d_v   (for weights @ V)
      + n^2             (softmax, approximate)

Write a function that computes the total attention FLOPs for a full model (all heads, all layers). Then fill in this table:

Sequence length	Attention FLOPs (Llama 2 7B)	Ratio vs 2K
2,048	?	1.0x
4,096	?	?
16,384	?	?
131,072 (128K)	?	?

Use Llama 2 7B parameters: 32 layers, 32 heads, d_head = 128.

Then answer:

1. By what factor do attention FLOPs increase when going from 4K to 128K context? Does this match your expectation from O(n²)?
2. If Flash Attention reduces memory from O(n²) to O(n) but does not change the FLOPs, why is it still faster?
3. Research: what is the approximate ratio of attention FLOPs to total model FLOPs for Llama 2 7B at 4K context? (Hint: the FFN uses approximately 2x the FLOPs of attention.)

Self-check

• ✓ The FLOPs ratio from 2K to 128K is approximately 4096x (since (128K/2K)² = 64² = 4096).
• ✓ You can explain the difference between compute-bound and memory-bound operations, and why attention is memory-bound.
• ✓ You understand that Flash Attention is faster because it reduces HBM reads/writes, not FLOPs.

Task 3: Architecture Trade-offs Essay

Goal: Articulate when you would choose each transformer architecture family and why.

Instructions

Write a short response (200–300 words) that addresses the following scenario:

Your team needs to build three different ML-powered features: (1) a semantic search engine over internal documents, (2) a customer-facing chatbot that answers questions in natural language, and (3) a system that translates user queries from English to SQL. For each feature, recommend a transformer architecture family (encoder-only, decoder-only, or encoder-decoder) and justify your choice. Consider factors like task structure, latency requirements, model size, and available pre-trained models.

Self-check

• ✓ You recommended a specific architecture for each of the three features.
• ✓ Your reasoning goes beyond "decoder-only is always best" — you articulate specific advantages of each architecture for each task.
• ✓ You mention at least one concrete model name for each recommendation (e.g., "a BERT-family model like..." or "a decoder model like...").
• ✓ You consider practical factors beyond raw capability (latency, cost, model size, availability of pre-trained weights).