KV-cache optimization: Exercises — Hinglish deep dive

6.1.13 · D4 · AI-ML › Scaling & Efficient Architectures › KV-cache optimization

Definition Is page par use kiye gaye symbols aur unit conventions

Kisi bhi problem se pehle, har symbol aur har unit ko pin down karte hain taaki kuch bhi undefined na rahe.

$L$ — sequence length (tokens ki sankhya).
$h$ — attention heads ki sankhya.
$d_{k}$ — head dimension: ek head ke key/value/query vector ki length.
$d$ — model width (a.k.a. hidden size): projection matrices mein feed ki jaane wali full embedding dimension. Yeh heads se $d = h \times d_{k}$ se relate hoti hai. Jab bhi aap $d^{2}$ dekhte ho toh yeh projection matrix $W \in R^{d \times d}$ se aata hai (Ex 2.3 / 4.1 dekho kyun).
precision — stored number per bytes (fp16 $= 2$ bytes).
Unit convention (neeche har jagah use hogi): hum binary units use karte hain. $1 KB = 2^{10} = 1024$ bytes, $1 MB = 2^{20}$ bytes, $1 GB = 2^{30}$ bytes. Toh "16 KB" yahan exactly $16 \times 1024 = 16, 384$ bytes matlab hai. (Strictly yeh KiB/MiB/GiB hain; hum KB/MB/GB readability ke liye likhte hain lekin arithmetic hamesha powers of two mein hoti hai.)

Shuru karne se pehle, do duniyaon ki ek shared picture — "sab kuch recompute karo" vs "yaad rakho aur append karo":

Figure padhna. Top row step 4 par cache ke bina generation hai: saare chaar key boxes $K_{1}, K_{2}, K_{3}, K_{4}$ black mein draw kiye gaye hain kyunki model unhe scratch se recompute karta hai. Bottom row cache ke saath generation hai: $K_{1}, K_{2}, K_{3}$ black hain (already stored, bas reuse ho rahe hain — "reuse" arrow follow karo) aur sirf $K_{4}$ red mein draw kiya gaya hai, kyunki yeh uss step mein cache ko compute karni wali ek nayi projection hai. Yahi contrast — teen reused vs ek nayi — cache ka poora idea hai, aur neeche ke har memory/FLOP problem mein bas inhi boxes ko count kiya ja raha hai.

Level 1 — Recognition

Exercise 1.1 (L1)

Batao ki neeche diye gaye mein se kaun sa autoregressive generation ke dauran KV-cache mein store hota hai, aur kaun sa har step par recompute hota hai: (a) key projections $K_{i}$ , (b) value projections $V_{i}$ , (c) query projections $Q_{t}$ , (d) final softmax attention weights.

Recall Solution

(a) $K_{i}$ — cached. Ek past token ki key ek baar compute hone ke baad kabhi nahi badalti, isliye hum ise store karte hain (Figure s01 mein black boxes).
(b) $V_{i}$ — cached, usi wajah se.
(c) $Q_{t}$ — recomputed har step par. Query yeh hoti hai ki "current token kya dhundh raha hai," aur har step par ek brand-new current token hota hai, isliye reuse karne ko kuch nahi hai (red box har step par nayi hoti hai).
(d) softmax weights — recomputed. Yeh nayi $Q_{t}$ par depend karte hain, isliye har step par change hote hain.

Cache exactly wahi cheezein store karta hai jo past positions ke liye fixed hain: $K$ aur $V$ . Isliye naam hai KV-cache.

Exercise 1.2 (L1)

Shape ki recognition. Ek single layer mein $h = 8$ heads aur head dimension $d_{k} = 64$ hai. 5 tokens generate karne ke baad, $Cache_{K}$ ki shape kya hai (tokens × heads × head-dim ke roop mein)?

Recall Solution

Har token ek head per ek key contribute karta hai: token per shape $(h, d_{k}) = (8, 64)$ hai. 5 tokens stack karne par milta hai $Cache_{K} \in R^{5 \times 8 \times 64} .$ $Cache_{V}$ ki shape identical hai. Pehla axis (yahan $5$ ) hi ek hai jo generate karte waqt badhta hai.

Level 2 — Application

Exercise 2.1 (L2)

Per-token cache memory. Ek layer mein $h = 32$ heads, $d_{k} = 128$ hai, fp16 mein store hai (2 bytes each). Ek token uss ek layer ke liye KV-cache mein kitne bytes add karta hai?

Recall Solution

Formula kahan se aata hai (short derivation). Ek token ke liye humein uski keys aur values store karni hoti hain. Ek single token ka key set ek $d_{k}$ -vector per head hota hai, isliye usmein $h \times d_{k}$ numbers hote hain; same value set ke liye bhi. Yeh do tensors hain $h \times d_{k}$ numbers ke, aur har number "precision" bytes occupy karta hai. Multiply karo: $bytes/token/layer = K and V 2 \times numbers per tensor h \times d_{k} \times bytes/number precision .$ Yahi reason hai har factor ka — kuch blindly cite nahi kiya gaya. Plug in karo: $= 2 \times 32 \times 128 \times 2 = 16, 384 bytes = 16 KB .$ (Unit convention yaad karo: $16 KB = 16 \times 1024 = 16, 384$ bytes exactly.)

Exercise 2.2 (L2)

Ek poori sequence ke liye full-model cache. Same layer specs ( $h = 32$ , $d_{k} = 128$ , fp16), ab $n_{layers} = 40$ aur $L = 1024$ tokens ki sequence. Total cache size gigabytes mein?

Recall Solution

Ek token ek layer mein Ex 2.1 se $2 h d_{k} \times precision$ cost karta hai. Ab hamare paas $L$ tokens hain (pehla cache axis badhta hai) aur $n_{layers}$ independent layers hain (har ek apna cache rakhti hai), isliye dono se multiply karte hain: $Memory = 2 \times n_{layers} \times L \times h \times d_{k} \times precision .$ Plug in karo: $= 2 \times 40 \times 1024 \times 32 \times 128 \times 2 bytes .$ Step karo: $2 \times 40 = 80$ ; $\times 1024 = 81, 920$ ; $\times 32 = 2, 621, 440$ ; $\times 128 = 335, 544, 320$ ; $\times 2 = 671, 088, 640$ bytes. Convert karo (binary GB): $671, 088, 640/ 2^{30} = 0.625$ GB. Answer: $0.625$ GB ( $= 640$ MB, kyunki $0.625 \times 1024 = 640$ ).

Exercise 2.3 (L2)

Speedup factor. Recall karo ki $d = h d_{k}$ model width hai, aur har $W_{Q}, W_{K}, W_{V} \in R^{d \times d}$ ek token ke $d$ -vector ko $d$ -vector mein project karta hai, har token ke liye $O (d^{2})$ multiply-adds lete hue. Caching ke bina, step $t$ par model saare $t$ tokens ko re-project karta hai, aur poori generation par sum karne se projection work $O (L^{2} d^{2})$ scale karti hai; caching ke saath har step par sirf ek nayi token project hoti hai, jo $O (L d^{2})$ deta hai. Pehle derive karo ki no-cache total $O (L^{2} d^{2})$ kyun hai, phir $L = 256$ ke liye projection speedup compute karo.

Recall Solution

No-cache cost $O (L^{2} d^{2})$ kyun hai (missing "why"). Ek token ki teen projections mein se har ek $O (d^{2})$ cost karti hai — kyunki ek $d$ -vector ko $d \times d$ matrix se multiply karna $d^{2}$ multiply-adds hai — isliye $O (d^{2})$ per token. Cache ke bina, generation step $t$ par model sequence ko fresh maanta hai aur ab tak dekhe gaye saare $t$ tokens ko re-project karta hai, us step mein $O (t d^{2})$ cost karke. Har step $t = 1 \dots L$ par sum karo: $\sum_{t = 1}^{L} O (t d^{2}) = O (d^{2}) \sum_{t = 1}^{L} t = O (d^{2} \cdot \frac{L ( L + 1 )}{2}) = O (L^{2} d^{2}) .$ $L$ mein quadratic purely triangular sum $\sum t = \frac{L ( L + 1 )}{2}$ se aata hai. Cache ke saath, har step par sirf ek nayi token project hoti hai: $\sum_{t = 1}^{L} O (d^{2}) = O (L d^{2})$ . Speedup: $\frac{L ^{2} d ^{2}}{L d ^{2}} = L = 256.$ $d^{2}$ cancel ho jata hai — speedup purely sequence-length factor $L$ hai. Lambi generations zyada benefit karti hain.

Level 3 — Analysis

Exercise 3.1 (L3)

MQA vs standard. Ek model $h = 32$ query heads, $d_{k} = 128$ use karta hai. Standard multi-head attention ke KV-cache size ko Multi-Query Attention (MQA) se compare karo, jo ek single K/V head share karta hai. Ratio do aur interpret karo.

Recall Solution

Standard cache (per token, per layer, constant factor ignore karke): $2 h d_{k} = 2 \times 32 \times 128.$ MQA ek K aur ek V head rakhta hai: $2 \times 1 \times d_{k} = 2 \times 128.$ Ratio: $\frac{MQA}{standard} = \frac{2 d _{k}}{2 h d _{k}} = \frac{1}{h} = \frac{1}{32} .$ MQA $32 \times$ kam KV-cache use karta hai. Interpretation: query heads ab bhi specialize karte hain (har ek ka apna $W_{Q}$ hai), lekin sab same keys aur values se padhte hain, isliye hum sirf ek set store karte hain. Cost: thodi si quality drop (~1–2% perplexity).

Exercise 3.2 (L3)

GQA middle ground. Llama-2 70B $h = 64$ query heads use karta hai jo $g = 8$ K/V groups mein grouped hain. (a) Standard attention ke versus cache reduction kya hai? (b) GQA standard aur MQA ke beech kahan hai?

Figure padhna. Teen head-sharing schemes side by side draw ki gayi hain. Top boxes query heads hain; bottom boxes stored KV heads hain; har line dikhati hai ki kaun si queries kaun sa KV padhti hain. Left ("standard $g = h$ "): har query ka apna private KV hai — koi sharing nahi, maximal cache. Right ("MQA $g = 1$ "): saari queries ek shared KV mein funnel hoti hain — minimal cache. Middle ("GQA $g = 2$ ", red mein drawn): queries groups mein bundle hoti hain, har group ek KV share karta hai — tunable middle ground. Jaise aap $g$ ko $h$ se $1$ tak slide karte ho, cache "standard" se "MQA" tak slide hota hai; red middle exactly woh case hai jo part (b) ke baare mein poochhta hai.

Recall Solution

(a) GQA $h$ ki jagah $g$ K/V heads store karta hai: $\frac{GQA}{standard} = \frac{2 g d _{k}}{2 h d _{k}} = \frac{g}{h} = \frac{8}{64} = \frac{1}{8} .$ Toh 8× reduction. (b) Figure ko left-to-right padho: standard mein har query head ke liye ek K/V head hota hai ( $g = h$ , koi reduction nahi); MQA ek single shared K/V head mein collapse ho jaata hai ( $g = 1$ , maximal $h \times$ reduction); GQA (beech mein red grouping) $1 < g < h$ ke saath beech mein hai. Yeh MQA ki zyaadatar memory savings recover karta hai jabki standard attention ki zyaada quality rakhta hai har group of queries ko apni keys/values dekar.

Exercise 3.3 (L3)

Sliding-window memory. Ek 100K-token context ke liye, full-cache memory ko $W = 4096$ width ke sliding-window cache se compare karo. Assume karo $n_{layers} = 32$ , $h = 32$ , $d_{k} = 128$ , fp16. Dono sizes aur ratio do.

Recall Solution

Per-token-per-layer bytes: $2 \times 32 \times 128 \times 2 = 16, 384$ B $= 16$ KB. Full cache ( $L = 100, 000$ ): $32 \times 100, 000 \times 16 KB = 51, 200, 000 KB = 48.83 GB .$ (Binary conversion: $51, 200, 000/ 2^{20} = 48.83$ GB.) Sliding-window ( $W = 4096$ ): $32 \times 4096 \times 16 KB = 2, 097, 152 KB = 2.0 GB .$ Ratio: $100000/4096 \approx 24.4 \times$ chhota. Trade-off: window memory ko context length se regardless constant $O (W)$ par rakhta hai, lekin model literally $W$ tokens se purani kisi bhi cheez ko attend nahi kar sakta — window ke paar ki long-range dependencies lost ho jaati hain.

Level 4 — Synthesis

Exercise 4.1 (L4)

FLOP accounting, cache vs no-cache. GPT-2 small ke liye ( $n_{layers} = 12$ , model width $d = 768$ ), caching ke saath aur bina $L = 100$ tokens generate karne ke liye projection FLOPs estimate karo. Dono aur speedup report karo.

Recall Solution

Per-step cost kahan se aata hai. Jaise Ex 2.3 mein establish hua, ek token ko $d \times d$ matrix se project karna $\approx d^{2}$ multiply-adds cost karta hai; hum teen projections $W_{Q}, W_{K}, W_{V}$ ko $O (d^{2})$ maante hain aur small constant factor drop karte hain ("3" aur koi bhi $+ d$ terms) kyunki hum sirf $L$ aur ratio ke saath scaling chahte hain — constants ratio mein cancel ho jaate hain. Toh:

No-cache, step $t$ , per layer: saare $t$ tokens re-project karo $\Rightarrow\approx t d^{2}$ .
Cache, step $t$ , per layer: sirf 1 nayi token project karo $\Rightarrow\approx d^{2}$ .

Without cache — saare steps par triangular series sum karo, times layers: $12 \times (\sum_{t = 1}^{100} t) \times 76 8^{2} = 12 \times 5050 \times 589, 824 \approx 3.57 \times 1 0^{10} FLOPs .$ (Use kiya $\sum_{t = 1}^{100} t = \frac{100 \cdot 101}{2} = 5050$ — Ex 2.3 jaisi hi triangular sum.) With cache — har step par ek nayi K/V/Q, toh $d^{2}$ per step: $12 \times 100 \times 76 8^{2} = 12 \times 100 \times 589, 824 \approx 7.08 \times 1 0^{8} FLOPs .$ Speedup: $\frac{3.57 \times 1 0 ^{10}}{7.08 \times 1 0 ^{8}} \approx 50.5 \times .$ Yeh expected $\sim \frac{L + 1}{2}$ factor se match karta hai: sum-of-integers $\frac{L ( L + 1 )}{2}$ vs $L$ deta hai, yaani $\frac{L + 1}{2} = 50.5$ . Dropped constant factor kabhi matter nahi kiya kyunki woh numerator aur denominator dono mein identically appear hota hai.

Exercise 4.2 (L4)

Memory budget ke under design karo. Tumhare paas 24 GB GPU hai. Model weights 14 GB lete hain, 10 GB KV-cache ke liye bachte hain. Layer specs: $n_{layers} = 32$ , $h = 32$ , $d_{k} = 128$ , fp16. (a) Standard attention ke saath maximum context length $L$ kya hai jo aap cache kar sakte ho? (b) $g = 8$ par GQA ke saath?

Recall Solution

Per-token, all-layers, standard: $2 \times 32 \times 32 \times 128 \times 2 bytes = 524, 288 B = 512 KB/token .$ Budget $= 10$ GB $= 10 \times 2^{30} = 10, 737, 418, 240$ B (binary GB). (a) $L_{max} = \frac{10 , 737 , 418 , 240}{524 , 288} = 20, 480$ tokens. (b) GQA with $g = 8$ head count $32$ ko K/V mein $8$ se replace karta hai → $\frac{8}{32} = \frac{1}{4}$ per-token size $= 128$ KB/token: $L_{max} = 20, 480 \times 4 = 81, 920 tokens .$ Design conclusion: GQA( $g = 8$ ) par switch karna same memory budget ke under affordable context length ko quadruple kar deta hai.

Level 5 — Mastery

Exercise 5.1 (L5)

Crossover analysis. MQA memory save karta hai lekin aap quality ke baare mein worried ho. Maano standard attention perplexity $p_{0}$ deta hai aur MQA $p_{0} (1 + 0.015)$ deta hai (1.5% increase). Meanwhile MQA tumhe $h = 32 \times$ zyada tokens cache karne deta hai. Agar lambi context perplexity ko roughly $2%$ reduce karti hai jab context $4 \times$ grow hoti hai (empirical), toh kya $4 \times$ context kharidne ke liye standard→MQA trade karna perplexity ko net help karta hai ya hurt karta hai?

Recall Solution

Do effects perplexity par multiplicatively compose karte hain:

MQA head-sharing penalty: factor $1.015$ (worse).
$4 \times$ lambi usable context: factor $0.98$ (better). Net factor: $1.015 \times 0.98 = 0.9947.$ Kyunki $0.9947 < 1$ hai, net perplexity ~0.53% drop hoti hai — context gain sharing penalty se outweigh karta hai. Mastery point: memory optimizations "free lunch vs pure cost" nahi hain; yeh longer context enable karte hain, jo khud quality improve karta hai. Tumhe direct penalty ko uske dwara unlock ki gayi capability ke against weigh karna hoga.

Exercise 5.2 (L5)

Full stack estimate. Sab kuch combine karo: $n_{layers} = 80$ , $h = 64$ , $d_{k} = 128$ , fp16, aur aap GQA( $g = 8$ ) sliding window $W = 8192$ ke saath deploy karte ho ek nominal 200K-token context ke liye. (a) 200K tokens par full standard cache kya cost karti? (b) GQA + sliding-window cache actually kya cost karta hai? (c) Combined reduction factor?

Recall Solution

Standard per-token-all-layers: $2 \times 80 \times 64 \times 128 \times 2 = 2, 621, 440 B = 2.5 MB/token .$ (Binary: $2, 621, 440/ 2^{20} = 2.5$ MB.) (a) Full standard at $L = 200, 000$ : $200, 000 \times 2.5 MB = 500, 000 MB = 488.28 GB .$ (b) GQA K/V heads $64 \to 8$ cut karta hai, factor $8/64 = 1/8$ → per-token $= 2.5 MB /8 = 320$ KB. Sliding window tokens ko $W = 8192$ par cap karta hai: $8192 \times 320 KB = 2, 621, 440 KB = 2.5 GB .$ (c) Reduction: $\frac{488.28}{2.5} \approx 195 \times .$ Interpretation: do orthogonal levers multiply karte hain — GQA $8 \times$ deta hai (head sharing) aur window $200000/8192 \approx 24.4 \times$ deta hai (bounded length), aur $8 \times 24.4 \approx 195$ . Isliye production long-context serving (7.2.1-llm-deployment, 6.2.3-inference-optimization) multiple cache tricks stack karta hai ek par rely karne ki jagah.

Recall Quick self-check summary

Cache sirf K aur V rakhta hai ::: queries har step fresh hoti hain Memory formula ka leading constant ::: 2, K aur V ke liye Heads ke terms mein model width ::: $d = h \times d_{k}$ Standard→MQA cache ratio ::: $1/ h$ GQA( $g$ ) cache ratio ::: $g / h$ Caching se projection speedup ::: factor $L$ Sliding window memory order ::: $O (W)$ , $L$ mein constant Combining independent cache savings ::: factors ko multiply karo

Related: 6.1.11-sparse-attention, 6.3.1-model-quantization, 5.4.2-beam-search, 6.1.13 KV-cache optimization (Hinglish).