Long-context LLMs: training and serving techniques

Summary

Long-context LLMs combine position-encoding extension (so the model generalizes past its training length), I/O-aware attention kernels (so attention fits in memory), and KV-cache management (so serving stays affordable at long inputs).

Frontier models in 2026 advertise 128K–2M-token context windows. The headline number hides three independent engineering problems, each with its own state of the art. Knowing which technique addresses which problem is the senior-level test.

The three problems

1. Position encoding has to extrapolate

A learned position table has a finite index range; other mechanisms may be defined beyond training positions without working reliably there. For RoPE checkpoints, candidate extension methods include:

• Position interpolation (Chen et al., 2023): linearly compress positions.
• NTK-aware scaling: alter the rotary frequency grid, with implementation-specific changes to different bands.
• YaRN (Peng et al., 2023): per-frequency interpolation tuned by training-length statistics.

These mapping changes may be paired with continued training. Evaluate actual long-context behavior; a defined rotation formula is not evidence that the checkpoint uses distant information correctly.

2. Attention must fit in memory

Naively materialized dense attention is $O(n^2)$ in memory; at $n=128,000$, one $n\times n$ score matrix alone is about 61 GiB (65.5 GB) at FP32. Solutions:

• FlashAttention: exact, tiled streaming softmax in SRAM. Memory drops to $O(n)$. Standard for both training and serving.
• Sparse attention (BigBird, Longformer): mask is sparse. Used for some encoder long-context models.
• Linear attention: methods change or approximate attention with different quality, arithmetic, and state trade-offs. Do not equate every such method with one low-rank construction.

Dense, sparse, and alternative attention mechanisms have different contracts. Choose them using a matched workload and task evaluation rather than an asserted universal production recipe.

3. KV cache becomes the cost driver

Under a fixed cache layout, retained KV bytes scale linearly with stored positions; the KV-cache lesson derives the ledger. For an illustrative 80-layer model with 8 KV heads, head width 128, two-byte entries, and 131,072 positions, one request needs 40 GiB of logical KV tensors. That excludes weights, allocator slack, temporary buffers, replication, and effects of sharding or compression.

• GQA / MQA: share K/V heads. Relative to otherwise matched multi-head attention, the logical KV reduction is the query-head count divided by the KV-head count, not a universal multiplier.
• PagedAttention: eliminate cache fragmentation and enable block reuse, including shared prefixes; it does not reduce the logical K/V bytes for a unique sequence.
• Quantization: int8 or int4 KV cache. 2–4× cache reduction.
• Sliding-window attention (Mistral): keep only the last $w$ KV positions per layer; lose strict global attention.
• KV cache eviction: heuristics like H2O (Zhang et al., 2023) keep only “heavy hitter” tokens.

A typical 2026 long-context production stack

• Llama-class architecture with GQA-8 and RoPE (YaRN-extended for long context).
• Training: BF16 mixed precision, FlashAttention-2 kernels, sequence packing, FSDP sharding.
• Serving: vLLM with PagedAttention + continuous batching, int8 weights, optional int8 KV cache.

Common pitfalls

• Quoting context length without measuring quality. A model can run at 128K but degrade rapidly past 32K. Use needle-in-a-haystack and long-doc QA evals.
• Confusing training length with usable context. Models often degrade on inputs longer than the longest examples seen during training (or RoPE extension).
• Ignoring serving cost. A 1M context window is feasible to compute but may cost $10+ per request at frontier prices.

Related

• Context parallelism and ring attention. Split exact long-context attention across devices.
• FlashAttention. Compute exact attention without writing the full score matrix to HBM.
• KV cache. Account for long-context serving memory.