Long-context training — including pretraining and post-training, has become a foundation of frontier AI models. Longer native sequence lengths give large language models (LLMs) the working memory to read a whole code repository or ingest hours of video, and reason through multi-step agentic tasks. In three years, context windows have grown from 128K–200K tokens to 1M–2M, and models trained toward 10M tokens are arriving. Training at those lengths, not just serving them, is where the system limits show up.
In particular, efficient training at multi-million-token sequence lengths runs into the following bottlenecks:
Today, we are excited to announce long-context training support in MaxText, Google Cloud's open-source JAX training framework. MaxText enables this by supporting four context-parallel strategies on Cloud TPU: All-gather, A2A (DeepSpeed Ulysses), P2P (Ring Attention), and the A2A + P2P hybrid (Unified Sequence Parallelism, or USP), combined with load balancing and sequence packing. All four use the Tokamax Splash Attention kernel. Combined with the 3D inter-chip interconnect (ICI) ( supported on TPUs such as TPU 5p or TPU 7 (Ironwood) ), this lets one sequence span an entire slice while communication overlaps compute and memory per chip stays flat. In our benchmarks, Llama 3.1 8B (25 head dim, 16 query heads) demonstrates exceptional long-context efficiency, achieving 70% Model FLOPs Utilization (MFU) on TPU 5p for sequence lengths up to 4 million tokens. Across both TPU 5p and TPU 7, context can be seamlessly extended to 14 million tokens—with TPU 7 consistently maintaining 50% MFU across all sequence scales below 6 million tokens.
This functionality is available today in MaxText v0.2.4(doc).
Context parallelism splits a long sequence across chips, so no chip holds the activations of the whole sequence. Context parallelism can be combined with data and tensor parallelism, and the smallest context-parallel size that fits your sequence in HBM is usually the right one. This also holds for sparser attention variants which cut attention compute but not the activations of every other layer.
Attention is the one global operation where every token needs the keys and values of every preceding token, so each context parallel strategy has to bring those keys and values back together. Each makes a different tradeoff between memory (how much KV a chip holds), communication (how much moves per layer, and whether it overlaps compute), and model shape (how many attention heads to split).
All-gather: Each chip keeps its chunk of the sequence and, at every layer, gathers the keys and values from every other chip. It is the simplest option. However, because causal masking causes later tokens to attend to larger prefixes, All-gather requires balancing compute load across devices.
All-to-All (Ulysses): Gathering keys and values makes every chip hold the whole sequence, so attention memory never shrinks no matter how many chips you add. All-to-All changes what each chip owns instead. Attention is independent across heads but not across tokens, so the head axis is free to shard and the token axis is not. An All-to-All before attention rotates the sharding from one onto the other: each chip gives up the heads it is not responsible for and receives the tokens it was missing for the heads it keeps, ending up with the full sequence for a few heads instead of a few tokens for all heads. This is also why Ulysses does not need token load balancing. Attention then runs as ordinary single-chip attention with standard kernels, and a second All-to-All rotates the output back. KV per chip shrinks by the context-parallel factor, but you cannot split across more chips than the model has heads — and tensor parallelism shards that same dimension, so TP × CP is capped by the head count (under GQA, by the KV heads: eight for Llama 3 8B). The all-to-all itself, though, does not overlap with compute — attention cannot start on a partial sequence, so every chip waits for data from every other chip before any of them begins.
P2P (ring attention): Ulysses' ceiling is the head count; ring attention removes it by never leaving the token axis. Keys and values never materialize in full. KV chunks travel around a ring of chips one hop at a time, neighbor to neighbor. Each chip attends its queries to the chunk it holds while the next one is in flight, and merges the partial results as it goes. A chip holds its own chunk and one visiting chunk, never the full sequence, and the communication overlaps compute. However, dividing attention computation into smaller chunks can reduce local kernel efficiency, and unlike All-to-All Ulysses, ring attention requires explicit load balancing.
A2A + P2P (USP): A2A + P2P (USP): Neither strategy scales on its own — All-to-All stops at the head count, and a long ring accumulates hops while thinning the overlap that makes it work. Unified sequence parallelism composes them on two axes of the chip mesh, so the context-parallel degree becomes the product of the two: All-to-All along one axis, sized at or below the KV-head count, and a ring along the other to extend past it. Because each All-to-All group has already pooled its tokens, a ring step now carries a chunk larger by the All-to-All factor, and the ring is correspondingly shorter — more compute per hop to hide the transfer behind, and fewer hops to accumulate latency and causal imbalance. On Cloud TPU the two map onto two axes of the slice's chip mesh.
These strategies cover the attention variants in today's major models, from multi-head and grouped-query attention to DeepSeek-style multi-head latent attention.
Splitting the sequence only pays off if the attention compute on each chip stays efficient. All four strategies run on Splash Attention, the block-sparse flash attention kernel from Tokamax, Google's open-source library of TPU kernels. It skips fully masked blocks and handles packed sequences without padding. P2P calls it once per ring step, on the local queries and the KV chunk that has just arrived, while the transfer for the next step runs underneath, in the forward and backward pass alike. Each chip does the same fused, block-level work it would do without context parallelism, with the communication hidden behind it.
P2P (Ring Attention) is the primary workhorse for very long-context training, whereas All-Gather is the standard baseline for most other cases where the KV can fit. These are the primary CP strategies to use in practice. We compare how far each can push context length on a fixed 64-chip TPU 5p slice with Llama 3.1 8B (256 head dim, 16 query heads) , and measure their MFUs. In ideal settings where both strategies can have large enough chunks (S/CP) and local batch size - the two major variables driving MFU on TPU 5p, we observe that P2P can outperform All-Gather almost every context length from 64K to 14M tokens, where communication becomes mostly hidden with the Ring circular transport schedule.
TPU v7x (Ironwood) takes the same recipe further. With twice the HBM per chip, P2P training of 14M tokens on 64 TPU 7 chips reached 48.7% Model FLOPs Utilization (MFU) compared to 37.1% on TPU 5p. We observed that ring attention performs better than other context parallelism strategies -- e.g. Ulysses and USP –- on TPU 7 chips.
The choice of context parallelism strategy depends mostly on whether the gathered KV fits in HBM, and on how many KV heads the model has:
For all CP strategies, when you run low on memory: Offload before you recompute. MaxText can move saved activations, optimizer states, and parameters to host memory tensor by tensor. At long context lengths, these transfers are hidden behind compute, making offloading effectively free. Offload attention outputs first, followed by optimizer states, layer inputs, parameters, and finally remat as a last resort. MFU stays flat through the offload steps while all-gather falls off and eventually stops fitting. For more details, see the remat tuning guide.
Looking ahead, we see two natural extensions. The first is broader model coverage. Models that mix sliding-window and global layers, such as Gemma, would run the window layers with a short exchange between neighboring chips while the global layers use the ring. We also plan to support linear attention mechanisms such as Gated Delta Net (GDN) and Kimi Delta Attention (KDA). The second is performance, with further tuning of the Tokamax kernels and the P2P backward pass to narrow the remaining gap to all-gather at short context, and recipes by model and context length for TPU 5p and 7 chips.
Context parallelism is available today in MaxText v0.2.4 on PyPI, the corresponding Docker images and on GitHub main. Sample recipes are available here for TPU 7 and TPU 5p. Pick a strategy with context_parallel_strategy (ring, ulysses, or usp), set the context-parallel size, and keep the rest of your run the same. The sharding guide has the flags and constraints for each strategy, and the release notes link to the pull requests behind this work. We can't wait to see what you train.
