The Attention Bottleneck in Decode-Phase Inference

Transformer inference has two phases: the prefill (processing the entire prompt in parallel) and decode (generating tokens one at a time autoregressively). What most tutorials gloss over is that decode is fundamentally memory-bound, not compute-bound.

During decode, each new token requires attending to every previous token in the sequence. The naive approach recomputes all key-value projections from scratch at every step, making generation O(n²) in sequence length. The KV-cache solves this by storing the K and V tensors from all previous positions:

# Simplified KV-cache during decode
# For each layer l, head h, and position t:
# K_cache[l][h][t] = W_k @ hidden_states[t]
# V_cache[l][h][t] = W_v @ hidden_states[t]

# At decode step T, attention for new token:
q = W_q @ hidden_states[T] # (1, d_head)
k = K_cache[l][h][:T+1] # (T+1, d_head)
v = V_cache[l][h][:T+1] # (T+1, d_head)
attn = softmax(q @ k.T / sqrt(d_head)) @ v

This reduces generation to O(n) per step. But the memory cost is brutal. For a 70B model (80 layers, 64 heads, 128 head dim, FP16):

KV-cache per token = 2 * n_layers * n_kv_heads * d_head * 2 bytes
 = 2 * 80 * 8 * 128 * 2 (GQA with 8 KV heads)
 = 327,680 bytes ≈ 320 KB per token

For a 4096-token sequence:
320 KB * 4096 = 1.28 GB per request

For 100 concurrent requests:
1.28 GB * 100 = 128 GB just for KV-cache

The model weights are ~140 GB in FP16. The KV-cache can rival the model itself in memory consumption, which creates a significant resource management challenge.

The Fragmentation Problem

Traditional serving frameworks pre-allocate a contiguous block of GPU memory for each request's KV-cache, sized for the maximum possible sequence length. If your max is 4096 tokens but the average completion is 500 tokens, you're wasting ~87% of allocated KV-cache memory.

Measurements on production clusters show average actual KV-cache utilization at 22%. This means 78% of the most expensive resource is being discarded.

# What naive allocation looks like:
def allocate_kv_cache(max_seq_len, batch_size, n_layers, n_heads, d_head):
 # Pre-allocate for worst case
 k_cache = torch.zeros(
 batch_size, n_layers, n_heads, max_seq_len, d_head,
 dtype=torch.float16, device='cuda'
 )
 v_cache = torch.zeros_like(k_cache)
 return k_cache, v_cache

# For batch=64, max_seq=4096 on a 70B model:
# This immediately claims ~82 GB of VRAM
# Most of which sits at zero throughout the request lifecycle

PagedAttention: OS Concepts Meet GPU Memory

PagedAttention from the vLLM team addressed this problem with a fundamentally different approach. The insight is deceptively simple: treat KV-cache memory the same way operating systems treat virtual memory.

Instead of contiguous per-request allocations, PagedAttention divides KV-cache into fixed-size blocks (typically 16 tokens each). Each request maintains a block table mapping logical positions to physical block locations in GPU memory, exactly like a page table maps virtual to physical addresses.

# Conceptual PagedAttention block table
# Each block stores KV for 16 tokens across all layers and heads

class BlockTable:
 def __init__(self, block_size=16):
 self.block_size = block_size
 self.table = [] # List of physical block indices

 def get_physical_block(self, logical_position):
 block_idx = logical_position // self.block_size
 block_offset = logical_position % self.block_size
 physical_block = self.table[block_idx]
 return physical_block, block_offset

# Kernel pseudocode for paged attention
# For each query token:
# For each block in request's block_table:
# Load K, V from physical block
# Compute partial attention scores
# Accumulate with online softmax (like FlashAttention)

The benefits cascade:

  1. No internal fragmentation: Blocks are allocated on demand. A 500-token response uses exactly ceil(500/16) = 32 blocks, not the full max_seq_len worth.

  2. Copy-on-write sharing: For beam search or parallel sampling, multiple sequences that share a common prefix can point to the same physical blocks. When a sequence diverges, only the diverging block gets copied.

  3. Dynamic memory: Blocks are freed immediately when a request completes, returned to the pool for the next request. No reserved-but-idle memory.

In a representative production cluster, switching from HuggingFace TGI to vLLM with PagedAttention increased concurrent request capacity from ~40 to ~110 on the same hardware. Throughput went from 480 to 1,340 tokens/second aggregate.

Continuous Batching: The Other Half of the Equation

Static batching waits for an entire batch of requests, processes them together, and doesn't start the next batch until all sequences in the current batch have finished generating. This means short completions sit idle while long ones keep going.

Continuous batching (also called iteration-level scheduling) is different. At each decode step, the scheduler can:

  • Evict completed sequences
  • Admit new sequences
  • Preempt lower-priority sequences if memory runs tight
# Simplified continuous batching scheduler
class ContinuousBatchScheduler:
 def __init__(self, max_batch_tokens, block_manager):
 self.waiting = deque() # Requests waiting to start
 self.running = [] # Currently generating
 self.max_batch_tokens = max_batch_tokens
 self.block_manager = block_manager

 def schedule(self):
 # Remove completed sequences
 self.running = [s for s in self.running if not s.is_finished()]

 # Try to admit new sequences
 while self.waiting:
 next_req = self.waiting[0]
 total_tokens = sum(s.get_len() for s in self.running)
 total_tokens += next_req.prompt_len

 if total_tokens > self.max_batch_tokens:
 break
 if not self.block_manager.can_allocate(next_req):
 break # Not enough KV-cache blocks

 self.running.append(self.waiting.popleft())
 self.block_manager.allocate(self.running[-1])

 return self.running

The practical impact: on a workload with mixed short (50-token) and long (2000-token) completions, continuous batching yields 2.8x higher throughput compared to static batching. The short requests weren't hostage to the long ones anymore.

FlashAttention-3: Because Memory Bandwidth Is the Real Enemy

FlashAttention rewrites the attention computation to be IO-aware, minimizing reads/writes to GPU HBM (High Bandwidth Memory) by keeping data in SRAM (the fast on-chip memory) as long as possible. Version 3 adds several tricks specific to Hopper architecture GPUs (H100):

Standard Attention:
 1. Compute S = Q @ K.T → write S to HBM
 2. Compute P = softmax(S) → write P to HBM
 3. Compute O = P @ V → write O to HBM
 Total HBM accesses: O(N² * d)

FlashAttention-3:
 1. Tile Q, K, V into SRAM-sized blocks
 2. For each tile: compute local S, P, O in SRAM
 3. Accumulate with online softmax rescaling
 4. Write final O to HBM once
 Total HBM accesses: O(N * d), linear, not quadratic

FA3 on H100s specifically exploits:

  • Warp specialization: Separate warps for MMA (matrix multiply-accumulate) and softmax, overlapping compute with data movement
  • FP8 support: Native FP8 tensor cores for the QK^T matmul, with FP32 accumulation
  • Asynchronous block-wise softmax: Using the TMA (Tensor Memory Accelerator) for async global-to-shared memory copies

For reproducible results, use the configurations and measurements reported in the FlashAttention 3 paper and benchmark the target workload. This guide does not provide an independently reproducible measurement for the specific H100 SXM configuration previously listed here.

The integration with vLLM's PagedAttention is where it gets interesting, the flash_attn_with_kvcache kernel handles the block table indirection directly, so you get both paged memory management and IO-aware attention in a single fused kernel.

Speculative Decoding: Trading Compute for Latency

The idea: use a small "draft" model to generate N candidate tokens quickly, then verify them all at once with the large "target" model in a single forward pass. If the target model agrees with the draft tokens (checked via a modified rejection sampling scheme), you've effectively generated N tokens in the time it takes the target to do one forward pass.

# Speculative decoding loop
def speculative_decode(target_model, draft_model, prompt, gamma=5):
 tokens = prompt
 while not is_done(tokens):
 # Draft phase: generate gamma candidates cheaply
 draft_tokens = []
 draft_probs = []
 for _ in range(gamma):
 logits = draft_model(tokens + draft_tokens)
 p = softmax(logits[-1])
 t = sample(p)
 draft_tokens.append(t)
 draft_probs.append(p)

 # Verify phase: single forward pass on target
 all_logits = target_model(tokens + draft_tokens)

 # Acceptance: modified rejection sampling
 accepted = 0
 for i in range(gamma):
 q = softmax(all_logits[len(tokens) + i])
 p = draft_probs[i]
 r = random()
 if r < min(1, q[draft_tokens[i]] / p[draft_tokens[i]]):
 accepted += 1
 else:
 # Resample from adjusted distribution
 adjusted = normalize(max(0, q - p))
 tokens.append(sample(adjusted))
 break
 else:
 # All accepted, sample one more from target
 tokens.extend(draft_tokens)
 q = softmax(all_logits[-1])
 tokens.append(sample(q))
 continue
 tokens.extend(draft_tokens[:accepted])

The math works out so that the output distribution is provably identical to sampling from the target model alone. You're not sacrificing quality.

In practice, the acceptance rate depends heavily on how well the draft model matches the target. Observed results include:

  • Llama-3-8B drafting for Llama-3-70B: ~75% acceptance rate, ~2.3x speedup
  • Tiny custom 100M drafter fine-tuned on target outputs: ~65%, ~1.8x
  • Code completion tasks (high predictability): ~85%, ~3.1x

The sweet spot for gamma (speculation length) is usually 4-6. Go higher and the probability of at least one rejection per round approaches 1, wasting the extra draft compute.

Putting It All Together: Production Configuration

Here is a vLLM configuration that can bring a cluster from 12 tok/s/user to 45 tok/s/user:

python -m vllm.entrypoints.openai.api_server \
 --model meta-llama/Llama-3-70B-Instruct \
 --tensor-parallel-size 8 \
 --max-num-batched-tokens 32768 \
 --max-num-seqs 256 \
 --block-size 16 \
 --swap-space 4 \
 --gpu-memory-utilization 0.92 \
 --enable-prefix-caching \
 --use-v2-block-manager \
 --speculative-model meta-llama/Llama-3-8B-Instruct \
 --num-speculative-tokens 5 \
 --quantization awq

Key decisions:

  • --max-num-batched-tokens 32768: Allows up to 32K tokens in flight per iteration. Tuned to just under the OOM threshold.
  • --enable-prefix-caching: Reuses KV-cache blocks for shared system prompts. With our 1,500-token system prompt, this saves ~480 KB per request.
  • --gpu-memory-utilization 0.92: Leave 8% headroom for activation memory spikes during prefill of long prompts.
  • --swap-space 4: 4 GB of CPU swap space per GPU for preempted requests, better than rejecting them outright.

The Numbers That Actually Matter

On 8xA100 80GB, serving Llama-3-70B-Instruct (AWQ 4-bit):

Metric Naive (TGI static) Optimized (vLLM) Improvement
Throughput (tok/s) 480 1,850 3.85x
Concurrent requests 40 180 4.5x
TTFT p50 820 ms 210 ms 3.9x
TTFT p99 4,200 ms 890 ms 4.7x
ITL p50 38 ms 22 ms 1.7x
Cost per 1M tokens $2.40 $0.62 3.87x

The cost reduction alone can pay for the engineering effort in about four days.

What's Coming Next

Several areas are worth watching for the next round of optimizations:

  • Disaggregated prefill/decode: Running prefill on different GPUs than decode, since they have completely different compute profiles. Splitwise and DistServe papers look promising.
  • Quantized KV-cache: Storing KV-cache in FP8 or even INT4, separate from model weight quantization. Early results show minimal quality loss with 2x memory reduction.
  • Ring attention: For ultra-long contexts (100K+), distributing the attention computation across multiple GPUs along the sequence dimension.

The field is moving incredibly fast. Six months ago PagedAttention was novel; now it's table stakes. The infrastructure advantage goes to whoever can integrate these techniques fastest into their serving stack.

For anyone running inference in production, examining KV-cache utilization is a high-value activity. That is likely where the majority of costs concentrate.