RingAttention and Context Parallelism: Mathematical Foundations, Distributed Blockwise Attention, Circular Communication Topologies, and Million-Token Context Scaling

Standard Transformer architectures face a quadratic memory and computational barrier in their self-attention mechanism. While FlashAttention solved the high-bandwidth memory (HBM) IO bottleneck on single devices by tiling matrices within SRAM, scaling sequence lengths beyond hundreds of thousands or millions of tokens quickly exceeds the physical memory capacity of any single GPU or TPU.

Context Parallelism (CP) distributes the sequence dimension across multiple accelerators. Among context-parallel approaches, RingAttention (Liu et al., 2023) provides an elegant distributed architecture that eliminates sequence-length limits without introducing communication overhead. By organizing compute devices into a logical ring, overlapping peer-to-peer key-value transfers with blockwise attention computation, and updating running online softmax statistics, RingAttention reduces per-device memory from $O(N)$ to $O(N/P)$ and enables linear context scaling up to millions of tokens.

The Long-Context Scaling Problem

In standard Multi-Head Attention (MHA), given an input sequence of length $N$, hidden dimension $d$, and number of heads $H$, the projections produce query, key, and value tensors:

$$Q, K, V \in \mathbb{R}^{N \times d}$$

The scaled dot-product attention computes:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

This formulation presents two primary physical constraints:

1. Quadratic Computation: Computing $QK^T$ requires $2N^2 d$ FLOPs per attention head.
2. Linear-to-Quadratic Memory: Storing raw attention weights requires $O(N^2)$ memory. While IO-aware tiling algorithms like FlashAttention avoid materializing the $N \times N$ matrix in HBM, the inputs $Q, K, V$ and output activations still require $O(N \cdot d)$ storage per layer. Across $L$ layers during training, storing activations for backpropagation scales as $O(L \cdot N \cdot d)$, quickly causing Out-of-Memory (OOM) errors at sequence lengths of $N \ge 64\text{K}$ on standard 80GB GPUs.

Standard Sequence Processing (Single Device):
Sequence (Length N) ───► [ OOM at N >= 64k-128k Tokens ]

Context Parallelism (P Devices):
Device 0: Tokens [0 ... N/P - 1]
Device 1: Tokens [N/P ... 2N/P - 1]
...
Device P-1: Tokens [(P-1)N/P ... N - 1]

Limitations of Traditional Parallelism Strategies

Existing distributed training paradigms fail to solve sequence-length scaling efficiently:

• Data Parallelism (DP / FSDP): Replicates the model across devices and partitions the batch dimension. It cannot process a single sequence whose activations exceed single-device memory.
• Tensor Parallelism (TP): Shards weight matrices across hidden dimensions ($d / P$). However, tensor parallelism requires two All-Reduce collectives per Transformer layer across the high-speed intra-node interconnect (NVLink). Scaling TP beyond a single 8-GPU node degrades throughput due to cross-node network latency, and TP cannot scale beyond the number of attention heads ($P \le H$).
• Pipeline Parallelism (PP): Partitions layers across stages. While PP distributes model weights, each stage must still hold the entire sequence's activations for its subset of layers, suffering from pipeline bubble overheads.

Blockwise Attention and Online Softmax Decomposition

RingAttention builds upon the mathematical framework of Online Softmax (Milakov & Gimelshein, 2018; Dao et al., 2022), extending it from local SRAM blocks to distributed accelerators.

The Online Softmax Formulation

Consider computing attention for a query block $Q_i \in \mathbb{R}^{B_r \times d}$ against two key-value blocks $(K_1, V_1)$ and $(K_2, V_2)$ of size $B_c \times d$.

For block $j \in \{1, 2\}$, the unnormalized attention scores are:

$$S_{ij} = \frac{Q_i K_j^T}{\sqrt{d_k}} \in \mathbb{R}^{B_r \times B_c}$$

The row-wise maximum of block $j$ is:

$$m_{ij} = \max_{\text{row}}(S_{ij}) \in \mathbb{R}^{B_r}$$

The unnormalized exponentiated scores and local row sum are:

$$P_{ij} = \exp(S_{ij} - m_{ij}), \quad \ell_{ij} = \sum_{\text{cols}} P_{ij} \in \mathbb{R}^{B_r}$$

When aggregating across blocks 1 and 2, the global maximum is updated dynamically:

$$m_i^{\text{new}} = \max(m_{i1}, m_{i2})$$

The updated normalization denominator $\ell_i^{\text{new}}$ is:

$$\ell_i^{\text{new}} = \ell_{i1} \cdot \exp(m_{i1} - m_i^{\text{new}}) + \ell_{i2} \cdot \exp(m_{i2} - m_i^{\text{new}})$$

The accumulated output block $O_i$ is updated via rescaled linear combination:

$$O_i^{\text{new}} = \frac{O_{i1} \cdot \ell_{i1} \cdot \exp(m_{i1} - m_i^{\text{new}}) + (P_{i2} V_2) \cdot \exp(m_{i2} - m_i^{\text{new}})}{\ell_i^{\text{new}}}$$

Because this update rule is associative and numerically exact, the attention output can be computed incrementally across arbitrary chunks of key-value pairs without ever materializing the full attention matrix.

The RingAttention Distributed Mechanism

RingAttention leverages online softmax across a ring of $P$ distributed devices.

Sequence Partitioning

Let the total sequence length be $N$. The sequence is partitioned evenly across $P$ devices, such that each device $p \in \{0, 1, \dots, P-1\}$ holds a local block of size $B = N / P$:

$$Q_p = Q[pB : (p+1)B], \quad K_p = K[pB : (p+1)B], \quad V_p = V[pB : (p+1)B]$$

Each device maintains:

• Its invariant local query block $Q_p$
• An active key-value block $(K_{\text{curr}}, V_{\text{curr}})$, initialized to $(K_p, V_p)$
• Running online softmax statistics: row maximum $m_p \in \mathbb{R}^{B}$ (initialized to $-\infty$), running sum $\ell_p \in \mathbb{R}^{B}$ (initialized to 0), and accumulated output accumulator $O_p \in \mathbb{R}^{B \times d}$ (initialized to 0)

Execution Algorithm

The computation proceeds over $P$ discrete steps ($s = 0, 1, \dots, P-1$):

Ring Step Progression (for Device p at step s):
1. Async Send: Send (K_curr, V_curr) to device (p + 1) mod P
2. Async Recv: Receive (K_next, V_next) from device (p - 1) mod P
3. Compute:    Local_Attention(Q_p, K_curr, V_curr)
4. Update:     Update running stats (m_p, \ell_p, O_p) using Online Softmax
5. Barrier:    Wait for communication to complete; set (K_curr, V_curr) <- (K_next, V_next)

By the end of step $P-1$, the key-value blocks have made a full circle around the ring. Every query block $Q_p$ has attended to every key-value block $K_j, V_j$ for all $j \in \{0, \dots, P-1\}$, yielding the exact global attention output without approximations.

Device 0          Device 1          Device 2          Device 3
 [Q0, K0, V0]      [Q1, K1, V1]      [Q2, K2, V2]      [Q3, K3, V3]
      │                 │                 │                 │
      ▼                 ▼                 ▼                 ▼
Step 0: Attn(Q0,K0)   Attn(Q1,K1)       Attn(Q2,K2)       Attn(Q3,K3)
      │  K0,V0 ─►       │  K1,V1 ─►       │  K2,V2 ─►       │  K3,V3 ──┐
      └─────────────────┴─────────────────┴─────────────────┴───────┘
      ┌─────────────────────────────────────────────────────────────┘
      ▼                 ▼                 ▼                 ▼
Step 1: Attn(Q0,K3)   Attn(Q1,K0)       Attn(Q2,K1)       Attn(Q3,K2)
      │  K3,V3 ─►       │  K0,V0 ─►       │  K1,V1 ─►       │  K2,V2 ──┐
      ...

Communication-Computation Overlap Arithmetic

The efficiency of RingAttention depends on hiding inter-device peer-to-peer (P2P) communication latency behind blockwise matrix multiplications.

Computation Cost per Step

At each step $s$, device $p$ performs two dense matrix multiplications:

1. $S = Q_p K_{\text{curr}}^T$ requiring $2 \cdot B^2 \cdot d$ FLOPs per head
2. $O = P V_{\text{curr}}$ requiring $2 \cdot B^2 \cdot d$ FLOPs per head

For $H$ attention heads, the total floating-point operations per ring step are:

$$\text{FLOPs}_{\text{step}} = 4 \cdot H \cdot B^2 \cdot d_k = 4 \cdot B^2 \cdot d_{\text{model}} = 4 \cdot \left(\frac{N}{P}\right)^2 \cdot d_{\text{model}}$$

Given device compute throughput $C_{\text{device}}$ (in FLOPS/s), the computation time per step is:

$$T_{\text{comp}} = \frac{4 \cdot (N/P)^2 \cdot d_{\text{model}}}{C_{\text{device}}}$$

Communication Cost per Step

At each step $s$, device $p$ sends its current $(K, V)$ tensors to device $(p+1) \bmod P$. Assuming 16-bit precision (2 bytes per element):

$$\text{Bytes}_{\text{transferred}} = 2 \times 2 \times B \times d_{\text{model}} = 4 \cdot \left(\frac{N}{P}\right) \cdot d_{\text{model}}$$

Given bidirectional interconnect bandwidth $W$ (in Bytes/s), the communication time is:

$$T_{\text{comm}} = \frac{4 \cdot (N/P) \cdot d_{\text{model}}}{W}$$

Perfect Overlap Condition

Communication is completely hidden ($T_{\text{comm}} \le T_{\text{comp}}$) when:

$$\frac{4 \cdot (N/P) \cdot d_{\text{model}}}{W} \le \frac{4 \cdot (N/P)^2 \cdot d_{\text{model}}}{C_{\text{device}}}$$

Simplifying yields the minimum local block size required for zero communication overhead:

$$\frac{N}{P} \ge \frac{C_{\text{device}}}{W}$$

Numerical Example (NVIDIA H100 SXM5 Cluster)

• Dense BF16 Tensor Core Peak: $C_{\text{device}} \approx 989 \times 10^{12} \text{ FLOPS/s}$ (at 50% MFU: $\approx 500 \text{ TFLOPS}$)
• Inter-Node Network Bandwidth (InfiniBand NDR 400 Gbps): $W \approx 50 \times 10^9 \text{ Bytes/s}$

$$\text{Block Size Threshold } B = \frac{N}{P} \ge \frac{500 \times 10^{12}}{50 \times 10^9} = 10,000 \text{ tokens}$$

If each GPU holds at least 10,000 tokens (e.g., an 8-GPU node running an 80K sequence, or a 64-GPU cluster running a 640K sequence), the communication is 100% overlapped with computation, introducing zero communication overhead.

Causal Masking and Load Balancing

In causal autoregressive language modeling, token $i$ can only attend to tokens $j \le i$. This creates an upper-triangular attention mask where computation is non-uniform across blocks.

Causal Attention Matrix (4 Devices):
      Dev 0    Dev 1    Dev 2    Dev 3
Dev 0 [  X  ]  [     ]  [     ]  [     ]   <- 1 block
Dev 1 [  X  ]  [  X  ]  [     ]  [     ]   <- 2 blocks
Dev 2 [  X  ]  [  X  ]  [  X  ]  [     ]   <- 3 blocks
Dev 3 [  X  ]  [  X  ]  [  X  ]  [  X  ]   <- 4 blocks

In a naive ring schedule:

• Device 0 only computes its diagonal block (Step 0) and remains idle for the remaining $P-1$ steps.
• Device $P-1$ computes in all $P$ steps.
• This causes severe device idle time, reducing overall cluster compute efficiency to ~50%.

Striped / Zigzag Ring Scheduling

To restore balanced computation across all devices, modern context parallelism implementations (USP; LoongTrain) employ zigzag token allocation or striped chunking.

Instead of assigning contiguous chunks $[0, N/P)$, each device receives an interleaved pair of chunks from both the first half and second half of the sequence:

$$\text{Device } p \text{ receives tokens: } \left[ p \cdot \frac{N}{2P}, (p+1) \cdot \frac{N}{2P} \right) \cup \left[ (2P - 1 - p) \cdot \frac{N}{2P}, (2P - p) \cdot \frac{N}{2P} \right)$$

Zigzag Indexing Example (P = 4 devices, 8 sequence chunks 0..7):
Device 0: Chunks 0 and 7
Device 1: Chunks 1 and 6
Device 2: Chunks 2 and 5
Device 3: Chunks 3 and 4

Under this partitioning:

• Total causal blocks for Device 0: Chunk 0 attends to 1 block; Chunk 7 attends to 8 blocks $\rightarrow 1 + 8 = 9$ chunk-attentions.
• Total causal blocks for Device 1: Chunk 1 attends to 2 blocks; Chunk 6 attends to 7 blocks $\rightarrow 2 + 7 = 9$ chunk-attentions.
• Total causal blocks for Device 2: Chunk 2 attends to 3 blocks; Chunk 5 attends to 6 blocks $\rightarrow 3 + 6 = 9$ chunk-attentions.
• Total causal blocks for Device 3: Chunk 3 attends to 4 blocks; Chunk 4 attends to 5 blocks $\rightarrow 4 + 5 = 9$ chunk-attentions.

Every device computes the exact same number of block-attentions, restoring full 100% hardware utilization and eliminating bubble latency under causal masks.

RingAttention vs. DeepSpeed Ulysses vs. Megatron Sequence Parallelism

Context Parallelism has evolved into several competing architectural paradigms:

Context Parallelism Taxonomy:

1. RingAttention:
   Sequence Partitioning: Along token dimension [N/P]
   Communication: Ring P2P (Send/Recv) overlapped with GEMM
   Constraint: N/P >= C/W (requires moderate local batch/sequence)
   Head Limit: Independent of head count (supports P > H)

2. DeepSpeed Ulysses:
   Sequence Partitioning: Along token dimension [N/P] -> All-to-All -> Head dimension [H/P]
   Communication: All-to-All collective before and after local FlashAttention
   Constraint: P <= H (devices cannot exceed head count)
   Latency: Non-overlapped All-to-All, but leverages ultra-fast full FlashAttention kernels

3. Hybrid 2D Sequence Parallelism (USP):
   Combines Ulysses within intra-node NVLink (fast All-to-All)
   and RingAttention across inter-node InfiniBand (overlapped P2P ring)

Architectural Comparison

• Communication Collective: RingAttention uses asynchronous point-to-point transfers (P2P ISend / IRecv), whereas DeepSpeed Ulysses utilizes collective All-to-All.
• Scaling Limits: Ulysses cannot scale beyond the number of attention heads $H$ (typically 32 to 128). In Grouped-Query Attention (GQA) architectures like Llama 3 (where KV heads $H_{KV} = 8$), Ulysses is strictly limited to $P \le 8$. RingAttention has no head-count constraint and can scale across thousands of GPUs ($P \gg H$).
• Memory Footprint: Both RingAttention and Ulysses achieve $O(N/P)$ memory scaling for KV caches and activation tensors.
• Hardware Topology Suitability: Ulysses excels within NVLink nodes where All-to-All latency is negligible ($< 1\text{ms}$). RingAttention excels across inter-node network topologies where bandwidth is limited but P2P ring communication can be fully overlapped with compute.

Python / PyTorch Algorithmic Implementation

Below is a minimal, self-contained PyTorch reference implementation demonstrating the RingAttention forward pass with Online Softmax accumulation and P2P communication simulation:

import torch
import torch.distributed as dist

def ring_attention_forward(
    q_local: torch.Tensor,       # [batch, seq_len_p, heads, head_dim]
    k_local: torch.Tensor,       # [batch, seq_len_p, heads, head_dim]
    v_local: torch.Tensor,       # [batch, seq_len_p, heads, head_dim]
    group: dist.ProcessGroup,
) -> torch.Tensor:
    rank = dist.get_rank(group)
    world_size = dist.get_world_size(group)
    scale = 1.0 / (q_local.shape[-1] ** 0.5)

    # Initialize online softmax statistics
    # m_i: running max, l_i: running sum of exp, out: accumulated output
    batch, seq_p, heads, d = q_local.shape
    m_i = torch.full((batch, heads, seq_p, 1), float("-inf"), device=q_local.device)
    l_i = torch.zeros((batch, heads, seq_p, 1), device=q_local.device)
    out = torch.zeros_like(q_local).permute(0, 2, 1, 3) # [B, H, S_p, D]

    # Rearrange Q for batched matmul: [B, H, S_p, D]
    q = q_local.permute(0, 2, 1, 3)

    k_curr = k_local.clone()
    v_curr = v_local.clone()

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

    for step in range(world_size):
        # Asynchronously send current KV to next rank, receive from prev rank
        if world_size > 1:
            k_next = torch.empty_like(k_curr)
            v_next = torch.empty_like(v_curr)
            
            send_k = dist.isend(k_curr, dst=next_rank, group=group)
            send_v = dist.isend(v_curr, dst=next_rank, group=group)
            recv_k = dist.irecv(k_next, src=prev_rank, group=group)
            recv_v = dist.irecv(v_next, src=prev_rank, group=group)

        # 1. Compute local attention scores: S_ij = Q_p * K_curr^T
        k_t = k_curr.permute(0, 2, 3, 1) # [B, H, D, S_p]
        scores = torch.matmul(q, k_t) * scale # [B, H, S_p, S_p]

        # 2. Local block reduction
        m_curr = torch.max(scores, dim=-1, keepdim=True)[0]
        p_curr = torch.exp(scores - m_curr)
        l_curr = torch.sum(p_curr, dim=-1, keepdim=True)

        # 3. Online Softmax update across distributed blocks
        m_new = torch.maximum(m_i, m_curr)
        alpha = torch.exp(m_i - m_new)
        beta = torch.exp(m_curr - m_new)
        
        l_new = l_i * alpha + l_curr * beta

        # 4. Rescale and accumulate output
        v_mat = v_curr.permute(0, 2, 1, 3) # [B, H, S_p, D]
        out = (out * l_i * alpha + torch.matmul(p_curr, v_mat) * beta) / (l_new + 1e-8)

        # Update running state
        m_i = m_new
        l_i = l_new

        # Wait for async communication to finish before next iteration
        if world_size > 1:
            send_k.wait()
            send_v.wait()
            recv_k.wait()
            recv_v.wait()
            k_curr = k_next
            v_curr = v_next

    return out.permute(0, 2, 1, 3) # Return [B, S_p, H, D]

Practical Deployment and 2D/3D Parallelism

In large-scale production training clusters (e.g., training 1M+ token foundation models like Llama 3.1 or Gemini 1.5), RingAttention is rarely used in isolation. Instead, it is combined with other parallelism dimensions in a 4D Parallelism hierarchy:

                    ┌───────────────────────────────┐
                    │      4D Parallelism Model     │
                    └──────────────┬────────────────┘
                                   │
         ┌─────────────────────────┼─────────────────────────┐
         ▼                         ▼                         ▼
┌──────────────────┐      ┌──────────────────┐      ┌──────────────────┐
│ Data Parallelism │      │ Pipeline Par.    │      │ Context Par.     │
│ (FSDP / ZeRO-3)  │      │ (Inter-Layer)    │      │ (Sequence Length)│
└──────────────────┘      └──────────────────┘      └────────┬─────────┘
                                                             │
                                        ┌────────────────────┴────────────────────┐
                                        ▼                                         ▼
                             ┌──────────────────────┐                  ┌──────────────────────┐
                             │ DeepSpeed Ulysses    │                  │ RingAttention        │
                             │ (Intra-Node NVLink)  │                  │ (Inter-Node IB/RoCE) │
                             └──────────────────────┘                  └──────────────────────┘

1. Intra-Node Sequence Parallelism (Ulysses): Within an 8-GPU node connected via 900 GB/s NVLink, sequence chunks are redistributed via All-to-All to parallelize across attention heads.
2. Inter-Node Context Parallelism (RingAttention): Across nodes connected via InfiniBand, key-value blocks circulate over a ring topology, overlapping network latency with local FlashAttention execution.
3. Fully Sharded Data Parallelism (FSDP): Model parameters and optimizer states are sharded across orthogonal data-parallel groups.

By decoupling sequence length from single-GPU memory limits, RingAttention transforms long-context training from an intractable memory bottleneck into a scalable distributed systems problem.

Sources

• Liu, H., Zaharia, M., & Abbeel, P. (2023). Ring Attention with Blockwise Transformers for Near-Infinite Context. arXiv:2310.01889
• Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135
• Jacobs, S. A., Tanaka, M., Zhang, C., et al. (2023). DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models. arXiv:2309.14509
• Fang, J., et al. (2024). USP: A Unified Sequence Parallelism Approach for Long Context Generative AI. arXiv:2405.07719
• Milakov, M., & Gimelshein, N. (2018). Online normalizer calculation for softmax. arXiv:1805.02867
• Liu, H., & Abbeel, P. (2023). Blockwise Parallel Transformer for Large Context Models. arXiv:2305.19370
• GitHub Repository: haoliuhl/ringattention