---
title: 'Ring Attention: Blockwise Transformers'
url: https://www.emergentmind.com/topics/ring-attention-with-blockwise-transformers
type: topic
---

# Ring Attention: Blockwise Transformers

Ring Attention with Blockwise Transformers enables exact Transformer self-attention and feedforward on sequences that are orders of magnitude longer than the memory limits of conventional architectures, without resorting to approximations or sacrificing compute/communication efficiency. The core principle is partitioning the sequence dimension across devices so that each processes local query blocks and rotates key-value blocks (or queries) in a ring topology. This achieves scalable context length, minimal activation memory, and high hardware utilization. Recent innovations, such as TokenRing, further address the communication bottlenecks inherent in earlier ring attention approaches by employing bidirectional, fully overlapping peer-to-peer data transfers, realizing substantial improvements in throughput and latency [2412.20501][2310.01889].

## 1. Blockwise Attention Formulation and Partitioning

Let a Transformer layer receive inputs $Q, K, V \in \mathbb{R}^{H \times L \times D}$, where $H$ is the number of attention heads, $L$ is the sequence length, and $D$ is the per-head hidden dimension. The sequence is partitioned into $N$ contiguous blocks of length $L_{\rm block} = L/N$, each assigned to a rank (device) $j=0,\ldots,N-1$:

- $Q = [Q^{(0)}, Q^{(1)}, ..., Q^{(N-1)}]$, with $Q^{(j)}\in \mathbb{R}^{H \times L_{\rm block} \times D}$.
- Similarly for $K, V$.

Each rank $j$ retains $K^{(j)}$ and $V^{(j)}$ locally. Iteratively, $Q$ blocks circulate through the logical ring so that, at each step $i$, rank $j$ receives $Q_{\rm in} = Q^{((j-i)\bmod N)}$.

Blockwise attention is computed as $\mathrm{FlashAttn}(Q_{\rm in}, K^{(j)}, V^{(j)})$, yielding:

- $\text{block\_out}^{(i,j)} \in \mathbb{R}^{H \times L_{\rm block} \times D}$
- $\text{block\_lse}^{(i,j)} \in \mathbb{R}^{H \times L_{\rm block}}$

Partial results are merged in a numerically stable log-sum-exp manner into global accumulators $\mathrm{Out}^{(i,j)}$ and $\mathrm{Lse}^{(i,j)}$, using the following update for head $h$ and token $t$:

\[
\begin{aligned}
\Delta &= \max\bigl({\rm Lse}^{(i-1,j)}_{h,t}, \text{block\_lse}^{(i,j)}_{h,t}\bigr) \\
{\rm Lse}^{(i,j)}_{h,t} &= \Delta + \log\bigl(\exp({\rm Lse}^{(i-1,j)}_{h,t}-\Delta) + \exp(\text{block\_lse}^{(i,j)}_{h,t}-\Delta)\bigr) \\
{\rm Out}^{(i,j)}_{h,t,d} &= \exp({\rm Lse}^{(i-1,j)}_{h,t} - {\rm Lse}^{(i,j)}_{h,t})~{\rm Out}^{(i-1,j)}_{h,t,d} \\
&~~~~~~~~~+ \exp(\text{block\_lse}^{(i,j)}_{h,t} - {\rm Lse}^{(i,j)}_{h,t})~\text{block\_out}^{(i,j)}_{h,t,d}
\end{aligned}
\]

After $N$ steps, each GPU owns the final outputs for a unique sequence segment [2412.20501].

## 2. Peer-to-Peer (P2P) Ring Communication Patterns

The underlying communication topology is a logical ring where each rank communicates only with its neighbors. In baseline Ring Attention [2310.01889], key-value ($K$, $V$) blocks are rotated clockwise; at each step, a rank computes attention for its fixed $Q$ block against the received ($K$, $V$) and sends out its local $K,V$ block to the next neighbor.

TokenRing introduces bidirectional communication to achieve near-perfect overlap between compute and communication:

- **Forward ring**: $Q$ blocks are sent clockwise.
- **Backward ring**: $\text{block\_out}$ and $\text{block\_lse}$ results are sent counter-clockwise.

This bidirectional strategy exploits full-duplex links typical in OAM full-mesh, NVLink mesh, NVSwitch, or Huawei Ascend HCCS, enabling both directions to transmit simultaneously, thereby halving the effective per-step communication latency.

### Example Pseudocode Structure

A simplified iteration (schematic):

```python
for i in range(N):
    # Forward: send Q to right neighbor, receive Q from left
    # Backward: send previous block_out, block_lse to left, receive from right
    # Compute local blockwise attention
    # Merge results using LogSumExpUpdate
    # Synchronize before next iteration
```

This concurrent messaging ensures no rank sits idle, and computation on the next block can begin as soon as data is available [2412.20501].

## 3. Asymptotic Complexity and Communication Overhead

Let $N$ = number of GPUs, $B = L/N$ = per-GPU block length, $D$ = head dim, $H = \#$ heads.

- **Compute per step per GPU**: $\mathcal O(H D B^2)$ (blockwise Flash Attention)
- **Communication per step**:
    - TokenRing:  send $Q$ block forward ($H B D$),  $\text{block\_out}$ + $\text{block\_lse}$ backward ($H(BD + B) \approx H B D$).
    - Baseline Ring-Transformers: transmit $2 H B D$ ($K,V$) unidirectionally.
- **Per-step latency**:
    - TokenRing: $T_{\rm comm}(TR) \approx H B D / \text{BW}_{\rm bidi}$
    - Baseline: $T_{\rm comm}(ring) \approx 2 H B D / \text{BW}_{\rm uni}$

Because computation scales as $O(B^2)$ but communication as $O(B)$, increasing $N$ eventually makes the system communication-bound. TokenRing's bidirectional overlap delays this scaling bottleneck by a factor of 2 [2412.20501].

### Activation Memory Scaling

| Scheme        | Per-device Memory         | Bottleneck      |
|---------------|--------------------------|-----------------|
| Vanilla       | $O(b h s^2)$             | Quadratic in $s$|
| BPT           | $O(b s h)$               | Linear in $s$   |
| Ring Attention| $6 b c h$ per device     | Linear in $c$   |

This enables effective "infinite context" with memory usage independent of the full sequence length $s = N c$ [2310.01889].

## 4. Empirical Results and Performance Benchmarks

On 4×NVIDIA A10, LLaMA2-7B with $H=32$, $D=128$, $L=24$K tokens, TokenRing achieves:

| Method              | Comm/Step (ms) | Compute/Step (ms) | Speedup |
|---------------------|----------------|-------------------|---------|
| RingAttention (PXB) | 7.6            | 2.5               | 1×      |
| TokenRing (PXB)     | 3.5–4.6        | —                 | 1.7×    |
| TokenRing (PIX)     | 3.2–4.0        | —                 | 2.0×    |

Maximum context lengths, measured on various accelerators and models:

| Accelerator         | Model  | Vanilla | MemEff Attn | BPT | Ring Attn |
|---------------------|--------|---------|-------------|-----|-----------|
| 8×A100 NVLink       | 7B     | 2K      | 16K         | 32K | 256K      |
| 32×A100 InfiniBand  | 13B    | 4K      | 32K         | 64K | 2048K     |
| TPUv4-1024          | 65B    | 4K      | 8K          | 16K | 4096K     |

Model FLOPs utilization suffers less than 5 percentage points degradation compared to blockwise partitioning at standard context length, maintaining $>70\%$ utilization in large-scale LLMs with $>256$K tokens [2310.01889].

On line retrieval tasks (LLaMA-13B, 512K context) Ring Attention outperforms GPT3.5-16K and Claude-100K by absolute margins up to 50 percentage points at maximal context, indicating that exact attention is retained even at extreme lengths.

## 5. Load Balancing and Topological Adaptability

Uniform sequence partitioning and zigzag ordering guarantee identical workload per rank and synchronized operation across the ring, avoiding load imbalance. For causal LLMs, the zigzag partition drops head-context pairs with no dependencies, reducing communication further.

TokenRing requires only a full-duplex, non-blocking P2P primitive—commodity interconnects (NVSwitch, OAM mesh, Huawei HCCS) suffice. On heterogenous or partially connected systems, hybrid schemes (TokenRing intra-node + All-to-All or higher-level RingAttention inter-node) preserve algorithmic efficiency [2412.20501].

## 6. Context, Generalization, and Practical Implications

Ring Attention and TokenRing extend blockwise Transformers into the regime of "near-infinite context," enabling context lengths scaling as $O(Nc)$, with exact self-attention semantics, numerically stable accumulation, and minimal per-device activation.

Unlike prior memory-efficient Transformers (linear, local, or approximative attention), this approach retains full attention expressiveness, allows training/inference with millions of tokens, and is broadly compatible with distributed training and inference frameworks (e.g., xDIT, FSDP pipelines). The only prerequisite is sufficient collective bandwidth to maintain compute-communication overlap for the target block size.

TokenRing's bidirectional, fine-grained ring communication pattern generalizes efficiently to clusters with arbitrary GPU counts and interconnects, sustaining high utilization across a diverse set of cloud and on-premise topologies. This design notably achieves substantial reductions in per-step network latency, high throughput, and enables training and inference at contexts unattainable by previous approaches [2412.20501][2310.01889].

Source: https://www.emergentmind.com/topics/ring-attention-with-blockwise-transformers