HAVE AI NEWS HAVE AI NEWS
tg subscribe
Models

The Evolution of Attention in LLMs: How Engineers Scale Context to Millions of Tokens

The Evolution of Attention in LLMs: How Engineers Scale Context to Millions of Tokens

We explore why the era of pure Multi-Head Attention is over and how a combination of MLA, learned sparsity, history compression, and FlashAttention 4 enables models to efficiently process ultra-long contexts.

In the modern large language model industry, the search for a "single ideal attention formula" has definitively come to an end. Today, LLM performance on long contexts is determined not by mathematical abstraction, but by a multi-layered engineering stack: KV cache storage formats, history retrieval algorithms, and low-level compute optimizations tailored to specific hardware accelerators.

Three Main Bottlenecks of Classic Attention

The fully connected Self-Attention mechanism requires comparing each token with all preceding ones. When scaling context, architects face three fundamental constraints:

  • Memory Bandwidth Deficit: During the generation (decode) phase, the GPU spends more time reading accumulated keys and values (KV cache) from HBM memory than performing actual computations.
  • Quadratic Prefill Complexity: Calculating attention matrices during the initial prompt processing phase (prefill) scales as O(n²), making work with long texts computationally prohibitive.
  • On-Chip Data Transfer Overhead: Classic Softmax requires repeatedly transferring intermediate matrices between compute core registers and accelerator memory.

Disentangling Concepts: The 4 Layers of the Attention Mechanism

Terminology confusion often arises from conflating different architectural layers. The practical attention stack consists of four independent components:

LevelPurposeExample TechnologiesArea of Impact
Head ArchitectureMethod of generating and structuring Q, K, V vectorsMHA, MQA, GQA, MLAKV cache size, memory bus load during generation
Connection PatternDefining the subset of tokens accessible to the current stepFull, Sliding Window, Block Sparse, Learned Top-kComputational complexity of prefill and long context
GPU KernelsHardware implementation of mathematical operationsFlashAttention, cuDNN SDPA, FlashInferActual TFLOPS, latency (TTFT/TPOT), HBM utilization
Positional EncodingInjecting token order informationRoPE, ALiBi, NoPEExtrapolation capability beyond the training context window

The Evolution of KV Cache: From MHA to MLA

Classic Multi-Head Attention (MHA) allocates unique sets of Q, K, and V for each attention head. This guarantees maximum model expressiveness but leads to catastrophic memory consumption. For a 128K token sequence on a 32-layer model in BF16 format, the cache size reaches dozens of gigabytes per single request.

To solve this problem, the industry progressively adopted several approaches:

  • MQA (Multi-Query Attention): Unifies all query heads around a single K and V pair. This dramatically reduces memory requirements but often leads to quality degradation in complex reasoning tasks.
  • GQA (Grouped-Query Attention): Splits query heads into groups sharing KV projections (e.g., 4 queries per 1 KV). GQA has become the industry standard, retaining MHA's expressiveness while achieving memory consumption close to MQA.
  • MLA (Multi-Head Latent Attention): An approach popularized by the DeepSeek-V2 architecture. Instead of storing raw projections, the model compresses the KV state into a compact low-rank latent vector, decoupling the RoPE positional component from the content component. In DeepSeek tests, this reduced KV cache size by 93.3% and boosted generation throughput by 5.76x.

Sparsity and Compression: Handling Millions of Tokens

While reducing KV heads addresses memory load during decoding, optimizing attention patterns reduces the O(n²) computational complexity during prefill.

Learned Sparse Attention

Unlike a fixed sliding window, adaptive sparse attention uses a lightweight indexer to identify semantically important context fragments:

  • MiniMax Sparse Attention (MSA): The algorithm divides the incoming stream into 128-token blocks and selects the 16 most relevant blocks for each GQA group. The core attention calculation is performed across only 2048 tokens regardless of prompt length. In an experiment on a 109B MoE model with a 1M context, this yielded a 28.4x reduction in attention computations, speeding up prefill by 14.2x and decode by 7.6x on H800 chips.

Hybrid Compression and Filtering (DeepSeek-V4)

The next step was combining compression and sparsity. The architecture uses two complementary mechanisms:

  1. CSA (Compressed Sparse Attention): Pre-aggregates neighboring tokens into consolidated representations, reducing overall sequence length (e.g., compressing 1M tokens into 250K elements), after which an indexer selects the top-k required blocks.
  2. HCA (Hierarchical Compressed Attention): Applies aggressive compression (up to 128 tokens into a single vector), running dense attention over the global summary while preserving an exact local window of recent tokens.

This approach allows the DeepSeek-V4-Pro model on a 1M context to consume only 10% of the KV cache and 27% of FLOPs per token compared to DeepSeek-V3.2.

Linear Attention and Recurrent Hybrids

A complete departure from quadratic complexity is offered by models with linear attention and selective state spaces (Mamba, Kimi Delta Attention). They compress the entire history into a fixed-size matrix (O(1) memory complexity per generation step).

However, purely recurrent architectures suffer from the progressive loss of fine details ("needle in a haystack"). The most viable direction has become hybrid design: the majority of layers use linear or recurrent state updates, while periodic dense- or MLA-attention layers restore exact contextual retrieval capabilities.

Hardware Kernels: The Evolution of FlashAttention

GPU kernel-level optimization (FlashAttention) does not alter the mathematics of token connections. Its purpose is to minimize reads and writes to slow HBM memory via tiled processing in the chip's fast SRAM.

  • FlashAttention-2: Optimized thread scheduling and parallelism within Streaming Multiprocessors (SMs).
  • FlashAttention-3: Tailored for NVIDIA Hopper architecture with an emphasis on asynchronous compute and FP8 operations.
  • FlashAttention-4: Redesigned for NVIDIA Blackwell (B200/GB200). Utilizing Tensor Memory and enlarged tile sizes in BF16, it achieves performance of up to 1613 TFLOPS (~71% of theoretical peak), outperforming cuDNN 9.13 by 1.3x and Triton by 2.7x.

Summary Table of Attention Architectures

MechanismKey PrincipleKV Cache SizePrefill ComplexityIndustry Status
MHAIndividual Q, K, V per headMaximumO(n²)Legacy standard
GQAGrouping query heads around shared KVModerateO(n²)Industry baseline standard
MLALow-rank latent KV compressionMinimalO(n²) (with dense pattern)Production solution (DeepSeek)
Learned Sparse (MSA)Indexer selects relevant history blocksStandard (same as GQA)Reduced (down to O(nk))Active adoption
CSA + HCAMulti-level token compression + sparse retrievalExtremely lowHybrid (sparse + summary)Frontier stack
Linear / SSM HybridFixed hidden state + sparse attention layersO(1) per stepLinear O(n)Promising direction

Practical Stack Selection Guide

Use CaseRecommended StackRationale
Chatbots and assistants (context up to 128K)GQA + FlashAttention/SDPAMinimal engineering risk, widespread support in inference engines (vLLM, SGLang).
High load and long generationMLAEliminates memory bandwidth bottlenecks under high batching concurrency.
Codebase and document analysis (256K+)MSA / Learned SparseLimits the active attention compute budget without degrading retrieval accuracy.
Ultra-long context (1M+ tokens)CSA + HCA hybridsReduces both physical memory footprint and the volume of actively accessed tokens.

Modern LLM inference is always a balancing act between context retention quality, hardware constraints of the target chip, and overall stack throughput. The winner is not the one with the best theoretical asymptotic formula on paper, but the one whose stack operates cohesively from the algorithm down to the GPU register.

Author: AquaDll50 минут назад

Source: habr.com