Skip to main content

Chapter 1: Inference Optimization

KV Cache​

Store computed key-value pairs to avoid recomputation:

# Without KV cache: recompute all KV for each new token
# With KV cache: only compute KV for new token

class KVCache:
    def __init__(self, max_batch, max_seq, num_heads, head_dim):
        self.k_cache = torch.zeros(max_batch, num_heads, max_seq, head_dim)
        self.v_cache = torch.zeros(max_batch, num_heads, max_seq, head_dim)
    
    def update(self, k, v, position):
        self.k_cache[:, :, position] = k
        self.v_cache[:, :, position] = v

Speculative Decoding​

Use small model to draft, large model to verify:

# 1. Draft: Small model generates N tokens quickly
# 2. Verify: Large model checks all N tokens in parallel
# 3. Accept matching tokens, regenerate from first mismatch

Continuous Batching​

Process multiple requests simultaneously:

StrategyThroughputLatency
StaticLowConsistent
ContinuousHighVariable