Long context: RoPE scaling, context parallelism

When Llama 3.1 shipped with a 128K context window, my first reaction was excitement. My second reaction, as someone who thinks about inference, was dread. A 128K context at 70B parameters means a KV cache of roughly 40 GB per request in FP16. That is more than the entire memory of many GPUs, for a single request. Serving long context is not just about making the model work at long sequences. It is about making the economics work.

RoPE: how position encoding scales

Most modern LLMs use Rotary Position Embeddings (RoPE) to encode token positions. RoPE works by rotating the query and key vectors in 2D subspaces at frequencies that depend on the position. The rotation angle for position m at frequency index i is:

# RoPE rotation angle
theta_i = base ** (-2i / d)  # base is typically 10,000
angle_m_i = m * theta_i      # rotation angle for position m

The key property: RoPE encodes relative positions through the angle difference between two tokens. The attention score between tokens at positions m and n depends on (m - n), not on the absolute positions.

The problem is that RoPE is trained with a maximum sequence length. If you try to use positions beyond the training range, the rotation angles become values the model has never seen, and attention patterns break down. The model starts hallucinating or losing coherence.

RoPE scaling methods

Three main approaches have emerged to extend RoPE beyond the training length:

Linear scaling (Position Interpolation). Instead of extrapolating to new positions, you interpolate: divide all positions by a scale factor so they fit within the original training range. If the model was trained on 4K tokens and you want 32K, set scale = 8 and use position m/8 instead of m.

# Position Interpolation
scale_factor = target_length / train_length  # e.g., 32768 / 4096 = 8
angle_m_i = (m / scale_factor) * theta_i

This works surprisingly well with minimal fine-tuning (a few hundred steps). The downside is that you compress all position information into the original range, which can reduce the model's ability to distinguish nearby positions.

NTK-aware scaling. Instead of scaling all frequencies equally, NTK-aware scaling modifies the RoPE base frequency. This preserves high-frequency (local) position information while stretching only the low-frequency (long-range) components:

# NTK-aware scaling
alpha = scale_factor
new_base = base * alpha ** (d / (d - 2))
theta_i = new_base ** (-2i / d)

The intuition is that local attention patterns (which depend on high-frequency rotations) should not change, while long-range patterns (low-frequency) can stretch. This gives better perplexity at long contexts without fine-tuning.

YaRN (Yet another RoPE extensioN). YaRN combines NTK-aware scaling with an attention temperature correction. It divides the frequency dimensions into three groups: low frequencies get NTK scaling, high frequencies get no scaling, and middle frequencies get interpolated. It also adds a temperature scaling factor to the attention logits to compensate for the changed magnitude of position encodings.

The real bottleneck: KV cache memory

RoPE scaling is the easy part. The hard part is memory. The KV cache stores the key and value tensors for every token in the context, at every layer, for every attention head. For a 70B model with 80 layers, 8 KV heads (GQA), and head dimension 128:

# KV cache size for one request at 128K context
layers = 80
kv_heads = 8
head_dim = 128
seq_len = 128_000
bytes_per_param = 2  # FP16

kv_bytes = 2 * layers * kv_heads * head_dim * seq_len * bytes_per_param
# 2 (K and V) * 80 * 8 * 128 * 128000 * 2
# = 2 * 80 * 8 * 128 * 128000 * 2
# = 41,943,040,000 bytes ≈ 42 GB per request

42 GB for a single request. An H100 has 80 GB of HBM. After loading the model weights (roughly 140 GB for 70B in FP16, so you need at least TP=2), you have maybe 5 to 10 GB per GPU for KV cache. That means you can serve zero concurrent 128K requests on a 2-GPU setup, and maybe one on a 4-GPU setup.

This is why long-context serving requires a fundamentally different approach.

KV cache compression techniques

Several strategies reduce the KV cache footprint:

• FP8 or INT8 KV cache. Quantize the cached keys and values from FP16 to FP8. Halves the memory with minimal quality loss. vLLM and TensorRT-LLM both support this. Our 42 GB drops to 21 GB.
• GQA (Grouped Query Attention). Already used by most modern models. Llama 3.1 70B uses 8 KV heads instead of 64 attention heads, an 8x reduction in KV cache size compared to MHA.
• Sparse attention patterns. Models like Gemini 1.5 use a mixture of local and global attention. Local attention layers only cache a sliding window of recent tokens, drastically reducing KV cache size for those layers.
• KV cache eviction. Drop KV entries for tokens the model is unlikely to attend to. H2O (Heavy Hitter Oracle) keeps only the tokens with the highest cumulative attention scores. StreamingLLM keeps the first few tokens (attention sinks) plus a sliding window.

Context parallelism for long prefill

Even if you can fit the KV cache, the prefill computation for 128K tokens is enormous. Prefill time scales quadratically with sequence length for standard attention (linearly with FlashAttention, but with a large constant). A 128K prefill on a single A100 with FlashAttention-2 takes roughly 30 to 60 seconds for a 70B model. That is an unacceptable TTFT.

Context parallelism, which I explored for video generation, applies here too. Split the 128K sequence across multiple GPUs, run Ring Attention, and reduce the prefill time by roughly the number of GPUs. With CP=8, that 30-second prefill drops to about 4 seconds.

The key difference from video CP is that for LLM long context, you typically combine CP with TP. A common setup for Llama 3.1 70B at 128K context:

• TP=4 to shard the model weights across 4 GPUs
• CP=2 to split the sequence across 2 groups
• Total: 8 GPUs, each holding 1/4 of the weights and 1/2 of the sequence

Chunked prefill: the practical middle ground

If you do not have enough GPUs for context parallelism, chunked prefill is the fallback. Instead of processing all 128K tokens in one monolithic prefill, break the prompt into chunks of 8K to 16K tokens and process them sequentially, interleaving decode steps for other requests between chunks.

# Chunked prefill scheduling
def chunked_prefill(prompt_tokens, chunk_size=8192):
    chunks = [prompt_tokens[i:i+chunk_size]
              for i in range(0, len(prompt_tokens), chunk_size)]
    for chunk in chunks:
        prefill_chunk(chunk)  # ~2-4 seconds per chunk
        # Between chunks, run decode iterations for other requests
        run_pending_decodes()

This prevents a single long-context request from blocking all other requests. The tradeoff is higher total prefill time (no parallelism benefit) but better system-level fairness and throughput.

Long context is a systems problem, not just a modeling problem. The model may handle 128K tokens perfectly, but if your serving system cannot fit the KV cache, parallelize the prefill, and manage memory, that context window is a marketing number, not a usable feature.

Tomorrow we shift to EAGLE speculative decoding, a clever approach that drafts at the feature level rather than the token level to achieve higher acceptance rates.