Data Movement Is All You Need: Optimizing Transformers

TL;DR

Most transformer operations are not bottlenecked by arithmetic — they are bottlenecked by data movement. This paper studies BERT-large training, classifies operators by their computation and dataflow, and shows that statistical normalization and element-wise operations consume 39% of an encoder layer’s runtime despite only 0.20% of its FLOPs. Moving tensors between HBM, caches, and registers is therefore a substantial optimization target. The authors then demonstrate that operator fusion and data layout optimizations can recover much of this wasted time, providing a principled framework for understanding where GPU cycles actually go in transformer workloads.

The Core Problem: Arithmetic Is Not the Bottleneck

As later hardware context, an A100 can offer 312 TFLOPS for eligible FP16 tensor-core matrix operations and around 2 TB/s of HBM bandwidth. That gives a tensor-core roofline ridge around 156 FLOPs/byte, not a compute ceiling for arbitrary pointwise operations. The paper’s measured setup is BERT-large training on Lassen V100 GPUs.

A useful supporting model is operational intensity, the ratio of floating-point operations to bytes moved:

An operation is compute-bound when I exceeds the machine’s compute-to-bandwidth ratio (the “ridge point” on the roofline model), and memory-bound when it falls below. The key finding: the majority of transformer operations — layer normalization, dropout, softmax, GELU activations, residual additions, and bias terms — are elementwise or reduction operations with I \approx O(1), placing them firmly in the memory-bound regime.

Large tensor contractions can be compute-bound, but shape and layout matter: the paper also finds small contractions such as QK^T with poor compute and memory utilization. A low FLOP count alone does not determine achieved bandwidth.

The Roofline Model Applied to Transformers

The roofline model provides supporting intuition: achievable performance is bounded by the applicable compute ceiling and memory bandwidth. This is an upper bound, not a measured runtime:

For a V100 SXM2, 125 TFLOPS of tensor-core throughput and 900 GB/s give a ridge near 139 FLOPs/byte for eligible matrix operations. General FP32 arithmetic instead has a 15.7 TFLOPS ceiling, giving a ridge near 17.4 FLOPs/byte. The figure below uses schematic operator intensities and selects the applicable ceiling; its values are not measurements from the paper.

In contrast, the large GEMM operations in the feed-forward layers (with dimensions [n, d] \times [d, 4d]) achieve I \approx O(d), which for typical hidden dimensions of 768–1024 places them well above the ridge point, allowing them to saturate the GPU’s compute capability.

Dataflow analysis and memory efficiency

The authors construct Stateful Dataflow multiGraphs (SDFGs) with DaCe. Nodes specify operators and edges specify data movement, exposing iteration spaces, dependencies and reuse. Sections 3.1–3.4 use BERT-large training on Lassen nodes with four 16 GB V100 GPUs, a batch of 8, sequence length 512, embedding dimension 1024 and 16 attention heads. They identify three operator classes:

1. Tensor contractions — the matrix multiplications in Q = XW_Q, K = XW_K, V = XW_V, attention scores A = QK^T / \sqrt{d_k}, and the feed-forward network layers. These have high arithmetic intensity and achieve good hardware utilization.

2. Statistical normalizations — softmax and layer normalization combine reductions with maps. Their access patterns in HBM, reduction strategy and vectorization affect performance.

3. Element-wise operations — biases, dropout, activations and residual connections do relatively little arithmetic per element and offer opportunities to reuse data through fusion.

The paper also evaluates memory usage efficiency (MUE), which combines avoidable traffic with achieved bandwidth:

Here Q is the theoretical I/O lower bound, D the bytes actually moved, and B/\hat{B} the achieved-to-peak bandwidth ratio. Reducing traffic and using the available bandwidth are separate goals; a kernel can fall short on either.

Optimization Strategies

The recipe combines compatible operator fusion, empirical layout benchmarking and end-to-end configuration selection.

Operator Fusion is the most impactful technique. Instead of launching separate GPU kernels for each elementwise operation (each requiring a full HBM read-write round-trip), fused kernels chain multiple operations together. The schematic below uses bias, GELU, dropout and a residual addition to illustrate intermediate traffic; it is not a measured paper kernel. The paper fuses operators with compatible iteration spaces and reduction dimensions, rather than every adjacent operation. The data stays in registers or shared memory between operations. For a fused sequence of k elementwise operations on a tensor of n elements, the memory traffic drops from O(kn) to O(n).

Data Layout Optimization benchmarks feasible permutations of transformer tensor dimensions, together with contraction algorithms and vectorization choices (Section 5). A layout fast for one contraction can be poor for the next reduction, so a generic NCHW-to-NHWC image-layout conversion is not the mechanism evaluated here.

Configuration Selection builds a directed graph whose edges represent measured runtimes for compatible input/output layouts and operator settings. A shortest path chooses an end-to-end configuration (Section 6.1). The search omits forward/backward dataflow connections, so it does not guarantee a globally optimal training layout. Communication overlap is not an evaluated contribution of this paper.

Why This Gets Worse Over Time

A crucial observation in the paper is that hardware trends are making this problem worse, not better. GPU compute throughput has been growing faster than memory bandwidth across successive generations:

These hardware figures are context, including generations after the paper; they are not its experiment configurations. They apply to eligible matrix operations. As the ratio increases, the corresponding ridge point moves rightward, meaning more operations fall into the memory-bound regime. An operation that was marginally compute-bound on a V100 may become memory-bound on an H100. This trend means that the data movement problem identified in the paper becomes more acute with each hardware generation, not less — making the proposed optimizations increasingly relevant.

Key Results

The profiling results are quantitative. On BERT-Large training with a V100:

• Tensor contractions: 99.80% of FLOPs, 61.0% of runtime
• Statistical normalizations: 0.17% of FLOPs, 25.5% of runtime
• Element-wise operations: 0.03% of FLOPs, 13.5% of runtime

These are the PyTorch BERT encoder-layer shares in Table 1.

The combined optimizations reduce data movement by up to 22.91%, achieve 1.30× speedup for a BERT encoder layer over general-purpose frameworks and 1.08× over DeepSpeed, and 1.19× for full BERT training. These are measured V100 results, not A100 or GPT-2 inference results.

For attention specifically, the softmax computation is a bottleneck because it requires a reduction across the sequence dimension (to compute the max for numerical stability and the normalization sum), followed by an elementwise exponentiation. This pattern — reduction followed by elementwise — is difficult to fuse naively but can benefit from online softmax algorithms that compute the result in a single pass.

Critical Analysis

Strengths:

• The paper provides one of the first rigorous, operation-level breakdowns of where time goes in transformer training. Prior work tended to profile at the layer level, obscuring the dominance of memory-bound operations.
• SDFG analysis and MUE separate redundant traffic from poor bandwidth use; the measured layout search then connects local kernel choices to the full layer.
• The optimization recipe can apply to other networks, while the reported performance gains belong to the evaluated BERT/V100 configurations.

Limitations:

• The optimizations proposed are relatively standard systems techniques (fusion and layout optimization). The paper’s contribution is more in the analysis than in the solutions.
• The evaluation was conducted on Lassen V100 hardware. Newer architectures with different compute-to-bandwidth ratios (e.g., H100 with its higher memory bandwidth from HBM3) may shift the bottleneck balance, though the qualitative conclusion — that non-GEMM operations are memory-bound — is unlikely to change.
• The paper does not address attention-specific optimizations like FlashAttention (Dao et al. 2022), which later demonstrated that even the “compute-bound” attention mechanism benefits from memory-aware tiling that avoids materializing the full N \times N attention matrix in HBM.
• Sparse attention patterns and mixture-of-experts architectures, which change the operational intensity profile, are not covered.
• The paper focuses on training workloads but does not deeply analyze autoregressive inference, where the operational intensity is qualitatively different: each token generation step involves matrix-vector (not matrix-matrix) multiplications, making even the GEMM operations memory-bound.

Impact and Legacy

This paper helped establish a critical shift in how the ML systems community thinks about transformer optimization: from “reduce FLOPs” to “reduce data movement.” The insight that most transformer operations are memory-bound directly motivated subsequent work.

Later work such as FlashAttention (Dao et al. 2022) shares the focus on data reuse and memory-aware tiling, but extends it to the attention mechanism itself, achieving 2–4x wall-clock speedups by avoiding HBM materialization of attention matrices. FlashAttention tiles the Q, K, and V matrices into blocks that fit in SRAM, computes partial attention outputs within each tile, and accumulates results without ever writing the full N \times N attention matrix to HBM. This is operator fusion taken to its logical conclusion for the attention operator.

The broader trend toward fused CUDA kernels in libraries like NVIDIA’s Transformer Engine, xFormers, and Triton-based custom kernels all reflect the core message of this paper: optimize for bytes moved, not FLOPs computed. PyTorch’s torch.compile and NVIDIA’s TensorRT both implement kernel fusion passes that automate the kind of manual optimizations the paper describes.

The supporting roofline model has become a standard tool for evaluating whether a proposed architectural change (e.g., grouped-query attention, multi-query attention) will translate to actual wall-clock improvements on real hardware. When researchers propose attention variants that reduce FLOPs but increase memory traffic (or vice versa), the roofline model provides the analytical framework to predict which will actually be faster on a given device.

For LLM inference specifically, the data movement analysis explains why autoregressive token generation is so slow: generating each token requires loading the entire model’s weights from HBM (a memory-bound operation), but performs only a single matrix-vector multiply per layer. This insight directly motivates techniques like KV caching, speculative decoding, and batching strategies that amortize weight loading across multiple tokens or sequences.

Related Reading

• Attention Is All You Need — the original transformer architecture whose operations this paper profiles
• Optimizing Transformer Inference — survey covering complementary optimization techniques including pruning, quantization, and knowledge distillation
• Vision Transformer — ViT extends transformers to vision, introducing the same data movement challenges in a new domain
• FlashAttention — applies this paper’s tiling-and-fusion philosophy directly to the attention operator