← Part 3: Sequence Sharding (Ulysses)

Part 4: Ring Attention

Distributed attention without gathering the full sequence

By Darshan Fofadiya

Let's take stock of where we are. We started in Part 1 with a seemingly impossible problem: running Llama-70B with a 1 million token context window on A100 GPUs. The model weights alone need 140 GB, the KV cache needs 328 GB, and the naive attention matrix would require 128 TB. None of this fits on a single 80 GB GPU.

In Part 2, we tackled the first bottleneck — the 140 GB of model weights. FSDP splits the weights across 8 GPUs, so each GPU stores only 17.5 GB. When a layer needs the full weights for computation, it gathers them temporarily via Ring AllGather, computes, and then frees them.

In Part 3, we tackled the second bottleneck — the activation memory. Each layer's Q, K, V tensors consume 20.5 GB for 1M tokens. Ulysses splits the sequence across 8 GPUs so each GPU only processes 125K tokens, reducing activations to 2.6 GB per GPU. The clever part is the All-to-All operation that reshapes data from "partial tokens, all heads" to "all tokens, partial heads" so each GPU can compute complete attention for its assigned heads using FlashAttention at full efficiency.

But two problems remain. First, the KV cache for inference consumes 41 GB per GPU — more than half the A100's memory for just a single request. In production, we'd want to serve multiple requests simultaneously, which would multiply this number. Second, Ulysses is limited to 8 GPUs because Llama-70B only has 8 KV heads, and each GPU needs at least one head after the All-to-All.

Ring Attention solves the scaling problem. It takes a fundamentally different approach: instead of redistributing data across all GPUs at once (like Ulysses's All-to-All), it passes KV blocks around a ring of GPUs one step at a time. This means it has no head count limit, and it only communicates with neighboring GPUs. In this part, we explain Ring Attention as a standalone technique. In Part 5, we'll show how Ring Attention and Ulysses combine in the full USP system.

Before we dive in, a quick note on terminology. When we say "node," we mean a single physical server containing multiple GPUs connected by fast NVLink interconnects (e.g., 8 A100 GPUs in one machine with 600 GB/s NVLink bandwidth). When we say "multiple nodes," we mean separate servers connected by slower network links like InfiniBand (typically 25-50 GB/s). This distinction matters because different parallelism strategies work better at different levels of the network hierarchy.

4.1 The Scaling Problem

Here's the memory budget after applying FSDP and Ulysses with 8 GPUs. All numbers assume a single request (batch_size=1):

MEMORY BUDGET PER GPU (8 GPUs, FSDP + Ulysses, batch_size=1):
═════════════════════════════════════════════════════════════
Weight shard (FSDP):              17.5 GB
KV cache (all heads, 1/8 seq):    41.0 GB  ← biggest consumer!
Activations (Q,K,V for 125K):      2.0 GB
Buffers & overhead:                ~2.0 GB
─────────────────────────────────────────────
Total:                             62.5 GB
A100 capacity:                     80.0 GB
Headroom:                          17.5 GB

It fits for a single request, but just barely. And here's the critical point: these numbers assume we're serving exactly one request at a time. In a real production system, we'd want to serve multiple requests concurrently to maximize GPU utilization. The KV cache scales linearly with batch size:

KV CACHE SCALES WITH BATCH SIZE:
════════════════════════════════
batch_size=1:  KV cache = 41 GB per GPU  → Total: 62.5 GB  ✓ (fits)
batch_size=2:  KV cache = 82 GB per GPU  → Total: 103.5 GB ❌ (exceeds 80 GB!)

Even serving just 2 concurrent 1M-token requests overflows the GPU.
And we're stuck at 8 GPUs because of GQA's 8 KV heads.

4.2 Ring Attention: The Core Idea

Ring Attention takes a completely different approach from Ulysses. Ulysses redistributes data so each GPU has the full sequence for a subset of heads, using an All-to-All operation that requires every GPU to communicate with every other GPU simultaneously. Ring Attention instead keeps each GPU's data local and passes KV blocks around a ring, one neighbor at a time.

The key insight is that attention can be computed incrementally. We don't need all K and V values at once. We can process one KV block at a time and accumulate the result using online softmax — the same mathematical trick that FlashAttention uses to process tiles within a single GPU (see Part 1, Section 1.4.6). The only difference is the granularity: FlashAttention tiles across SRAM blocks, Ring Attention tiles across GPUs.

THE RING ATTENTION IDEA:
════════════════════════
Arrange 8 GPUs in a ring: GPU0 → GPU1 → GPU2 → ... → GPU7 → GPU0

Each GPU holds:
  - Its local Q chunk (fixed — never moves)
  - Its local KV chunk (gets passed to the next GPU each step)

Step 0: Each GPU computes attention of its Q against its LOCAL KV block
Step 1: Each GPU passes its KV to the next neighbor, receives a new KV
        block from the previous neighbor, and computes attention against it
Step 2-7: Same pattern — pass KV forward, receive new KV, compute attention

After 8 steps, every KV block has visited every GPU. Each GPU has computed
full attention for its Q tokens against ALL 1M tokens of K and V — without
ever holding more than 125K tokens of KV at any point.

Notice that Ring Attention doesn't split by heads at all. It splits purely by sequence position. This means it works regardless of how many KV heads the model has — whether it's 8 (Llama-70B with GQA), 1 (multi-query attention), or 64 (standard MHA). There is no head count limit on the number of GPUs you can use.

4.3 Step-by-Step Ring Attention

Let's trace through Ring Attention with our 8 GPU setup. Each GPU starts with 125K tokens (1M ÷ 8).

4.3.1 Initial State

INITIAL STATE:
══════════════
GPU 0: Q₀ [125K tokens], KV₀ [125K tokens]  ← tokens 0-124,999
GPU 1: Q₁ [125K tokens], KV₁ [125K tokens]  ← tokens 125,000-249,999
GPU 2: Q₂ [125K tokens], KV₂ [125K tokens]  ← tokens 250,000-374,999
...
GPU 7: Q₇ [125K tokens], KV₇ [125K tokens]  ← tokens 875K-999K

Q stays fixed on each GPU throughout. Only KV blocks rotate.

4.3.2 Step 0: Local Attention

STEP 0 — Each GPU computes attention with its LOCAL KV:
═══════════════════════════════════════════════════════
GPU 0: Attention(Q₀, KV₀) → partial output O₀, running stats (m₀, l₀)
GPU 1: Attention(Q₁, KV₁) → partial output O₁, running stats (m₁, l₁)
...
GPU 7: Attention(Q₇, KV₇) → partial output O₇, running stats (m₇, l₇)

Each GPU uses FlashAttention internally. The "running stats" are the
online softmax values: running max (m) and running sum of exponentials (l).

THEN: each GPU sends its KV block to the next GPU in the ring.
  GPU 0 sends KV₀ → GPU 1,  and receives KV₇ from GPU 7
  GPU 1 sends KV₁ → GPU 2,  and receives KV₀ from GPU 0
  ...
  GPU 7 sends KV₇ → GPU 0,  and receives KV₆ from GPU 6

4.3.3 Step 1: First Rotation

STEP 1 — Each GPU now has a DIFFERENT KV block:
════════════════════════════════════════════════
GPU 0: Q₀ (fixed) + received KV₇ (from GPU 7)
GPU 1: Q₁ (fixed) + received KV₀ (from GPU 0)
...
GPU 7: Q₇ (fixed) + received KV₆ (from GPU 6)

Each GPU computes attention against the NEW KV block and merges the
result with its existing partial output using online softmax correction.

THEN: pass KV to the next GPU again.

4.3.4 Steps 2 through 7: Completing the Ring

Each subsequent step follows the same pattern:
  1. Receive a new KV block from the previous neighbor
  2. Compute attention of local Q against the new KV block
  3. Merge the result into the running output using online softmax
  4. Send the KV block to the next neighbor

AFTER ALL 8 STEPS:
══════════════════
GPU 0 has computed: Attention(Q₀, KV_full) — all 1M tokens
GPU 1 has computed: Attention(Q₁, KV_full) — all 1M tokens
...
GPU 7 has computed: Attention(Q₇, KV_full) — all 1M tokens

Every GPU has computed FULL attention for its Q chunk against the
entire 1M token sequence — without ever holding more than 125K
tokens of KV in memory at any point.

4.4 Online Softmax Across KV Blocks

The mathematical trick that makes Ring Attention work is the same online softmax we saw in FlashAttention (Part 1, Section 1.4.6). Standard softmax requires knowing the maximum value and the sum of exponentials across the entire row of attention scores. But in Ring Attention, at step 0 we only have scores for 125K out of 1M tokens.

The solution is to maintain three running values for each query row — a running maximum (m), a running sum of exponentials (l), and a running weighted output (o). As each new KV block arrives, we update these values and correct the accumulated output:

STEP 0 — Process local KV₀ (125K tokens):
  Compute scores: s₀ = Q_row @ K₀ᵀ
  m₀ = max(s₀)
  l₀ = Σ exp(s₀ - m₀)
  o₀ = [exp(s₀ - m₀) / l₀] @ V₀

STEP 1 — Process received KV₇ (next 125K tokens):
  Compute scores: s₇ = Q_row @ K₇ᵀ
  m_new = max(s₇)
  
  Update global max:  m₁ = max(m₀, m_new)
  Correction factor:  α = exp(m₀ - m₁)   ← rescales old values
  
  Update running sum:    l₁ = α × l₀ + Σ exp(s₇ - m₁)
  Update running output: o₁ = α × o₀ + [exp(s₇ - m₁)] @ V₇

STEPS 2-7 — Same update pattern for each new KV block.

FINAL — After all 8 KV blocks:
  output = o_final / l_final
  
  This divides the accumulated weighted sum (o_final) by the total
  sum of exponentials (l_final), producing the exact softmax-weighted
  average of all V values across the entire 1M token sequence.

This is not an approximation. The correction factor α = exp(m_old - m_new) perfectly rescales all previous partial results whenever a new maximum is discovered. After processing all blocks, the final normalization by l_final produces the exact same result as computing attention over all 1M tokens at once.

4.5 Communication-Computation Overlap

Ring Attention has a useful property: while a GPU is computing attention on the current KV block, it can simultaneously send that block to the next GPU and receive the next block from the previous GPU. The compute and communication happen in parallel on separate hardware (tensor cores vs network interface), using CUDA streams.

4.5.1 The Numbers for Llama-70B (8 GPUs)

KV BLOCK SIZE (what gets transferred each step):
  K: [1, 125K, 8 heads, 128 dim] × 2 bytes = 0.26 GB
  V: [1, 125K, 8 heads, 128 dim] × 2 bytes = 0.26 GB
  Total KV per block: 0.51 GB

COMMUNICATION TIME (point-to-point, one neighbor):
  Within node (NVLink, 300 GB/s per direction):
    0.51 GB ÷ 300 GB/s = 1.7 ms per step

  Across nodes (InfiniBand HDR, 25 GB/s per direction):
    0.51 GB ÷ 25 GB/s = 20.4 ms per step

COMPUTE TIME (FlashAttention per step):
  Each GPU has all 64 query heads for 125K tokens.
  Per head, attention involves two matrix multiplications:
    QK^T: 2 × 125K × 125K × 128 FLOPs (multiply-add)
    score @ V: 2 × 125K × 125K × 128 FLOPs
    Total per head: 4 × 125K × 125K × 128 = 8.0e12 FLOPs

  All 64 heads: 64 × 8.0e12 = 5.12e14 FLOPs per step

  At ~175 TFLOPS (FlashAttention-2 realistic utilization on A100):
    5.12e14 ÷ 175e12 ≈ 2.9 seconds per step

  (We also omit FSDP weight gathering and Ulysses All-to-All overhead
  in these calculations — see note below.)

OVERLAP ANALYSIS:
  Within node:  compute (2.9 s) vs comm (1.7 ms)  → 1700× → fully hidden
  Across nodes: compute (2.9 s) vs comm (20.4 ms) → 142×  → fully hidden

4.5.2 Total Time for Ring Attention

RING ATTENTION TIMING (8 A100 GPUs, 1M tokens):
════════════════════════════════════════════════

PER LAYER:
  Steps: 8 (one per KV block)
  Compute per step: ~2.9 seconds
  Communication per step: 1.7 ms (completely hidden behind compute)
  Total per layer: 8 × 2.9 s ≈ 23 seconds

FOR 80 LAYERS:
  Attention: 23 s × 80 ≈ 31 minutes
  FFN + projections: ~4% additional (FFN is O(n), negligible vs O(n²) attention)
  Total forward pass: ~32 minutes on 8 A100 GPUs

FOR COMPARISON — at 128K tokens (more common in production):
  Compute per step: 128K²/1M² = 0.016× → ~47 ms per step
  Per layer: 8 × 47 ms = 376 ms
  80 layers: ~30 seconds  ← much more practical!

4.6 Ring Attention vs Ulysses: Why We Need Both

Now we have two approaches for distributing attention. A natural question is: if Ring Attention can scale to any number of GPUs and hides its communication behind compute, why would we ever use Ulysses?

The answer is compute efficiency. Ring Attention breaks the attention computation into N sequential steps, each processing a smaller block. FlashAttention's fused kernel is less efficient on these smaller blocks — the GPU cores are underutilized because smaller matrix multiplications don't generate enough parallel work to saturate all the streaming multiprocessors. Ulysses, by contrast, gives each GPU the full sequence for its assigned heads, so FlashAttention operates on the full-size attention matrix and runs at peak efficiency.

This isn't theoretical — the USP paper benchmarked both approaches on 8xA100-SXM4 NVLink GPUs and found that Ulysses achieves 1.4-1.8× higher throughput than Ring Attention on the same hardware:

BENCHMARK: Attention throughput on 8xA100 NVLink (from USP paper, Table 3)
═══════════════════════════════════════════════════════════════════════════
                        32K seq len      128K seq len
Ulysses-only (8×1):    7.36 iters/sec   1.01 iters/sec
Ring-only (1×8):        5.14 iters/sec   0.56 iters/sec
Ratio:                  1.43×            1.80×

Ulysses is significantly faster on the same hardware.

4.6.1 Why Ring is Slower: The Block Size Effect

The performance gap comes from how GPUs handle matrix multiplications of different sizes. An A100 has 108 streaming multiprocessors (SMs), each capable of running multiple warps in parallel. To fully utilize all SMs, the matrix multiplication needs to be large enough to generate sufficient parallel work:

ULYSSES (each GPU processes full sequence for its heads):
  FlashAttention processes: [1M tokens, 128 dim] × [1M tokens, 128 dim]
  This is a massive computation that fully saturates all 108 SMs.
  GPU utilization: high (large matrices, lots of parallel work)

RING ATTENTION (each GPU processes one block per step):
  FlashAttention processes: [125K tokens, 128 dim] × [125K tokens, 128 dim]
  This is 64× less work per step (1M² vs 125K² = 64× reduction)
  GPU utilization: lower (smaller matrices, SMs may be underutilized)
  Plus: 8 sequential steps with kernel launch overhead between each

The total FLOPs are the same, but the GPU executes them less efficiently
when they're broken into smaller sequential chunks.

4.6.2 Side-by-Side Comparison

4.7 Scaling Beyond 8 GPUs

Ring Attention's biggest advantage over Ulysses isn't speed — it's the ability to scale to any number of GPUs. At 8 GPUs, both approaches use the same amount of memory per GPU. But Ring Attention can go to 16, 32, or 64 GPUs, reducing per-GPU memory linearly.

4.7.1 KV Cache: Same at 8 GPUs, Different at Scale

At 8 GPUs, the KV cache per GPU is identical for both approaches. Let's verify the math:

ULYSSES (8 GPUs):
  After All-to-All, each GPU has the full 1M sequence for 1 KV head.
  KV cache per GPU: 1,000,000 × 1 head × 128 dim × 2 (K+V) × 80 layers × 2 bytes
                  = 1,000,000 × 256 × 80 × 2
                  = 40,960,000,000 bytes = 40.96 GB

RING ATTENTION (8 GPUs):
  Each GPU caches KV for its local 125K tokens, all 8 heads.
  KV cache per GPU: 125,000 × 8 heads × 128 dim × 2 (K+V) × 80 layers × 2 bytes
                  = 125,000 × 2,048 × 80 × 2
                  = 40,960,000,000 bytes = 40.96 GB

Same! 1M × 1 head = 125K × 8 heads = 1M head-token pairs per GPU.

The difference appears when we scale beyond 8 GPUs — something only Ring Attention can do:

4.7.2 Memory at Different Scales

Everything scales linearly. Double the GPUs, halve the memory per GPU.

4.7.3 The Compute Tradeoff

More GPUs means more ring steps, but each step has less work because each GPU holds fewer Q tokens and each KV block is smaller. The total compute per GPU decreases linearly:

COMPUTE PER GPU (attention, 80 layers, at ~175 TFLOPS):
  Total attention FLOPs: 64 heads × 4 × 1M × 1M × 128 × 80 layers = 2.62e18
  Per GPU = Total / N_GPUs

  8 GPUs:  3.28e17 FLOPs/GPU → ~31 minutes
  16 GPUs: 1.64e17 FLOPs/GPU → ~16 minutes
  32 GPUs: 8.19e16 FLOPs/GPU → ~8 minutes
  64 GPUs: 4.10e16 FLOPs/GPU → ~4 minutes

Each doubling of GPUs halves the compute time per GPU.

However, there's a subtlety: with more GPUs, each step processes a smaller matrix (fewer queries × fewer keys). Smaller matrix multiplications have lower GPU utilization because there's less parallel work to saturate all 108 streaming multiprocessors on the A100. This is why the USP paper found that Ulysses (which keeps full-size matrices) achieves 1.4-1.8× higher throughput than Ring on the same hardware. The theoretical linear speedup may only translate to ~70-85% efficiency in practice.

4.8 Causal Masking and Load Balancing

For decoder-only models like Llama, causal attention means token i can only attend to tokens 0 through i. This creates an opportunity to skip computation, but also a load balancing challenge.

4.8.1 Skipping Fully-Masked Blocks

CAUSAL MASKING IN RING ATTENTION (8 GPUs):
══════════════════════════════════════════
GPU 0 has Q for tokens 0 through 124,999 (the first 125K tokens).
When it receives KV₇ (tokens 875K-999K):
  All of GPU 0's tokens are EARLIER than GPU 7's tokens.
  Causal mask blocks all attention. This entire block can be SKIPPED.

Blocks each GPU actually needs to process:
  GPU 0 (first 125K tokens):   only KV₀                    → 1 block
  GPU 1 (tokens 125K-249K):    KV₀, KV₁                    → 2 blocks
  GPU 2 (tokens 250K-374K):    KV₀, KV₁, KV₂               → 3 blocks
  ...
  GPU 7 (tokens 875K-999K):  all 8 blocks                 → 8 blocks

Total blocks: 1+2+3+...+8 = 36 (vs 64 without masking)
Savings: 36/64 = 56% → 44% compute saved overall

4.8.2 The Load Imbalance Problem

With contiguous assignment, GPU 7 processes 8 blocks while GPU 0 processes only 1. The wall-clock time is determined by the slowest GPU (GPU 7), so GPUs 0-6 sit idle waiting. The effective throughput is limited by GPU 7's workload of 8 blocks, even though the average is only 4.5 blocks per GPU.

4.8.3 Zigzag (Striped) Load Balancing

The standard solution, used in both the Striped Attention paper and the USP paper (and implemented in Megatron-LM), is to reorder how blocks are assigned to GPUs. Instead of giving each GPU a contiguous chunk, we pair early blocks with late blocks.

Here's how it works concretely. First, we split the 1M sequence into 16 blocks (2× the number of GPUs):

16 blocks of 62.5K tokens each:
  B0  B1  B2  B3  B4  B5  B6  B7  B8  B9  B10 B11 B12 B13 B14 B15
  ←── early tokens ──────────────────────────────── late tokens ──→

Then each GPU gets one early block paired with one late block. Since each GPU gets 2 blocks of 62.5K tokens each, every GPU still has 125K tokens total — the same as the contiguous assignment:

ZIGZAG ASSIGNMENT (each GPU: 1 early block + 1 late block = 125K tokens):
═════════════════════════════════════════════════════════════════════════
GPU 0: B0  + B15   (tokens 0-62.4K      + tokens 937.5K-999K)
GPU 1: B1  + B14   (tokens 62.5K-124.9K + tokens 875K-937.4K)
GPU 2: B2  + B13   (tokens 125K-187.4K  + tokens 812.5K-874.9K)
GPU 3: B3  + B12   (tokens 187.5K-249.9K + tokens 750K-812.4K)
GPU 4: B4  + B11   (tokens 250K-312.4K  + tokens 687.5K-749.9K)
GPU 5: B5  + B10   (tokens 312.5K-374.9K + tokens 625K-687.4K)
GPU 6: B6  + B9    (tokens 375K-437.4K  + tokens 562.5K-624.9K)
GPU 7: B7  + B8    (tokens 437.5K-499.9K + tokens 500K-562.4K)

Now let's count the compute load carefully. With 16 small blocks, each block's Q tokens need to attend to all blocks at earlier or equal positions (causal mask). We count each block interaction as one unit of work:

COMPUTE LOAD WITH ZIGZAG (counting small block interactions):
═════════════════════════════════════════════════════════════
GPU 0: B0 attends to blocks ≤ B0 → 1 small block
        B15 attends to blocks ≤ B15 → 16 small blocks
        Total: 1 + 16 = 17 small blocks

GPU 1: B1 attends to ≤ B1 → 2 small blocks
        B14 attends to ≤ B14 → 15 small blocks
        Total: 2 + 15 = 17 small blocks

GPU 2: B2 → 3,  B13 → 14  → Total: 17 small blocks
GPU 3: B3 → 4,  B12 → 13  → Total: 17 small blocks
GPU 4: B4 → 5,  B11 → 12  → Total: 17 small blocks
GPU 5: B5 → 6,  B10 → 11  → Total: 17 small blocks
GPU 6: B6 → 7,  B9 → 10   → Total: 17 small blocks
GPU 7: B7 → 8,  B8 → 9    → Total: 17 small blocks

Every GPU processes exactly 17 small block interactions. Perfectly balanced!

To understand the speedup, we need to convert small blocks to big-block equivalents. A small block has half the tokens of a big block, so the attention computation for one small-block pair (62.5K × 62.5K) is 1/4 the FLOPs of a big-block pair (125K × 125K) — because FLOPs scale as queries × keys.

CONVERTING TO BIG-BLOCK EQUIVALENTS:
════════════════════════════════════
1 small-block interaction = (62.5K × 62.5K) = 1/4 of a big-block (125K × 125K)
17 small blocks per GPU = 17/4 = 4.25 big-block equivalents per GPU

WALL-CLOCK TIME COMPARISON:
═══════════════════════════
WITHOUT zigzag (contiguous, big blocks):
  GPU 0: 1 block, GPU 1: 2, ..., GPU 7: 8 blocks
  Total across all GPUs: 1+2+...+8 = 36 big blocks
  Wall-clock = 8 big blocks (bottleneck is GPU 7)

WITH zigzag:
  Every GPU: 17 small blocks = 4.25 big-block equivalents
  Total across all GPUs: 8 × 17 = 136 small = 34 big-block equivalents
  Wall-clock = 4.25 big blocks (all GPUs finish together)

Speedup: 8 / 4.25 = 1.88×

Note: the total work is 36 vs 34 big-block equivalents — the small
difference comes from the diagonal blocks being slightly overestimated
in the big-block counting (each diagonal block is partially masked by
causal attention, but we count it as a full block in both cases).
The key point: zigzag distributes work evenly, cutting wall-clock
time nearly in half.

This works because attention computation doesn't depend on token order — only on positional encoding (RoPE). The tokens on each GPU are non-contiguous in the original sequence, but RoPE is applied based on the token's original position, not its physical location on the GPU. FlashAttention handles the resulting non-triangular causal mask pattern without issues.

4.8.4 Zigzag Speedup Converges to 2×

The zigzag speedup is nearly the same regardless of GPU count, and converges to exactly 2× as N grows. Here's why:

GENERAL FORMULA (N GPUs):
═════════════════════════
Split into 2N small blocks. Each small block has seq/(2N) tokens.
GPU i gets block i and block (2N-1-i).

Work for GPU i (counting each block interaction as 1 unit):
  Block i needs blocks 0..i → (i+1) small blocks
  Block (2N-1-i) needs blocks 0..(2N-1-i) → (2N-i) small blocks
  Total: (i+1) + (2N-i) = 2N + 1 small blocks  ← same for ALL GPUs!

Converting to big-block equivalents:
  1 small block = (seq/(2N))² FLOPs
  1 big block   = (seq/N)² FLOPs = 4 × small block FLOPs
  So: (2N+1) small blocks = (2N+1)/4 big-block equivalents

Without zigzag: wall-clock = N big blocks (bottleneck GPU)
With zigzag:    wall-clock = (2N+1)/4 big-block equivalents

Speedup = N / ((2N+1)/4) = 4N / (2N+1)

As N → ∞:  4N / (2N+1) → 4N / 2N = 2

  N=8:   4×8  / 17 = 1.88×
  N=16:  4×16 / 33 = 1.94×
  N=32:  4×32 / 65 = 1.97×
  N=64:  4×64 / 129 = 1.98×
  N→∞:                 2.00×

4.9 The Complete Memory Picture

Let's bring together everything from Parts 1 through 4:

4.9.1 What the Headroom Means: Batch Capacity

All the numbers above assume batch_size=1. In production, the headroom is what allows us to serve multiple requests concurrently. Weights are shared across requests — only KV cache and activations scale with batch size:

BATCH CAPACITY AT 32 GPUs (Llama-70B, 1M context):
═══════════════════════════════════════════════════
Fixed cost:  weights (4.4 GB) + overhead (2.0 GB) = 6.4 GB
Per request: KV cache (10.2 GB) + activations (0.5 GB) = 10.7 GB

batch=1:  6.4 + 10.7 = 17.1 GB  ✓
batch=2:  6.4 + 21.4 = 27.8 GB  ✓
batch=3:  6.4 + 32.1 = 38.5 GB  ✓
batch=4:  6.4 + 42.8 = 49.2 GB  ✓
batch=5:  6.4 + 53.5 = 59.9 GB  ✓
batch=6:  6.4 + 64.2 = 70.6 GB  ✓
batch=7:  6.4 + 74.9 = 81.3 GB  ❌ (exceeds 80 GB)

Max batch size: 6 concurrent 1M-token requests per GPU.

4.9.2 A Throughput Primer

LLM inference has two distinct phases with very different compute profiles. Let's walk through both.

4.9.3 Prefill Phase: Processing the Input

The prefill phase processes all input tokens through the model to build the KV cache. This is exactly the forward pass we've been calculating — all 1M tokens go through all 80 layers:

PREFILL TIME (Llama-70B, 1M tokens, A100 GPUs, without causal optimization):
════════════════════════════════════════════════════════════════════════════
  8 GPUs:  ~31 minutes  (23 s/layer × 80 layers)
  32 GPUs: ~8 minutes   (5.9 s/layer × 80 layers)
  64 GPUs: ~4 minutes   (2.9 s/layer × 80 layers)
  128 GPUs: ~2 minutes  (1.5 s/layer × 80 layers)

With zigzag scheduling (causal models, ~1.9× speedup):
  8 GPUs:  ~16.5 minutes  (1.88× speedup)
  32 GPUs: ~4.1 minutes   (1.97× speedup)
  64 GPUs: ~2.0 minutes   (1.98× speedup)

  The zigzag speedup approaches 2× as GPU count grows: 4N/(2N+1) → 2.

For comparison — 128K tokens (more common production context):
  8 GPUs:  ~30-50 seconds (without zigzag), ~16-26 seconds (with zigzag)
  32 GPUs: ~8-12 seconds (without zigzag), ~4-6 seconds (with zigzag)

Yes, 1M token prefill takes minutes, not seconds. This is the reality of O(n²) attention at extreme sequence lengths. It's why most production deployments use 128K or shorter contexts, and why 1M-token contexts are reserved for specialized use cases.

4.9.4 Decode Phase: Generating Output Tokens

After prefill, the model generates output tokens one at a time. Each new token requires a forward pass through all 80 layers, but with a crucial difference: only 1 new query token is processed, and the KV cache (already built during prefill) is reused.

4.9.4.1 What happens for each new token

For each of the 80 layers:
  1. FSDP AllGather: reconstruct full layer weights from shards
  2. Compute Q, K, V for the new token (tiny matmul: [1, 8192] × [8192, 8192])
  3. Append new K, V to the KV cache
  4. Attention: Q (1 token) attends to entire KV cache (1M+ tokens)
  5. FFN: process 1 token through the feed-forward network
  6. Free gathered weights, move to next layer

4.9.4.2 Why decode is memory-bandwidth bound

With only 1 query token, the matrix multiplications are tiny — the GPU spends most of its time loading data from HBM, not computing. Let's break down the time per layer:

DECODE TIME BREAKDOWN PER LAYER (32 A100 GPUs, 1M context):
════════════════════════════════════════════════════════════

1. Weight loading (dominates!):
   Full layer weights: ~856M params × 2 bytes = 1.71 GB
   Reading 1.71 GB from HBM at 2 TB/s = 0.86 ms per layer
   
2. KV cache reading:
   Local KV per layer: 31.25K tokens × 8 heads × 128 dim × 2 (K+V) × 2 bytes = 128 MB
   Reading 128 MB at 2 TB/s = 0.064 ms per layer
   
3. Attention compute:
   FLOPs: 64 heads × 4 × 1 × 31.25K × 128 = 1.02e9
   At 175 TFLOPS: 0.006 ms  ← negligible!

4. FFN compute:
   Also memory-bound (loading FFN weights for 1 token)
   Included in weight loading above.

Total per layer: ~0.92 ms (dominated by weight loading)
Total for 80 layers: ~74 ms per output token

4.9.4.3 Decode at different configurations

DECODE LATENCY PER TOKEN — FSDP vs TP (Llama-70B, 1M context, A100):
════════════════════════════════════════════════════════════════════

With FSDP (our setup — must AllGather full layer weights each token):
  Weight load per layer: 1.71 GB at 2 TB/s = 0.86 ms
  80 layers: ~69 ms just for weight loading
  Plus KV cache reads + communication: ~74-89 ms total per token

With Tensor Parallelism (each GPU already has its weight shard):
  Weight load per layer: 214 MB at 2 TB/s = 0.11 ms  (8× less!)
  80 layers: ~9 ms for weight loading
  Plus KV cache reads: ~9-15 ms total per token

4.9.4.4 How attention works during decode with distributed KV cache

During decode, the KV cache is already distributed across GPUs from the prefill phase. Each GPU holds KV for its chunk of the sequence. Meta's Context Parallelism paper describes a "pass-Q" variant for decode: the new token's Q embedding is passed around the ring (Q is tiny — just 16 KB for one token). Each GPU forwards Q to the next neighbor while simultaneously computing partial attention against its local KV cache. Partial outputs are then gathered back via All-to-All to produce the final result.

4.9.5 End-to-End Example

For causal (decoder-only) models, zigzag scheduling reduces the prefill attention time by ~1.88× (from Section 4.8.3). Let's apply this to our estimates:

EXAMPLE: 1M input + 500 output tokens, 32 GPUs, batch_size=6:
══════════════════════════════════════════════════════════════
Prefill (without zigzag): ~8 minutes
Prefill (with zigzag, 1.97× speedup): ~8 / 1.97 ≈ 4.1 minutes
Decode:  500 tokens × ~15 ms/token (with TP decode) = ~7.5 seconds

Total per batch:  4.1 min + 0.1 min = ~4.2 minutes
Throughput: 6 requests / 4.2 min ≈ 1.4 requests/minute

For 128K context + 500 output tokens, 32 GPUs, batch_size=6:
Prefill (with zigzag): ~4-6 seconds
Decode:  500 × ~15 ms/token = ~7.5 seconds

Total per batch:  ~12-14 seconds
Throughput: 6 / 13 s ≈ 28 requests/minute  ← much more practical!

4.10 Summary

What's Next

Weights are sharded. Sequence is sharded. Attention is ringed. Now let's orchestrate them all.

@misc{fofadiya2026ringattention,
  author       = {Darshan Fofadiya},
  title        = {Ring Attention — The Illustrated Guide to LLM Inference at Scale, Part 4},
  year         = {2026},
  url          = {https://darshanfofadiya.com/llm-inference/ring-attention.html}
}

Comments