GPU Parallel Primitives Lab

Eight numbers sit in GPU memory: 3, -1, 4, 1, 5, -9, 2, 6. An elementwise ReLU can compute each output independently: load a value, compare it with zero, and store the result. A sum cannot. How does one output learn about all eight inputs?

Attention softmax, LayerNorm, RMSNorm, cross-entropy loss, and top-k sampling need information from multiple elements. The exchange may stay within a warp, use shared memory within a block, or cross blocks through global memory. A row that fits in one block needn't use every scope.

These building blocks are parallel primitives. Reduction folds a tensor or tile into one summary. Prefix scan keeps a cumulative total at every boundary. Arg reduction tracks the winning value and its source position. Online softmax computes its maximum and normalizer together; producing every probability still requires access to the original logits after the final normalizer is known.

The GPU Kernel Performance Engineering Lab examined memory coalescing and shared-memory bank conflicts through matrix transpose. Here, the layout is simple; the challenge is coordinating thread dependencies: which values must be written, synchronized, and published before another lane can proceed.

One total or every boundary?

Pair adjacent values in the running array [3, -1, 4, 1, 5, -9, 2, 6]. Adding neighboring pairs yields [2, 5, -4, 8], then [7, 4], then 11. A balanced binary tree with $n\geq1$ leaves has $n-1$ additions and $\lceil\log_2 n\rceil$ combining levels. The eight-leaf tree has exactly three levels; padding a hardware tile with identities can add work beyond those seven useful additions.

A prefix scan asks for a completely different contract: it keeps every intermediate boundary. An exclusive sum at index $i$ accumulates elements strictly before index $i$, placing the additive identity 0 at index 0. An inclusive sum also includes $x_i$. For the first four elements [3, -1, 4, 1], the exclusive prefix sum is [0, 3, 2, 6] and the inclusive prefix sum is [3, 2, 6, 7].

For exact real addition, the two scan conventions connect through a direct identity:

Arg reduction carries a structured payload: (value=6, index=7). Equal maxima need a tie rule. Nonfinite values need a policy too: this reference rejects them and chooses the lowest index among equal finite values.

from math import isfinite

x = [3, -1, 4, 1, 5, -9, 2, 6]
exclusive, inclusive, total = [], [], 0
for value in x:
    exclusive.append(total)
    total += value
    inclusive.append(total)

def argmax(values):
    if not values:
        raise ValueError("Argmax requires a nonempty sequence")
    if any(not isfinite(value) for value in values):
        raise ValueError("Argmax inputs must be finite")
    index = max(range(len(values)), key=lambda i: (values[i], -i))
    return values[index], index

assert total == 11
assert exclusive == [0, 3, 2, 6, 7, 12, 3, 5]
assert inclusive == [3, 2, 6, 7, 12, 3, 5, 11]
assert argmax(x) == (6, 7)
assert argmax([6, -1, 6]) == (6, 0)
print("sum:", total)
print("exclusive:", exclusive)
print("inclusive:", inclusive)
print("argmax (value, index):", argmax(x))

sum: 11
exclusive: [0, 3, 2, 6, 7, 12, 3, 5]
inclusive: [3, 2, 6, 7, 12, 3, 5, 11]
argmax (value, index): (6, 7)

Scan also turns decisions into addresses. Keep only the positive values: scan their 0/1 flags to obtain a unique destination for each retained item. A rejected item must skip its write, even if its prefix shares a destination with the next item.

flags = [int(value > 0) for value in x]
offsets, retained = [], 0
for flag in flags:
    offsets.append(retained)
    retained += flag
compacted = [None] * retained
for value, flag, destination in zip(x, flags, offsets):
    if flag:
        compacted[destination] = value
assert offsets == [0, 1, 1, 2, 3, 4, 4, 5]
assert compacted == [3, 4, 1, 5, 2, 6]
boundaries = [0, 4, 4, 8]
segments = [sum(x[a:b]) for a, b in zip(boundaries, boundaries[1:])]
assert segments == [7, 0, 4]  # The empty segment has sum identity zero.
for values in ([float('nan'), 6], [float('inf'), 6], [6, float('-inf')]):
    try:
        argmax(values)
    except ValueError:
        continue
    raise AssertionError("Nonfinite argmax input accepted")
print("Keep flags:", flags)
print("Destinations:", offsets)
print("Compacted:", compacted)
print("Segment sums (including an empty segment):", segments)
print("Nonfinite argmax inputs rejected")

Keep flags: [1, 0, 1, 1, 1, 0, 1, 1]
Destinations: [0, 1, 1, 2, 3, 4, 4, 5]
Compacted: [3, 4, 1, 5, 2, 6]
Segment sums (including an empty segment): [7, 0, 4]
Nonfinite argmax inputs rejected

The segment boundaries specify half-open input ranges. A GPU implementation must preserve that grouping; one whole-array reduction would erase it. This CPU scatter uses a separate output buffer and doesn't establish that an in-place GPU compaction is safe.

Reduce within a warp using shuffle intrinsics

Early GPU reduction kernels relied heavily on shared memory: threads wrote inputs into a __shared__ array, synchronized via __syncthreads(), and accumulated across stride intervals. That approach cost shared memory capacity, suffered from bank conflicts when strided accesses hit identical memory banks, and stalled on block barriers at every tree level.

Within a CUDA warp (32 threads), __shfl_down_sync exchanges register values without shared memory. It removes the shared-memory accesses for that exchange, but doesn't promise single-cycle latency. Type width and generated instructions matter; the lab exchanges 64-bit values.^{[1]Reference 1C/C++ Language Extensionshttps://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/cpp-language-extensions.html}

Offsets 16, 8, 4, 2, 1 construct a five-stage reduction tree across all 32 lanes. Lane 0 accumulates the warp total, while higher lanes hold intermediate partials:

__device__ long long warp_sum(long long value) {
    constexpr unsigned kFullMask = 0xffffffffu;
    const unsigned lane = threadIdx.x % warpSize;
    for (int offset = warpSize / 2; offset > 0; offset /= 2) {
        const long long other = __shfl_down_sync(kFullMask, value, offset);
        if (lane + offset < warpSize) {
            value += other;
        }
    }
    return value;  // Full aggregate is valid in lane 0.
}

This helper assumes a full, participating 32-thread warp, with representable signed 64-bit intermediate sums. The downloaded kernel launches 256 threads and keeps tail lanes alive. Volta introduced Independent Thread Scheduling; old assumptions that neighboring lanes implicitly synchronize are unsafe. The shuffle mask names participants: every named, non-exited lane must reach the corresponding collective with the same mask. Shuffles exchange registers; they aren't shared-memory fences.^{[1]Reference 1C/C++ Language Extensionshttps://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/cpp-language-extensions.html[2]Reference 2CUDA Programming Guide.https://docs.nvidia.com/cuda/cuda-programming-guide/}

Participation differs from owning an input. With 20 inputs, lanes 20 through 31 carry zero and still execute every shuffle. If lane 20 exits and lane 4 requests its value, the result is undefined; it isn't guaranteed to double lane 4. A full mask cannot recreate an exited lane. The separate down-shuffle rule for a source beyond the warp's upper boundary leaves the caller's value unchanged; the conditional addition discards that result. Keep the collective outside the conditional so all 32 lanes participate.

The CPU model verifies this behavior by checking lane participation, missing source rejections, and signed 64-bit integer overflow protection:

from assets.primitives_model import warp_sum, I64_MAX, expect_error

lanes = [1] * 20 + [0] * 12
assert warp_sum(lanes) == 20
expect_error(ValueError, lambda: warp_sum(lanes, active=range(20)))
assert warp_sum([0] * 31 + [I64_MAX]) == I64_MAX
expect_error(OverflowError, lambda: warp_sum([I64_MAX, 1] + [0] * 30))
print("20 values + 12 participating zero lanes: sum 20")
print("Absent source rejected; representable and overflowing int64 sums distinguished")

20 values + 12 participating zero lanes: sum 20
Absent source rejected; representable and overflowing int64 sums distinguished

To aggregate across a 256-thread block (8 warps), lane 0 of each warp deposits its warp total into __shared__ long long warp_sums[8]. A block-wide barrier __syncthreads() guarantees all 8 warps finish their writes before warp 0 reads the shared array. Warp 0's first 8 lanes load the 8 partial sums, while lanes 8 through 31 load zero. A second warp_sum folds those 8 partials, leaving the entire block total in lane 0:

sum = warp_sum(sum);
if (lane == 0) {
    warp_sums[warp] = sum;
}
__syncthreads();
if (warp == 0) {
    long long block_sum = lane < 8 ? warp_sums[lane] : 0;
    block_sum = warp_sum(block_sum);
    if (lane == 0) {
        block_sums[blockIdx.x] = block_sum;
    }
}

Notice the barrier placement: __syncthreads() sits outside all warp-divergent branches. Placing a block barrier inside if (lane == 0) or if (warp == 0) triggers undefined behavior because only a fraction of the block's threads reach the barrier.

Aggregate partials across blocks: multi-pass kernels versus global atomics

Each thread loads up to two elements: one at block_start + threadIdx.x and one 256 elements later. Full warps access adjacent, aligned elements; the final partial tile still requires bounds guards. A 256-thread block consumes up to 512 inputs and emits one 64-bit partial. Later passes read 64-bit partials rather than the original int32 inputs.^{[3]Reference 3Optimizing Parallel Reduction in CUDAhttps://developer.download.nvidia.com/compute/cuda/1.1-Beta/x86_website/projects/reduction/doc/reduction.pdf}

Why not insert a grid barrier into an ordinary launch? Resident blocks can wait forever for predecessors that cannot be scheduled until those blocks retire. __syncthreads() synchronizes one block. Cooperative Groups does support grid synchronization with a cooperative launch, supported hardware, and a grid sized for co-residency; that is a different launch contract from this lab.^{[4]Reference 4CUDA: Device-Callable APIs and Intrinsicshttps://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/device-callable-apis.html}

Two architectural strategies solve this grid-scale aggregation:

1. Multi-pass hierarchical kernel launches: Each kernel pass reduces elements by a factor of 512. For 4,194,304 inputs:

4,194,304 inputs -> 8,192 partials -> 16 partials -> 1 scalar
                  pass 1          pass 2         pass 3

Same-stream ordering lets a later pass consume completed earlier-pass outputs. It doesn't repair races inside either kernel. Blocks write separate partials; subsequent passes still have reduction dependencies, smaller grids, and launch costs.

1. Global atomics: Another implementation can atomically accumulate each block partial into an initialized output. Many updates to one address can become a throughput limit; separate bins distribute updates, but a skewed histogram can still concentrate them. Compare the complete algorithm rather than assuming atomics are always slower.

The accumulator type and API must match. The lab uses signed long long, whereas legacy atomicAdd's 64-bit integer overload is unsigned. An atomic variant isn't a drop-in replacement without a compatible signed operation and overflow contract.^{[1]Reference 1C/C++ Language Extensionshttps://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/cpp-language-extensions.html} Hierarchical passes avoid one shared atomic hotspot, at the cost of partial buffers and more launches. They don't eliminate serialization or guarantee bandwidth scaling.

Floating-point trees can disagree without a race

Integer addition preserves grouping and ordering when every intermediate fits. Checking only the final sum is insufficient. Predict what the scan model will do when its first pair overflows but a later negative value brings the final sum back into range.

from assets.primitives_model import block_scan, I64_MAX

values = [I64_MAX, 1, -1, 0]
assert sum(values) == I64_MAX
print("Final mathematical sum fits:", sum(values) == I64_MAX)
try:
    block_scan(values, width=4)
except OverflowError as error:
    print("First tree pair rejected:", str(error))
else:
    raise AssertionError("Overflowing partial accepted")

Final mathematical sum fits: True
First tree pair rejected: Signed 64-bit intermediate overflow

The up-sweep first combines I64_MAX + 1; a valid final total can't repair that invalid intermediate. The CUDA fixture guards the absolute sum bound 9N, which bounds every partial for its inputs in [-9, 9]. An arbitrary-input kernel needs its own accumulator-width or overflow contract.

Exact real addition is associative; floating-point addition rounds. Changing the grouping can therefore change the result:

The standard library example below forces 32-bit float truncation to reveal how operand ordering changes low bits on identical inputs:

from struct import pack, unpack
from math import fsum

def f32(value):
    return unpack("f", pack("f", value))[0]

a, b, c = f32(100_000_000), f32(1), f32(-100_000_000)
left_grouped = f32(f32(a + b) + c)
cancel_first = f32(f32(a + c) + b)
assert left_grouped == 0.0
assert cancel_first == 1.0
assert fsum([a, b, c]) == 1.0
print("(a + b) + c:", left_grouped)
print("(a + c) + b:", cancel_first)
print("Higher-precision reference:", fsum([a, b, c]))

(a + b) + c: 0.0
(a + c) + b: 1.0
Higher-precision reference: 1.0

Blocks needn't finish in index order. That alone doesn't change a fixed tree: a later pass can combine indexed partials in the same order every run. An atomic or schedule-dependent combining order can change rounding, sometimes substantially for ill-conditioned sums such as this example.

Different valid combining trees can disagree without a race. Unexpected disagreement still needs investigation; non-associativity isn't permission to dismiss a correctness failure.

Repeatability and numerical accuracy are two distinct properties:

• Repeatability concerns identical output bits under specified conditions. Current CUB documents run_to_run as DeviceReduce's default, with a fixed tree for the same GPU/build/configuration. Supported environment APIs can request other guarantees. This doesn't promise identical results across toolkit/library releases or tuning changes.^{[5]Reference 5cub::DeviceReducehttps://nvidia.github.io/cccl/unstable/cub/api/structcub_1_1DeviceReduce.html[6]Reference 6CCCL: Determinismhttps://nvidia.github.io/cccl/unstable/cccl/determinism.html}
• Accuracy concerns distance from the mathematical result. Pairwise summation can reduce rounding error, but cancellation and input conditioning still matter. Test against a suitable reference with declared tolerances; if bitwise repeatability is part of the contract, test that separately. A fixed tree can repeat an inaccurate answer.

Parallel scan: Hillis-Steele step efficiency versus Blelloch work efficiency

Scan must preserve input order, even when its operator isn't commutative. For power-of-two $n$, compare two tree structures:

• Hillis-Steele (Step-efficient): At stride $2^k$, position $i\geq2^k$ combines its value with position $i-2^k$ from the previous stage. There are $\log_2 n$ stages and $n-2^k$ additions per stage: $n\log_2 n-(n-1)$ total additions. A 32-lane shuffle scan has five stages and 129 additions. Its short dependency depth can be useful, but the extra arithmetic isn't free. Every stage must read a consistent previous state; an unsynchronized in-place shared-memory update isn't equivalent.

• Blelloch (Work-efficient): Blelloch scan takes $2 \log_2 n$ steps, but performs only $O(n)$ additions ($2n - 2$ operations total). It splits the scan into two distinct phases in shared memory:

1. Up-Sweep (Reduce): A binary reduction tree where stride doubles each round (1, 2, 4, ...). Threads add values upward until the root (last array position) holds the grand total.
2. Down-Sweep (Distribute): Save the root total, overwrite the root with 0 (the additive identity), and reverse the tree with stride halving (... 4, 2, 1). At each step, a node sends its current value to its left child, while its right child receives the sum of its old left child and its current value.

For 512 items, the counts are 4,097 versus 1,022 additions, about 4.01x. That's an arithmetic-work ratio, not a measured shared-memory traffic or kernel-speed ratio. Blelloch also performs swaps, identity placement, and synchronization; storage layout and the generated instructions affect cost.^{[7]Reference 7Single-pass Parallel Prefix Scan with Decoupled Look-backhttps://research.nvidia.com/sites/default/files/pubs/2016-03_Single-pass-Parallel-Prefix/nvr-2016-002.pdf}

The trace below captures the exact array mutations during Blelloch exclusive scan on [3, -1, 4, 1]:

from assets.primitives_model import block_scan, device_scan

trace = []
local, total = block_scan([3, -1, 4, 1], width=4, trace=trace)
for phase, values in trace:
    print(f"{phase:8}: {values}")
assert local == [0, 3, 2, 6] and total == 7
x = [3, -1, 4, 1, 5, -9, 2, 6]
assert block_scan(x[4:], width=4)[0] == [0, 5, -4, -2]
assert device_scan(x, width=4) == [0, 3, 2, 6, 7, 12, 3, 5]
print("Device-wide:", device_scan(x, width=4))

up 1    : [3, 2, 4, 5]
up 2    : [3, 2, 4, 7]
root = 0: [3, 2, 4, 0]
down 2  : [3, 0, 4, 2]
down 1  : [0, 3, 2, 6]
Device-wide: [0, 3, 2, 6, 7, 12, 3, 5]

for n in (32, 512):
    levels = n.bit_length() - 1
    hillis = sum(n - (1 << k) for k in range(levels))
    blelloch = 2 * n - 2
    print(f"N={n}: Hillis-Steele {hillis}, Blelloch {blelloch}, ratio {hillis / blelloch:.2f}x")
assert hillis == 4097 and blelloch == 1022

N=32: Hillis-Steele 129, Blelloch 62, ratio 2.08x
N=512: Hillis-Steele 4097, Blelloch 1022, ratio 4.01x

Propagate block offsets: multi-pass hierarchies versus decoupled look-back

A local Blelloch scan computes correct prefixes within each block, but every block starts at local offset 0. To stitch them into a device-wide prefix scan, each block must receive the cumulative sum of all preceding blocks.

In a hierarchical multi-pass scan:

1. Each block scans its 512 elements, emits local prefixes, and writes its block total to a global block_sums buffer. For blocks with inputs [3,-1,4,1] and [5,-9,2,6], the block totals are [7, 4].
2. A second kernel scans block_sums to compute block offsets: exclusive_scan([7, 4]) = [0, 7].
3. A third kernel broadcasts each block's offset across its threads: Block 0 adds +0, while Block 1 adds +7 to all its local prefixes, converting [0, 5, -4, -2] into [7, 12, 3, 5].

Testing scan hierarchies requires exercising recursion boundaries. If your block size is 512, an input of $512^2 = 262,144$ elements produces exactly 512 block totals, which fit in a single metadata block. But $512^2 + 1 = 262,145$ elements produce 513 block totals, spilling into two metadata blocks and triggering another recursion level. Testing only small powers of two misses that structural threshold.

Decoupled Look-Back (Single-Pass Scan): The three-stage example writes local outputs, scans block totals, then updates those outputs. Larger inputs need recursive metadata scans, so the downloadable lab can launch more than three kernels. Decoupled look-back avoids a separate full-array offset-add pass. CUB's DeviceScan uses this approach, with additional implementation/setup work beyond the input/output stream.^{[8]Reference 8cub::DeviceScanhttps://nvidia.github.io/cccl/unstable/cub/api/structcub_1_1DeviceScan.html}

• Tile assignment and forward progress are implementation concerns. The original report describes an atomic-counter adaptation; the current look-back agent instead uses start_tile + blockIdx.x. Neither is a universal definition of the algorithm.^{[7]Reference 7Single-pass Parallel Prefix Scan with Decoupled Look-backhttps://research.nvidia.com/sites/default/files/pubs/2016-03_Single-pass-Parallel-Prefix/nvr-2016-002.pdf[9]Reference 9CUB Look-Back Scan Agent Sourcehttps://github.com/NVIDIA/cccl/blob/main/cub/cub/agent/agent_scan.cuh}
• Each block computes its local reduction aggregate and publishes a status flag in global memory:
  • X: Uninitialized.
  • A: Local aggregate computed and published.
  • P: Inclusive prefix through this tile computed and published.
• When block $k$ finishes its local aggregate, it looks back at tile $k-1$:
  • If that tile is still X, wait for valid published data; reading its uninitialized payload is incorrect.
  • If tile $k-1$ is marked P, block $k$ reads the prefix directly, adds its local aggregate, and publishes its own status P.
  • If tile $k-1$ is marked A, block $k$ adds tile $k-1$'s aggregate and looks further back to tile $k-2$, accumulating aggregates until it reaches a tile marked P.

The payload must be visible before a reader accepts its status, using a correct publication protocol. Plain flag polling doesn't establish memory ordering or forward progress. The design targets roughly one input read and one output write per element, plus metadata and redundant look-back work. For this lab's int32-to-int64 contract, those useful bytes are $4N+8N$, not $2N$ equal-width bytes. Current CUB scan defaults to not_guaranteed determinism; supported environment APIs can request run-to-run float addition. Consult the installed release rather than assuming the reduction default also applies to scan.^{[8]Reference 8cub::DeviceScanhttps://nvidia.github.io/cccl/unstable/cub/api/structcub_1_1DeviceScan.html}

Online softmax carries a scale with its sum

Softmax converts an array of unnormalized logits $x_i$ into a probability distribution:

For finite unmasked logits, $x_i-m\leq0$, so their exponentials cannot overflow. Small terms can underflow to zero. Masks, NaNs, infinities, and all-masked rows need explicit policies; this model uses None for masks, rejects other nonfinite values, and rejects a row with no unmasked logits.

A simple implementation makes three logical passes over the input:

1. Reduction pass: find the global maximum $m = \max_j x_j$.
2. Reduction pass: compute the normalizer denominator $d = \sum_j e^{x_j - m}$.
3. Elementwise pass: compute each probability $p_i = e^{x_i - m} / d$.

That's three input reads and one output write per element before extra reduction traffic. It doesn't prove three HBM reads: cache service matters, and a row retained in registers/shared memory can reuse its values without these separate global passes.

The online normalizer (Milakov & Gimelshein 2018) merges maximum search and summation into one input pass, maintaining $(m,d)$: the running maximum and exponential sum relative to it. Its standalone softmax algorithm then rereads logits to emit probabilities, reducing logical input reads from three to two.^{[10]Reference 10Online normalizer calculation for softmax.https://arxiv.org/abs/1805.02867} The companion softmax rejects one-shot iterators; supply a reusable sequence whose values stay unchanged between passes, or store values for the output stage.

When merging two partial states $(m_a, d_a)$ and $(m_b, d_b)$, the combined maximum is $m = \max(m_a, m_b)$. Before adding the sums, each denominator must be rescaled to match the new reference point:

For the left half [3, -1, 4, 1], $m_a=4$ and $d_a\approx1.424404$. For the right half [5, -9, 2, 6], $m_b=6$ and $d_b\approx1.386195$. At the common maximum 6, the left contribution becomes about $0.192772$ after multiplying by $e^{-2}$; the right stays unchanged. The combined denominator is approximately $1.578968$. These displayed decimals are rounded; the merge identity is exact over real arithmetic, with rounding in a floating-point implementation.

from math import exp, inf, isclose, fsum
from assets.primitives_model import softmax_state, merge_states, softmax, expect_error

x = [3, -1, 4, 1, 5, -9, 2, 6]
left, right = softmax_state(x[:4]), softmax_state(x[4:])
merged = merge_states(left, right)
assert merged[0] == 6
assert isclose(merged[1], fsum(exp(v - 6) for v in x), rel_tol=1e-15)
assert isclose(fsum(softmax(x)), 1.0, abs_tol=1e-15)
assert softmax([None, 5, None]) == [0, 1, 0]
assert softmax([-10000, 0, 10000]) == [0, 0, 1]
assert merge_states((-inf, 0), (-inf, 0)) == (-inf, 0)
expect_error(ValueError, lambda: softmax([None, None]))
expect_error(ValueError, lambda: softmax([inf, 1]))
for a, b in zip(softmax(x), softmax([v + 1000 for v in x])):
    assert isclose(a, b, rel_tol=1e-14, abs_tol=1e-15)
print(f"left: m={left[0]}, d={left[1]:.6f}")
print(f"right: m={right[0]}, d={right[1]:.6f}")
print(f"merged: m={merged[0]}, d={merged[1]:.6f}")
print("Shift invariance, extreme logits, empty identity and masking checked")

left: m=4, d=1.424404
right: m=6, d=1.386195
merged: m=6, d=1.578968
Shift invariance, extreme logits, empty identity and masking checked

class CountReads:
    def __init__(self, values):
        self.values = values
        self.reads = 0

    def __iter__(self):
        for value in self.values:
            self.reads += 1
            yield value

row = CountReads(x)
assert softmax_state(row)[0] == 6
assert row.reads == 8
print("Normalizer input reads:", row.reads)
row.reads = 0
probabilities = softmax(row)
assert len(probabilities) == 8 and row.reads == 16
assert isclose(fsum(probabilities), 1.0, abs_tol=1e-15)
expect_error(ValueError, lambda: softmax(iter(x)))
print("Complete softmax input reads:", row.reads)
print("Eight probabilities need another input pass; one-shot iterator rejected")

Normalizer input reads: 8
Complete softmax input reads: 16
Eight probabilities need another input pass; one-shot iterator rejected

This counts Python model reads, not CUDA transactions. If a later logit changes the maximum or denominator, a probability emitted earlier would be wrong unless it were updated. Keeping logits on chip or fusing a different output contract changes the storage problem.

How FlashAttention builds on the online normalizer: This online state update is the foundation behind FlashAttention (Dao et al., 2022).^{[11]Reference 11FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.https://arxiv.org/abs/2205.14135} Standard attention computes $S = QK^T$, takes the full softmax $P = \operatorname{softmax}(S)$, and multiplies by $V$, materializing an $N \times N$ attention matrix in HBM ($O(N^2)$ memory traffic).

FlashAttention loads blocks of $Q$, $K$, and $V$ into fast SRAM (shared memory/registers). For each block of keys and values, it computes local dot products $S_{\text{tile}}$, updates the running online softmax state $(m, d)$, and rescales the accumulated output accumulator $O$ in registers:

Here $O$ denotes the unnormalized weighted-value accumulator. At the end, divide each row by its final $d$. FlashAttention emits weighted values, so it can avoid storing every probability. For fixed head dimension, it avoids quadratic score/probability storage and retains linear-size outputs/statistics. It still does quadratic attention arithmetic and has shape-dependent HBM traffic; avoiding materialization doesn't guarantee a compute-bound kernel.^{[11]Reference 11FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.https://arxiv.org/abs/2205.14135}

Race conditions and verification: sanitizing shared-memory hazards

The scan exchanges partials through shared arrays. Diagnose three access-order hazards:

• RAW (Read-After-Write): Thread A writes to a shared index, but Thread B reads before the write finishes.
• WAR (Write-After-Read): Thread A reads from a shared location, but Thread B overwrites that location before Thread A finishes loading.
• WAW (Write-After-Write): Multiple threads attempt to write to the same shared address without a deterministic sequence.

NVIDIA's Compute Sanitizer provides specialized tools to detect these defects at the machine-instruction level:^{[12]Reference 12Compute Sanitizerhttps://docs.nvidia.com/compute-sanitizer/ComputeSanitizer/}

compute-sanitizer --tool memcheck --error-exitcode=99 ./primitives_lab 4099 1
compute-sanitizer --tool racecheck --error-exitcode=99 ./primitives_lab 4099 1
compute-sanitizer --tool synccheck --error-exitcode=99 ./primitives_lab 4099 1

Each tool addresses a specific failure mode:

1. synccheck diagnoses supported synchronization violations, including divergent block/warp barriers. A block barrier in a nonuniform branch can hang or behave incorrectly; keep the lab's barriers outside its per-index guards.
2. racecheck detects supported shared-memory hazards on executed paths and reports conflicting accesses. Two unordered reads aren't a race; a conflicting write is necessary. A warp barrier only orders its participants, so it doesn't replace a block barrier when another warp consumes the data. This tool isn't a general global-memory race detector.
3. memcheck detects out-of-bounds reads and writes in global and shared memory, plus unaligned memory accesses.

Keep two verification rules firmly separated:

• A clean sanitizer run doesn't prove arithmetic correctness. Forgetting block offsets can leave the kernel free of shared-memory hazards but numerically wrong. In the pictured example, all four second-block outputs miss 7; if an earlier block's total were zero, the omission could be hidden.
• A passing numerical run doesn't prove race-freedom. A missing barrier can happen to produce the right answer for one schedule and fail on another. Run reference checks and the relevant sanitizer tools, then retain the shape/source/settings for each result. Clean reports cover their exercised paths, not all possible executions.

Select the primitive by its output contract

Different workloads demand different output contracts:

Use CUB's DeviceReduce and DeviceScan as strong baselines for standalone whole-array operations. BlockReduce and BlockScan are different: threads invoke them cooperatively inside your kernel, so they can participate in fusion. CUB is part of CUDA Core Compute Libraries (CCCL) and dispatches architecture/type-specific policies; that isn't a promise of runtime benchmarking or optimal end-to-end performance.^{[13]Reference 13CUB Developer Overviewhttps://nvidia.github.io/cccl/unstable/cub/developer_overview.html}

The explicit-temporary-storage overloads query bytes, allocate storage, then execute. Current documentation also exposes environment-based overloads that manage temporary allocation, so the query pattern isn't mandatory for every API. Check your installed CCCL version. For an int32-to-int64 sum comparison, explicitly widen inputs; recording only the output type is insufficient to describe the complete accumulator contract.

#include <cub/cub.cuh>
#include <thrust/iterator/transform_iterator.h>
// Declare this conversion at namespace scope.
struct Widen {
    __host__ __device__ long long operator()(int value) const {
        return static_cast<long long>(value);
    }
};
// Remaining fragment: inside a function with device pointers, count, and stream.
// CHECK_CUDA is the lab's error-checking macro.
auto input64 = thrust::make_transform_iterator(input, Widen{});
void* temporary = nullptr;
std::size_t temporary_bytes = 0;
CHECK_CUDA(cub::DeviceReduce::Sum(
    temporary, temporary_bytes, input64, output, count, stream));
CHECK_CUDA(cudaMalloc(&temporary, temporary_bytes));
CHECK_CUDA(cub::DeviceReduce::Sum(
    temporary, temporary_bytes, input64, output, count, stream));
CHECK_CUDA(cudaStreamSynchronize(stream));  // Complete before releasing storage.
CHECK_CUDA(cudaFree(temporary));

CCCL 3.0 removed CUB's legacy TransformInputIterator; use the Thrust transform iterator above to widen each loaded element before reduction.^{[14]Reference 14CCCL 2.x - CCCL 3.0 migration guidehttps://nvidia.github.io/cccl/unstable/cccl/3.0_migration_guide.html}

Here input is an initialized device int*, output a device long long*, and all intermediate sums must fit. The fragment illustrates API use; the supplied benchmark doesn't include a CUB variant. Reuse scratch storage outside measured repetitions rather than timing this allocation/free sequence.^{[5]Reference 5cub::DeviceReducehttps://nvidia.github.io/cccl/unstable/cub/api/structcub_1_1DeviceReduce.html}

When might a custom kernel help?

1. Fusion: A device-wide reduction followed by normalization may write intermediates and require separate launches. A fused LayerNorm or attention kernel can reuse data on chip, possibly with CUB block/warp primitives. RoPE's pairwise rotation isn't itself a reduction.
2. Layout or specialization: Strides aren't automatically beyond CUB: its iterators can adapt layouts. Register-tiled fragments or a specialized surrounding computation may still favor an integrated implementation. Compare compatible types, operators, outputs, determinism, and full-pipeline timing before replacing the baseline.

Run the CUDA lab and verify correctness boundaries

Download these files while preserving the relative assets/ paths used by the commands and Python imports:

• The CUDA lab source
• Host validation header
• CPU algorithm model
• Historical benchmark receipt

The CUDA download implements integer reduction and exclusive scan, each with a deliberately serial baseline. Arg reduction and online softmax are CPU examples and extension tasks, not existing CUDA kernels. The receipt belongs to an older source hash, not the revised download. Compile with optimization and line information:

: "${GPU_ARCH:?Set GPU_ARCH to your installed GPU target, for example sm_80}"
nvcc -O3 -std=c++17 -lineinfo -arch="$GPU_ARCH" \
  assets/primitives_lab.cu -o primitives_lab
./primitives_lab 4194304 30

Before benchmarking, run correctness sweeps across warp boundaries (31, 32, 33), block boundaries (255, 256, 257, 511, 512, 513), and recursive scan boundaries ($512^2 + 1 = 262,145$):

for n in 1 31 32 33 255 256 257 511 512 513 4099 262145; do
  ./primitives_lab "$n" 1 || exit 1
done

CPU checks you can run immediately

The companion Python model verifies these same index boundaries, checking 17 hierarchy sizes, 260 recursive scan cases, signed overflow handling, and online softmax partition merges:

from assets.primitives_model import run_checks
run_checks()

CPU models: 17 hierarchy sizes and 260 recursive scan cases pass
Overflow, absent shuffle sources, invalid widths and masked/nonfinite rows checked
No CUDA compilation, GPU synchronization, sanitizer or performance claim

Interpret timing without confusing traffic and bandwidth

With the displayed 4194304 30 arguments, the benchmark uses CUDA events around 30 complete variant invocations after five warmups each. Running with no arguments instead selects $2^{24}$ inputs and 50 iterations. Allocation and initial host-to-device transfers are outside the event interval. Hierarchical reduction includes its final device-to-device scalar copy; scan includes every recursive scan/offset launch. The event interval can include submission gaps between these operations.

The benchmark reports logical byte rate:

• For reduction: $4N$ input bytes divided by elapsed time.
• For scan: $4N$ input bytes plus $8N$ output bytes ($12N$ bytes total) divided by elapsed time.

This numerator counts useful algorithmic data, not physical DRAM transactions. Timing includes the intermediate work, but the numerator omits its bytes and the reduction scalar output. Cache hits change where bytes are served, not this logical byte count. Buffers are reused after warmup; residency isn't measured. Correctness is checked before timing, not for every timed invocation.

Record the source hash, compiler/toolkit, GPU, shape, iteration count, cache policy, and the exact timed composition. The one-thread kernels explain dependency and serve as correctness baselines; their speedups don't establish competitiveness with CUB. For a library comparison, add a matched CUB variant and rerun actual GPU checks/timing.

Review questions

1. Why does warp shuffle avoid shared memory bank conflicts during intra-warp reduction?

The exchange uses registers rather than shared-memory addresses, so it doesn't incur shared-memory bank conflicts. It still requires valid participants and source lanes, with target-dependent instruction cost. The block reduction retains a block-wide barrier when other warps publish shared partials.

2. What is the fundamental algorithmic tradeoff between Hillis-Steele and Blelloch parallel scan?

Hillis-Steele executes $\log_2 N$ steps with $O(N \log_2 N)$ total additions, prioritizing step efficiency at the expense of work efficiency. Blelloch splits execution into Up-Sweep reduction and Down-Sweep distribution phases taking $2 \log_2 N$ steps but only $O(N)$ operations, matching serial work complexity.

3. Why do floating-point parallel reductions yield non-identical results across different tile configurations?

IEEE 754 floating-point addition is non-associative: $(a + b) + c \neq a + (b + c)$. Altering thread block dimensions, tile sizes, or reduction tree hierarchies shifts intermediate rounding boundaries, producing slight bitwise discrepancies even in strictly race-free code.

4. How does decoupled look-back eliminate multi-pass kernel launches in global prefix scans?

Tiles publish aggregates and completed prefixes with a correct visibility protocol. A tile can accumulate predecessor aggregates until it finds a completed prefix, incorporating that offset before writing its own outputs. This avoids a separate full-array offset pass; metadata, initialization, and forward-progress requirements remain.

5. Why does online softmax rescale running accumulators when updating the maximum?

Changing the reference maximum changes every previous exponential's scale: $e^{x-m_{\text{new}}}=e^{x-m_{\text{old}}}e^{m_{\text{old}}-m_{\text{new}}}$. Rescale the existing sum before adding a new contribution. This yields the normalizer in one input pass, but a standalone full probability vector still needs stored logits or another input pass. An empty state must be handled separately to avoid evaluating $-\infty-(-\infty)$.