The Atomic Price of Saving Memory

TL;DR for operators

A transformer kernel can minimize theoretical memory movement and still execute more slowly. In this study, combining attention’s forward and backward passes initially produced worse GPU performance than running them separately. Hardware profiling revealed excessive atomic accumulation; removing one source of contention recovered up to 2.49 times the naive baseline’s throughput on an NVIDIA H200. Removing another atomic operation, however, increased memory traffic and slowed larger fused workloads.

For AI infrastructure teams, three acceptance criteria should remain separate: computational correctness, theoretical memory efficiency, and measured execution performance. Benchmark representative workloads, inspect hardware counters, and test each optimization across relevant sequence lengths and devices. The reported gains concern the authors’ own implementations, not comparisons with production libraries such as FlashAttention.

The memory-saving fusion that initially ran slower

Consider a training system that combines two computational passes to avoid writing a large intermediate array to memory and reading it back. The expected benefit is straightforward: fewer transfers should reduce pressure on the memory system.

Now suppose the combined implementation runs slower on both deployment GPUs. The immediate question is whether the computational design is fundamentally inefficient or whether its execution introduces another bottleneck.

Lenore M. Mullin and Gaetan Hains encounter precisely this problem in Validating Memory-Optimal Transformer Kernels on Real Hardware. Their initial fused forward-plus-backward attention implementation is slower than separate execution at every tested sequence length on both NVIDIA A100 and H200 GPUs.¹

The comparison uses an appropriately constructed baseline:

This matters because a fused kernel performs both computations. Comparing its runtime directly with backward-only execution would charge fusion for additional work without crediting the forward pass it replaces.

The authors’ result initially challenges the practical value of fusion. Yet the benchmark establishes only that this particular implementation is slower. It does not establish that eliminating the intermediate array was a poor computational decision.

Understanding that difference requires separating the mathematical specification from the hardware execution.

The atomic bottleneck was in the execution arrangement

The paper validates transformer kernels derived in four companion studies using the Mathematics of Arrays (MoA) formalism.

MoA separates two descriptions of a computation:

• Denotational Normal Form (DNF): the hardware-independent mathematical specification, expressed through array shapes, indexing, and operations.
• Operational Normal Form (ONF): the machine-specific arrangement of computation, including execution order, intermediate storage, memory access, and parallel work assignment.

A transformation denoted by $\gamma$ connects the specification to its operational realization.

The distinction creates an opportunity. A kernel can preserve its specified mathematical computation while changing how GPU workers calculate and accumulate its results.

The profiler identifies where the original implementation went wrong.

During parallel gradient computation, different GPU workers may update the same destination. Atomic operations prevent unsafe concurrent updates, but large numbers of contending operations can create substantial execution overhead.

At sequence length $N=2048$, profiling on an A100 records:

The atomic count is exactly twice as large in the fused implementation. This comparison diagnoses synchronization overhead; it is not the full naive-versus-fused performance benchmark, because backward alone performs less work.

The authors then restructure the GPU implementation into four passes, eliminating atomic accumulation into the GV gradient array.

The change reverses the initial regression.

At $N=8192$, the corrected fused-to-naive execution-time ratios are 0.517 on A100 and 0.401 on H200, equivalent to approximately 1.93-fold and 2.49-fold speedups. The revised implementation outperforms the naive baseline across all tested sequence lengths on both GPUs.

This is the paper’s strongest mechanistic evidence: hardware counters identify a specific bottleneck, a targeted implementation change removes it, and subsequent measurements show the expected improvement.

The mathematical computation remains the same. Its operational arrangement changes.

Removing another atomic operation brought the slowdown back

After the successful GV correction, the authors investigate whether eliminating the remaining GK atomic accumulation can deliver further improvements.

This follow-up is an implementation ablation rather than a separate validation of fusion’s theoretical memory advantage.

Removing GK atomics improves backward-only execution by approximately 19–44%. For fused execution, the results are conditional on workload size: performance improves through $N=2048$, but deteriorates at $N=4096$ and $N=8192$.

Profiling explains the trade-off.

At $N=4096$, the atomic-free fused implementation transfers 89.83 GB, compared with 45.60 GB for the version retaining GK atomic accumulation. The revised implementation generates approximately 1.97 times as much DRAM traffic, largely because a separate GK pass rereads the GS intermediate array.

Eliminating synchronization therefore exchanges one resource cost for another. At larger sequence lengths, the additional memory traffic outweighs the benefit of avoiding GK atomics.

The two correction stages must be distinguished. Removing GV atomic accumulation fixes the initial fusion regression; removing GK atomic accumulation subsequently introduces a workload-dependent trade-off.

For kernel developers, the relevant optimization target is total execution cost, not the count of atomic operations considered independently.

The cost model predicts scaling better than physical memory traffic

The study also tests whether formally derived memory-cost predictions survive execution on real hardware.

For forward attention, observed scaling approaches the predicted quadratic behavior. At $N=8192$, local scaling exponents reach 1.97 on A100 and 2.04 on H200, close to the theoretical exponent of 2.

However, predicting the growth rate of computation is different from predicting exact memory transactions.

For forward attention with batch size 2, sequence length 2048, and embedding dimension 64, the formal element-count model predicts 75.497 MB of memory traffic. Nsight Compute measures 151.378 MB on the A100, a ratio of 2.01.

The discrepancy arises partly because logical array accesses do not map perfectly onto physical memory transfers. GPU memory coalescing combines suitable accesses into efficient transactions; poorly coalesced accesses can transfer more data than idealized element accounting suggests.

A complete-transformer-block experiment shows a similar, smaller discrepancy: 15.74 MB measured against 10.49 MB predicted, or 1.50 times the idealized estimate. Coalescing contributes, but the authors do not fully resolve the excess traffic, which may also involve intermediate rereads across kernel launches.

These experiments validate the broad scaling predictions without establishing exact byte-level predictions for every implementation.

They also limit simplistic hardware comparisons. Although the H200 outperforms the A100 in the tested forward-attention configurations, its advantage contracts to just 1.06 times at $N=8192$.

CPU thread allocation can matter more than available core count

The paper’s inference-decode experiments reveal an even more dramatic hardware-specific effect.

On NCSA Delta, a dual-socket AMD EPYC 7763 node has 128 physical cores distributed across eight NUMA domains. At $N=1024$, decode execution takes approximately 0.000083 seconds with 64 threads but 0.044417 seconds with 128 threads.

Increasing the thread count produces a 535.1-fold slowdown.

The authors interpret this anomaly through NUMA topology, memory-access locality, and scheduling overhead. The comparison with Purdue Anvil reinforces the importance of machine structure: Anvil’s tested single-socket, 32-core configuration experiences an oversubscription penalty below 2.9 times.

These are not otherwise identical hardware experiments, so they do not isolate a universal causal effect of NUMA crossings. The extreme Delta result is specific to its tested configuration and workload. At larger decode sizes, additional threads can again become beneficial.

The implication for CPU inference deployment is concrete: select thread counts using representative sequence lengths and actual node topology rather than setting parallelism to the maximum available core count.

A further compiler comparison supports the same hardware-dependent interpretation. C outperforms Fortran by factors of 1.2–1.85 for the tested CPU MLP implementations, while Fortran outperforms C by 1.45–2.81 for the corresponding GPU fused norm-plus-MLP implementations. Profiling narrows the GPU discrepancy to memory-latency stalls in the gated-MLP backward kernel, but the underlying compiler-code-generation cause remains unresolved.

Source-language rankings alone cannot substitute for measurements of compiled kernels on the intended hardware.

Kernel acceptance needs three separate tests

For AI infrastructure teams, the paper supports a validation workflow that treats formal derivation as the beginning of performance engineering rather than its final approval.

The third test deserves particular attention because benchmark procedures can themselves mislead.

The authors initially observe GPU decode timings dominated by approximately 0.30–0.40 seconds of cold-start overhead. Introducing an untimed warmup and measuring multiple problem sizes within a single initialized process corrects the contamination and reveals the expected scaling behavior.

For production evaluation, Cognaptus would extend this workflow into explicit implementation-version tracking and acceptance thresholds for the intended workload distribution. Every structural modification should be compared against the correct baseline, with hardware profiling used to explain significant regressions. Repeated measurements are especially necessary when unusual timing discontinuities appear.

This recommendation is an operational inference from the study, not a production deployment protocol experimentally validated by the authors.

What the benchmarks establish—and what they leave open

The evidence is strongest for the measured GPU atomic-contention mechanism and the recovery achieved through targeted restructuring. The CPU topology results, memory-traffic comparisons, and compiler experiments broaden the case for hardware-specific validation.

Several boundaries remain material.

The benchmarks cover two HPC clusters, AMD EPYC CPU configurations, and NVIDIA GPUs using the nvc/OpenACC toolchain. They compare the authors’ MoA-derived implementations with their formal predictions and internal baselines, not with FlashAttention, vendor libraries, or other independently optimized production kernels.

Statistical repetition is uneven, and confidence intervals are not reported. The complete-block CPU timing increase of 12.5 times between $N=2048$ and $N=4096$ persists after five-run averaging, but its proposed cache-capacity explanation lacks direct profiling confirmation. Some numerical and compiler-level issues also remain unresolved, and the paper describes portions of its reproducibility artifacts as still being prepared for archival release.

Consequently, the results support targeted implementation improvement, not universal performance superiority, demonstrated financial savings, or an automated hardware-portability system.

Preserve the computation, revise its execution

The central contribution is a practical connection between formal computation and empirical performance.

A formal specification makes memory costs and computational dependencies explicit. Hardware measurements expose costs that simplified models omit. Targeted rewrites can then change the operational implementation without redesigning the underlying mathematical computation.

The successful GV correction demonstrates that sequence directly. The later GK experiment establishes its necessary qualification: every improvement must be evaluated against the other costs it introduces.

For organizations maintaining transformer infrastructure across devices and compiler versions, the promising long-term asset is a stable computational specification paired with replaceable, hardware-specific implementations. This paper demonstrates that approach for selected kernels and machines. Extending it into a general optimization framework remains a research agenda.

Cognaptus: Automate the Present, Incubate the Future.