Attention: Everything You Need

Introduction

Since Vaswani et al. dropped the Transformer paper in 2017, we've basically been optimizing the same few equations for nearly a decade. That elegant $\text{softmax}(QK^T/\sqrt{d_k})V$ has spawned everything from GPT-4 to Llama to Claude, but getting from that formula to a system that actually serves millions of requests at reasonable cost is a different beast entirely.

This post is my attempt to go through the past decade. We'll move quickly through the fundamentals (you probably know them), then dig into what actually matters when you're trying to train on 100K token contexts or serve at scale without burning through your infrastructure budget.

What is Attention?

Attention determines relevance. Given a query, which parts of the input matter? Scaled dot-product attention:

Three matrices:

• Query (Q): What we're looking for
• Key (K): What we're looking at
• Value (V): What we retrieve

Compute dot products between queries and keys for relevance scores, then weight the values accordingly.

The $\sqrt{d_k}$ Scaling Factor

The scaling isn't arbitrary. Vaswani et al.: "We suspect that for large values of $d_k$, the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients".

If $Q$ and $K$ contain i.i.d. elements with mean 0 and variance 1, their dot product has variance $d_k$. Large dimensions → large dot products → softmax saturation → vanishing gradients.

Dividing by $\sqrt{d_k}$ normalizes variance back to ~1. Gradients stay healthy, training stays stable.

Multi-Head Attention

Single attention is useful, but multi-head attention is where things get interesting. Compute attention $h$ times in parallel, each with different learned projections:

where each head is:

Each head learns different weight matrices ($W_i^Q$, $W_i^K$, $W_i^V$), focusing on different relationships. One head might capture subject-verb dependencies, another tracks long-range references, another focuses on syntactic patterns.

Multi-Head Design Benefits

Three advantages for production systems:

1. Embarrassingly Parallel: Each head computes independently, which is perfect for GPU acceleration and essential for production throughput.
2. Multiple Representational Subspaces: The original paper notes multi-head attention "allows the model to jointly attend to information from different representation subspaces at different positions". Different heads specialize in different phenomena simultaneously.
3. Robustness Through Redundancy: No single attention pattern required. Research shows certain heads can be pruned post-training without significant degradation, suggesting built-in redundancy.

What Attention Heads Learn

Research shows attention heads develop distinct patterns:

1. Positional heads: Attend to adjacent tokens (position $n \pm 1$) for local syntax
2. Syntactic heads: Track grammatical relationships like subject-verb agreement
3. Semantic heads: Capture long-range dependencies and coreference
4. Rare-word heads: Attend to the least frequent tokens in the input

Voita et al. found the positional, syntactic, and rare-word heads in translation models. Clark et al. found BERT heads tracking syntax and coreference, and many heads parking most of their attention on [SEP] or punctuation, effectively a no-op when the head has nothing useful to do.

Tools like BertViz, and methods like attention rollout and attention flow, let you inspect these patterns.

Studies show 10-20% of heads can be pruned with <1% quality loss; useful for model compression. Some heads are redundant, others critical.

Attention entropy as diagnostic: Low entropy (peaked) = focused on specific tokens. High entropy (uniform) = either not learning useful patterns or attending broadly.

The Permutation Equivariance Property

Here's a subtle but important property: self-attention without a causal mask is permutation-equivariant. If you swap input positions, outputs swap accordingly. Mathematically, for any permutation matrix $P$:

This means the attention mechanism itself has no inherent bias toward position; it treats all positions equally. This is why positional encodings are necessary for sequence order. (A causal mask breaks the symmetry and leaks some positional information, which is why decoder-only models trained with no positional encoding at all still work, but explicit encodings remain the norm.) The mechanism is extremely flexible, but requires careful position encoding design (which led to innovations like RoPE and other alternatives).

Implementation Insights

Attention dropout (typically $p=0.1$) is applied to attention weights after softmax. This prevents overfitting to specific attention patterns and encourages robustness. During training, it randomly zeros out some attention connections, forcing the model to learn multiple paths to each value. Most large LLM pretraining runs (Llama, for example) set attention dropout to 0, since a single pass over trillions of tokens rarely overfits.

In practice, the per-head projections are fused: one $d_{model} \times d_{model}$ matrix each for $Q$, $K$, and $V$, whose outputs are reshaped into $h$ heads of size $d_k = d_{model}/h$. Multi-head attention therefore costs about the same as single-head attention at full width. Masks are applied to the scores before softmax by setting disallowed positions to $-\infty$.

Cross-Attention vs Self-Attention

While self-attention relates positions within a single sequence, cross-attention relates positions between two different sequences. This is fundamental for many architectures:

• In self-attention: $Q$, $K$, $V$ all come from the same sequence
• In cross-attention: $Q$ comes from one sequence (target), $K$ and $V$ come from another (source)

Applications:

• Encoder-decoder models: Decoder attends to encoder outputs (original Transformer, T5, BART)
• Vision-language models: Text attends to image features (Flamingo, BLIP-2's Q-Former). CLIP is a dual encoder and has no cross-attention
• Speech: A text decoder attends to audio encoder features (Whisper); speech-to-speech models attend to source speech features
• Retrieval-augmented generation: The generator attends to encoded retrieved chunks (RETRO, Fusion-in-Decoder). Most production RAG simply concatenates documents into the prompt and uses self-attention

The Transformer architecture uses both: self-attention in encoder and decoder, plus cross-attention connecting them.

Causal Attention: Preventing Information Leakage

Autoregressive models like GPT and Llama must generate tokens sequentially; each token can only depend on previous tokens, not future ones. Causal masking (also called autoregressive or look-ahead masking) enforces this constraint:

This ensures token $i$ can only attend to positions $\leq i$. The mask is applied before softmax, setting future positions to $-\infty$ so they have zero weight after softmax.

In production:

• Training: Can parallelize across sequence (teacher forcing with causal mask)
• Inference: The prompt is processed in parallel (prefill), but output tokens are generated one at a time (decode)
• KV caching: Essential for efficient inference (cache previous tokens' K, V to avoid recomputation)

Bidirectional vs Causal:

• BERT uses bidirectional attention (no mask) for understanding tasks
• GPT uses causal attention for generation tasks
• Encoder-decoder models use both: bidirectional in encoder, causal in decoder

Memory & Compute Optimizations

Modern transformer training faces significant memory bottlenecks. The following techniques address these challenges to enable efficient training at scale.

Mixed Precision Training: Foundation for Efficiency

Before diving into attention-specific optimizations, mixed precision training (FP16/BF16) is the foundational technique that enables all other optimizations. Instead of using 32-bit floating point (FP32) for all computations, mixed precision uses 16-bit formats for most operations while maintaining FP32 for critical accumulations.

Two main formats:

• FP16 (Float16): Standard half precision, halves activation memory, but prone to underflow/overflow
• BF16 (BFloat16): Same range as FP32 but lower precision, preferred for stability

BF16 advantages:

• Same dynamic range as FP32 (8-bit exponent) reduces underflow/overflow compared to FP16
• No loss scaling required (unlike FP16), which is the main practical benefit
• Native hardware support on A100, H100, and newer GPUs
• Pairs perfectly with FlashAttention and other optimizations
• Roughly halves activation memory (optimizers still keep FP32 master weights, so total savings are smaller)
• 2-3x training speedup from faster compute on Tensor Cores
• Negligible quality difference versus FP32 when done correctly

Mixed precision is now standard in all production training pipelines; it's the first optimization to enable, not an optional extra.

The Next Frontier: FP8 Training

Modern hardware (NVIDIA H100/B200) natively supports FP8 (8-bit floating point), doubling throughput over BF16. FP8 uses two formats: E4M3 (higher precision) for forward pass, E5M2 (wider range) for gradients.

Libraries like NVIDIA's Transformer Engine handle FP8 automatically with delayed scaling. For inference, FP8 is gaining ground for Llama-4 class models: 2x faster weight loading than BF16 without INT4's quantization noise. On older GPUs without FP8 support, INT8 and INT4 weight quantization remain the common choices.

Trade-off: Requires H100+ hardware. BF16 remains the safer default, but FP8 is the future for cutting-edge systems.

FlashAttention: The I/O Revolution

Understanding attention's computational complexity is one thing; understanding its memory access patterns is another. The real bottleneck in modern hardware isn't compute; it's memory bandwidth.

Important distinction: FlashAttention reduces memory usage from $O(N^2)$ to $O(N)$ by not materializing the full attention matrix, while maintaining $O(N^2)$ computational complexity (FLOPs). The 2-4x speedup comes from optimizing memory I/O between HBM and SRAM, not from reducing arithmetic operations. This is a common misconception worth clarifying.

FlashAttention reformulates attention to minimize memory reads and writes between GPU high-bandwidth memory (HBM) and on-chip SRAM. Key insight: instead of materializing the full $N \times N$ attention matrix in HBM, compute attention in blocks that fit in SRAM, using tiling and recomputation strategies.

The Memory Hierarchy Problem

Modern GPUs have a memory hierarchy (A100 numbers from the FlashAttention paper):

• HBM (High-Bandwidth Memory): Large (40GB+) but slow (~1.5 TB/s)
• SRAM (On-chip memory): Small (~20MB) but fast (~19 TB/s)

Standard attention implementations:

1. Load $Q$, $K$ from HBM → compute $QK^T$ → write to HBM
2. Load $QK^T$ from HBM → compute softmax → write to HBM
3. Load softmax output and $V$ from HBM → compute final output

Each step involves expensive HBM reads/writes. For a sequence of length 4096, the score matrix has 4096 × 4096 ≈ 16.8M entries, or 32MB in FP16, for a single head of a single sequence. Multiply by heads and batch size and it dwarfs SRAM.

FlashAttention's Solution

FlashAttention uses tiling and kernel fusion:

1. Tiling: Split $Q$, $K$, $V$ into blocks that fit in SRAM
2. Online Softmax: Compute softmax incrementally without materializing the full attention matrix
3. Recomputation: In the backward pass, recompute attention on-the-fly rather than storing it

Result: 2-4x speedups and the ability to handle much longer sequences. This isn't just an optimization; it's what makes training on sequences of 16K+ tokens practical.

By recomputing attention during backpropagation rather than storing it, FlashAttention trades minimal extra compute for massive memory savings, a favorable tradeoff since modern GPUs are compute-rich but memory-bound.

FlashAttention-2 and FlashAttention-3

FlashAttention-2 made three changes. It cut non-matmul FLOPs by deferring the softmax rescaling to the end of the loop, since GPUs run matmuls far faster than anything else. It made the outer loop run over query blocks so it could parallelize across the sequence dimension, not just batch and heads, which keeps the GPU busy for long sequences with small batches. And within each block it split work across warps by query rows instead of by keys, removing most of the shared-memory communication between warps. The result was ~2x over FA-1, reaching 50-73% of A100 peak FLOPs.

FlashAttention-3 targets H100's new hardware: the Tensor Memory Accelerator (TMA) for asynchronous memory transfers and WGMMA instructions for warpgroup-wide matmuls. It uses warp specialization, where producer warps load data while consumer warps compute, and it overlaps softmax with the matmuls both across warpgroups (ping-pong scheduling) and within a warpgroup. FA-3 is 1.5-2x faster than FA-2 in FP16, reaching ~75% of H100 peak (about 740 TFLOPs/s), and gets close to 1.2 PFLOPs/s in FP8, using block quantization and incoherent processing to keep FP8 error in check. The evolution from FA-1 → FA-2 → FA-3 demonstrates how algorithms must co-evolve with hardware for optimal performance.

Gradient Checkpointing: Trading Compute for Memory

Core Idea: During backpropagation, recompute intermediate activations on-demand instead of storing them all in memory.

Standard backpropagation:

• Forward: Compute and store ALL intermediate activations (attention scores, FFN activations, dropout masks, layer outputs)
• Backward: Use stored activations to compute gradients
• Memory: $O(L \cdot S \cdot H)$ where $L$ = num_layers, $S$ = seq_len, $H$ = hidden_dim

With checkpointing:

• Forward: Only save inputs to checkpointed segments, discard internal activations
• Backward: Recompute activations from saved checkpoints when needed
• Memory: $O(\sqrt{L} \cdot S \cdot H)$ with a checkpoint every $\sqrt{L}$ layers. Checkpointing every layer still stores one input per layer, $O(L \cdot S \cdot H)$, but with a much smaller constant since none of the layer's internal activations are kept

Tradeoffs:

• Memory reduction: 60-80% for large models
• Speed cost: ~20-33% slower (each checkpointed forward pass runs again during backward, and the forward pass is roughly a third of training compute)
• Gradient quality: Mathematically identical to standard backprop

When to use: Training large models that don't fit in GPU memory, maximizing batch size or sequence length.

When NOT to use: Inference (no backprop), when memory isn't the bottleneck.

Combined techniques: FlashAttention + checkpointing enables 10x+ longer sequences. Works well with mixed precision and ZeRO optimizer for massive models.

Inference Optimizations

In production, inference latency and throughput directly drive user experience and cost. LLMs face unique challenges during inference, particularly around memory efficiency and latency. The following techniques address these bottlenecks to enable fast, cost-effective serving at scale.

Grouped Query Attention (GQA): The KV Cache Solution

During autoregressive generation, models cache key-value pairs from previous tokens to avoid recomputation. For a model with 32 attention heads and 128-dimensional head size, each token requires storing $32 \times 128 \times 2 = 8{,}192$ values per layer (keys and values). At 100K context length, this becomes approximately 819M values per layer, a massive memory bottleneck.

Multi-Query Attention (MQA) proposed a radical solution: share a single key-value head across all query heads. This reduces KV cache by 32x but comes with potential quality degradation (~1-2% on benchmarks), which can be largely mitigated through uptraining (continuing training with MQA architecture).

Grouped Query Attention (GQA) found the sweet spot. Instead of 32 separate KV heads (MHA) or 1 shared KV head (MQA), use, for example, 8 KV heads with 4 query heads per group.

The results are remarkable: a 4x reduction in KV cache size in this configuration, with quality close to MHA and speed close to MQA. Llama 2 70B, all Llama 3 models, Mistral 7B, and Gemma 2 use GQA. Even better, existing MHA models can be "uptrained" to GQA using only 5% of original training compute.

Multi-Head Latent Attention (MLA): Extreme Compression

While GQA reduces memory by sharing KV heads, Multi-Head Latent Attention (MLA), introduced in DeepSeek-V2 and refined in DeepSeek-V3, takes a more aggressive approach: low-rank key-value joint compression.

Instead of storing full Key and Value matrices (even shared ones), MLA projects them into a low-dimensional latent vector $c_{KV}$ during inference.

The mechanism:

1. Compression: Inputs are projected down to a small latent dimension $d_c$ (e.g., 512) using a down-projection matrix $W_{DKV}$
2. Storage: Only this tiny compressed latent vector $c_{KV}$ is stored in the KV cache, drastically reducing memory footprint
3. Decompression: During attention computation, the latent vector is up-projected using $W_{UK}$ and $W_{UV}$ to generate the full heads on-the-fly. At inference, $W_{UK}$ can be absorbed into the query projection and $W_{UV}$ into the output projection, so the full keys and values are never materialized
4. Decoupled RoPE: RoPE's position-dependent rotation would block absorbing $W_{UK}$ into the query projection, so MLA carries position in a small separate key vector (64 dimensions in DeepSeek-V2) that bypasses the compression and is cached alongside $c_{KV}$

Impact: MLA enables extreme KV cache compression. Models with massive head counts (e.g., 128 heads) achieve KV cache sizes smaller than even GQA models, often compressing memory usage by 90%+ compared to standard MHA.

Real-world deployment: This breakthrough is why DeepSeek-V3 can efficiently serve a 671B parameter model at production scale. By compressing the KV cache to a fraction of what GQA would require, MLA makes serving massive MoE models economically viable, a critical enabler for the next generation of ultra-large language models.

Adoption note: MLA started as a DeepSeek innovation (V2 and V3). Kimi K2 reuses the DeepSeek-V3 architecture including MLA, but most other model families still use GQA as the standard KV cache optimization.

Beyond GQA: KV Cache Compression

While GQA reduces cache size through architectural changes, additional compression techniques push efficiency further:

1. Quantization: Store KV cache in 8-bit or 4-bit precision instead of 16-bit (2-4x savings with minimal quality loss)
2. Eviction policies: Drop least-important cached tokens based on attention scores
3. Streaming LLMs: Keep only recent tokens plus "attention sinks" (initial tokens that accumulate disproportionate attention)
4. H2O (Heavy Hitters Oracle): Cache only high-attention tokens, evict low-attention ones

Example impact: For a Llama-2-70B-shaped model (80 layers, 64 query heads of dimension 128), a single 4K-token sequence needs ~10.7GB of KV cache with full MHA in FP16. GQA with 8 KV heads cuts that to ~1.3GB, and INT8 quantization halves it again to ~0.67GB.

Attention sinks: Initial tokens accumulate disproportionate attention weight even when semantically irrelevant. Keeping these tokens in cache is critical for long-context generation quality; evicting them causes catastrophic degradation.

Rotary Position Embedding (RoPE): Encoding Position Elegantly

Remember the permutation equivariance property? We need positional information, but how we encode it matters enormously. The original Transformer added fixed sinusoidal position encodings to the input (learned absolute embeddings, later used by BERT and GPT-2, performed about the same). RoPE took a different approach: encode position by rotating query and key vectors in high-dimensional space.

Key insight: if we rotate $Q_m$ by angle $m\theta$ and $K_n$ by angle $n\theta$, their dot product naturally depends on the relative distance $(m-n)\theta$:

Here each pair of dimensions is treated as a complex number, and the bar is complex conjugation.

In practice, this complex rotation is implemented using sine and cosine operations on real-valued vectors. There are two common implementation variants:

Interleaved (Llama-style): Rotates adjacent pairs of features $(q_{2i}, q_{2i+1})$:

Rotary Half (GPT-NeoX-style): Splits the vector into two halves and rotates them against each other:

where $\theta_i = 10000^{-2i/d}$ for the $i$-th dimension pair.

Critical: Both variants are mathematically valid and produce equivalent relative position encoding, but they are incompatible when loading pre-trained weights. Meta's original Llama checkpoints expect interleaved, while GPT-NeoX expects rotary half. Hugging Face's Llama port uses rotary half and permutes the Q/K projection weights at conversion time to compensate, which is exactly the detail that bites when you mix checkpoints and code from different sources. Using the wrong variant produces garbage outputs.

Note: In practice, libraries like Hugging Face Transformers handle RoPE automatically for each model architecture, so you rarely implement it yourself.

RoPE advantages:

• Zero parameters: No learned embeddings to store
• Inherently relative: Attention depends on relative distance, not absolute position
• Extrapolation: Can handle sequences longer than training length with proper scaling
• Universal adoption: Llama, GPT-NeoX, PaLM, Mistral all use RoPE

For extending to ultra-long contexts (100K+ tokens), techniques either rescale positions (linear position interpolation) or rescale the rotation frequencies (NTK-aware scaling, YaRN) so that longer sequences map into the range seen during training.

ALiBi: Simpler Position Encoding

An alternative to RoPE is ALiBi (Attention with Linear Biases), which adds position-dependent bias directly to attention scores:

where $m$ is a fixed, head-specific slope. For $H$ heads, the slopes form a geometric sequence $m_h = 2^{-8h/H}$ for $h = 1, \ldots, H$.

Even simpler than RoPE:

• No position embeddings to compute or apply
• Just subtract distance penalty from attention scores
• Used in: BLOOM, MPT models

Trade-off: ALiBi's main selling point was length extrapolation: trained on short sequences, it degrades gracefully on longer ones. RoPE with scaling has since won out in most new models, but ALiBi offers a simpler implementation and competitive performance. The choice often depends on engineering constraints and specific use cases.

PagedAttention: Virtual Memory for KV Cache

Even with GQA reducing cache size, managing memory efficiently during inference remains challenging. Consider serving multiple requests with varying sequence lengths: traditional approaches pre-allocate maximum memory per request, leaving 60-80% of KV cache memory wasted on over-reservation and fragmentation.

PagedAttention, implemented in the vLLM serving framework, applies operating system virtual memory concepts to KV cache management. Instead of contiguous memory allocation, it stores KV cache in fixed-size blocks (e.g., 16 tokens) with a logical-to-physical mapping, much like OS page tables.

Advantages:

• Dynamic allocation: Allocate memory as sequences grow
• Memory sharing: Multiple sequences can share prefix blocks (critical for parallel sampling and beam search), with copy-on-write when they diverge
• No external fragmentation: Any free block can serve any sequence, and waste is limited to each sequence's partially filled last block
• Continuous batching: Efficiently batch requests of different lengths

Result: 2-4x higher throughput than FasterTransformer and Orca at the same latency. Earlier systems used only 20-40% of KV cache memory for actual token states; vLLM wastes under 4%. vLLM has become the de facto standard for LLM serving.

Continuous Batching: Maximizing GPU Utilization

Traditional batching waits for all requests in a batch to complete before starting new ones:

Batch 1: [Request A (1000 tokens), Request B (100 tokens), Request C (500 tokens)]
→ Wait for A (1000 tokens) to finish before starting Batch 2
→ GPU idle waiting for longest request

Continuous batching (Orca, implemented in vLLM) adds new requests as soon as capacity is available.

Advantages:

• Much higher GPU utilization than static batching, since slots freed by finished requests are refilled immediately
• Lower latency for short requests (don't wait for long requests)
• Better throughput overall

Enabler: Orca's iteration-level scheduling predates PagedAttention, but paged KV memory is what lets a server pack many requests of different lengths into a batch without reserving worst-case memory for each.

Speculative Decoding: Accelerating Inference 2-3x

Autoregressive generation is inherently sequential; each token requires a full model forward pass. Speculative Decoding achieves 2-3x speedup without changing outputs or retraining by using a small "draft" model to predict multiple tokens, then verifying them in parallel with the target model.

Core insight: Language modeling tasks often contain easier subtasks. A small, fast model can draft plausible continuations, and the large model can verify multiple candidates in a single forward pass.

Algorithm:

1. Draft phase: Small model generates $k$ candidate tokens (fast, lower quality)
2. Verification phase: Large model processes all $k$ candidates in parallel (single forward pass)
3. Acceptance: Accept each draft token with probability $\min(1, p(x)/q(x))$, where $p$ and $q$ are the target and draft probabilities
4. Correction: At the first rejection, sample a replacement from the normalized residual $\max(0, p - q)$. If all $k$ are accepted, sample one bonus token from the target. This keeps the output distribution exactly equal to the target model's

How this works:

• Draft model is 10-100x faster (smaller, quantized, or distilled)
• With a 50% per-token acceptance rate and $k=4$, each target pass yields ~1.9 tokens on average, before subtracting the draft model's cost
• Mathematically guaranteed to match target model's distribution exactly
• No retraining or architecture changes required

Production adoption: Google has reported using it in production, and vLLM, TensorRT-LLM, and SGLang all support it. Variants like Medusa (extra decoding heads instead of a separate draft model) and EAGLE push acceptance rates higher. Works best when:

• Draft model is 10x+ faster than target
• Acceptance rate > 60% (predictable text, code, structured output)
• Batch size is small (latency-sensitive applications)

Trade-offs: Less effective for creative/unpredictable generation where draft model struggles to match target distribution.

Prompt Caching: Eliminating Redundant Computation

Many LLM requests share common prefixes: system prompts, few-shot examples, or document context. Prompt caching (prefix caching, KV cache reuse) eliminates redundant computation by caching and reusing KV pairs from identical prefixes.

How it works:

1. Compute KV cache for prompt prefix during first request
2. Store cached KV pairs with hash of prefix tokens
3. For subsequent requests with same prefix, load cached KV and only compute new tokens
4. Entries expire after a period of inactivity (5 minutes by default on Anthropic, with an optional 1-hour TTL) or are evicted LRU-style in self-hosted servers

Example scenario:

Request 1: [System Prompt (1000 tokens)] + [User Query A (50 tokens)]
→ Compute full 1050 tokens, cache first 1000

Request 2: [System Prompt (1000 tokens)] + [User Query B (50 tokens)]
→ Load cached 1000 tokens, compute only 50 new tokens
→ ~20x less prefill work for this request; cached tokens are billed at a discount

Advantages:

• Latency: Anthropic reports up to 85% lower time-to-first-token on long cached prompts
• Cost: Anthropic bills cache reads at 10% of the base input price (cache writes cost 25% more than base); OpenAI also discounts cached input tokens, by an amount that varies by model
• Longer context: Enables longer system prompts without latency penalty
• RAG optimization: Cache document context across multiple queries

Implementation details:

• Cache granularity: Minimum cacheable prefix of around 1024 tokens on both Anthropic (higher for some models) and OpenAI, with OpenAI matching in 128-token increments
• Cache points: Up to 4 breakpoints per prompt
• Scope: Caching applies to entire prefix up to cache_control block
• Automatic: vLLM, TensorRT-LLM detect and cache common prefixes

Design pattern: Put the stable content (system prompt, tool definitions, few-shot examples, shared documents) first and mark the end of it as a cache breakpoint; put per-request content after it. Anything variable placed before the stable content changes the prefix and invalidates the cache for everything after it.

Note: Anthropic requires explicit cache_control breakpoints; OpenAI caches eligible prefixes automatically, and vLLM's automatic prefix caching matches shared prefix blocks without any markers.

Production considerations:

• Structure prompts with stable prefixes first, variable content last
• Cache hit rate depends on exact token sequence matching
• Particularly effective for: RAG systems, multi-turn conversations with system prompts, code generation with examples
• Provider support: Anthropic Claude, OpenAI GPT-4, Google Gemini, AWS Bedrock

Real-world impact: Anthropic reports up to 90% cost reduction and 85% latency reduction for long-context applications with stable prefixes.

Long-Context Solutions

As context windows extend beyond 10K-20K tokens, full $O(N^2)$ attention becomes prohibitive. These techniques enable efficient processing of ultra-long sequences.

Sparse and Sliding Window Attention

Full $O(N^2)$ attention becomes prohibitive beyond 10K-20K tokens. Sparse attention patterns offer linear or near-linear complexity by restricting which positions can attend to each other.

Sliding Window Attention restricts each token to attend only to the previous $W$ tokens.

Mistral 7B (v0.1) uses a 4096-token sliding window. Key insight: stacking layers provides global context through composition, since layer $L$ can see information from up to $L \times W$ tokens away. With a 4K window and 32 layers, the theoretical receptive field is ~131K tokens, although information degrades over that many hops, and later Mistral releases dropped the sliding window in favor of full attention.

Longformer combines sliding window with global attention tokens: certain positions (e.g., $\text{[CLS]}$) attend to all positions and are attended to by all positions. This hybrid approach balances efficiency with global reasoning.

Ring Attention: Scaling Beyond Single-Device Limits

While FlashAttention optimizes memory on a single device, Ring Attention tackles a different problem: what if your sequence is too long to fit on any single device? Even with FlashAttention, processing 100M tokens requires over 1TB of memory, far exceeding typical GPU capacity.

Core mechanism: Ring Attention distributes the sequence across multiple devices arranged in a ring. Each device holds one query block and iterates through all key-value blocks by passing them around the ring. The key insight is overlapping communication with computation: while device $i$ computes attention, it simultaneously sends KV blocks to device $(i+1)$ and receives from device $(i-1)$.

This works because blockwise attention computation is permutation-invariant: blocks can be processed in any order as long as statistics are combined correctly. When block size is large enough (≥ FLOPS/Bandwidth), communication is fully hidden by computation, resulting in near-zero overhead.

Memory per device: $O(b \cdot c \cdot h)$ per block, in the paper's notation:

• $b$ = batch size
• $c$ = block size (tokens per device, e.g., 1024)
• $h$ = hidden dimension

Each device stores:

1. Local query block ($b \cdot c \cdot h$)
2. Current KV block being processed ($2 \cdot b \cdot c \cdot h$)
3. Next KV block being received ($2 \cdot b \cdot c \cdot h$, overlapped with computation)
4. Output accumulator ($b \cdot c \cdot h$)

Total: ~$6 \cdot b \cdot c \cdot h$ bytes, independent of total sequence length $N$. Whether processing 1M or 100M tokens, each device uses the same memory.

Scaling: Context length scales linearly with device count. With 8 GPUs, train 8× longer sequences; with 512 TPUs, train 512× longer. Experiments show training sequences exceeding 100M tokens, enabling whole-codebase understanding, long document processing, and high-resolution multimodal models.

Comparison with alternatives:

• vs FlashAttention: FlashAttention optimizes single-device memory; Ring Attention distributes across devices. They're complementary: Ring Attention uses FlashAttention for local computation.
• vs Sparse attention: Ring Attention computes exact attention without approximations, while sparse methods sacrifice some interactions for efficiency.
• vs Traditional sequence parallelism: Ring Attention only communicates with neighbors (not all-to-all), overlaps communication with computation, and can hide communication entirely when blocks are large enough.

Ring Attention demonstrates that when single-device memory is the bottleneck, clever distributed algorithms can unlock orders of magnitude improvements without sacrificing accuracy.

The Modern LLM Stack

Production systems combine these innovations:

The Complete Stack:

• Training: FlashAttention for memory efficiency, RoPE for position encoding
• Architecture: GQA for reduced KV cache, optional sliding window for long context
• Serving: PagedAttention (vLLM) for memory management, continuous batching for throughput
• Inference acceleration: Speculative decoding for 2-3x speedup, prompt caching for up to 90% cost reduction

Open-weight models like Llama 3 and Mistral document most of this stack publicly; closed labs don't publish theirs, but the same bottlenecks apply. Each innovation solves a specific bottleneck:

*Note: FlashAttention reduces memory usage from $O(N^2)$ to $O(N)$ but maintains $O(N^2)$ computational complexity (FLOPs). The speedup comes from optimizing memory I/O (HBM ↔ SRAM), not reducing arithmetic operations.

Techniques Simplified: What to use?

Training from scratch: Mixed Precision (BF16, or FP8 on H100+), FlashAttention, Gradient Checkpointing (if memory-constrained), GQA in architecture, RoPE for position encoding

Fine-tuning: Mixed Precision (BF16, or FP8 on H100+), FlashAttention, Gradient Checkpointing, LoRA/QLoRA for parameter-efficient fine-tuning

Serving (latency-critical, chatbots): FP8 inference (H100+) or BF16, GQA or MQA, PagedAttention (vLLM), Continuous Batching, Speculative Decoding (with draft model), Prompt Caching (with common prefixes), KV cache quantization (INT8)

Serving (throughput-critical, batch jobs): FP8 inference (H100+) or BF16, PagedAttention, Continuous Batching, Larger static batches, Prompt Caching. Skip Speculative Decoding, which is less effective for large batches.

Long context (100K+ tokens): FlashAttention, Sliding Window Attention, RoPE scaling (NTK-aware, YaRN), Sparse patterns

Ultra-long context (1M+ tokens, research): Ring Attention, Sparse attention patterns, Efficient position encoding. Consider alternatives like SSMs or linear attention.

Hardware-aware algorithm design matters as much as the math. Understanding memory hierarchy, parallelism opportunities, and hardware capabilities is essential.

Attention vs Alternative Sequence Models

While attention + modern optimizations dominates production systems, it's worth understanding alternatives to appreciate why attention won, and when you might consider something else.

Attention isn't the only way to model sequences. Understanding the trade-offs helps explain why attention dominates:

Where $n$ = sequence length, $d$ = hidden dimension, $k$ = kernel size

Attention won because: Parallelization + long-range modeling > complexity cost. Modern hardware (GPUs/TPUs) is optimized for parallel matrix operations, making attention's $O(n^2)$ complexity acceptable for practical sequence lengths.

Emerging alternatives:

• Linear Attention (Linformer, Performer): Approximate attention with $O(n)$ complexity using low-rank projections or random features
• State Space Models (Mamba, S4): Recurrent models with efficient parallelization during training
• RWKV: Hybrid approach combining RNN efficiency with Transformer-like parallelization

These alternatives show promise for ultra-long contexts (1M+ tokens) where $O(n^2)$ becomes prohibitive, but standard attention remains dominant for most applications.

When to Consider Alternatives

Attention alternatives make sense when:

1. Ultra-long context (1M+ tokens): SSMs or linear attention may be viable when $O(n^2)$ becomes truly prohibitive, though Ring Attention + sparse patterns often work better
2. Streaming inference: RNNs/SSMs have constant memory regardless of history length, useful for infinite-length streaming applications
3. Edge deployment: Simpler models (RNNs, small CNNs) may fit resource constraints better than attention-based models
4. Specific inductive biases: CNNs for local patterns (images, audio), RNNs for true sequential processing with hidden state

For most production LLMs, standard attention + modern optimizations (FlashAttention, GQA, PagedAttention) remains optimal. The ecosystem, tooling, and empirical results strongly favor attention-based architectures for the 1K-100K token range that covers most real-world applications.

Key Takeaways

1. Scaled dot-product attention uses $\sqrt{d_k}$ scaling to maintain gradient health as dimensionality increases, a simple fix for variance growth.
2. Multi-head attention provides parallel representational subspaces, allowing the model to capture diverse relationships simultaneously. Research shows heads develop distinct patterns (positional, syntactic, semantic), and 10-20% can be pruned with minimal loss.
3. Attention dropout regularizes against overfitting to specific attention patterns by randomly zeroing connections during training. It matters most for smaller models and datasets; large LLM pretraining runs usually disable it.
4. Cross-attention vs self-attention: Self-attention relates positions within a sequence; cross-attention relates positions between sequences (critical for translation, multimodal models, retrieval-augmented generation).
5. Causal masking prevents information leakage in autoregressive models by masking future positions. This enables parallel training but requires sequential generation at inference.
6. Permutation equivariance means attention has no inherent position bias, necessitating explicit positional encodings (RoPE, ALiBi) but providing maximum flexibility.
7. Mixed precision training (BF16) is the foundational optimization that enables all others: about half the activation memory, 2-3x speedup, with minimal quality loss. Always enable this first.
8. FlashAttention solved the "how" of efficient computation through IO-aware tiling, achieving 2-4x speedups and enabling long-context training by minimizing memory bandwidth bottlenecks.
9. GQA solved the "what to cache" problem by reducing KV cache size (4x in a typical configuration) with minimal quality loss, critical for inference at scale. Further compression via quantization and eviction policies pushes efficiency even further.
10. RoPE solved the "where" of position encoding with zero parameters and inherent relative positioning, becoming the de facto standard. ALiBi offers a simpler alternative with competitive performance.
11. PagedAttention solved the "where to store" problem by applying virtual memory concepts to KV cache management, achieving 2-4x serving throughput and making continuous batching far more memory-efficient.
12. Sparse patterns (sliding window, block-sparse) solve the "when to attend" problem for ultra-long contexts, enabling 100K+ token processing with linear complexity through layer composition.
13. Ring Attention solves the "how to distribute" problem by enabling context length to scale linearly with device count. By organizing devices in a ring and overlapping communication with computation, it achieves near-zero overhead and supports sequences up to device-count times longer than single-device memory-efficient Transformers, enabling 100M+ token contexts.
14. Speculative Decoding achieves 2-3x lower latency in small-batch serving by using a fast draft model to generate candidate tokens, then verifying them in parallel with the target model. Maintains exact output distribution without retraining.
15. Prompt Caching eliminates redundant computation by caching and reusing KV pairs from identical prefixes, with Anthropic reporting up to 90% cost reduction and 85% latency reduction for requests with stable system prompts or document context.

References

1. Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS 2017.
2. Alammar, J. (2018). The Illustrated Transformer.
3. UvA Deep Learning Notebooks. Tutorial 6: Transformers and Multi-Head Attention.
4. Raschka, S. (2024). Understanding and Coding Self-Attention, Multi-Head Attention.
5. Dao, T., et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022.
6. Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. arXiv preprint.
7. Ainslie, J., et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. arXiv preprint.
8. Su, J., et al. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv preprint.
9. Peng, B., et al. (2023). YaRN: Efficient Context Window Extension of Large Language Models. arXiv preprint.
10. Kwon, W., et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP 2023.
11. Yu, G., et al. (2022). Orca: A Distributed Serving System for Transformer-Based Generative Models. OSDI 2022.
12. Jiang, A. Q., et al. (2023). Mistral 7B. arXiv preprint.
13. Beltagy, I., et al. (2020). Longformer: The Long-Document Transformer. arXiv preprint.
14. Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv preprint.
15. Shah, J., et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision.
16. Liu, H., Zaharia, M., & Abbeel, P. (2023). Ring Attention with Blockwise Transformers for Near-Infinite Context. ICLR 2024.
17. Leviathan, Y., Kalman, M., & Matias, Y. (2023). Fast Inference from Transformers via Speculative Decoding. ICML 2023.
18. DeepSeek-AI, et al. (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. arXiv preprint.
19. DeepSeek-AI, et al. (2024). DeepSeek-V3 Technical Report. arXiv preprint.
20. Micikevicius, P., et al. (2022). FP8 Formats for Deep Learning. arXiv preprint.
21. NVIDIA. (2023). Transformer Engine: FP8 Training for Transformers.
22. Voita, E., et al. (2019). Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned. ACL 2019.
23. Clark, K., et al. (2019). What Does BERT Look At? An Analysis of BERT's Attention. BlackboxNLP 2019.