Accelerating Language Model Generation Speed — Li Hong-Yi 2026 Course

In LLM applications, the inference phase's text continuation (decoding) process often becomes a performance bottleneck. To make models generate faster, we must optimize the core of Transformers — the Self-Attention mechanism. This tutorial introduces two key optimization techniques.

Part One: Flash Attention — Overcoming Hardware Data Transfer Bottlenecks¶

Flash Attention's core innovation isn't changing computation results but optimizing GPU-level data access logic.

1. Hardware Background: Warehouse vs. Workbench¶

• HBM (Warehouse): GPU's large storage space (e.g., A100's 80GB memory), but relatively slow access.
• SRAM (Workbench): Small, high-speed cache near compute units. Ultra-fast computation, but very limited capacity (only ~10s of MB).
• Bottleneck: Computation itself is fast, but frequently shuttling data between warehouse and workbench (like Attention Weights) slows everything down.

2. Flash Attention's Solution¶

• Tiling/Chunking: Cut long sequences into small chunks that fit in SRAM for computation.
• Online Softmax: Through math tricks (multiplying exponential correction factors), correct previous values while processing each chunk, no need to wait for the full sequence.
• Advantages:
• Unchanged Results: Not approximate; output is identical to original Attention.
• Plug-and-Play: No model retraining needed.
• Significant Speedup: On long sequences (e.g., 4096 tokens), achieves 8–9× speedup.

Part Two: KV Cache — Avoiding Redundant Computation¶

During text generation (decoding), each new token must consider all past information.

1. Core Concept¶

• Redundant Computation Problem: Without optimization, generating the next token requires recomputing Key (K) and Value (V) vectors for all old tokens—wasteful.
• Solution: Store computed K and V in cache; only compute new token's QKV and operate with cached data.

2. Space Challenge¶

• KV Cache grows rapidly with sequence length. Example: Gemma 2 27B model requires ~0.72MB KV per token.
• Oversized sequences can exhaust even massive HBM, causing Out of Memory (OOM).

Part Three: Advanced KV Cache Optimization¶

To reduce KV Cache memory footprint, researchers developed several variants:

1. Change K, V Group Counts (Architecture Optimization)¶

• MQA (Multi-Query Attention): All Queries share one K, V pair. Saves space but hurts model capability.
• GQA (Grouped-Query Attention): Multiple Query groups share one K, V. Balance between performance and space; widely used in Llama and Gemma.
• MLA (Multi-Head Latent Attention): Compress K, V into low-dimensional vectors. Magic: models can operate directly in compressed space without decompression, saving resources.

2. Change Attention Range or Purge Cache¶

• Sliding Window Attention: Focus only on fixed-length past information, set cache limits.
• Streaming LLM: Models heavily rely on first few tokens (Attention Sink); keep them in sliding window to handle ultra-long sequences without training.
• KV Cache Pruning: Discard rarely-attended K, V vectors (e.g., H2O algorithm). Experiments show retaining 20% of content still performs well on many tasks.

Conclusion: Developer Cost-Saving Tips¶

When designing AI Agents, using Prefix Caching dramatically cuts costs:

• Rule: Place stable, unchanging content (like System Prompt, tool descriptions) at the front; dynamic content (dates, variables) at the end.
• Benefit: Cached portions often get heavy discounts (e.g., 90% off), saving over 50% of API fees.