Transformers rarely fail loudly. They stall—one saturated lane on a multi-lane highway—while everything else idles. The art is knowing which lane is clogged: quadratic attention, KV cache memory, kernel launch overhead, CPU–GPU transfers, or network collectives. Once you see it, fixes are mechanical. Until you see it, you waste weeks turning the wrong knobs.

What is the Transformer bottleneck, in practice?

Short answer: it’s the slowest component dominating end-to-end step or token time—often attention math, memory movement (KV cache), or communication. You don’t guess it; you profile it.

• Algorithmic: O(n²) attention during prefill; softmax and matmuls dominate.
• Memory-bound: KV cache size/bandwidth limits decoding; allocator fragmentation hurts reuse.
• Kernel/pathology: too many small kernels, unfused ops, poor SM occupancy.
• I/O-bound: dataloader stalls, CPU tokenization, CPU–GPU copies without pinning/overlap.
• Distributed: NCCL all-reduce/all-gather stalls; pipeline bubbles; imbalance between ranks.
• Sampler/decoding: temperature/top‑p/beam logic on CPU; Python overhead per token.

How do you locate the bottleneck quickly?

Short answer: record a timeline and match symptoms to causes. If GPU SMs are low but HBM is high, you’re memory‑bound; if both are low, you’re I/O or Python‑bound; if both are high, you’re compute‑bound (often attention).

• Tools: PyTorch Profiler + TensorBoard, Nsight Systems, Nsight Compute. For distributed, add NCCL debug/timeline.
• Signals to read:
  • GPU utilization vs HBM bandwidth: compute‑bound vs memory‑bound.
  • Kernel count and gaps: many micro‑kernels and idle gaps imply fusion/graphs needed.
  • CPU time: tokenizer, dataset transforms, Python callbacks during decode.
  • Comm bars: long all‑reduce/all‑gather vs compute overlap indicates synchronization hotspots.
• Break down by phase: prefill (context ingest), decode (token‑by‑token), or training step (fwd+bwd+optimizer).

Is attention the problem, or is it memory bandwidth?

Short answer: prefill is usually attention‑compute bound; decode is often KV cache bandwidth/size bound.

• Prefill symptoms: O(n²) attention time grows fast with context length; FlashAttention or chunked attention helps.
• Decode symptoms: token time flatlines as sequence grows; KV reads dominate; paged attention, GQA/MQA, or cache quantization helps.
• Quality guardrails: some speedups (e.g., linearized attention) change model behavior; others (FlashAttention) are numerically safe within training configs.

Context length and the bottleneck: what actually scales?

Short answer: prefill cost scales with sequence length squared; decode scales roughly linearly with sequence length through KV cache reads.

• Prefill (O(n²)): doubling context ~quadruples attention work; batching amortizes cost across requests.
• Decode (O(n)): each new token attends over n keys/values; KV cache memory grows with n.
• Mitigations: FlashAttention/tiling for prefill; paged attention, GQA/MQA, KV quantization, sliding windows for decode.

How much memory does the KV cache use?

Short answer: approximately 2 × L × H × n × d_head × bytes, for keys and values across layers.

Back‑of‑envelope (batch=1):

• Formula: memory_bytes ≈ 2 × L × H_kv × n × d_head × bytes_per_elem.
• Example (LLaMA‑like 7B: L=32, H_kv=32, d_head=128, n=8,000, FP16=2 bytes):
  • 2 × 32 × 32 × 8,000 × 128 × 2 ≈ 4.19e9 bytes ≈ 3.9 GiB.
• With GQA (H_kv=8): ~4× smaller, ≈ 1.0 GiB for the same n.
• Multiply by batch size to approximate multi‑request decoding memory.

Training vs inference: where are the typical choke points?

Short answer: training hits attention compute, activation memory, and comms; inference hits KV cache, kernel overhead, and sampler logic.

Training bottlenecks and practical fixes

• Use BF16 or FP16 with loss scaling; prefer BF16 on newer GPUs for stability.
• Enable FlashAttention in forward/backward; confirm kernels are actually used in profiler.
• Activation checkpointing to trade compute for memory; combine with sequence parallelism to raise batch/seq length.
• Fused optimizers and norm layers; reduce kernel count. Consider CUDA Graphs for stable shapes.
• Distributed: match parallelism to model shape (tensor parallel for wide layers; pipeline for deep stacks). Overlap all‑reduce with backward when possible.
• Dataloader: pinned memory, prefetch, CPU tokenization off critical path, shard files, and cache frequently accessed samples.

Inference bottlenecks and practical fixes

• Batching: continuous batching to amortize prefill; prefix caching for shared prompts.
• KV cache strategy: GQA/MQA to reduce K/V heads, paged attention to limit fragmentation, and cache quantization (e.g., INT8/INT4) where quality permits.
• Decoding: speculative decoding with a draft model; keep sampling on GPU; avoid Python per‑token overhead.
• Kernel/runtime: FlashAttention for prefill, fused sampling kernels, CUDA Graphs to remove launch overhead, ensure large enough batch/sequence tiles.
• I/O: stream inputs, avoid sync points, and ensure tokenization is batched and off the hot path.

Which fixes actually move the needle first?

Short answer: apply the fix that addresses your phase‑dominant cost—FlashAttention for prefill; cache/batching/quantization for decode; checkpointing/parallelism for training.

• If prefill dominates: enable FlashAttention/FlashAttention‑2; increase batch size; ensure fused softmax and matmuls; consider short‑chunked attention when quality allows.
• If decode dominates: adopt paged attention and GQA/MQA; enable cache quantization; implement continuous batching and GPU‑resident sampling.
• If memory limits batch: activation checkpointing, ZeRO optimizer states, gradient accumulation, sequence parallelism.
• If CPU shows up: move tokenization/filters to separate workers; use pinned memory; prefetch aggressively.
• If comms block: rebalance pipeline; adjust micro‑batch count to avoid bubbles; overlap compute/collectives; verify NCCL bandwidth and topology.

Long‑context transformers: fix the right problem

Short answer: decide if you need exact softmax attention at long ranges or an approximation; then combine positional strategy with memory‑aware kernels.

• Positional handling: RoPE scaling or interpolation; ALiBi to bias attention with minimal overhead; verify retention quality on target tasks.
• Approximations: Performer, Linformer, BigBird, Longformer provide sub‑quadratic attention with quality tradeoffs—validate on your data.
• Windows/recency: sliding window or block sparse attention can bound KV size and speed decode without retraining some models.
• Retrieval: retrieve‑then‑read often beats brute‑force long context for factual tasks and reduces n² prefill cost.

Diagnose like an operator: a 30‑minute profiling plan

Short answer: capture one representative run, tag phases, and classify each minute by the dominant resource.

1. Warm‑up: run a short pass to stabilize kernels and allocator behavior.
2. Record: enable PyTorch Profiler with CUDA and stack traces for 60–90 seconds (or Nsight Systems for finer GPU timelines).
3. Mark phases: tag prefill vs decode (inference) or fwd/bwd/optimizer (training) in your logs.
4. Classify resources:
   • Compute‑bound: high SM occupancy, high FP/BF utilization, attention kernels top of list.
   • Memory‑bound: high HBM bandwidth, low SM occupancy, frequent large memory copies.
   • Kernel launch overhead: many tiny kernels with gaps; Python shows in call stacks.
   • Comm‑bound: visible all‑reduce/all‑gather bars with low overlap; ranks idle waiting.
5. Apply one fix per class:
   • Compute: FlashAttention, increase tile sizes, upgrade precision path (BF16).
   • Memory: paged attention, GQA/MQA, cache quantization, larger contiguous blocks.
   • Launch/Python: fused kernels, CUDA Graphs, move sampling to GPU, batch tokenizer.
   • Comm: overlap, topology‑aware placement, adjust micro‑batches to fill pipeline.
6. Re‑measure: confirm the top 3 kernels/time slices changed; don’t stack fixes blindly.

Common mistakes that keep the bottleneck hidden

Short answer: measuring averages only, optimizing the wrong phase, and trusting flags that aren’t actually active.

• Assuming FlashAttention is used because a config says so; verify kernel names in the timeline.
• Benchmarking short prompts only, then deploying long‑context workloads.
• Batching without continuous batching or prefix caching, underutilizing the GPU at peak times.
• Leaving sampling on CPU with per‑token Python loops.
• Ignoring allocator fragmentation; paged attention or pooling can recover large contiguous memory.
• Tuning beam search when top‑p is the production default (or vice versa).