KV-cache optimization: Foundations — Hinglish deep dive

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

Isse pehle ki tum samjho kyun hum K aur V cache karte hain (aur Q nahi), tumhe yeh jaanna zaroori hai ki parent note mein har letter aur picture ka matlab kya hai. Yeh page unhe order mein build karta hai, har cheez pichli cheez pe depend karti hai. Kuch bhi assume nahi kiya gaya.

Related maps: 6.1.13 KV-cache optimization (Hinglish) · parent KV-cache optimization · prerequisites 6.1.1-transformer-architecture, 6.1.4-multi-head-attention.

1. Ek token aur uski embedding — sab kuch ka atom

Picture: ek row of boxes imagine karo, ek per word: "The cat sat". Har box ko numbers ka ek arrow se replace kiya jaata hai jo "meaning space" mein point karta hai. Position $t = 1$ hai "The", $t = 2$ hai "cat", aur aise hi aage.

Topic ko yeh kyun chahiye: poori cost story yeh hai ki "kya hota hai per token". Agar tum yeh nahi jaante ki har token ek vector $x_{t}$ ban jaata hai, toh tum count nahi kar sakte ki uske liye kya compute hota hai.

Figure 1 — teen word-boxes "The", "cat", "sat" positions $t = 1, 2, 3$ par, har ek ke saath ek neecha arrow numbers ki ek bracketed list mein: length $d$ ka embedding vector $x_{t}$ .

Ek embedding arrow mein components ki sankhya model dimension $d$ hai. Bade models $d = 768$ (GPT-2 small) se lekar hazaron tak use karte hain. Woh akela number $d$ cost formulas mein har jagah aata hai, isliye use abhi pin kar lo.

Notation

x_{t}

position

t

par token ka embedding vector hai;

d

batata hai ki usme kitne numbers hain.

2. Q, K, V — har token se puchhe gaye teen sawaal

Attention (6.1.4-multi-head-attention mein built) har token ki embedding $x_{t}$ ko teen naye vectors mein badalta hai, use teen fixed number-grids (matrices) se multiply karke jinhe $W_{Q}$ , $W_{K}$ , $W_{V}$ kehte hain.

Picture — ek library: har past book ki ek spine label (uski Key) aur contents (uski Value) hoti hai. Naya token ek search slip (uski Query) lekar andar aata hai. Woh apni slip ko har spine label se compare karta hai, phir un books se contents uthata hai jo best match karti hain.

Figure 2 — do past-token "books" har ek mein ek green Key label aur ek blue Value content dikha rahi hain, plus ek orange new-token "search slip" $Q_{t}$ jisme dashed red arrows har Key se compare kar rahe hain.

Topic ko yeh kyun chahiye — aur yeh caching ka dil hai: ek past token ka label $K_{i}$ aur contents $V_{i}$ kabhi nahi badalte ek baar likhne ke baad. Lekin har naya token ek fresh search slip $Q_{t}$ laata hai. Yahi asymmetry exactly woh reason hai kyun hum K aur V cache karte hain lekin Q recompute karte hain.

Recall Hum Q cache kyun nahi kar sakte?

Kyunki Q naye token ka sawaal hai, jo har step par alag hota hai. K aur V purane tokens ke fixed jawab hain. ::: Bilkul sahi.

W_{K}

kya hai?

Shape

d_{k} \times d

ki ek fixed matrix (training mein seekhi gayi, generation ke dauran frozen) jo ek embedding

x_{t}

ko uske Key vector

K_{t}

mein map karti hai.

3. Books ko rank karne wala multiply — $Q K^{T}$ aur dot product

Yeh measure karne ke liye ki search slip spine label se kitna match karti hai, attention dot product use karta hai: matching components ko multiply karo aur jod lo. Bada dot product = strong match.

Shapes ko pin down karna (yahan beginners confuse ho jaate hain): generation step $t$ par ek naya query hai, ek row vector $Q_{t} \in R^{1 \times d_{k}}$ aur $t$ cached keys hain jo ek matrix mein stack hain, ek key per row: $K = K_{1} K_{2} ⋮ K_{t} \in R^{t \times d_{k}} .$

Query ko har key ke saath ek saath dot karne ke liye, hume har key ek column ke roop mein chahiye. Transpose $K^{T}$ exactly wahi karta hai — woh rows ka $t \times d_{k}$ stack ko columns ka $d_{k} \times t$ block mein badal deta hai: $K^{T} \in R^{d_{k} \times t} .$

Ab shapes matrix multiply ke liye line up hoti hain: $1 \times d_{k} Q_{t} d_{k} \times t K^{T} = 1 \times t scores,$ har past token ke liye ek score. Inner $d_{k}$ dimensions match karke cancel ho jaate hain — transpose karne ki poori wajah yehi hai.

Picture: query arrow ki shadow jo har key arrow par padti hai — jitni lambi shadow, utna zyada score.

Figure 3 — origin se ek orange query vector $Q$ aur teen key vectors $K_{1}, K_{2}, K_{3}$ ; dotted projection lines dikhate hain $K_{3}$ (green) $Q$ ke saath sabse zyada align karta hai, sabse zyada dot-product score deta hai.

Parent note ke single-layer worked example mein (3rd token generate karna, 4 heads, $d_{k} = 64$ ), $Q_{3} (K_{1}, K_{2}, K_{3})^{T}$ exactly 3 scores produce karta hai — ek per past position — jo upar derive ki gayi length $t = 3$ ki row se match karta hai.

d_{k}

se divide kyun karte hain?

Scores ko bahut bada hone se rokne ke liye jab

d_{k}

badhta hai, jo agla step (softmax) bahut spiky bana deta. Yeh ek stabiliser hai.

4. Softmax — scores ko ek choice mein badalna

Picture — attention ka ek pie chart: bada raw score → bada slice, aur saare slices ek poori pie mein sum hote hain. Exponential $e^{s}$ leader ko exaggerate karta hai taaki strongest match dominate kare lekin kuch fully ignore na ho.

Figure 4 — left: teen raw scores ka ek bar chart; right: wahi scores softmax ke baad ek pie ke roop mein dikhaye gaye jinke teen slices 1 mein sum hote hain, sabse bada score sabse bada slice le raha hai.

Topic ko yeh kyun chahiye — aur yahan V ki shape enter karti hai: yeh weights cached Values ko blend karke ek output banane ke liye multiply karte hain. $t$ Value vectors ko waise hi stack karo jaise humne Keys stack ki thi: $V = V_{1} V_{2} ⋮ V_{t} \in R^{t \times d_{k}} .$ Softmax output $t$ weights ki ek row hai, shape $1 \times t$ . Shapes multiply karna: $1 \times t softmax scores t \times d_{k} V = 1 \times d_{k} output .$ Inner $t$ cancel ho jaata hai, toh ek $1 \times t$ weight row times $t \times d_{k}$ Value stack ek ek length- $d_{k}$ output vector deta hai — saari Values ka weighted blend. Toh poora attention formula

Attention (Q, K, V) = softmax (\frac{Q K ^{T}}{d _{k}}) V

left-to-right padhta hai: labels ko score karo → pie mein badlo → contents mix karo.

5. Heads, $h$ , aur $d_{k}$ — kaam ko split karna

Ek attention kaafi nahi hai, toh model kaafi heads parallel mein run karta hai — har ek ek head hai.

Ek "layer" kya count hota hai? Ek transformer identical blocks ko ek ke upar stack karta hai; har block ek layer hai (6.1.1-transformer-architecture inhe build karta hai). Ek layer mein apna attention (apne $W_{Q}, W_{K}, W_{V}$ ke saath) plus ek feed-forward network hota hai. GPT-2 small mein $12$ aisi layers hain; GPT-3-style model mein $96$ hain. $N$ = layers ki sankhya. Har layer apna alag KV-cache rakhti hai, isliye full-model memory $N$ se multiply hoti hai.

Topic ko yeh kyun chahiye: cache memory per head, per layer, per token count ki jaati hai. Ab ki $h$ , $d_{k}$ aur $p$ sab define ho gaye hain, memory formula padhta hai:

Memory (per layer, per token) = 2 \times h \times d_{k} \times p

$2$ isliye kyunki hum do cheezein store karte hain: Key stack aur Value stack.

6. Multi-Query Attention — jahaan $h$ -fold saving aati hai

Memory derivation, explicitly:

\frac{MQA cache per token}{Standard cache per token} = \frac{2 \times d _{k}}{2 \times h \times d _{k}} = \frac{1}{h} .

$h = 12$ heads ke saath, MQA $12 \times$ kam KV-cache store karta hai, kyunki K aur V dono ek baar (length $d_{k}$ par) store hote hain $h$ times ki jagah. Quality typically sirf 1–2% drop hoti hai.

7. Autoregressive generation aur causal mask

Picture — ek lower-triangular grid: row $i$ (attending token) columns $j \leq i$ (khud aur past) fill kar sakta hai lekin upper-right triangle (future) blacked out hai.

Figure 5 — ek $L \times L$ grid attending position $i$ (rows) aur key position $j$ (columns) se indexed; green cells $j \leq i$ mein $0$ (allowed), red upper-right cells $j > i$ mein $- \infty$ (future blocked).

Topic ko yeh kyun chahiye: kyunki mask guarantee karta hai ki ek token sirf positions $\leq t$ attend karta hai, aur woh past K,V frozen hain — yahi exactly unhe cache karne ka license hai. Cache ke hote hue bhi tumhe mask add karna hi hoga, warna model illegally future positions read kar lega.

8. Big-O, $L$ , aur woh cost jisse hum lad rahe hain

Topic ko yeh kyun chahiye: poora motivation ek cost comparison hai. Per token projection cost per layer $O (d^{2})$ hai ( $d_{k} \times d$ matrix times length- $d$ vector, $h$ heads mein sum karne par $h d_{k} d = d \cdot d = d^{2}$ milta hai). Big-O constant factors $N$ (layers) aur per-head split chhupaata hai, toh neeche clean formulas per-layer projection costs hain; full model ke liye $N$ layers se multiply karo:

Without cache (per layer): With cache (per layer): Speedup: t = 1 \sum L O (t d^{2}) = O (L^{2} d^{2}) t = 1 \sum L O (d^{2}) = O (L d^{2}) \frac{L ^{2} d ^{2}}{L d ^{2}} = L .

100-token sequence ke liye, yeh roughly $100 \times$ kam projection work hai. Parent note ka GPT-2-small worked example (jisme dono sides par $N = 12$ layer factor include hai, toh woh cancel ho jaata hai) $L = 100$ ke liye $\approx 50 \times$ real speedup measure karta hai — ideal $L$ se thoda kam kyunki attention khud abhi bhi har step par thoda cost karta hai, lekin same order of magnitude hai.

O (L^{2} d^{2})

O (L d^{2})

ka matlab?

Per-layer projection cost: pehla length mein quadratic hai (naive), doosra linear (cached). Unka ratio

L

hai, speedup. Pure model ke liye kisi bhi ko

N

layers se multiply karo.

Do classic mistakes (ek line mein kyun galat hain)

Prerequisite map

Har foundation agla feed karta hai; saath mein woh justify karte hain kyun K aur V cache karna (Q nahi) safe aur sasta hai. Deeper follow-ups: 6.2.3-inference-optimization, 6.1.11-sparse-attention, 6.3.1-model-quantization, 7.2.1-llm-deployment, 5.4.2-beam-search.

Equipment checklist

Khud test karo — right side cover karo.

x_{t}

kya hai?

Position

t

par token ke liye embedding (numbers ka vector).

d

ka matlab kya hai?

Model dimension — ek embedding mein kitne numbers hain.

Q, K, V ek-ek word mein?

Query = main kya dhoondh raha hoon; Key = mera label; Value = mera content.

Har projection matrix

W_{Q}, W_{K}, W_{V}

ki shape kya hai?

d_{k} \times d

— yeh length-

d

embedding ko length-

d_{k}

vector mein map karta hai.

Q, K, V mein se kaunse cache hote hain, aur kyun?

K aur V — woh past tokens ke hain aur kabhi nahi badalte; Q har step par fresh hota hai.

Step

t

par

Q_{t} K^{T}

mein shapes kya hain?

Q_{t}

1 \times d_{k}

hai,

K

t \times d_{k}

hai, toh

K^{T}

d_{k} \times t

hai aur scores

1 \times t

hain.

Kyun softmax(

\cdot

)

V

length-

d_{k}

output deta hai?

1 \times t

weight row times

t \times d_{k}

Value stack inner

t

cancel karta hai,

1 \times d_{k}

chodta hai.

Softmax kya karta hai?

Raw scores ko positive weights mein badalta hai jo 1 mein sum hote hain (attention ki ek pie).

d

h

d_{k}

ka relationship?

d = h \times d_{k}

Memory formula mein precision

p

kya hai?

Stored number per bytes — fp16 mein

p = 2

, fp32 mein

p = 4

Total KV-cache footprint formula?

2 \times N \times L \times h \times d_{k} \times p

(K+V, saari layers, saare tokens).

Cache memory formula mein 2 ka factor kyun hai?

Do stacks store hoti hain: K aur V.

Ek "layer" kya hota hai aur cache ke liye yeh kyun matter karta hai?

Ek transformer block (attention + feed-forward);

N

layers mein se har ek apna KV-cache rakhti hai, toh memory

N

se scale hoti hai.

MQA structurally kya change karta hai, aur cache kitna shrink hota hai?

Saare heads ek K aur ek V share karte hain (

d_{k}

size par stored,

h \times d_{k}

nahi), cache

h

factor se shrink hota hai.

Causal mask ki dimensions kya hain, aur yeh kaise apply hota hai?

L \times L

matrix

M

jo softmax se pehle raw scores mein add hoti hai;

M [i, j] = 0

for

j \leq i

aur

- \infty

for

j > i

Naive vs cached per-layer generation cost, aur speedup?

O (L^{2} d^{2})

O (L d^{2})

; speedup

= L

(full model ke liye

N

se multiply karo).