KV-cache optimization — Hinglish notes

6.1.13 · AI-ML › Scaling & Efficient Architectures

Yeh Kyun Chahiye: First Principles se Derivation

Naive Attention Cost

Multi-head self-attention mein, sequence length $L$ ke liye, hum compute karte hain:

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

Step-by-step cost analysis:

Training ke dauran (teacher forcing): Hamare paas ek saath saare $L$ tokens hote hain.
- $Q, K, V \in R^{L \times d}$ ek baar compute hote hain
- Attention matrix: $O (L^{2} d)$ operations
- Total per layer: $O (L^{2} d + L d^{2})$
Autoregressive generation ke dauran (problem):
- Position $t = 1$ par token generate karo: $K_{1}, V_{1}$ compute karo ✓
- Position $t = 2$ par token generate karo: $Q_{2}$ compute karo (naya), lekin recompute $K_{1}, V_{1}, K_{2}, V_{2}$ ✗
- Position $t = L$ par token generate karo: recompute $K_{1}, \dots, K_{L}$ aur $V_{1}, \dots, V_{L}$ ✗

Yeh step kyun? Har generation step mein sirf ek naya token add hota hai, lekin caching ke bina model is poori sequence ko fresh treat karta hai aur saare previous tokens ke projections recompute karta hai.

$L$ tokens generate karne ka total naive cost:

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

Yeh sequence length mein quadratic hai sirf projection overhead ke liye, attention karne se pehle bhi!

KV-Cache Solution

Savings ki derivation:

Cache ke saath: Cache ke bina: Speedup: t = 1 \sum L O (d^{2}) = O (L d^{2}) t = 1 \sum L O (t d^{2}) = O (L^{2} d^{2}) \frac{L ^{2} d ^{2}}{L d ^{2}} = L

Yeh kyun kaam karta hai: Attention mein causality mask ensure karta hai ki token $t$ sirf positions $\leq t$ par attend kare. Kyunki $K_{i}, V_{i}$ for $i < t$ fixed hain (model generation ke dauran frozen hai), unki values kabhi nahi badalti — caching ke liye perfect.

Full-model cache (saare layers mein) = $2 \times n_{layers} \times L \times h \times d_{k} \times precision$ .

Example: GPT-3 style model ( $n_{layers} = 96$ , $h = 96$ heads, $d_{k} = 128$ , fp16):

Per token, per layer: $2 \times 96 \times 128 \times 2 = 49, 152 bytes \approx 48 KB$
Per token, saare 96 layers: $96 \times 48 KB = 2 \times 96 \times 96 \times 128 \times 2 bytes \approx 4.7 MB$
2048 tokens ke liye (saare layers): $2048 \times 4.7 MB \approx 9.6 GB$

Yeh step kyun? Hamare paas 2 matrices hain (K aur V), har ek ka size (num_tokens × num_heads × head_dim) hai, jo per layer store hoti hain. Full-model cost layers ki number se multiply hoti hai, isliye per-token figure ~48 KB (ek layer) se ~4.7 MB (saare 96 layers) tak jump karta hai.

Worked Examples

Step 1: Previous cache state (tokens 1-2 generate karne ke baad):

Cache_{K} = [K_{1} K_{2}] \in R^{2 \times 4 \times 64}, Cache_{V} = [V_{1} V_{2}] \in R^{2 \times 4 \times 64}

Yeh step kyun? Har $K_{i}$ ka shape (num_heads, head_dim) = (4, 64) hai. Ab tak 2 tokens generate ho chuke hain.

Step 2: Naye token ke liye projections compute karo:

K_{3} = W_{K} x_{3} \in R^{4 \times 64}, V_{3} = W_{V} x_{3} \in R^{4 \times 64}, Q_{3} = W_{Q} x_{3} \in R^{4 \times 64}

Step 3: Cache update karo (concatenation):

Cache_{K} \leftarrow K_{1} K_{2} K_{3} \in R^{3 \times 4 \times 64}

Step 4: Sirf $Q_{3}$ use karke saare cached keys ke saath attention compute karo:

scores = \frac{Q _{3} ( K _{1} , K _{2} , K _{3} ) ^{T}}{64} \in R^{4 \times 3}

Yeh step kyun? $Q_{3}$ (shape 4×64) cached K se multiply hota hai (shape 3×4×64 → transpose karke 4×64×3). Har head ko ek 1×3 score vector milta hai (3 positions par attending).

Cost: $O (d^{2})$ ek token ke projections ke liye + $O (t \cdot d)$ attention scores ke liye. $K_{1}, K_{2}$ ka koi recomputation nahi.

100-token sequence generate karna:

Cache ke bina:

Per layer, per token $t$ : $t$ key-value pairs compute karo → $O (t \times 76 8^{2})$
Total: $12 \times \sum_{t = 1}^{100} t \times 76 8^{2} \approx 3.5 \times 1 0^{10}$ FLOPs

Cache ke saath:

Per layer, per token: 1 naya K, V compute karo → $O (76 8^{2})$
Total: $12 \times 100 \times 76 8^{2} \approx 7.1 \times 1 0^{8}$ FLOPs
Speedup: 50× ( $L = 100$ ke liye approximately $L$ factor se match karta hai)

Memory overhead:

Cache size: $12 layers \times 2 (K+V) \times 100 tokens \times 12 heads \times 64 dim \times 2 bytes \approx 3.7 MB$

Yeh kyun important hai: 3.7 MB, 500 MB model weights ke comparison mein negligible hai, lekin 98% projection compute bachata hai.

Advanced Optimizations

1. Multi-Query Attention (MQA)

Cache reduction:

\frac{MQA cache}{Standard cache} = \frac{2 L d _{k}}{2 L h d _{k}} = \frac{1}{h}

Yeh step kyun? $h = 32$ heads ke saath, MQA 32× kam KV-cache memory use karta hai. Quality drop typically 1-2% perplexity hota hai.

2. Grouped-Query Attention (GQA)

Middle ground: $h$ query heads ko $g$ groups mein group karo ( $g < h$ ), har group ek K, V head share karta hai.

GQA cache = 2 Lg d_{k}, jahan g = h / k kisi k > 1 ke liye

Example: Llama-2 70B $h = 64$ , $g = 8$ use karta hai (standard ke against 8× cache reduction).

3. Sliding Window / Local Attention

Bahut lambi sequences ke liye, sirf last $W$ tokens cache karo:

Cache_{K}^{(t)} = [K_{t - W + 1}, \dots, K_{t}]

Tradeoff: Constant $O (W)$ memory, lekin $W$ tokens se aage long-range dependencies kho jaate hain.

Use case: 100K+ token contexts ke saath summarization jahan recent context sabse zyada important hota hai.

4. FlashAttention Integration

FlashAttention operations ko fuse karke aur SRAM tiling use karke, poora $O (L^{2})$ score matrix materialize kiye bina attention compute karta hai. KV-cache + FlashAttention mila ke kaam karte hain:

Cache redundant K, V projections eliminate karta hai
FlashAttention attention scores ka memory overhead eliminate karta hai

Combined speedup: Long sequences ke liye 10-20× (dono alag-alag bottlenecks reduce karte hain).

Common Mistakes

Yeh sahi kyun lagta hai: $Q$ ko bhi ek matrix multiply chahiye, toh caching logical lagti hai.

Steel-man: Ek typical forward pass mein, $Q$ cache karna help karta. Lekin autoregressive generation mein, har naya token ek fresh query generate karta hai jo saari previous positions par attend karta hai. Humein token $t$ ke liye naya $Q_{t}$ chahiye kyunki yeh encode karta hai "naya token kya dhundh raha hai."

Fix: Sirf K aur V reusable hain kyunki yeh represent karte hain "previous tokens kya offer karte hain." Query represent karta hai "current token ko kya chahiye" — ise recompute karna hi hoga.

Yeh sahi kyun lagta hai: Cache mein data hai, toh use kyun na karein?

Steel-man: Full sequences ke saath training ke dauran hum causal mask use karte hain. Yeh bhoolna asaan hai ki generation ke dauran bhi yeh zaruri hai.

Fix: Caching ke saath bhi, causal mask apply karo:

mask [i, j] = {0 - \infty if j \leq i if j > i

Yeh ensure karta hai ki token $t$ sirf positions $1, \dots, t$ dekhe (chahe physically autoregressive mode mein $K_{t + 1}, \dots$ abhi exist hi na karein, yeh batched generation ke liye matter karta hai).

Yeh sahi kyun lagta hai: Notation $O (L)$ memory kehta hai.

Steel-man: Notation constants hide karta hai. 32-layer, 32-head model ke liye $d_{k} = 128$ fp16 mein:

Cache = 2 \times 32 \times 32 \times 128 \times L \times 2 bytes = 1, 048, 576 L bytes \approx 1 MB \times L

$L = 100, 000$ tokens ke liye: sirf cache ke liye ~105 GB (zyaatar GPU memory se zyada).

Fix: Hybrid strategies use karo jaise sliding window cache, purane cache ko CPU par offload karna, ya ultra-long contexts ke liye sparse attention.

Implementation Details

Typical PyTorch pattern:

class CachedAttention:
    def __init__(self):
        self.kv_cache = None  # Will hold (batch, num_heads, seq_len, head_dim)
    def forward(self, x, use_cache=True):
        B, T, C = x.shape
        Q = self.W_q(x)  # Always compute fresh
        K_new = self.W_k(x)
        V_new = self.W_v(x)
        if use_cache and self.kv_cache is not None:
            K_cached, V_cached = self.kv_cache
            K = torch.cat([K_cached, K_new], dim=2)  # Concat along seq_len
            V = torch.cat([V_cached, V_new], dim=2)
        else:
            K, V = K_new, V_new
        # Standard attention computation
        scores = (Q @ K.transpose(-2, -1)) / math.sqrt(self.head_dim)
        atn = F.softmax(scores, dim=-1)
        out = atn @ V
        
        if use_cache:
            self.kv_cache = (K.detach(), V.detach())  # Detach to avoid gradient retention
        
        return out

Detach kyun? Poore cached history ke through backprop rokta hai (sirf training ke dauran zaruri hai; inference ke dauran waise bhi koi gradients nahi hote).

Kab Kya Use Karein

Context Length	Strategy	Example Use Case
< 2K tokens	Standard KV-cache	Chatbot turns, code completion
2K - 8K	KV-cache + MQA/GQA	Document Q&A
8K - 32K	Sliding window cache	Long-form writing
32K+	Sparse/local + offload	Book summarization, legal docs

Recall Ek 12-Saal Ke Bache Ko Samjhao

Socho tum ek story likh rahe ho ek word at a time. Har baar jab tum naya word likhte ho, tumhe yaad rakhna hota hai ki saare previous words kiske baare mein the taaki tumhara naya word sense kare.

Slow tarika: Har baar jab tum naya word likhte ho, tum poori story ki shuruaat se re-read karte ho aur har word ke baare mein notes lete ho. 100 words likhne ka matlab hai story 100 baar padhna!

KV-cache tarika: Pehli baar jab tum har word padhte ho, sticky notes par notes likhte ho (wahi "cache" hai). Jab tum word #50 likhte ho, tum sirf apne 49 sticky notes par ek nazar daalte ho instead of poori story re-read karne ke. Tum sirf word #50 ke liye ek naya sticky note add karte ho.

Tradeoff: Tumhe apni desk par saare sticky notes ke liye jagah chahiye (memory), lekin tum bahut sara time bachate ho poori story re-read na karke. Agar tumhari desk bahut bhar jaaye (bahut lambi story), toh shayad tum sirf last 20 sticky notes rakho aur puraane throw away karo (sliding window).

Connections

6.1.1-transformer-architecture - Self-attention mechanism jise KV-cache optimize karta hai
6.1.4-multi-head-attention - Kai heads cache requirements ko kaise multiply karte hain
6.2.3-inference-optimization - Deployment speedups ki broader category
6.3.1-model-quantization - Lower precision (int8, fp16) ke through cache memory reduce karna
5.4.2-beam-search - Generation algorithm jo multiple beams ke saath cache-intensive ho jaata hai
6.1.11-sparse-attention - $O (L^{2})$ attention cost reduce karne ka alternative approach
7.2.1-llm-deployment - Real-world systems jahan cache management critical hai

#flashcards/ai-ml

KV-cache kis primary problem ko solve karta hai? :: Autoregressive generation ke dauran, previously generated tokens ke key aur value projections ka redundant recomputation rokna. Caching ke bina, hum har step par saare prior tokens ke K aur V recompute karte hain, jisse $O (L^{2} d^{2})$ cost aati hai $O (L d^{2})$ ki jagah.

KV-cache memory formula per layer derive karo :: Cache size = $2 \times L \times h \times d_{k} \times bytes$ . 2 alag K aur V matrices ke liye hai, $L$ sequence length hai, $h$ attention heads ki number hai, $d_{k}$ head dimension hai, aur bytes precision hai (fp16 ke liye 2, fp32 ke liye 4). Full model ke liye $n_{layers}$ se multiply karo.

Hum Query (Q) projections cache kyun nahi karte?

Query represent karta hai "current token kya dhundh raha hai" aur har naye token ke liye recompute karna padta hai. Keys aur values represent karte hain "previous tokens kya offer karte hain," jo past tokens ke liye fixed hain, isliye unhe safely cache kiya ja sakta hai. Q cache karna naye token ko apna context-dependent query banane se rokta.

Multi-Query Attention (MQA) kya hai aur iska cache benefit kya hai?

MQA ek single shared K aur V projection saare query heads mein use karta hai, per-head K, V pairs ki jagah. Yeh cache size ko

h

(heads ki number) factor se reduce karta hai, kyunki cache

2 L d_{k}

ho jaata hai

2 h L d_{k}

ki jagah. Trade-off: massive memory savings ke liye ~1-2% quality drop.

GPT-2 Small ke liye 512 tokens generate karne par KV-cache size calculate karo :: GPT-2 Small: 12 layers, 12 heads, $d_{k} = 64$ , fp16. Cache = $12 \times 2 \times 512 \times 12 \times 64 \times 2$ bytes = 18.9 MB. Yeh per-sample hai; batch size ise multiply karta hai.

KV-cache ke saath bhi causal masking kyun zaruri hai?

Generation ke dauran cache mein sirf past tokens hone ke bawajood, masking kyun crucial hai: (1) batched generation jahan sequences ki alag-alag lengths hain, (2) training behavior ke saath consistency, (3) attention score indexing mein off-by-one errors rokna. Mask ensure karta hai ki position

i

sirf positions

\leq i

par attend kare.

Grouped-Query Attention (GQA) kya hai?

Standard multi-head attention aur MQA ke beech ka middle ground. Query heads ko

g

groups mein divide kiya jaata hai (jahan

g < h

), aur har group ek K, V head share karta hai. Partial cache reduction (

h / g

factor) deta hai full MQA ke comparison mein kam quality degradation ke saath. Example: Llama-2 70B

h = 64, g = 8

use karta hai.

KV-cache kab bottleneck ban jaata hai?

(1) Bahut lambi sequences (100K+ tokens) jahan

O (L)

memory phir bhi GPU capacity se zyada ho, (2) Production mein large batch sizes, (3) Kai layers aur heads wale models (GPT-3 scale), (4) Multi-beam search cache ko beam width se multiply karta hai. Solutions: sliding window, offloading, ya sparse attention.