Google announces Long-Context Training on TPUs with MaxText Context

Google Cloud announced long-context training support in MaxText, its open-source JAX training framework, scaling native sequence lengths up to 14 million tokens on Cloud TPUs, according to Google Developers Blog's report. The release implements four distinct context-parallel strategies alongside Tokamax Splash Attention kernels to bypass memory bottlenecks and prevent out-of-memory errors on accelerator chips.

This technical breakdown covers the underlying context parallelism algorithms, hardware performance benchmarks across Cloud TPU generations, and setup guidelines for distributed training. It is written for machine learning engineers, systems architects, and infrastructure teams training frontier foundation models across large accelerator clusters.

01Four context parallel strategies in MaxText

MaxText addresses quadratic attention scaling by dividing single long sequences across chips so no individual device stores full sequence activations. The framework introduces four context-parallel modes: All-gather, All-to-All (DeepSpeed Ulysses), P2P (Ring Attention), and Unified Sequence Parallelism (USP). All four methods run on top of Tokamax Splash Attention, Google's block-sparse TPU kernel that skips fully masked blocks and handles packed sequences without extra padding.

All-gather represents the simplest baseline, retaining token chunks locally and gathering key and value pairs across all chips at each layer. All-to-All rotates sharding from token space to head space before attention runs, allowing execution via single-chip kernels without token load balancing.

P2P circulates key and value chunks in a ring, overlapping computation with communication while chips store only visiting chunks. USP combines All-to-All within device groups and P2P across groups to eliminate head-count limits.

02Benchmark gains across Cloud TPU hardware

Context parallelism in MaxText v0.2.4 prevents activation growth from exhausting high-bandwidth memory on Cloud TPU slices. In benchmark evaluations using Llama 3.1 8B with 256 head dimension and 16 query heads, the system maintained strong Model FLOPs Utilization while scaling sequence lengths from standard ranges to multi-million token spans. The implementation pairs these algorithmic strategies with inter-chip interconnect fabrics across TPU v5p and TPU 7 hardware.

The table details performance throughput recorded across hardware generations. P2P serves as the primary engine for very long contexts, whereas All-gather provides the standard baseline when activations fit in device memory. On Cloud TPU 7, the framework sustained 50 percent MFU across sequence lengths below six million tokens, ultimately reaching 14 million tokens across full slices.

03Deployment commands and configuration switches

Engineers can access these context-parallel strategies directly in the MaxText v0.2.4 release using standard repository updates and environment switches. Upgrading existing local checkouts requires pulling the latest release branch from the open-source repository.

git clone https://github.com/google/maxtext.git
cd maxtext
git checkout tags/v0.2.4
pip install -r requirements.txt

Activating specific context-parallel modes occurs through model training flags during slice invocation. Teams specify the desired sequence parallelism degree and kernel configuration directly within the run command:

python3 MaxText/train.py MaxText/configs/base.yml \
  model_name=llama3.1-8b \
  context_parallel_size=8 \
  context_parallel_type=p2p \
  attention_kernel=tokamax_splash

Switching to Unified Sequence Parallelism pairs ring and all-to-all axes across the chip mesh. Setting default parameters in the configuration YAML pins cluster training runs:

python3 MaxText/train.py MaxText/configs/base.yml \
  model_name=llama3.1-8b \
  context_parallel_size=16 \
  context_parallel_type=usp \
  usp_a2a_size=2 \
  usp_p2p_size=8

04Architectural tradeoffs and hardware limits

Selecting a context-parallel configuration involves balancing memory storage, communication overhead, and model architecture. All-gather gathers full key-value sets on every chip, exhausting high-bandwidth memory if sequences expand without proportional batch adjustments. It also requires load balancing because causal masking assigns larger prefixes to later tokens.

All-to-All avoids token load balancing but cannot split sequences across more chips than available attention heads. Under grouped-query attention, the product of tensor parallelism and context parallelism cannot exceed key-value head counts, creating a strict ceiling of eight devices for Llama 3.1 8B. Ring attention bypasses head limits but risks reducing local kernel efficiency when dividing computation into smaller chunks.

05Topology mapping across accelerator slices

Unified Sequence Parallelism addresses single-strategy bottlenecks by mapping two communication patterns onto the multidimensional physical interconnect mesh of Cloud TPU hardware. Sizing All-to-All within key-value head limits pools tokens locally, allowing subsequent ring hops to carry larger chunks over shorter paths. This topology keeps the 3D inter-chip interconnect saturated while hiding data transfer behind block computations.

The unified context-parallel pipeline in MaxText v0.2.4 supports modern architectural variations including multi-head attention, grouped-query attention, and DeepSeek-style multi-head latent attention. As pretraining and post-training scale toward agentic tasks requiring full codebases or long video sequences, infrastructure teams can extend native context lengths across full Cloud TPU slices without rewriting underlying attention code.