Ring Attention with Blockwise Transformers: Mathematical Foundations, Circular KV Passing, Overlapped Communication-Computation, and Linear Context Scaling

Scaling the context length of transformer-based large language models has historically collided with two fundamental walls: the $O(S^2)$ computational and memory complexity of standard self-attention, and the high-bandwidth memory (HBM) capacity of individual accelerator devices. While IO-aware tiling algorithms such as FlashAttention eliminate intermediate quadratic activation storage by computing softmax within fast on-chip SRAM, the entire sequence key and value tensors must still reside within the memory pool accessible to a single device or tightly coupled tensor-parallel group.

Standard tensor parallelism (Megatron-LM TP) partitions attention heads across $N$ accelerators, but requires two costly All-Reduce operations per transformer layer and cannot scale beyond the number of query heads (often 32 or fewer in modern architectures). Furthermore, sequence parallelism approaches that gather the full sequence into each device recreate the memory bottleneck.

Ring Attention, introduced by Hao Liu, Matei Zaharia, and Pieter Abbeel, resolves this limitation by combining blockwise online softmax with a circular peer-to-peer communication topology. By distributing the sequence across $N$ devices and circulating key-value blocks in a ring while query blocks remain resident, Ring Attention eliminates the quadratic memory barrier per GPU, hides communication latency behind compute FLOPs, and achieves context scaling that is strictly linear in the number of participating accelerators.

The Memory and Communication Bottlenecks in Long Sequences

Standard scaled dot-product attention maps query tensor $Q \in \mathbb{R}^{S \times d}$ , key tensor $K \in \mathbb{R}^{S \times d}$ , and value tensor $V \in \mathbb{R}^{S \times d}$ to an output tensor $O \in \mathbb{R}^{S \times d}$ :

$O = \text{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)V$

For a sequence of length $S = 1,000,000$ tokens with hidden dimension $d = 8,192$ in 16-bit precision, storing the raw activation matrices $Q, K, V$ requires $3 \times 10^6 \times 8192 \times 2 \text{ bytes} \approx 49.15\text{ GB}$ . The intermediate attention matrix $A = QK^T \in \mathbb{R}^{S \times S}$ would require $10^{12} \times 2 \text{ bytes} = 2\text{ TB}$ of storage per attention layer, making single-device materialization physically impossible on modern 80GB or 141GB GPUs.

While FlashAttention reduces the memory complexity to $O(S \cdot d)$ by streaming blocks through SRAM via online softmax, training a 1-million-token sequence still requires storing the input activations, key-value states, and optimizer parameters across all layers. Traditional distributed approaches fail under extreme sequence lengths:

Tensor Parallelism (TP): Partitions weights and heads across $N$ devices. When $S$ is massive, activation memory per device remains proportional to $S$ , and intra-layer All-Reduce over high-speed NVLink becomes a primary latency bottleneck.
All-Gather Sequence Parallelism (DeepSpeed Ulysses): Partitions sequence dimension $S/N$ across devices, then executes an All-to-All collective to transform sequence-partitioned states into head-partitioned states for local full-context attention. While efficient within high-bandwidth clusters, Ulysses is bounded by the head count ( $N \le H$ ) and requires global all-to-all communication that degrades over inter-node InfiniBand networks.

Ring Attention eliminates both constraints by never collecting the full sequence on any single accelerator.

Blockwise Attention and Online Softmax Formulations

The algorithmic cornerstone of Ring Attention is the online normalizer calculation for softmax formulated by Milakov and Gimelshein, which allows an attention output to be computed incrementally over sequence blocks without storing intermediate attention weights.

Let a sequence of length $S$ be split into $N$ blocks of size $B = S / N$ . An accelerator holding query block $Q_i \in \mathbb{R}^{B \times d}$ receives a sequence of key blocks $K_j \in \mathbb{R}^{B \times d}$ and value blocks $V_j \in \mathbb{R}^{B \times d}$ for $j \in \{0, 1, \dots, N-1\}$ .

For each query token row, the attention block computation proceeds iteratively. Let $S_{ij} = \frac{Q_i K_j^T}{\sqrt{d}} \in \mathbb{R}^{B \times B}$ . To maintain numerical stability without computing the global maximum beforehand, the algorithm tracks running statistics:

Running Row Maximum: $m_i^{(k)} \in \mathbb{R}^B$ , tracking the maximum pre-softmax logit seen up to block step $k$ .
Running Normalizer (Partition Function): $l_i^{(k)} \in \mathbb{R}^B$ , tracking the running sum of exponentiated logits.
Unnormalized Running Output: $O_i^{(k)} \in \mathbb{R}^{B \times d}$ .

When encountering the $k$ -th key-value block $(K_{(k)}, V_{(k)})$ , the update rules are evaluated:

$\tilde{m}_i^{(k)} = \text{rowmax}(S_{i,(k)})$

$m_i^{(k)} = \max(m_i^{(k-1)}, \tilde{m}_i^{(k)})$

$P_{i,(k)} = \exp(S_{i,(k)} - m_i^{(k)})$

$l_i^{(k)} = \exp(m_i^{(k-1)} - m_i^{(k)}) \odot l_i^{(k-1)} + \text{rowsum}(P_{i,(k)})$

$O_i^{(k)} = \text{diag}\left(\exp(m_i^{(k-1)} - m_i^{(k)})\right) O_i^{(k-1)} + P_{i,(k)} V_{(k)}$

After iterating through all $N$ blocks ( $k = 0, \dots, N-1$ ), the final exact attention output for query block $Q_i$ is recovered by dividing by the accumulated normalizer:

$O_i^* = \text{diag}\left(l_i^{(N-1)}\right)^{-1} O_i^{(N-1)}$

This recurrence produces the identical mathematical result as standard full-sequence softmax attention, while requiring only $O(B \cdot d)$ active working memory on the accelerator.

The Ring Attention Protocol and Overlapped Communication

Ring Attention maps this blockwise recurrence onto a logical ring topology of $N$ devices. Each device $p \in \{0, 1, \dots, N-1\}$ is initialized with its local shard of input tokens:

Query block: $Q_p \in \mathbb{R}^{B \times d}$ (stays resident on device $p$ across all $N$ steps)
Key block: $K_p \in \mathbb{R}^{B \times d}$ (rotates around the ring)
Value block: $V_p \in \mathbb{R}^{B \times d}$ (rotates around the ring)

The forward pass executes across $N$ discrete iterations:

Step 0 (Local Initialization): Device $p$ computes blockwise attention between local $Q_p$ and local $(K_p, V_p)$ , initializing running statistics $(m_p^{(0)}, l_p^{(0)}, O_p^{(0)})$ . Simultaneously, device $p$ issues non-blocking asynchronous sends (ncclSend) of $(K_p, V_p)$ to its downstream neighbor $(p + 1) \pmod N$ , while posting non-blocking asynchronous receives (ncclRecv) from upstream neighbor $(p - 1) \pmod N$ .
Step $k$ ( $k = 1, \dots, N-1$ ): Device $p$ waits for the incoming key-value buffer $(K_{recv}, V_{recv})$ from step $k-1$ . While computing the next blockwise attention update between resident $Q_p$ and $(K_{recv}, V_{recv})$ , the runtime simultaneously dispatches $(K_{recv}, V_{recv})$ to device $(p + 1) \pmod N$ .
Completion: After $N$ iterations, the rotating key-value buffers complete a full $360^\circ$ traversal of the ring and return to their origin device. Device $p$ applies final division by $l_p^{(N-1)}$ , yielding exact attention output $O_p$ .

Overlapping Arithmetic and Communication Ratios

The primary operational advantage of Ring Attention is the zero-bubble overlap between matrix multiplication and inter-device communication.

Consider a block of size $B$ tokens and hidden dimension $d$ on an accelerator with peak computation throughput $F$ (in FLOPs/s) and bidirectional ring interconnect bandwidth $W$ (in bytes/s):

Computation Cost per Step: Computing $Q_i K_j^T$ takes $2 B^2 d$ FLOPs. Multiplying attention weights by $V_j$ takes another $2 B^2 d$ FLOPs. Total computation per step is $4 B^2 d$ FLOPs. Time required:

$T_{\text{comp}} = \frac{4 B^2 d}{F}$

Communication Volume per Step: Each device transfers one key block $K \in \mathbb{R}^{B \times d}$ and one value block $V \in \mathbb{R}^{B \times d}$ in FP16/BF16 (2 bytes per element). Total payload sent per step is $2 \times (2 B d) = 4 B d$ bytes. Time required:

$T_{\text{comm}} = \frac{4 B d}{W}$

Communication is fully masked behind computation when $T_{\text{comp}} \ge T_{\text{comm}}$ :

$\frac{4 B^2 d}{F} \ge \frac{4 B d}{W} \implies B \ge \frac{F}{W}$

The minimum block size $B_{\text{min}}$ required to achieve 100% communication overlap depends strictly on the arithmetic intensity ratio $F/W$ :

On an NVIDIA H100 SXM (FP16 Tensor Core throughput $F \approx 989\text{ TFLOPs}$ , NVLink bidirectional ring bandwidth $W \approx 450\text{ GB/s}$ per direction):

$B_{\text{min}} \approx \frac{989 \times 10^{12}}{450 \times 10^9} \approx 2,197\text{ tokens}$

For any per-device shard size $B \ge 2,048$ tokens, the peer-to-peer transmission of key-value blocks is completely hidden behind the tensor contractions. The communication adds zero overhead to the model execution time.

Causal Masking Dynamics: Striped Attention and Zig-Zag Topologies

In non-causal encoder settings (such as bidirectional embeddings or vision transformers), every query block attends to every key-value block, ensuring uniform computational load across all $N$ steps on all devices.

In autoregressive (causal) language modeling, token $t_i$ can only attend to tokens $t_j$ where $j \le i$ . If a sequence is partitioned sequentially (Device 0 holds tokens $[0, B-1]$ , Device 1 holds $[B, 2B-1]$ , etc.), a naive ring schedule leads to severe compute imbalance:

Device 0 only attends to its own block; for the remaining $N-1$ steps, all incoming key-value blocks are fully masked out ( $j > i$ ), leaving Device 0 completely idle.
Device $N-1$ must compute attention against all $N$ blocks.
Overall cluster utilization drops by 50%, as half of the theoretical block interactions are masked.

Two distinct architectural solutions eliminate this causal imbalance:

1. Striped Attention (Interleaved Token Assignment)

Proposed by Brandon et al., Striped Attention permutes token assignment across devices in a round-robin interleaved fashion rather than contiguous chunks. Token index $t$ is assigned to device $p = t \pmod N$ .

Under this permutation, every device holds a uniform distribution of prefix, middle, and suffix tokens. On every step of the ring exchange, each device processes an identical mix of valid and masked causal relations. Every device evaluates an upper-triangular causal mask within its local tiles, equalizing execution latency across all devices and recovering near-optimal compute efficiency without idle bubbles.

2. Zig-Zag Ring Attention (Block Reordering)

In Zig-Zag Ring Attention (adopted in modern sequence parallelism runtimes like DeepSpeed-Ulysses / USP), each device is assigned two half-blocks: one from the early sequence and one from the late sequence (for example, Device $p$ holds blocks $p$ and $2N - 1 - p$ ).

By interleaving the forward rotation of the second half-block with the first, the total number of non-zero attention tiles per device across the ring traversal is strictly equalized to $N/2$ full block equivalents. Causal computation is balanced without requiring token-level scattering and gathering.

The Backward Pass and Gradient Exchange

During backpropagation, Ring Attention must compute gradients with respect to queries ( $\nabla Q$ ), keys ( $\nabla K$ ), and values ( $\nabla V$ ).

Computing these gradients requires the attention probability matrix $P_{ij}$ , which is recomputed on the fly using saved forward statistics $(m_i, l_i)$ via FlashAttention recomputation:

Query Gradients ( $\nabla Q_i$ ): Accumulate locally on device $p$ as incoming key-value blocks $(K_j, V_j)$ and output gradients $\nabla O_i$ circulate through the ring.
Key and Value Gradients ( $\nabla K_j, \nabla V_j$ ): Gradients with respect to incoming blocks must either be accumulated locally into a temporary buffer and circulated in lockstep with $(K_j, V_j)$ , or output gradients $\nabla O_i$ can be circulated in reverse ring order while keys and values remain stationary.

In the standard reverse-ring implementation, $\nabla K$ and $\nabla V$ rotate in the same circular direction as during forward propagation, accumulating gradient contributions at each device before returning to their home device. Activation memory is bounded strictly to $O(B \cdot d)$ , allowing backpropagation through multi-million-token sequences without out-of-memory errors.

Ring Attention vs. Ulysses vs. Megatron Sequence Parallelism

Modern distributed long-context architectures select context parallelism paradigms based on cluster network hierarchy and model configurations:

Megatron Sequence Parallelism (Megatron-SP): Splits non-attention operators (LayerNorm, MLP) along sequence length and executes All-Gather / Reduce-Scatter around attention. It is bounded by single-node tensor parallel groups and cannot extend context beyond memory limits of a single node.
DeepSpeed Ulysses: Uses All-to-All collective communication across $N$ devices to convert sequence parallelism to head parallelism. Ulysses offers high kernel efficiency because each device executes a single local FlashAttention kernel on a subset of heads. However, Ulysses requires $N \le H_{\text{heads}}$ (preventing scaling to large cluster sizes when using Grouped-Query Attention with few KV heads) and requires high all-to-all cross-sectional network bandwidth.
Ring Attention: Uses peer-to-peer non-blocking point-to-point ring transfers. It scales independently of head count ( $N$ can exceed $H_{\text{heads}}$ by arbitrary factors), operates efficiently over lower-tier ring topologies across distributed nodes, and requires only $O(S/N)$ memory per device.
Unified Sequence Parallelism (USP): Hybridizes Ulysses within high-speed intra-node NVLink domains and Ring Attention across inter-node InfiniBand networks, optimizing both collective bandwidth and global scaling.

Practical Implementation Patterns

In PyTorch with NCCL, a minimal non-blocking Ring Attention step leverages asynchronous communication handles to overlap computation:

import torch
import torch.distributed as dist

def ring_attention_forward(
    q_local: torch.Tensor,     # Shape: (batch, seq_block, num_heads, head_dim)
    k_local: torch.Tensor,
    v_local: torch.Tensor,
    ring_group: dist.ProcessGroup,
) -> torch.Tensor:
    rank = dist.get_rank(ring_group)
    world_size = dist.get_world_size(ring_group)

    next_rank = (rank + 1) % world_size
    prev_rank = (rank - 1 + world_size) % world_size

    # Allocate double buffers for async P2P communication
    k_curr, v_curr = k_local.clone(), v_local.clone()
    k_recv = torch.empty_like(k_local)
    v_recv = torch.empty_like(v_local)

    # Initialize online softmax statistics
    out_accum = torch.zeros_like(q_local)
    m_i = torch.full((q_local.shape[0], q_local.shape[2], q_local.shape[1]), -float('inf'), device=q_local.device)
    l_i = torch.zeros((q_local.shape[0], q_local.shape[2], q_local.shape[1]), device=q_local.device)

    for step in range(world_size):
        # Asynchronously dispatch current KV to next rank and receive from prev rank
        reqs = []
        if step < world_size - 1:
            reqs.append(dist.isend(k_curr, dst=next_rank, group=ring_group))
            reqs.append(dist.isend(v_curr, dst=next_rank, group=ring_group))
            reqs.append(dist.irecv(k_recv, src=prev_rank, group=ring_group))
            reqs.append(dist.irecv(v_recv, src=prev_rank, group=ring_group))

        # Compute blockwise attention between resident Q and current KV block
        # (Executed concurrently while NCCL transfers execute in background streams)
        scale = 1.0 / (q_local.shape[-1] ** 0.5)
        scores = torch.einsum("bqhd,bkhd->bhqk", q_local, k_curr) * scale

        m_block = scores.max(dim=-1).values
        m_next = torch.maximum(m_i, m_block)
        p_block = torch.exp(scores - m_next.unsqueeze(-1))

        alpha = torch.exp(m_i - m_next)
        l_next = alpha * l_i + p_block.sum(dim=-1)

        # Update running output accumulator
        out_accum = out_accum * alpha.permute(0, 2, 1).unsqueeze(-1) + torch.einsum("bhqk,bkhd->bqhd", p_block, v_curr)

        m_i, l_i = m_next, l_next

        # Await communication completion before next iteration
        for req in reqs:
            req.wait()

        k_curr, k_recv = k_recv, k_curr
        v_curr, v_recv = v_recv, v_curr

    # Final normalizer division
    final_output = out_accum / l_i.permute(0, 2, 1).unsqueeze(-1)
    return final_output

Architectural Trade-Offs and Scaling Limits

While Ring Attention enables theoretically infinite context scaling as additional GPU nodes are introduced, production deployments must account for several structural trade-offs:

Kernel Launch Overhead on Fine Block Partitions: As the ring degree $N$ grows large for a fixed context length $S$ , block size $B = S/N$ decreases. When $B$ drops below the threshold required to saturate GPU Tensor Cores ( $B < 1,024$ ), GEMM computation time drops below kernel launch latency, destroying the computation-communication overlap.
Gradient Synchronization in Backpropagation: Backpropagation requires maintaining accurate causal masks and executing reverse rotations for gradient tensors. If network jitter occurs on any single node in the ring, the synchronous ring pipeline stalls across all $N$ devices.
Serving KV Cache Fragmentation: In autoregressive decoding, new tokens arrive sequentially. Using Ring Attention for single-token prefill or step-by-step decoding introduces high inter-GPU coordination latency. Consequently, production serving engines generally utilize Ring Attention exclusively for chunked long-context prefill, transitioning to paged attention mechanisms during single-token decoding phases.

By reformulating attention as a distributed blockwise recurrence overlapped with circular communication, Ring Attention provides the foundational systems architecture enabling frontier language models to train and reason across million-token sequences.