KV-cache optimization: Visual walkthrough — Hinglish deep dive

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

Shuru karne se pehle, teen words jo baar baar aayenge. Hum unhe ek baar plain language mein define karte hain, aur pehle kabhi nahi:

Step 1 — Ek word at a time ("autoregressive" ka matlab)

KYA. GPT jaisa model text ek token at a time likhta hai. Token 4 likhne ke liye pehle tokens 1, 2, 3 likh chuka hona chahiye. Yeh aage skip nahi kar sakta.

KYUN. Yeh left-to-right rule causality kehlata hai: ek word sirf apne pehle ke words dekh sakta hai, baad ke nahi. Yeh ek restriction hi is page ki har cheez ki seed hai — yahi cheez past ke kaam ko reusable banati hai.

PICTURE. Har nayi step dayi taraf ek square add karti hai. Arrows dikhate hain ki ek naya word sirf peeche dekhta hai — kabhi aage nahi.

token 1 x_{1} \to token 2 x_{2} \to token 3 x_{3} \to new x_{4}

$x_{t}$ — position $t$ pe word ki embedding (numbers ki list).
$\to$ arrows ka matlab hai "baad mein generate hua," strictly left to right.

Step 2 — Har naye word ko kya compute karna padta hai

KYA. Jab word $x_{t}$ aata hai, model use teen fixed weight-matrices se multiply karke uske teen characters banata hai:

Q_{t} = W_{Q} x_{t}, K_{t} = W_{K} x_{t}, V_{t} = W_{V} x_{t}

$W_{Q}, W_{K}, W_{V}$ — fixed number-tables (training mein seekhe, ab frozen). "Fixed" key word hai.
$x_{t}$ — aane wala word vector, length $d$ .
$Q_{t}, K_{t}, V_{t}$ — teen output vectors, har ek bhi length $d$ .

KYUN. Kyunki $W_{K}$ aur $W_{V}$ generation ke dauran kabhi nahi badlte, ek word jo key aur value produce karta hai woh ek baar stamp ho jaati hai aur hamesha ke liye rehti hai. Yahi deep reason hai ki caching possible hai — $K_{1}$ ke baare mein kuch bhi depend nahi karta ki baad mein kitne words aate hain.

PICTURE. Word bayi taraf box mein enter karta hai; teen arrows $Q, K, V$ lekar nikalte hain. Notice karo $K$ aur $V$ pe ek chhota padlock 🔒 hai — ek baar ban gaye, toh frozen.

Step 3 — Naive tarika: sab kuch recompute karo (badhta hua staircase)

KYA. Maan lo hum cache nahi karte. Token $t$ generate karne ke liye, ek plain transformer poore prefix ko brand-new sequence maanta hai aur $K_{1}, V_{1}, \dots, K_{t}, V_{t}$ scratch se recompute karta hai — chahe $K_{1}$ pehle $t - 1$ baar compute ho chuka ho.

KYUN. Ek raw forward pass ke steps ke beech koi memory nahi hoti. Use sirf pata hai "yeh raha $t$ tokens ka sequence, unke projections banao." Isliye step ke saath kaam badhta hai.

PICTURE. Har step ek bar hai; step $t$ ka bar $t$ blocks uncha hai. Sab milake ek triangle banaate hain — ek staircase jo baar baar unchi hoti jaati hai.

Ek word ke projection ka cost roughly $d^{2}$ multiply-adds hai (length- $d$ vector times $d \times d$ table). Step $t$ pe hum $t$ words ke liye yeh dobara karte hain:

work at step t = words redone t \times cost each d^{2}

$1$ se $L$ tak har step add karo:

t = 1 \sum L t d^{2} = d^{2} \frac{L ( L + 1 )}{2} = O (L^{2} d^{2})

$\sum_{t = 1}^{L} t$ — "1 + 2 + 3 + … + L," triangle ka area, $\frac{L ( L + 1 )}{2}$ ke barabar.
$L^{2}$ hi quadratic curse hai: text double karo, waste work chaar guna ho jaata hai.

Step 4 — Cached tarika: ek baar compute karo, yaad rakho (flat road)

KYA. Ek notebook rakho — KV-cache. Jab word $t$ aata hai, sirf apna $K_{t}, V_{t}$ compute karo (ek word ka), phir unhe notebook mein glue karo:

Cache_{K} \leftarrow [Cache_{K}; K_{t}], Cache_{V} \leftarrow [Cache_{V}; V_{t}]

Semicolon $;$ ka matlab hai "neeche ek nayi row stack karo." Notebook exactly ek row per step badhti hai.
Purani rows kabhi touch nahi hoti — yeh Step 2 se frozen 🔒 keys/values hain.

KYUN. Kyunki woh purani rows kabhi nahi badal sakti (frozen weights, past words fixed), unhe dobara derive karna bilkul waste hai. Hum $d^{2}$ cost ek baar per word, hamesha ke liye pay karte hain.

PICTURE. Ab har step ki height same hai (ek block). Staircase flat road mein collapse ho jaati hai jiska length $L$ hai.

t = 1 \sum L one word 1 \times d^{2} = L d^{2} = O (L d^{2})

Inner count ab $1$ hai, $t$ nahi — yahi poora trick hai.

Step 5 — Factor $L$ kahan se paida hota hai (triangle ÷ rectangle)

KYA. Dono totals ko divide karo.

Speedup = \frac{L ^{2} d ^{2} naive, triangle}{cached, road L d ^{2}} = L

KYUN. Saving exactly sequence length ke barabar hai. 2048 tokens likho → roughly 2048× kam projection work. Yeh koi approximation trick nahi hai; yeh triangle ke area aur rectangle ka ratio hai.

PICTURE. Unche triangle (naive) ko flat bar (cached) ke upar overlay karo. Triangle ka area bar ke area ke upar ka ratio $L /1$ hai — factor $L$ visible ho gaya.

Step 6 — Attention ko Query chahiye hoti hai ( $Q$ cache KYUN nahi hota)

KYA. Hum $K$ aur $V$ cache karte hain. $Q$ cache nahi karte. Step $t$ pe hum fresh $Q_{t}$ compute karte hain aur use saare cached keys ke against score karte hain:

scores_{t} = \frac{Q _{t} [ K _{1} , \dots , K _{t} ] ^{⊤}}{d _{k}}

$Q_{t}$ — naye word ke liye ek fresh query (uski curiosity uske liye unique hai).
$[K_{1}, \dots, K_{t}]^{⊤}$ — har cached key, columns mein lined up compare karne ke liye.
$d_{k}$ — ek shrink factor jo numbers ko blow up hone se rokta hai ( $d_{k}$ = ek head ka size); yeh counting argument ko affect nahi karta.

KYUN. Ek Key/Value "woh hai jo ek word offer karta hai" — yeh kabhi nahi badlta, isliye cache karo. Query "woh hai jo current word poochh raha hai" — har step pe brand-new question, isliye reuse nahi ho sakti. $Q$ cache karna galat sawaal ka jawab dena hoga.

PICTURE. Ek glowing nayi query arrow saare padlocked keys ke peeche fire karti hai. Query nayi hai; keys sab purani-aur-frozen hain.

Step 7 — Edge cases: woh corners jahan cheezein break ho sakti hain

KYA. Boundaries check karo taaki koi reader kabhi unshown case pe na atke.

Case A — pehla token ( $t = 1$ ). Notebook shuru mein empty hoti hai. Hum $K_{1}, V_{1}$ compute karte hain, cache ek row ban jaata hai, aur $Q_{1}$ sirf apne aap ko attend karta hai. Koi division problem nahi, koi "attend to nothing" nahi — flat road ki length 1 hai.

Case B — causal mask phir bhi apply hona chahiye. Cache ke saath bhi, ek word sirf rows $\leq t$ dekh sakta hai. Kyunki hum $K_{t}$ tab append karte hain jab word $t$ aata hai, cache mein future keys physically ho hi nahi sakti — append order mask ko free mein enforce karta hai. (Kisi given prompt ke prefill ke dauran, explicit mask phir bhi apply hoti hai.)

Case C — sliding window (bounded memory). Agar hum sirf aakhiri $W$ rows rakhein, toh memory badhna band ho jaati hai: $O (L)$ ki jagah $O (W)$ . Road sirf flat nahi balki fixed length ki ho jaati hai. Cost: $W$ steps se purane words bhool jaate hain. Dekho sparse attention.

Case D — memory, woh price jo hum pay karte hain. Speed per step mein extra memory free hai, lekin notebook itself linearly badhti hai: $2 \times n_{layers} \times L \times h \times d_{k} \times bytes$ . Yeh trade-off inference optimization mein explore hota hai aur aage quantization aur multi-/grouped-query attention se shrink hota hai.

PICTURE. Ek single strip: ek empty box (t=1), diagonal ke upar gray-out triangle (mask future block karta hai), aur ek fixed-width window ek lambe ribbon pe slide karti hui.

Ek-picture summary

Uncha triangle (sab kuch recompute karo) versus flat road (ek baar compute karo), padlocked notebook fresh query ko feed karti hui — poori derivation ek single frame mein.

Recall Feynman retelling — plain words mein vapas bolo

Socho tum ek essay likh rahe ho jahan, har naya sentence add karne se pehle, tum poora essay phir se shuru se padhte ho. Sentence 100 tumhe woh 99 sentences dobara padhwaata hai jo tum pehle se jaante ho. Yeh re-reading hi naive cost hai — badhta hua staircase, $L^{2}$ .

Ab socho tum har likhe hue sentence ki ek notebook mein ek summary line rakhte ho. Naya sentence add karne ke liye tum notebook dekho (woh pehle se wahan hai) aur ek nayi line jodo. Same story, lekin ab har sentence ka kaam constant hai — flat road, $L$ .

Do facts notebook ko trustworthy banate hain: (1) model ke weights frozen hain, isliye kisi past word ki Key/Value kabhi nahi badal sakti — store karna safe hai; (2) tum sirf aage likhte ho, isliye notebook kabhi future word leak nahi kar sakti — causal rule free mein hold karta hai. Staircase ko road se divide karo aur speedup nikalta hai: exactly sequence length $L$ . Ek cheez jo tum store nahi kar sakte woh hai Query, kyunki har naya word ek brand-new sawaal poochh raha hota hai.

Recall Quick self-check

Naive projection cost $O (L^{2} d^{2})$ kyun hai? ::: Har step $t$ pe $t$ words ke projections recompute hote hain, har ek $d^{2}$ pe; $1 + 2 + \dots + L = \frac{L ( L + 1 )}{2} = O (L^{2})$ sum karne se $O (L^{2} d^{2})$ milta hai. Exactly $L$ speedup kyun, kuch aur kyun nahi? ::: Yeh triangle area $\frac{L ^{2}}{2} d^{2}$ aur flat road $L d^{2}$ ka ratio hai, jo $L$ ke barabar hai. $K, V$ cache kyun karo lekin $Q$ kabhi nahi? ::: $K, V$ = "ek past word kya offer karta hai," frozen aur reusable; $Q$ = "current word kya poochh raha hai," har step pe naya. Cached generation mein causal mask automatically kya enforce karta hai? ::: Hum $K_{t}$ tabhi append karte hain jab word $t$ aata hai, isliye future keys cache mein physically kabhi present nahi hoti.

Yeh bhi dekho: Transformer architecture · Beam search (cache ko beam width se multiply karta hai) · LLM deployment.