Various notes on engineering with Transformers.

Index

Mixture of Experts

  • Helps scale model size without increasing FLOPs
  • Dipatch to expert works with
  • HF blog post on MoE
  • MoE is typically applied to MLP layers but can also be applied to attention heads

Faster Inference

Speculative decoding

  • Use a smaller model to predict next k-tokens and then run a batch of k trajectories with the larger model (in parallel)
  • Inference is memory bound, running 1 prediction takes the same time as k concurrent ones
  • Methods
    • Medusa: multiple light weight prediction heads that predict future tokens
    • Multi-token prediction: predict multiple tokens at each inference step
    • N-gram speculation: based off of current context
    • EAGLE: perdict next hidden state instead of next token and then uses the same head for prediction. Very strong!
    • DSpark: from DeepSeek V4; dynamic verification length, parallel backbone + sequential for dependencies when drafting + per-position confidence predictions.

Prefill decode disaggregation

  • Prefill: Compute bound; mostly computation of large matmuls to store in KV cache
  • Decode: Memry bound; retrieval and storage of KV cache + some compute
  • Longer multi-turn sessions = more resuse of KV cache
  • Mostly useful for heterogenous compute (high HBM v/s faster compute)
  • This only helps when communication (network, disk, memory etc) overhead is lesser than prefilling locally at the inference node
  • vLLM supports disaggregated prefilling with LMCache.

Paged attention

  • KV cache is normally store as a contiguous block of memory.
  • This is wasteful as not all sequences use up the entire context window.
  • This leads to lower effective batch size and memory fragmentation.
  • With paged attention, KV cache is stored as smaller pages and retrieved on the fly with a lookup table.

Prefix caching

  • Store KV cache of an exisiting query to reduce latency.
  • Helps with things such as system prompt, multi-turn conversation.
  • Essentially helps reduce prefill with stored caches across multiple requests.

KV cache offloading

  • KV cache size grows linearly with sequence length.
  • User interactions are mostly sparse, KV caches for all requests don’t need to stay in HBM all the time.
  • Unused KV cache can be paged out to host RAM or even persistent storage and then paged in when required.
  • vLLM supports offloading with LMCache.