KV-cache optimization: Worked examples — Hinglish deep dive

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

Shuru karne se pehle, un symbols ki ek choti si reminder jo hum use karte rahenge, taaki koi peeche na rahe:

Definition Char characters

$L$ — sequence length, yani kitne tokens abhi tak aa chuke hain. Ek row of boxes imagine karo, ek box per word.
$d$ — model width (jise hidden size bhi kehte hain): ek token ko describe karne ke liye kitne numbers hote hain. Har box ki height imagine karo.
$h$ — number of heads: model parallel mein kitne independent attention "viewpoints" chalata hai. Har token ke box ko $h$ patli vertical strips mein slice kiya hua imagine karo, ek strip per head.
$d_{k}$ — head dimension: ek attention head kitne numbers dekhta hai (ek strip ki height). Usually $d = h \times d_{k}$ , toh strips wapas poore box mein jod di jaati hain.

Ek token ke liye "K" ya "V" ek chota rectangle hota hai jiska shape $h \times d_{k}$ hai (saari strips, side by side). $L$ ko stack karo aur tumhe cache milta hai: ek $L \times h \times d_{k}$ block. Yahi block neeche diye gaye har example ka star hai.

Scenario matrix

Har KV-cache question in mein se kisi ek cell ka hota hai. Neeche diye gaye examples shuru mein declare karte hain ki woh kis cell(s) ko cover karte hain.

#	Case class	Concrete trigger	Covered by
A	Empty / first token ( $L = 1$ )	pehle token ka prefill, cache empty se shuru	Ex 1
B	Growing cache, one step	token $t$ append karo, $t$ positions pe attend karo	Ex 2
C	Prefill vs decode asymmetry	poora prompt ek saath jaata hai, phir ek-ek karke decode	Ex 3
D	Zero / degenerate input	$L = 0$ , ya single-head ( $h = 1$ ), ya $d_{k}$ bahut chhota	Ex 4
E	Limiting behaviour ( $L \to \infty$ )	cache memory kab model weights se aage nikal jaati hai?	Ex 5
F	Head-layout variants	standard vs MQA vs GQA cache size	Ex 6
G	Real-world word problem	"kya yeh 32 GB GPU pe serve hoga?" — deployment sizing	Ex 7
H	Exam twist	sliding window: cache barhna band, break-even token dhundo	Ex 8

Hum inhe order mein attack karte hain. Cells A–C mechanics hain, D–E edge aur limit, F–H engineering reality.

Example 1 — Cell A: empty cache, pehla token

Forecast: Abhi compute karne se pehle andaza lagao — cache mein pehle kitne rows hain? Aur ek token ke baad?

Steps:

Generation se pehle, cache empty hai. Iska shape hai $0 \times 4 \times 64$ — literally zero rows. Yeh step kyun? Caching incremental hai; pehle token ke exist hone se pehle kuch bhi remember karne ke liye nahi hota. Yeh woh base case hai jis se har loop shuru hoti hai.
Pehle token ke projections compute karo: $K_{1} = W_{K} x_{1}$ aur $V_{1} = W_{V} x_{1}$ , dono ka shape $4 \times 64$ hai. Yeh step kyun? $x_{1}$ token 1 ki embedding hai (ek length- $d$ vector). Fixed matrix $W_{K}$ se multiply karne par is token ki key saare 4 heads ke liye milti hai. Value ke liye bhi yahi. Dekho projection step.
Append karo. Cache $0 \times 4 \times 64$ se $1 \times 4 \times 64$ ho jaata hai. Yeh step kyun? Appending hi poori trick hai: hum exactly ek row per token badhate hain.
Store hue numbers count karo. Do matrices (K aur V), har ek $4 \times 64 = 256$ numbers: $2 \times 1 \times 4 \times 64 = 512 numbers = 512 \times 2 = 1024 bytes = 1 KB .$ Yeh step kyun? Yeh parent note ka per-token memory formula hai, $L = 1$ par evaluate kiya hua.

Verify: Token 1 ka attention sirf position 1 par attend karna chahiye — ek $4 \times 1$ score (ek score per head, ek position). Exactly 1 row wale cache ke saath exactly yahi milta hai. Units check: $numbers \times bytes/number = bytes$ . ✓

Example 2 — Cell B: ek growth step, teesra token

Forecast: Token 3 ke score row ki width (kitne columns) kitni hogi?

Steps:

Cache pehle: $2 \times 4 \times 64$ (tokens 1,2). Yeh step kyun? Hum sirf re-read karte hain; $K_{1}, K_{2}$ kabhi recompute nahi karte. Cache ka yahi point hai.
Naye token ki key aur value compute karo: $K_{3} = W_{K} x_{3}$ aur $V_{3} = W_{V} x_{3}$ , dono ka shape $4 \times 64$ , jahan $x_{3}$ token 3 ki embedding hai. Yeh step kyun? Sirf token 3 genuinely naya hai, isliye sirf iske K aur V compute karne hain — woh reusable "what previous tokens offer" pieces jo hum cache mein append karenge.
Naya query alag compute karo: $Q_{3} = W_{Q} x_{3}$ , yeh bhi $4 \times 64$ — aur hum ise cache nahi karte. Yeh step kyun? $Q_{3}$ encode karta hai "token 3 kya dhundh raha hai". Har future token apna fresh query lekar aata hai, aur purane queries par kabhi re-attend nahi hota, isliye $Q$ cache karna memory waste hoga (parent mein yeh Mistake 1 hai). Keys aur values yaad rakhe jaate hain; queries hamesha ek use ke baad discard kar diye jaate hain.
K aur V append karo: cache $3 \times 4 \times 64$ ho jaata hai. ( $Q_{3}$ append nahi hota.)
Score shape. Har head ke liye, $Q_{3}$ (ek $1 \times 64$ vector) 3 cached keys ( $3 \times 64$ ) se dot product karta hai: $scores = \frac{Q _{3} K _{1 : 3}^{⊤}}{64} \in R^{4 \times 3} .$ Yeh step kyun? Ek row per head (4), ek column per attended position (3). $d_{k} = 64 = 8$ scaling dot products ko blow up hone se rokti hai (scaled dot-product attention se).

Figure 1 — cache append aur score fan-out. Left mein, do lavender rows frozen cache hain ( $K_{1}, K_{2}$ ) jo hum kabhi recompute nahi karte; coral row ( $K_{3}$ , marked NEW) woh single row hai jo humne abhi append ki, cache ko $3 \times 4 \times 64$ banate hue. Right mein, mint box fresh query $Q_{3}$ hai — teen coral arrows trace karo: $Q_{3}$ 3 cached keys mein se har ek tak wapas pahunchta hai, aur vertical slate arrow (labelled "/ 8") $64$ se divide karta hai butter box se pehle jo result collect karta hai: ek $4 \times 3$ score, ek row per head, teen columns teen attended positions ke liye. Notice karo $Q_{3}$ ka koi arrow cache mein nahi jaata — yeh ek baar use hota hai aur discard ho jaata hai, exactly jaisa step 3 ne argue kiya.

Verify: Figure dekho — naya row (coral) height-1 hai, aur score arrow exactly 3 cached columns tak fan out karta hai. Column count = current token index = 3. ✓ Aur $64 = 8$ , ek clean integer. ✓

Example 3 — Cell C: prefill vs decode

Forecast: Kya cached run roughly $L$ times sasta hoga, ya beeche mein kuch, bade prefill chunk ki wajah se?

Steps:

Unit fix karo. Ek token project karna matlab iska $d$ -vector $x_{t}$ ko $W_{K}, W_{Q}, W_{V}$ mein se har ek se multiply karna (har ek $d \times d$ ), jiska cost $O (d^{2})$ multiply-adds hai. Hum isse ek block kehte hain. $n$ tokens project karna $n$ independent aisi products hain, isliye exactly $n$ blocks cost hogi — chahe ek-ek karke karo ya ek single $(n \times d) \cdot (d \times d)$ batched matmul mein fuse karo. Yeh step kyun? Ek batched matmul $n$ alag waalon se kam arithmetic nahi karta — woh same $n$ token-products ko ek call mein pack karta hai hardware efficiency ke liye. Isliye " $P$ tokens ek matmul mein" phir bhi $P$ blocks kaam hai, constant nahi. Yehi wajah hai ki $P$ tokens ka prefill $P$ blocks cost karta hai, constant nahi.
Naive projection cost. Har token position $t$ par naive model saare $t$ positions ke liye K,V recompute karta hai: $\sum_{t = 1}^{15} t = \frac{15 \cdot 16}{2} = 120$ blocks. Yeh step kyun? Cache ke bina, token $t$ generate karna sequence ko fresh treat karta hai — har baar $t$ projections.
Cached projection cost. Prefill saare $P = 10$ prompt tokens ko ek baar project karta hai (10 blocks, step 1 se), phir $G = 5$ decode steps mein se har ek exactly 1 naya token project karta hai (5 blocks): $10 + 5 = 15$ blocks. Yeh step kyun? Cache ka matlab hai hum prompt ka K,V kabhi redo nahi karte, aur har decode step sirf apna single naya token add karta hai.
Speed ratio. $120/15 = 8$ . Yeh step kyun? Yeh theoretical " $\times L$ " speed-up ki concrete cash value hai — lekin yeh 15 nahi, 8 aata hai, kyunki naive sum $L (L + 1) /2$ hai, $L^{2}$ nahi, aur cached cost $L$ ke barabar hai.

Verify: General check: naive $= \frac{L ( L + 1 )}{2}$ , cached $= L$ , ratio $= \frac{L + 1}{2}$ . $L = 15$ ke liye: $\frac{16}{2} = 8$ . ✓ Step 4 se match karta hai. Parent ka " $\times L$ " big-O scaling hai; exact ratio $\frac{L + 1}{2}$ hai.

Example 4 — Cell D: degenerate inputs

Forecast: In mein se kaunsa zero bytes deta hai, aur kaunsa sabse chhota non-zero cache deta hai?

Steps:

(a) $L = 0$ . Total cache $= 2 \times L \times h \times d_{k} \times bytes = 0$ . Yeh step kyun? Empty sequence → kuch store nahi karna. Confirm karta hai ki formula gracefully degrade karta hai; koi divide-by-zero nahi, koi negative memory nahi.
(b) $h = 1$ . Per token $= 2 \times 1 \times 64 \times 2 = 256$ bytes. Yeh step kyun? Single-head model $h \times d_{k}$ ko sirf $d_{k}$ mein collapse karta hai. Yeh MQA per-token figure bhi hai (ek shared K/V) — Ex 6 ka preview.
(c) $d_{k} = 1$ , $h = 8$ . Per token $= 2 \times 8 \times 1 \times 2 = 32$ bytes. Yeh step kyun? Head dimension ko 1 tak shrink karne se har key/value ek single scalar ban jaata hai. Phir bhi 8 heads hain, toh K ke liye 8 scalars + V ke liye 8 = 16 numbers = 32 bytes.

Verify: (a) $= 0$ . ✓ (b) $256$ bytes, non-zero aur "real" cases mein heads ki number ke hisaab se sabse chhota. (c) $32$ bytes yahan actually sabse chhota total hai. Ordering: $0 < 32 < 256$ . ✓ Har formula degenerate inputs ke neeche finite aur non-negative rehta hai.

Example 5 — Cell E: crossover limit ( $L \to \infty$ )

Forecast: Hazaron tokens? Laakhon? Order of magnitude guess karo.

Steps:

Per-token full-model cache. $2 \times n_{layers} \times h \times d_{k} \times bytes = 2 \times 12 \times 12 \times 64 \times 2 = 36, 864$ bytes $= 36$ KB. Yeh step kyun? Yeh parent ka formula layers par summed hai, per token. Har naya token yeh fixed slab add karta hai.
Cache = weights set karo. $L \times 36, 864 = 500 \times 1 0^{6}$ solve karo: $L = \frac{500 , 000 , 000}{36 , 864} \approx 13, 563 tokens .$ Yeh step kyun? Cache linearly $L$ mein badhta hai jabki weights constant hain — isliye hamesha ek crossover hoti hai. Iske baad, memory cache se dominated hoti hai, model se nahi.
Limit interpret karo. Jaise $L \to \infty$ cache sab kuch dwarf kar deta hai: $lim_{L \to \infty} \frac{cache}{weights} = \infty$ . Yehi wajah hai ki MQA/GQA aur sliding windows exist karte hain (Ex 6, 8). Yeh step kyun? Parent ka poora "advanced optimizations" section is unbounded growth ka response hai.

Verify: $36, 864 \times 13, 563 \approx 5.0 \times 1 0^{8} = 500$ MB. ✓ Linear-in- $L$ growth ka matlab hai ratio ka koi finite bound nahi hai. ✓ (Yahan MB = $1 0^{6}$ bytes, decimal — units clarification ke liye Ex 7 dekho.)

Example 6 — Cell F: standard vs MQA vs GQA

Forecast: MQA koi bada divisor chhota hona chahiye. Exactly kitne factor se, aur GQA kahan land karta hai?

Steps:

(a) Standard. $2 \times h \times d_{k} \times 2 = 2 \times 64 \times 128 \times 2 = 32, 768$ bytes $= 32$ KB. Yeh step kyun? Har query head apna K aur V rakhta hai. Yeh woh baseline hai jo parent ka formula deta hai.
(b) MQA. Ek shared K/V head: $2 \times 1 \times 128 \times 2 = 512$ bytes. Yeh step kyun? MQA $h$ K/V heads ki jagah ek single rakh deta hai. Reduction $= 64/1 = 64 \times$ .
(c) GQA, $g = 8$ . $2 \times 8 \times 128 \times 2 = 4096$ bytes $= 4$ KB. Yeh step kyun? GQA ek K/V per group rakhta hai. 8 groups ke saath reduction $64/8 = 8 \times$ hai — parent ka stated Llama-2 figure.

Figure 2 — cache size staircase. Teen bars, sabse lambe se sabse chhote: lavender bar standard multi-head hai 32,768 bytes/token/layer par (saare 64 K/V heads rakhe hue); mint bar GQA hai 8 groups ke saath 4,096 bytes par (labelled "8x smaller"); coral bar MQA hai 512 bytes par (labelled "64x smaller"). Left-to-right padhne par staircase har baar girta hai jab hum zyada queries mein K/V heads share karte hain — height literally memory bill hai, isliye chhota bar sasta cache hai. GQA deliberately dono extremes ke beech mein baith hai: MQA ki bahut zyada saving, lekin uski quality loss bahut kam.

Verify: Standard $32768$ B; MQA $512$ B → ratio $32768/512 = 64$ . ✓ GQA $4096$ B → ratio $32768/4096 = 8$ . ✓ GQA MQA aur standard ke beech mein hai, exactly jaisa figure ka staircase dikhata hai. ✓

Example 7 — Cell G: real-world deployment sizing

Forecast: GPU ka kaafi headroom hai, toh kya cache yahan problem bhi hai — haan ya nahi?

Steps:

Per-user cache (standard). Ex 5 se, per-token full-model cache $= 36, 864$ bytes. 2048 tokens ke liye: $2048 \times 36, 864 = 75, 497, 472$ bytes. Binary MiB mein ( $/ 2^{20}$ ): $75, 497, 472/1, 048, 576 = 72.0$ MiB exactly. Yeh step kyun? Ek user ka cache = per-token slab $\times$ context length. Bytes $2^{20}$ se cleanly divide hoti hain, isliye 72 MiB exact hai, rounded nahi.
8 users + weights. $8 \times 72 MiB + 500 MiB = 576 + 500 = 1076$ MiB $\approx 1.05$ GiB. Yeh step kyun? Concurrent requests mein se har ek ko apna khud ka cache chahiye; weights ek baar shared hain. Total 32 GiB budget mein fit hona chahiye.
Fit check. $1076 MiB ≪ 32 \times 1024 = 32, 768 MiB$ → haan, easily fit ho jaata hai. Weights reserve karne ke baad, caches ka budget $= 32, 268 - 500 = 32, 268$ MiB, isliye $max users = ⌊ \frac{32 , 268}{72} ⌋ = ⌊ 448.2 ⌋ = 448 users (standard) .$ Yeh step kyun? Cache, weights nahi, concurrency ke liye scaling bottleneck hai — yeh core inference-optimization insight hai.
MQA switch. MQA K/V ko head factor $h = 12$ se shrink karta hai, isliye per-user cache $= 72/12 = 6$ MiB. Max users $= ⌊ 32, 268/6 ⌋ = 5378$ . Yeh step kyun? Cache $12 \times$ cut karne se serviceable concurrency roughly $12 \times$ multiply hoti hai (448 → ~5378). Yehi wajah hai production LLMs MQA/GQA use karte hain — dekho deployment.

Verify: Standard: $8 \times 72 + 500 = 1076$ MiB, aur $1076 < 32768$ MiB. ✓ Max standard users $= 448$ ; MQA $= 5378$ ; ratio $\approx 12$ . ✓ Units: MiB / (MiB/user) = users, poora binary throughout. ✓

Example 8 — Cell H: exam twist, sliding window break-even

Forecast: Kya window cache token 1 se hi memory bachata hai, ya kisi threshold ke baad?

Steps:

Growth laws. Full cache rows $= t$ (hamesha badhta hai). Window cache rows $= min (t, W) = min (t, 512)$ . Yeh step kyun? Sliding window $W$ exceed hone par purana token drop kar deta hai, isliye iska size 512 par plateau ho jaata hai — constant memory, parent ke "Sliding Window" section se.
Pehla farq. Har $t \leq 512$ ke liye dono caches $t$ rows rakhte hain (abhi tak kuch drop nahi hua), isliye woh identical hain. $t = 513$ par full cache ke 513 rows hain lekin window ne token 1 already evict kar diya hai aur sirf $min (513, 512) = 512$ rakhta hai. Toh woh pehli baar $t = 513 = W + 1$ par alag hote hain. Yeh step kyun? Window size se neeche koi saving nahi hoti; saving exactly $W$ se ek token aage shuru hoti hai. Woh threshold $t = W + 1$ hi woh "break-even" hai jo question pooch raha hai.
$t = 4096$ par ratio. Full $= 4096$ rows; window $= min (4096, 512) = 512$ rows. $ratio = \frac{4096}{512} = 8.$ Yeh step kyun? Long context par window yahan $8 \times$ chhota hai — tradeoff hai 512 tokens se purane dependencies ka loss (relevant to long-horizon decoding jahan distant context matter kar sakta hai).
Limit behaviour. Window memory $O (W)$ constant hai jabki full $O (t)$ linear hai, isliye jaise $t \to \infty$ ratio $t / W \to \infty$ : saving unbounded badhti hai. Yeh step kyun? 100K+ token contexts ke liye sliding-window attention ka yahi poora point hai — memory barhna bilkul band ho jaata hai.

Verify: $min (513, 512) = 512 \neq = 513$ aur $min (512, 512) = 512 = 512$ → pehla divergence $t = 513 = W + 1$ par. ✓ $4096/512 = 8$ . ✓ Ratio $t / W$ $t \to \infty$ par unbounded hai. ✓

Recall Self-check

With-cache projection cost vs naive, exact ratio for length $L$ ? ::: $\frac{L + 1}{2}$ (nahi $L$ — woh big-O hai). KV-cache pehli baar length- $W$ window cache se kab differ karta hai? ::: Token $t = W + 1$ par. MQA cache reduction factor vs standard multi-head with $h$ heads? ::: $h \times$ (ek shared K/V head). $Q$ kabhi cache kyun nahi hota? ::: Har naye token ko ek fresh query chahiye ("woh kya dhundh raha hai"); purane queries par kabhi re-attend nahi hota. Per-token full-model cache formula? ::: $2 \times n_{layers} \times h \times d_{k} \times bytes$ .

Related: 6.3.1-model-quantization upar diye gaye har formula mein bytes-per-number factor ko shrink karta hai — head-count tricks se orthogonal, aur unke saath stack hota hai.