Paper Methods
- Sequence-Parallel Blockwise Q, K, V Partitioning Across N Ring Hosts
- Asynchronous Peer-to-Peer (P2P) Ring KV Block Send/Recv Double Buffering
- Cross-Device Online Softmax Accumulator Rescaling (Exact FlashAttention)
- Arithmetic Intensity & Interconnect Bandwidth Lower-Bound Analysis (c >= F / B)
Engineering Limitations
- •Standard causal masking causes 50% compute load imbalance across ring ranks unless paired with Striped Attention or ZigZag token partitioning
- •When block size c falls below the hardware FLOP-to-bandwidth ratio F / B, NVLink/InfiniBand communication latency becomes exposed
- •Total cluster FLOPs still scale quadratically O(N_seq^2) with total sequence length
Ring Attention Paper Breakdown
A distributed-systems and hardware-interconnect breakdown of Ring Attention with Blockwise Transformers for Near-Infinite Context (Liu, Zaharia, & Abbeel, ICLR 2024 / arXiv:2310.01889), the foundational sequence-parallelism architecture that scales exact Transformer context lengths linearly with the number of GPUs (L_total = N_hosts * c) with zero communication overhead when compute dominates interconnect transfer.
1. The Single-GPU HBM Context Wall
Even though FlashAttention eliminates the O(L^2) intermediate attention matrix in GPU High Bandwidth Memory (HBM), a single GPU's HBM capacity still bounds maximum sequence length L because the activations, Queries (Q), Keys (K), Values (V), and feed-forward intermediates scale as O(L * d) per layer:
Tensor Parallelism (Megatron-LM style) splits head dimensions d across up to 8 intra-node GPUs, but requires synchronous AllReduce barriers twice per layer and cannot scale sequence length across hundreds of multi-node hosts.
2. Ring Topology & Asynchronous Double-Buffered KV Circulation
Ring Attention partitions a sequence of length L = N_h * c into N_h contiguous token blocks of length c, assigning block i (Q_i, K_i, V_i in R^{c x d}) to device rank_i arranged in a logical ring topology:
The N_h-Step Ring Execution Schedule
Notice a key algebraic property of blockwise attention: to compute the output block O_i for local queries Q_i, rank_i never needs to move Q_i or O_i off-chip. It only needs to stream each remote (K_j, V_j) block past Q_i in N_h ring stages (s = 0, 1, ..., N_h - 1):
Because O_i, running row-max m_i in R^c, and running normalizer l_i in R^c stay pinned in rank_i's local SRAM/HBM, each (K_j, V_j) block is discarded immediately after its inner FlashAttention tile pass completes!
3. Cross-Device Online Softmax Accumulator Math
Ring Attention relies on the associativity and commutativity of Online Softmax Rescaling. Across ring steps s = 0 ... N_h - 1, device i maintains three state tensors for its c query tokens:
m_i^{(s)} in R^{c x 1}: Running row-wise maximum logit seen across all KV blocks visited so far.l_i^{(s)} in R^{c x 1}: Running row-wise exponential sum normalizer.O_i^{(s)} in R^{c x d}: Running unnormalized attention output accumulator.
When remote block (K_j, V_j) arrives at step s:
After all N_h ring steps complete (s = N_h - 1), device i performs a single final division O_i = O_i^{(N_h - 1)} / l_i^{(N_h - 1)}, producing bit-for-bit identical mathematical outputs to monolithic single-GPU attention over length L = N_h * c without ever materializing more than one (K_j, V_j) block in memory.
4. Zero-Overhead Overlapping Condition & Roofline Analysis
When is the P2P ring communication latency T_comm completely hidden behind the blockwise FlashAttention compute time T_comp?
Let:
c= block length (tokens per host)d= hidden dimensionF= effective Tensor Core FLOP/s per GPU (e.g.,600 TFLOPSin BF16 on H100)B= unidirectional P2P interconnect bandwidth in Bytes/s (e.g.,450 GB/sNVLink or50 GB/sInfiniBand)
Compute vs. Communication Time per Ring Step
- Compute Time (
T_comp): ComputingQ_i K_j^T(2 * c^2 * dFLOPs) andP_{i,j} V_j(2 * c^2 * dFLOPs) requires4 * c^2 * dFLOPs (plus causal/backward passes):
- Communication Time (
T_comm): Sending bothK_jandV_j(2 * c * delements in 2-byte FP16/BF16 =4 * c * dBytes) over bandwidthB:
- Zero-Overhead Condition (
T_comp >= T_comm):
Notice that d cancels out completely! As long as the per-GPU block length c exceeds the hardware ridge ratio F / B, communication is 100% hidden behind computation:
| Interconnect Tier | Effective F (BF16) | Link Bandwidth B | Critical Block Size c_min = F / B | Practical Block Size c | Communication Overhead |
|---|---|---|---|---|---|
| Intra-Node NVLink 4.0 (H100) | 600 TFLOPS | 450 GB/s | 1,334 tokens | c = 4,096 | 0.0% (Compute-Bound) |
| Multi-Node 400GbE / NDR IB | 600 TFLOPS | 50 GB/s | 12,000 tokens | c = 16,384 | 0.0% (Compute-Bound) |
| Standard 100GbE RoCEv2 | 312 TFLOPS (A100) | 12.5 GB/s | 24,960 tokens | c = 32,768 | 0.0% (Compute-Bound) |
5. Causal Masking Imbalance & Striped / ZigZag Attention
In autoregressive causal attention (m >= n), a naive contiguous block assignment causes severe compute load imbalance:
- At step
s = 0, allN_hranks compute lower-triangular diagonal blocks (50%active FLOPs). - At step
s > 0, ranks where block indexi < j(Q_iprecedesK_jin the sequence) are 100% masked out (0%active FLOPs) while ranks withi > jmust compute full densec x cblocks (100%active FLOPs).
Production Ring Attention deployments (such as Striped Attention and Megatron Context Parallelism ZigZag Partitioning) assign tokens to rank i either modulo-striped (token t -> rank (t mod N_h)) or symmetrically paired from the front and back of the sequence ({block_i, block_{2*N_h - 1 - i}}), guaranteeing that every GPU computes the exact same number of unmasked causal tokens at every single ring step.
