First kernels: elementwise, reductions, softmax

Everything in this part has an arithmetic intensity between 0.2 and 3. For a large enough tensor, Part 3 therefore gives the first target: compulsory bytes divided by achievable bandwidth. Reaching it still requires coalesced traffic, enough work in flight, and little unnecessary movement. A naive elementwise kernel gets maybe 60% of peak; a naive reduction gets 5%.

Start with a copy. Its job is to read n bytes and write n bytes, so on an H100 the floor for a 1 GB tensor is 2×10⁹ / 3.35×10¹² = 597 µs. Whether you get near it depends entirely on how much traffic you keep in flight, which is Little's Law from Part 1:

Two knobs get you there: how many warps are resident, and how many independent loads each thread has outstanding. They multiply.

In code, the two knobs are the vector width and the unroll factor:

// A grid-stride copy with 16-byte accesses and 4 loads in flight per thread.
__global__ void scale4(int n4, float a, const float4* __restrict__ x, float4* __restrict__ y)
{
    int stride = blockDim.x * gridDim.x;
    for (int i = blockIdx.x*blockDim.x + threadIdx.x; i < n4; i += stride*4) {
        float4 a0, a1, a2, a3;                // four INDEPENDENT loads —
        a0 = x[i];                            // the compiler issues all four
        a1 = (i+stride   < n4) ? x[i+stride  ] : a0;   // before waiting on any
        a2 = (i+stride*2 < n4) ? x[i+stride*2] : a0;
        a3 = (i+stride*3 < n4) ? x[i+stride*3] : a0;
        a0.x*=a; a0.y*=a; a0.z*=a; a0.w*=a;   // scale every loaded vector
        a1.x*=a; a1.y*=a; a1.z*=a; a1.w*=a;
        a2.x*=a; a2.y*=a; a2.z*=a; a2.w*=a;
        a3.x*=a; a3.y*=a; a3.z*=a; a3.w*=a;
        y[i] = a0;
        if (i+stride   < n4) y[i+stride  ] = a1;
        if (i+stride*2 < n4) y[i+stride*2] = a2;
        if (i+stride*3 < n4) y[i+stride*3] = a3;
    }
}

# Triton gets the same effect from a large BLOCK_SIZE: the block is
# lowered to vectorized loads and the loads are issued together.
@triton.jit
def scale(x_ptr, y_ptr, a, n, BLOCK: tl.constexpr):
    pid  = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    x = tl.load(x_ptr + offs, mask=mask)
    tl.store(y_ptr + offs, a * x, mask=mask)

# BLOCK=4096 with num_warps=8 gives each thread 16 elements in flight.
# This is what triton.autotune searches over — see Part 9.

y = a * x     # dispatches to a hand-tuned elementwise kernel
              # torch.compile will fuse this with its neighbours

A reduction turns n values into one. It is the shape underneath sum, max, norm, softmax, and every loss function, and it is the standard worked example of GPU optimization because the naive version is so bad.

The difficulty is that a reduction is inherently sequential in its dependency structure, so you have to build a tree. How you index that tree determines whether your warps diverge and whether your shared-memory accesses conflict — both of which you now know how to reason about.

The modern answer is the third one. Since Kepler, a warp can exchange registers directly:

// Sum a contiguous set of active lanes starting at lane 0.
// A full warp takes 5 instructions; the result lands in its first lane.
__inline__ __device__ float warpReduceSum(float v, unsigned mask) {
    int lane = threadIdx.x & 31;
    int count = __popc(mask);
    int off = 1;
    while (off < count) off <<= 1;
    for (off >>= 1; off > 0; off >>= 1) {
        float other = __shfl_down_sync(mask, v, off);
        if (lane + off < count) v += other;
    }
    return v;
}

// Two-level: reduce within each warp, stage up to 32 partials in SMEM,
// then have warp 0 reduce those. Works for any block size up to 1024.
__inline__ __device__ float blockReduceSum(float v) {
    __shared__ float s[32];
    int lane = threadIdx.x & 31, wid = threadIdx.x >> 5;
    int lanesThisWarp = min(32, blockDim.x - wid*32);
    unsigned warpMask = lanesThisWarp == 32
        ? 0xffffffffu : ((1u << lanesThisWarp) - 1u);
    v = warpReduceSum(v, warpMask);
    if (lane == 0) s[wid] = v;
    __syncthreads();
    int numWarps = (blockDim.x + 31) >> 5;
    float blockSum = 0.0f;
    if (wid == 0 && lane < numWarps) {
        blockSum = s[lane];
        unsigned finalMask = numWarps == 32
            ? 0xffffffffu : ((1u << numWarps) - 1u);
        blockSum = warpReduceSum(blockSum, finalMask);
    }
    return blockSum;                  // valid in thread 0
}

# In Triton the reduction is a single call; the compiler emits the
# shuffle tree, the shared-memory staging, and the barriers for you.
total   = tl.sum(x, axis=0)
maximum = tl.max(x, axis=0)

# Associative custom reductions are available too:
@triton.jit
def _combine(m0, d0, m1, d1):
    m = tl.maximum(m0, m1)
    return m, d0 * tl.exp(m0 - m) + d1 * tl.exp(m1 - m)

The step count is worth stating exactly, because it is the whole complexity story. Within a full warp: log₂32 = 5 shuffle steps. Within a block: those 5, plus one shared-memory round trip, one barrier, and up to 5 more shuffles for the warp partials. Across the grid: each block writes its single partial to global memory, and a second tiny kernel over the B partials finishes the job — or each block does one FP32 atomicAdd, which saves the launch but can make the result run-to-run nondeterministic because floating-point addition is not associative and block arrival order can change. The two-kernel path adds a logical 8B bytes for writing and rereading the partials; atomics instead perform contended read-modify-writes whose physical traffic depends on the cache and memory system. Tree depth is O(log N) at every level, but only the first level touches the full input; that is the difference between a tuned reduction sitting on bandwidth and a naive one sitting at 5% (Harris 2007).

One variant to know: built as a butterfly with __shfl_xor_sync, the same 5 steps leave the sum in every lane instead of only lane 0 — which is what you want when all threads consume the result, as in the norms of §4.5 (those broadcast through shared memory instead). Same depth either way; the choice is only about where the answer has to end up.

Softmax needs a maximum for numerical stability, a sum of exponentials for normalization, and then a division. Written naively that is three passes over the data:

Milakov and Gimelshein's observation is that the first two passes can be fused, if you are willing to retroactively correct the running sum every time the maximum increases:

That bolded factor is the whole trick. When the running max jumps from m_j−1 to a larger m_j, every term already accumulated in d was scaled by the wrong exponent; multiplying by exp(m_j−1 − m_j) fixes all of them at once. And because the exponent is always ≤ 0, the correction never overflows.

The state is a pair (m, d), and the combining operator is associative and commutative:

which means it drops straight into the shuffle tree from §4.2. Every thread keeps its own partial (m, d), and one warp reduction combines them. That is the entire parallel online softmax.

The payoff is memory accesses per element: 4 for the safe 3-pass version, 3 for the online version, and the paper measures a 1.15–1.3× end-to-end speedup for softmax alone. That is a modest win here — and a decisive one in attention, where the "array" is a score matrix you cannot afford to store at all.

From Part 1's instruction table: the special-function unit produces 16 results per clock per SM. On an H100 that works out to about 3.9 trillion exp results/s, against 989.5 TFLOP/s of BF16 matmul — unlike units, but a useful throughput ratio for the operation mix that FlashAttention-3 rounds to 256×.

Practical consequences:

• Use exp2, not exp. The hardware instruction is MUFU.EX2. Computing e^x costs an extra multiply by log₂e — so fold that constant into the scale you were already applying. Attention kernels multiply QK^T by 1/√d · log₂e and then call exp2.
• Prefer __expf/-use_fast_math semantics where accuracy allows; the accurate expf is a multi-instruction sequence.
• Overlap it with matmul. This is exactly what FlashAttention-3's ping-pong schedule does: while one warpgroup runs its GEMM on the tensor cores, the other runs its softmax on the SFU. Two different units, so they genuinely run concurrently. Part 7.

RMSNorm is the modern default (Llama, Mistral, Gemma, DeepSeek) because it drops the mean-centering entirely:

One reduction instead of two, no mean subtraction, same normalization effect in practice. It is entirely memory-bound: for hidden size H it does O(H) FLOPs on 2H·s bytes, an intensity around 0.75.

// One block per row. H must fit in the block's registers + a reduction.
template<int BLOCK>
__global__ void rmsnorm(const __nv_bfloat16* __restrict__ x,
                        const __nv_bfloat16* __restrict__ g,
                        __nv_bfloat16* __restrict__ y, int H, float eps)
{
    const int row = blockIdx.x;
    x += (size_t)row * H;  y += (size_t)row * H;

    float acc = 0.f;
    for (int i = threadIdx.x; i < H; i += BLOCK) {
        float v = __bfloat162float(x[i]);
        acc = fmaf(v, v, acc);                 // accumulate in FP32 — always
    }
    acc = blockReduceSum(acc);

    __shared__ float inv;
    if (threadIdx.x == 0) inv = rsqrtf(acc / H + eps);
    __syncthreads();

    for (int i = threadIdx.x; i < H; i += BLOCK)   // second pass hits L1/L2
        y[i] = __float2bfloat16(__bfloat162float(x[i]) * inv * __bfloat162float(g[i]));
}

@triton.jit
def rmsnorm(X, G, Y, stride, H, eps, BLOCK: tl.constexpr):
    row  = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    mask = cols < H
    x = tl.load(X + row*stride + cols, mask=mask, other=0.).to(tl.float32)
    inv = tl.math.rsqrt(tl.sum(x*x, axis=0) / H + eps)
    g = tl.load(G + cols, mask=mask, other=0.).to(tl.float32)
    tl.store(Y + row*stride + cols, (x*inv*g).to(tl.bfloat16), mask=mask)

rms = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + eps)
y   = (x.float() * rms * g.float()).to(x.dtype)

4.5.1 If you need the mean too: Welford

LayerNorm needs variance, and the textbook formula Var = E[x²] − E[x]² is numerically dangerous: for data with a large mean and small spread it subtracts two nearly equal large numbers. Welford's algorithm avoids that by updating incrementally:

And it merges in parallel — the merge below is Chan, Golub and LeVeque's update — which is what makes it usable on a GPU: the merge is associative in exact arithmetic, so it goes through the same shuffle tree:

For the large-tensor cases modelled in this part, bandwidth is the limiting roof, so the main structural optimization is to do more per transferred byte. Three levels of it, in increasing order of payoff:

1. Fuse elementwise chains. RMSNorm → SiLU → multiply → residual as one kernel instead of four is a 4× win, as computed in Part 3. This is what torch.compile does automatically, and it is why it helps so much on small models.
2. Fuse into the GEMM epilogue. The accumulator is already in registers when the mainloop ends. Bias, activation, residual add, and quantization all cost a handful of instructions there, versus two full passes over the output tensor as separate kernels. Part 5.
3. Fuse across the algorithm. This is FlashAttention: not "fuse these three elementwise ops" but "restructure so the intermediate never exists". Part 7.

4.6.1 What Triton does and does not do for you

Triton handles shared-memory allocation, swizzled layouts, software pipelines, reduction trees, and tensor-core instruction selection. For a masked tl.load, the mask makes out-of-bounds lanes safe and the compiler can vectorize when the pointer layout, alignment and divisibility permit; coalescing still depends on the layout you supplied and should be checked. You choose block sizes, num_warps, num_stages, and the algorithm. Current Triton APIs can express persistent scheduling and explicit warp-specialized producer/consumer pipelines, but the compiler does not invent that algorithmic restructuring for you. Those schedules remain programmer-visible, which is why the fastest attention kernels still spend so much code on them.

For everything in this part, that gap does not matter. A Triton elementwise, reduction or norm kernel will match a good CUDA one, because the ceiling is memory bandwidth and both reach it. Write these in Triton.

Next: the op at the other end of the roofline. Matmul, from a naive triple loop to something within 10% of cuBLAS, one measured step at a time.