Atomic Operations, Warp Shuffle, and the Register Cache: A Deep Dive with Online Softmax

Image of the Tesla V100 Accelerator with Volta GV100 GPU. SXM2 Form Factor.

There are multiple posts out there comparing atomic operations and warp shuffle, and for the most part they are right that warp shuffle tends to be more efficient. But these posts rarely go deep enough on why, multiple times, they stop at the benchmark number and leave the reader without a mental model of what is actually happening in hardware. This article tries to fill that gap, building from the primitives up through the profiler output, and ending with how these ideas quietly shape the design of modern AI inference systems.

We cover the mechanics of atomic operations mostly atomicAdd as the case study, the privatization techniques built on top of them, the register cache abstraction that warp shuffle enables, and a hands-on experiment comparing both approaches on an online softmax kernel profiled on a Turing Arch. The Nsight Compute screenshots from this experiment, memory workload charts, instruction mix histograms, occupancy analysis, and L2 distribution tables, provides hardware-level evidence for every claim made in the theoretical sections.

While reading the article, it is important we state the foundation of these forms of reduction, as they sprung from SIMT (Single instruction Multiple Thread), which is a subset of Flynn’s Taxonomy of computer architecture. The advantage of this is that, rather than a single thread issuing vector based instructions, we have multiple threads do this in parallel. Parallel programming often employs the idea of communication across threads in the same thread block and operations such as parallel reduction and scans.

Atomic Operations: The Problem They Solve and The Problem They Create

Thread Interference in Global Memory

To understand why atomic operations exist, you first need to understand the limitations they were designed to combat. When multiple threads attempt to read, modify, and write back to the same global memory location concurrently, the results are undefined. This is the classic read-modify-write hazard:

thread A reads  x = 5
thread B reads  x = 5
thread A writes x = 6   (5 + 1)
thread B writes x = 6   (5 + 1)  ← B overwrote A's update

Correct result: 7.   Actual result without coordination: 6.

This is not a bug in program logic, it is a fundamental consequence of parallel execution without guaranteed ordering on concurrent memory accesses. Atomic operations solve this by making the read-modify-write sequence indivisible at the hardware level. The L2 cache line containing the target address is locked, the operation executes, and the lock releases before any other thread can access that address. Threads are forced into a serialized queue purely by arrival order at the memory controller, not by examining the values being computed.

It is very important to note that atomic operations do not reason about data dependencies. They enforce ordering at the memory level, and that ordering carries a latency cost that scales directly with the number of concurrent threads targeting the same location.

The Contention Scaling Problem

In performing atomic operations on a particular memory location, the highest throughput that one can achieve is one atomic operation every 400 cycles (200 cycles for the read and 200 cycles for the write). This translates into a time-based throughput of 1/400 atomics/clock 1 G (clocks/second) = 2.5 M atomics/second. This is dramatically lower than most users expect from a GPU memory system. Furthermore, the long latency of the sequence of atomic opera tions will likely dominate the kernel execution time and can dramatically lower the execution speed of the kernel.

Privatization: Reducing Contention Through Indirection

The Core Idea

The natural response to atomicAdd contention is privatization. Instead of having all threads write to one shared accumulator, you create private copies, one per thread block or per warp and have threads accumulate into their private copy first. At the end of the kernel, private copies are merged back into the single shared result with one atomicAdd per copy. The contention reduction is proportional to the number of private copies. With one private copy per block of 256 threads, global atomic operations drop from N to N/256. The contention on the global accumulator drops by the same factor. The advantage of this is that the programmer can adopt shared memory and thread synchronization.

Contiguous vs. Interleaved Partitioning

Privatization comes with two natural partitioning strategies, and the choice between them matters.

Contiguous partitioning: each block owns a contiguous chunk of the input. This maximizes memory coalescing during the private accumulation phase — consecutive threads access consecutive addresses. This is the right strategy when the merge is a bijective write or a serial reduction, because there is no contention to spread and you want maximum bandwidth.

Interleaved partitioning: threads stride across the input thread 0, handles elements 0, T, 2T; thread 1 handles 1, T+1, 2T+1. For workloads like histograms, where spatially adjacent inputs tend to map to the same output bucket, contiguous assignment causes intra-block collisions at merge time. Interleaved assignment spreads those collisions across different buckets.

The decision rule is precise:

merge == atomicAdd (many-to-few)   →  interleaved wins   (spread collision probability)
merge == bijective write            →  contiguous wins    (maximize cache coalescing)

They solve different problems. Interleaved partitioning solves contention. Contiguous partitioning solves memory throughput. Applying the wrong one in the wrong place produces suboptimal results in both directions.

The Register Cache: Warp Shuffle as a Memory Layer

The Missing Layer in the Hierarchy

The GPU memory hierarchy has a structural gap. Global memory is slow and large, meaning it is shared across all threads. Shared memory is fast, small and shared within a thread block. Registers, being the smallest, are fastest and private to each thread. But there is no hardware cache layer scoped to a single warp which is the natural granularity of lockstep execution.

This matters because many algorithms are warp-centric. A warp of 32 threads executes in lockstep on a single set of functional units. Computations that require threads within a warp to share intermediate values are forced through shared memory, which requires explicit synchronization via __syncthreads() and carries SRAM access latency.

Hamilis and Silberstein in 2017 at the Technion introduced a software abstraction called the register cache that fills this gap. Distributing warp-level data across thread registers and using the shuffle instruction to communicate when any thread needs data held by another. The result is a virtual caching layer at warp granularity, backed entirely by registers and inter-thread register exchange, with no shared memory involved.

The Shuffle Primitive

NVIDIA introduced the SHFL instruction in the Kepler microarchitecture in 2012. The primitive __shfl_sync(mask, r, t) allows an issuing thread to read register r from thread t within the same warp, while simultaneously publishing its own value for others to read. Mask is defined as the set of thread in the warp that participates in a collective operation. They determine the program logic.

Register-to-register communication operates at warp speed with no memory system involvement that is approximately 5 cycles per round, no barrier required (provided no intra-warp divergence). The register file on modern GPUs is also significantly larger than shared memory from 256KB versus 64KB meaning spare registers can serve as a caching resource that would otherwise sit idle.

However, starting with CUDA 9, the original __shfl() was deprecated in favor of __shfl_sync(), which takes an explicit thread mask. This was required for correctness on Volta and later architectures, where independent thread scheduling gives each thread its own program counter, making implicit warp-level synchronization unsafe. The explicit mask makes the synchronization contract visible in source code and safe across all architectures.

Applying the Register Cache to Online Softmax

Softmax requires computing a denominator Z = Σ exp(xᵢ) , a global reduction where every element contributes to one scalar. The naive approach fires atomicAdd directly from every thread:

__global__ void onlineSoftmaxDenom_Atomic(
    const float* input, float* denom, int n)
{
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid < n) {
        float val = expf(input[tid]);
        atomicAdd(denom, val);  // N = 16M threads, one address — maximum contention
    }
}

The warp shuffle version replaces the global atomic with a three-level hierarchy:

Level 1 — each thread holds exp(input[tid]) in a register

Level 2 — warp butterfly (5 × __shfl_down_sync, zero memory traffic):
    val += __shfl_down_sync(0xffffffff, val, 16);
    val += __shfl_down_sync(0xffffffff, val, 8);
    val += __shfl_down_sync(0xffffffff, val, 4);
    val += __shfl_down_sync(0xffffffff, val, 2);
    val += __shfl_down_sync(0xffffffff, val, 1);
    // lane 0 holds warp partial sum — entirely in registers

Level 3 — one warp_partial[warp_id] write to shared memory
          + __syncthreads()
          + second butterfly over WARPS_PER_BLOCK = 8 values
          + one atomicAdd per block  (N/256 operations, not N)

__global__ void onlineSoftmaxDenom_WarpShuffle(
    const float* input, float* denom, int n)
{
    __shared__ float warp_partial[WARPS_PER_BLOCK];
    int tid     = blockIdx.x * blockDim.x + threadIdx.x;
    int lane    = threadIdx.x & 31;
    int warp_id = threadIdx.x >> 5;

    float val = (tid < n) ? expf(input[tid]) : 0.0f;

    val += __shfl_down_sync(0xffffffff, val, 16);
    val += __shfl_down_sync(0xffffffff, val, 8);
    val += __shfl_down_sync(0xffffffff, val, 4);
    val += __shfl_down_sync(0xffffffff, val, 2);
    val += __shfl_down_sync(0xffffffff, val, 1);

    if (lane == 0) warp_partial[warp_id] = val;
    __syncthreads();

    if (warp_id == 0) {
        float block_val = (lane < WARPS_PER_BLOCK) ? warp_partial[lane] : 0.0f;
        block_val += __shfl_down_sync(0xffffffff, block_val, 4);
        block_val += __shfl_down_sync(0xffffffff, block_val, 2);
        block_val += __shfl_down_sync(0xffffffff, block_val, 1);
        if (lane == 0) atomicAdd(denom, block_val);
    }
}

Experimental Results: Turing architecture, 16M Elements

Both kernels were profiled on an NVIDIA GeForce Turing based GPU (sm_75, 14 SMs) with N = 16,777,216 elements over 100 benchmark iterations. Timing was measured with CUDA events (device-side, excluding CPU overhead). Profiling used Nsight Systems for the full GPU timeline and Nsight Compute for per-instruction hardware counters. What follows draws directly from the Nsight Compute screenshots.

Kernel Duration

KernelTimeCycles:

onlineSoftmaxDenom_Atomic34.95ms

onlineSoftmaxDenom_WarpShuffle930.82 µs

Warp shuffle is 37× faster as measured inside Nsight Compute at the same kernel invocation. The atomic kernel consumed 53,053,297 hardware cycles; the warp shuffle kernel consumed 1,407,499 cycles, a 37.7× cycle count difference on identical hardware running identical logic. That ratio is the cost of serializing 16 million atomic requests through a single L2 cache line lock.

The Memory Throughput Paradox

The memory workload screenshots reveal the single most counterintuitive result in the entire experiment:

Figure 1

Figure 2

The atomic kernel has an L2 hit rate of 88.91% (figure 1) yet achieves only 1.94 GB/s of memory throughput. The warp shuffle kernel has an L2 hit rate of 3.18% (figure 2) yet achieves 72.62 GB/s. This appears completely backwards. The explanation is that a high L2 hit rate does not mean fast execution when the bottleneck is the L2 atomic serialization unit rather than bandwidth.

The single accumulator address fits permanently in L2 , every atomic request hits cache at 88.91% rate. But hitting L2 does not skip the lock. The hardware atomic unit on that one L2 slice still serializes every request one at a time. The cache line is hot but contested, and contested cache lines serialize even when cached. The 1.94 GB/s throughput reflects how little actual data movement happens when all you are doing is serialized single-word updates.

The warp shuffle kernel (figure 2), by contrast, streams 16 million distinct float values sequentially from device memory. Each value is read once and discarded while L2 cannot help because there is no reuse. The 3.18% hit rate is expected and correct for a streaming access pattern. The 72.62 GB/s is approaching the Turing based GPU device memory bandwidth ceiling, confirming this kernel is memory-bandwidth-bound rather than contention-bound.

Key insight: High L2 hit rate + low throughput = contention bottleneck, not bandwidth bottleneck. The memory system is not the constraint, the atomic serialization unit is.

Symmetric Atomics: The Read-Modify-Write Signature

Figure 3

The Memory Chart for the atomic kernel in figure 3 shows an exact symmetry: 524.29K global load requests and 524.29K global store requests. These numbers are equal because every atomicAdd is physically a read-modify-write, the hardware must read the current value, add the new value, and write the result back. Every atomic operation generates both a load and a store at the L2 level.

Figure 4

The warp shuffle kernel in figure 4 shows a different pattern: 524.29K global loads (reading the input array) but only 65.54K global stores. That store count is N/256 = 65,536 exactly one global write per thread block, which is the single atomicAdd per block designed into the kernel. The asymmetry is the architectural fingerprint of the privatization strategy: load everything, reduce privately, write once.

Occupancy: The Cost of the Synchronization Barrier

The occupancy analysis for the warp shuffle kernel shows a gap between theoretical and achieved occupancy:

Figure 5

Theoretical occupancy is 100%, the register and shared memory footprint of the warp shuffle kernel are small enough that the SM could theoretically run 32 warps simultaneously. Achieved occupancy is 76.33%, meaning only 24.42 out of 32 theoretical warps are active on average. Nsight Compute attributes this to warp scheduling overhead and load imbalance and the single __syncthreads() between the shared memory write and the second butterfly is the structural source.

When warp 0 writes to warp_partial and calls __syncthreads(), all 8 warps in the block must arrive at the barrier before any can proceed. Warps that finish the first butterfly earlier sit idle at the barrier. During those idle cycles, the SM slots those warps occupy are unavailable for other work. The 23.67% estimated speedup is the performance left on the table by this barrier stall.

The Register Count chart confirms registers are not the occupancy limiter while occupancy stays at 100% up to 16 registers per thread and the current kernel sits comfortably in that flat region. The barrier is the constraint, not register pressure.

Figure 6

Interestingly, this occupancy gap does exist for the atomic kernel.

Pipeline Utilization Differences

The Estimated Instructions per Pipeline charts reveal a structural difference in how the two kernels use SM resources.

Figure 7

In the atomic kernel, the OTHER pipeline (which handles atomicAdd and transcendentals like expf) and the ALU pipeline carry most of the load, with LSU in third place. The CBU (Convergence Barrier Unit) is minimal meaning no barriers exist in the kernel.

Figure 8

In the warp shuffle kernel, the LSU carries significantly more load relative to the atomic version because all shared memory reads and writes in the block-level reduction route through the LSU. The CBU is small but non-zero, capturing the one __syncthreads(). The ADU (Address Datapath Unit) appears in the warp shuffle chart but not the atomic chart, the shuffle butterfly accesses thread-indexed registers in a pattern that exercises the address computation logic more than a scalar atomic.

The overall pipeline balance in the warp shuffle kernel is healthier due to the load being distributed across more pipeline stages, resulting to fewer bottlenecks. The atomic kernel concentrates everything on the path from the load unit through the atomic unit in L2, leaving most pipeline stages underutilized while that one path is overloaded.

Where Atomic Operation Logic Lives Today

While atomic operations may seem primitive, their philosophy is being adopted and extended in some of the most performance-sensitive systems in modern AI infrastructure. Serialized access to shared state only at the moment of commitment, maximize parallelism everywhere else appears repeatedly at higher levels of abstraction.

Even in warp level programming atomic operation is still being used in cases where a function has to be passed to the program flow. Such action can prove to be difficult because at that level you can’t change the function interface. Hence we employ atomic operations such as atomicAggInc()andatomicAdd() . This is called Opportunistic warp-level programming.

// increment the value at ptr by 1 and return the old value
__device__ int atomicAggInc(int *ptr) {
    int mask = __match_any_sync(__activemask(), (unsigned long long)ptr);
    int leader = __ffs(mask) – 1;    // select a leader
    int res;
    if(lane_id() == leader)                  // leader does the update
        res = atomicAdd(ptr, __popc(mask));
    res = __shfl_sync(mask, res, leader);    // get leader’s old value
    return res + __popc(mask & ((1 << lane_id()) – 1)); //compute old value
}

atomicAggInc() atomically increments the value pointed to by ptr by 1 and returns the old value. It uses the atomicAdd() function, which may incur contention. To reduce contention, atomicAggInc replaces the per-thread atomicAdd() operation with a per-warp atomicAdd(). The __activemask() in line 4 finds the set of threads in the warp that are about to perform the atomic operation. __match_any_sync() returns the bit mask of the threads that have the same value ptr, partitioning the incoming threads into groups whose members have the same ptr value. Each group elects a leader thread (line 5), which performs the atomicAdd() (line 8) for the whole group. Every thread gets the old value from the leader (line 9) returned by the atomicAdd(). Line 10 computes and returns the old value the current thread would get from atomicInc() if it were to call the function instead of atomicAggInc.

KV Cache Management in vLLM

vLLM’s paged attention system manages GPU memory the way an operating system manages virtual memory. The KV cache for each sequence is divided into fixed-size physical blocks (pages), and a page table maps logical sequence positions to physical block addresses. When a new token is generated, the page table must be updated atomically two sequences cannot simultaneously claim the same physical block.

The reference counting layer that tracks block ownership is built on atomic increments and decrements. When prefix sharing is enabled (multiple sequences sharing the same KV prefix), every sequence that accesses a shared block increments its reference count on arrival and decrements it on departure. Under high concurrency, these updates become a contention hotspot structurally identical to our softmax atomicAdd many concurrent writers targeting a small number of counters.

The L2 slice imbalance we observed (one slice 74.55% above average) would manifest here as the L2 slice responsible for the page table entries becoming overloaded relative to the rest of the cache. The mitigation is the same: interleaved sequence assignment to SMs spreads simultaneously active sequences across different physical block regions, reducing the probability that two sequences target the same page table entry at the same time.

FlashDecoding’s Cross-Block Reduction

FlashAttention computes attention entirely within a single thread block, the KV data fits in SRAM and requires no cross-block synchronization. For long sequences during decoding, the KV sequence is too long for one block to handle. FlashDecoding splits the KV sequence across multiple blocks, each computing partial attention output and partial softmax statistics (max, sum, output).

The merge step mirrors the structure of our warp shuffle kernel’s block-level reduction, lifted to the grid level. Each block writes its partial result to a pre-allocated global buffer; a bijective write, one writer per output slot, no atomics needed. A second kernel reads all partial results and performs the final reduction. This two-pass structure is privatization: blocks accumulate privately, merge once.

The cross-block merge cannot use atomicAdd directly because online softmax's update rule is non-linear and the rescaling factor requires knowing both the current and incoming maximum simultaneously, which no additive atomic can express. The two pass approach sidesteps this by making the merge an explicit serial reduction over a small array of partial statistics, following the same design principle the softmax denominator comparison demonstrates: push the global synchronization point as late and as rarely as possible.

Speculative Decoding and Private KV Blocks

Speculative decoding generates multiple candidate tokens in parallel using a draft model, then verifies them with the target model in a single pass. During the draft phase, candidate tokens are generated into private KV blocks physical memory allocations not yet committed to the main sequence’s page table. If verification accepts a candidate, its private block is committed (page table updated atomically). If rejected, the block is freed.

This is privatization at the KV block level. The draft phase is the private accumulation step. The verification and page table update is the merge. The atomic operation appears only at the commit point, that is, the same structure the warp shuffle kernel demonstrates that they do the most work possible before touching shared state, then merge once. The L2 slice imbalance lesson applies here too, if many speculative sequences commit simultaneously, their page table updates all target the same region, creating a hotspot in whichever L2 slice holds those entries. Staggering commit times or interleaving page table layout reduces this.

Conclusion

Atomic operations were a foundational primitive for enabling parallel accumulation on GPUs. They solve a real problem, thread interference on shared mutable state however they solve it by serializing access, which at scale becomes the bottleneck they were meant to avoid. The Nsight Compute data from this experiment quantifies that bottleneck precisely 53 million cycles for 16 million serialized atomic requests on a kernel that requires only 1.4 million cycles when the serialization is removed.

The profiler data revealed three things that go beyond what most benchmark comparisons show. First, high L2 cache hit rate does not imply fast execution instead it implies the atomic kernel’s 88.91% L2 hit rate coexisted with 1.94 GB/s of effective throughput because cache hits still pass through the serialized atomic unit. Second, the L2 slice imbalance (74.55% above average on the hot slice) shows exactly where contention concentrates in hardware where one physical piece of the cache becomes a serial bottleneck while everything else waits. Third, the instruction count paradox shows that the warp shuffle kernel executes roughly 2× more instructions than the atomic kernel yet runs in 1/37th the time, because instruction count and wall-clock time completely decouple when the bottleneck is stall cycles.

Warp shuffle breaks the dependency on shared mutable state within the warp entirely. By distributing data across registers and using intra-warp register exchange, it eliminates the memory round trip that atomic operations require. The register cache abstraction formalizes this into a principled design pattern. The deeper lesson is that the same serialization versus parallelism tradeoff appears at every level of the software stack, vLLM’s page table, FlashDecoding’s two-pass merge, speculative decoding’s private block commit all embody the same answer the warp shuffle kernel demonstrates accumulate privately, merge once, make the merge cheap.

Reference

Programming Massively Parallel Processors: A Hands-on Approach

Book by David Kirk and Wen-mei Hwu

[embed]Using CUDA Warp-Level Primitives | NVIDIA Technical Blog NVIDIA GPUs execute groups of threads known as warps in SIMT (Single Instruction, Multiple Thread) fashion. Many CUDA…developer.nvidia.com&text=atomicAggInc()%20atomically%20increments%20the,the%20function%20instead%20of%20atomicAggInc%20.)

[embed]Register Cache: Caching for Warp-Centric CUDA Programs | NVIDIA Technical Blog In this post we introduce the "register cache", an optimization technique that develops a virtual caching layer for…developer.nvidia.com

[embed]Voting and Shuffling to Optimize Atomic Operations | NVIDIA Technical Blog 2iSome years ago I started work on my first CUDA implementation of the Multiparticle Collision Dynamics (MPC)…developer.nvidia.com