Highly optimized CUDA snippets

Optimization is not a grab-bag of tricks. Every fast kernel is fast for one reason: it stopped wasting the resource that was holding it back. Learn to find that wall, attack it, and re-measure.

The roofline mental model

A GPU can do two things: move bytes (between HBM and the chip) and do arithmetic (FLOPs in the SMs). These have separate ceilings. For a given kernel define its arithmetic intensity:

intensity I = total FLOPs / total bytes moved to/from HBM   (FLOP per byte)

The achievable rate is the lower of two limits: what compute can sustain (a flat ceiling, peak FLOP/s) and what bandwidth can feed (a slope, bandwidth × intensity). Plot attainable FLOP/s against intensity and you get the roofline:

The ridge point I* = peak_FLOPs / peak_BW is the crossover. On a modern data-center GPU peak fp32 might be ~20 TFLOP/s and HBM bandwidth ~2 TB/s, so I* ≈ 10 FLOP/byte. Read that number carefully:

• I < I* → you are on the slope → memory-bound. You will never hit peak FLOP/s no matter how clever your math is. The only thing that helps is moving fewer bytes (or moving them more efficiently).
• I > I* → you are under the flat roof → compute-bound. Bandwidth is not your problem; reducing instructions, raising ILP, or using tensor cores helps.

Classify before you optimize. Count the bytes a kernel must touch and the FLOPs it must do, divide, and compare to I*:

Memory coalescing from first principles

HBM is not byte-addressable for the SM. The memory system serves aligned 32-byte sectors, grouped into 128-byte cache lines. When a warp's 32 lanes issue a load, the hardware looks at the 32 addresses and figures out the minimum set of sectors that covers them. If the 32 lanes touch one contiguous, aligned 128-byte region, the hardware coalesces the request into a handful of sector transactions and every fetched byte is used. If the addresses are scattered, each lane may drag in its own 32-byte sector of which only 4 bytes are wanted — and you pay for the rest.

// Each lane strides by blockDim.x -> neighbouring lanes are far apart.
// Lane i and lane i+1 are 'blockDim.x' floats apart: scattered access,
// many sectors fetched, most bytes wasted.
__global__ void bad_transpose_read(const float* in, float* out, int n) {
    int t = threadIdx.x;                                  // 0..blockDim-1
    for (int i = 0; i < n; ++i)
        out[i * blockDim.x + t] = in[i * blockDim.x + t]; // stride blockDim.x
}

// Consecutive global thread indices touch consecutive addresses.
// Lane i and lane i+1 are adjacent floats -> one aligned line per warp.
__global__ void good_copy(const float* in, float* out, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) out[i] = in[i];
}

The rule reduces to: index global memory with blockIdx.x * blockDim.x + threadIdx.x so that threadIdx.x (the fastest-varying lane index) maps to the fastest-varying address. A stride-2 access wastes half your bandwidth; a stride-32 access can waste ~87%. On a memory-bound kernel that is a direct multiplier on runtime.

The optimization decision table

This is the table to think through when a kernel is slow. Symptom → the wall it implies → the first lever to pull. Profilers (Nsight Compute) report most of these symptoms directly.

Snippet 1 — grid-stride loop with vectorized memory

For a bandwidth-bound elementwise kernel, two things matter: every byte fetched should be used (coalescing) and you should issue the widest load the hardware supports so fewer instructions move the same bytes. A float4 load moves 16 bytes per lane in one instruction — a warp then moves 512 bytes in one coalesced transaction set. The grid-stride loop makes one launch handle any input size and lets you tune the grid independently of n.

// Bandwidth-bound: y = a*x + y  (SAXPY).
// Precondition for the vectorized path: x and y are 16-byte aligned and
// contiguous. reinterpret_cast to float4 is only valid if the base pointer
// is aligned to alignof(float4) == 16. cudaMalloc guarantees >= 256-byte
// alignment, so a base pointer is fine; a mid-array offset may not be.
__global__ void saxpy_vec4(int n, float a, const float* x, float* y) {
    int gtid   = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;

    int n4 = n / 4;                                   // number of full float4s
    const float4* x4 = reinterpret_cast<const float4*>(x);
    float4*       y4 = reinterpret_cast<float4*>(y);

    // Vectorized body: each iteration moves 4 floats per thread, coalesced.
    for (int i = gtid; i < n4; i += stride) {
        float4 xv = x4[i];
        float4 yv = y4[i];
        yv.x = a * xv.x + yv.x;
        yv.y = a * xv.y + yv.y;
        yv.z = a * xv.z + yv.z;
        yv.w = a * xv.w + yv.w;
        y4[i] = yv;
    }

    // Scalar tail: handle the up-to-3 leftover elements when n % 4 != 0.
    for (int i = n4 * 4 + gtid; i < n; i += stride)
        y[i] = a * x[i] + y[i];
}

The grid-stride loop also means a thread block stays resident and reused rather than launching one thread per element; you size the grid to fill the GPU (e.g. a few blocks per SM) and the loop covers the rest. Vectorization does not increase intensity — it is still ~0.17 FLOP/byte — but it reduces instruction overhead and maximizes per-transaction efficiency, which is exactly what a memory-bound kernel needs.

Snippet 2 — warp reduction via __shfl_down_sync

A reduction (sum, max, …) collapses many values to one. The naive version writes partial sums to shared memory and barriers between every tree level. But the 32 lanes of a warp can exchange registers directly with shuffle instructions — no shared memory, no __syncthreads. The data lives in 32 lane registers; each step halves the active distance, folding the upper half into the lower half.

// Sum the 32 values held one-per-lane across a warp. Returns the full
// sum in lane 0 (shfl_down leaves upper lanes stale, so only lane 0's
// result is meaningful unless you broadcast it afterward).
// Defined here cleanly: lesson 03 reuses warp_reduce_sum by name.
__inline__ __device__ float warp_reduce_sum(float v) {
    // Full warp participates; mask names all 32 lanes.
    for (int offset = 16; offset > 0; offset >>= 1)
        v += __shfl_down_sync(0xffffffffu, v, offset);
    return v;   // valid in lane 0
}

The mask 0xffffffffu declares all 32 lanes participate; on modern GPUs warps execute independently and the _sync intrinsics require you to name the participating lanes so the compiler can reason about reconvergence. Use a partial mask only when you know some lanes are inactive (e.g. a tail warp), and compute it with __activemask() rather than guessing.

Snippet 3 — block reduction (two-level tree)

To reduce a whole block (up to 1024 threads = up to 32 warps) we compose Snippet 2 with one short shared-memory stage. Each warp reduces itself to one value (no barriers), writes that single value to a 32-slot shared array, and then warp 0 reduces those ≤32 partial sums with one more warp-reduce. This minimizes both barriers (exactly one) and shared traffic (32 writes + 32 reads per block).

__inline__ __device__ float block_reduce_sum(float v) {
    static __shared__ float warp_sums[32];    // one slot per warp (max 32 warps)

    int lane = threadIdx.x & 31;               // lane index within warp (0..31)
    int wid  = threadIdx.x >> 5;               // warp index within block

    v = warp_reduce_sum(v);                    // level 1: each warp -> its lane 0
    if (lane == 0) warp_sums[wid] = v;         // each warp publishes its sum

    __syncthreads();                           // the single barrier we need

    // level 2: warp 0 loads the per-warp sums and reduces them.
    int num_warps = (blockDim.x + 31) / 32;
    v = (threadIdx.x < num_warps) ? warp_sums[lane] : 0.0f;
    if (wid == 0) v = warp_reduce_sum(v);      // block total ends up in thread 0
    return v;                                  // valid in threadIdx.x == 0
}

Compare to the textbook shared-memory tree, which barriers log2(blockDim) times and reads/writes shared memory at every level. The two-level version does the bulk of the work in registers via shuffles and touches shared memory only at the warp boundary. Thread 0 of each block then does one atomicAdd to a global accumulator (one atomic per block — see Snippet 5 for why that matters).

Snippet 4 — shared-memory tiled matmul

Naive matmul reloads a full row of A and column of B from HBM for every output element: O(N³) bytes for O(N³) FLOPs → intensity ≈ O(1), hopelessly memory-bound. The fix is reuse: load a TILE×TILE block of A and of B into shared memory once, then every thread in the block reads those tiles many times from on-chip memory. An A tile is reused across all output columns in the tile; a B tile across all output rows. This raises intensity by roughly a factor of TILE, pushing the kernel toward the compute roof.

#define TILE 32
// C = A * B, all row-major, square N x N for clarity.
__global__ void matmul_tiled(const float* A, const float* B, float* C, int N) {
    // +1 padding on the leading dim of As avoids shared-memory bank conflicts
    // on the column-wise read of As below (see "bank conflicts" section).
    __shared__ float As[TILE][TILE + 1];
    __shared__ float Bs[TILE][TILE];

    int ty = threadIdx.y, tx = threadIdx.x;
    int row = blockIdx.y * TILE + ty;
    int col = blockIdx.x * TILE + tx;

    float acc = 0.0f;                        // accumulator lives in a register

    for (int t = 0; t < N; t += TILE) {
        // Coalesced loads: consecutive tx -> consecutive global columns.
        As[ty][tx] = (row < N && t + tx < N) ? A[row * N + (t + tx)] : 0.0f;
        Bs[ty][tx] = (col < N && t + ty < N) ? B[(t + ty) * N + col] : 0.0f;
        __syncthreads();                     // tiles fully loaded before use

        #pragma unroll
        for (int k = 0; k < TILE; ++k)
            acc += As[ty][k] * Bs[k][tx];    // both reads hit shared memory

        __syncthreads();                     // done with tiles before reload
    }
    if (row < N && col < N) C[row * N + col] = acc;
}

#pragma unroll on the inner loop lets the compiler issue the multiply-adds back to back with no loop overhead and good ILP, while acc stays in a register the whole time. Each output element is touched once in HBM (the final write); every reuse comes from shared memory.

Shared-memory bank conflicts

Shared memory is split into 32 banks, each 4 bytes wide, interleaved: address a lives in bank (a/4) % 32. A warp can read 32 different banks in one cycle. But if two lanes of a warp hit two different addresses in the same bank, the accesses serialize — an N-way conflict costs N cycles. (All lanes reading the same address is fine; that is a broadcast, not a conflict.)

Consider a __shared__ float tile[32][32]. Reading tile[ty][k] across the warp (varying tx, fixed k) — a column access — means lane i touches address (i*32 + k)*4. The bank is (i*32 + k) % 32 = k % 32 for every lane: all 32 lanes land in the same bank → a 32-way conflict, serialized.

That is the +1 in As[TILE][TILE + 1]. Padding the leading dimension by one element makes consecutive rows fall in consecutive banks, so a column read spreads across all 32 banks with no conflict. It costs a few unused bytes of shared memory and removes the serialization entirely.

Snippet 5 — privatized histogram

A histogram increments a bin per input element. If every thread does atomicAdd(&global_bins[v], 1) directly, and the data is skewed, thousands of threads serialize on the same hot bin — atomic contention dominates and the kernel crawls. With N elements you issue N global atomics.

Privatize: give each block its own copy of the bins in shared memory, accumulate there (fast on-chip atomics with far less contention — only the threads of one block compete), then have the block flush its partial histogram to global with one atomic per nonzero bin.

#define NBINS 256
__global__ void histogram_privatized(const unsigned char* data, int n,
                                      unsigned int* global_bins) {
    __shared__ unsigned int local_bins[NBINS];

    // Cooperatively zero the shared bins.
    for (int b = threadIdx.x; b < NBINS; b += blockDim.x)
        local_bins[b] = 0u;
    __syncthreads();

    // Grid-stride over the input; atomics now hit shared memory, not HBM.
    int gtid   = blockIdx.x * blockDim.x + threadIdx.x;
    int stride = blockDim.x * gridDim.x;
    for (int i = gtid; i < n; i += stride)
        atomicAdd(&local_bins[data[i]], 1u);
    __syncthreads();

    // Flush: one global atomic per nonzero bin per block.
    for (int b = threadIdx.x; b < NBINS; b += blockDim.x) {
        unsigned int c = local_bins[b];
        if (c != 0u) atomicAdd(&global_bins[b], c);
    }
}

The traffic against the global atomic unit drops from N atomics to at most blocks × 256. For a 100M-element image with ~1000 blocks that is ~256k global atomics instead of 100M — a ~400× reduction in contention, plus the per-element atomics now run on much faster shared memory. The wall here was contention, not bandwidth; privatization is the lever aimed at it.

Occupancy and register pressure

Occupancy is resident warps / maximum warps the SM can hold. The GPU hides memory latency by switching to another ready warp while one stalls; more resident warps means more latency to hide. Occupancy is capped by whichever per-block resource runs out first:

• Registers: an SM has a fixed register file. If each thread needs 64 registers and the SM has 65536, at most ~1024 threads fit regardless of other limits. Use too many and either occupancy drops or registers spill to local memory (which is really HBM) — a latency disaster the profiler flags as spills.
• Shared memory: a block's shared-memory request is subtracted from the SM's fixed pool; large tiles cap how many blocks co-reside.
• Block/warp slots: hard hardware limits on resident blocks and warps per SM.

// Hint the compiler: at most 256 threads/block, aim for >= 4 blocks/SM.
// The compiler then budgets registers so this occupancy is achievable,
// trading register count for spills as needed.
__global__ void __launch_bounds__(256, 4)
my_kernel(const float* in, float* out, int n) { /* ... */ }

Snippet 6 — pinned memory + CUDA event timing

Kernel launches are asynchronous, so wrapping a launch in CPU wall-clock timing measures almost nothing — the call returns before the GPU finishes. CUDA events are timestamps recorded in the stream; the elapsed time between two events is the actual GPU work done between those two points. Pinned (page-locked) host memory additionally lets cudaMemcpyAsync run by DMA concurrently with compute, instead of the slower pageable path.

float* h_buf;
cudaMallocHost(&h_buf, bytes);           // pinned host memory (page-locked)

cudaEvent_t start, stop;
cudaEventCreate(&start);
cudaEventCreate(&stop);

cudaEventRecord(start);                   // timestamp in the (default) stream
my_kernel<<<grid, block>>>(/* ... */);     // enqueued after start
cudaEventRecord(stop);                    // timestamp after the kernel

cudaEventSynchronize(stop);               // wait until 'stop' has occurred
float ms = 0.0f;
cudaEventElapsedTime(&ms, start, stop);   // GPU time between the two events

cudaEventDestroy(start);
cudaEventDestroy(stop);
cudaFreeHost(h_buf);

The key point: cudaEventElapsedTime measures the GPU timeline between two stream points, immune to host-side scheduling. CPU timers around an async launch measure launch overhead, not the kernel. Always synchronize on the stop event before reading the elapsed time.

Snippet 7 — CUDA Graphs for repeated small-kernel sequences

Each kernel launch costs a few microseconds of host-side work (argument marshaling, driver bookkeeping). For a few big kernels that is invisible; for a loop that fires dozens of tiny kernels per iteration — an autoregressive decode step, say — those microseconds dominate and the GPU sits idle in the gaps. CUDA Graphs let you capture the whole sequence once and replay it with a single launch, collapsing the per-kernel host overhead into one.

cudaStream_t stream;
cudaStreamCreate(&stream);

cudaGraph_t graph;
cudaGraphExec_t exec;
bool captured = false;

for (int step = 0; step < num_steps; ++step) {
    if (!captured) {
        // Capture the sequence of small kernels once.
        cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);
        decode_attention<<<g1, b1, 0, stream>>>(/* ... */);
        decode_mlp      <<<g2, b2, 0, stream>>>(/* ... */);
        decode_sample   <<<g3, b3, 0, stream>>>(/* ... */);
        cudaStreamEndCapture(stream, &graph);
        cudaGraphInstantiate(&exec, graph, nullptr, nullptr, 0);
        captured = true;
    }
    cudaGraphLaunch(exec, stream);    // one launch replays all three kernels
}
cudaStreamSynchronize(stream);

Graphs attack launch overhead only. They do not move fewer bytes or do less math — a bandwidth-bound kernel is exactly as bandwidth-bound inside a graph. Use them when the timeline shows the GPU idling between many short kernels. (Fusing the kernels into one is the other lever for the same wall; graphs are preferred when fusion is impractical or the kernels are heterogeneous.)

Async copies and double buffering (Ampere+)

In the tiled matmul, the block loads a tile, barriers, computes, barriers, loads the next tile — load and compute are serialized. On Ampere and later, cp.async (exposed via cuda::memcpy_async / the __pipeline API) issues a global→shared copy that proceeds in the background without occupying registers as a staging area. Combined with double buffering — two shared-memory tile buffers, computing on buffer A while the copy fills buffer B — this overlaps the next tile's load with the current tile's math, hiding memory latency behind compute. It is the modern path to keeping the SMs fed in a tiled kernel and a core ingredient of library-grade matmul.

Optimization answer template

When asked "how would you make this kernel faster?", do not list tricks. Walk the loop, aiming each step at the confirmed wall:

1. Classify. Count bytes moved and FLOPs done; compute intensity and compare to the ridge point. Memory-bound or compute-bound? This decides everything that follows.
2. Coalesce. Ensure threadIdx.x maps to consecutive addresses so the warp's accesses collapse into minimal, fully-used transactions. Vectorize (float4) where alignment allows.
3. Reuse. If a value is read more than once, stage it on-chip — shared memory tiles, then registers / micro-tiles — to raise intensity and stop re-fetching from HBM.
4. Fuse. Combine adjacent kernels so an intermediate is produced and consumed while still on-chip, eliminating a full HBM round trip. For many tiny kernels, fuse or capture as a CUDA Graph to kill launch overhead.
5. Reduce contention. Replace global atomics with privatized accumulation; remove unnecessary barriers and fix bank conflicts (pad the leading dimension); restructure branches to keep warps uniform.
6. Measure. Time with CUDA events, profile with Nsight, and confirm the bottleneck moved. Re-classify and repeat from step 1. The loop ends when the kernel is bound by a resource you cannot cheaply reduce.