Implementing and optimizing GPU Kernels - Softmax
GPU Kernels are mathematical functions (or programs) that are executed on the GPU. The implementation of these programs, can go really deep. Not only can the math be complicated, but you also need to optimize these programs as much as possible, to make sure they don’t become a bottleneck in scenarios like AI inference / training.
In this post, my goal was to dive deep in how kernels are actually implemented and optimized. To get a better sense of this “world”. I implemented different versions of the same kernel, and benchmarked each against the same inputs.
All the code for this post lives in a small lab repo, gpu_fun, where I implemented the same row-wise softmax six different ways — from a naive PyTorch baseline to a hand-tuned CUDA kernel — and benchmarked each one against the same swept inputs (row widths from 32 to 16384, ~200MB per tensor). The metric I care about throughout is achieved bandwidth: softmax is a memory-bound op (barely any math per byte moved), so the “roofline” ceiling is simply the GPU’s peak HBM bandwidth. The closer a kernel’s curve hugs that ceiling, the less time it’s wasting on redundant memory traffic.
Triton Softmax (Python)
Before touching raw CUDA, it’s worth seeing what a naive, unfused softmax actually costs. Written as separate PyTorch ops, each line is its own kernel launch that round-trips through HBM:
# ======================================================================
# BASELINE: naive softmax, written as separate PyTorch ops on purpose.
# Each line below is roughly its own GPU kernel launch that reads the
# tensor from HBM and writes a result back. Count the HBM round-trips:
# read x (max) -> read x + write (subtract) -> read + write (exp)
# -> read (sum) -> read + write (divide). That's a LOT of wasted traffic.
# The fused kernel does it all in ONE pass.
# Each operation, is an independent GPU operation (slow!)
# ======================================================================
def naive_softmax(x: torch.Tensor) -> torch.Tensor:
#fyi: a tensor is basically an array. it can be 1d, 2d, etc..
# dim=0 # down the rows, column-wise
# dim=1 # across the columns, row-wise
# keepdim -> keeps the result as a col-shaped tensor
# [0] is because max returns (values, indices), where values are the max elements
x_max = x.max(dim=1, keepdim=True)[0] # read x
# subtract each row's max, from that row. x_max is automatically stretched across cols (broadcasting)
z = x - x_max # read x, write z
num = torch.exp(z) # read z, write num
# sums each row, across columns
denom = num.sum(dim=1, keepdim=True)
# divide each value in each row, by that sum row's sum
return num / denom
Five separate reads/writes to global memory for one softmax. Triton lets us fuse all of that into a single kernel: load the row once into fast on-chip SRAM, do the max/exp/sum/divide entirely on-chip, write the result back once.
# ======================================================================
# PHASE 1 KERNEL: fused softmax. One program handles one row.
# The whole row is loaded into fast on-chip SRAM, normalized there, and
# written back once. No intermediate tensors touch HBM.
#
# Limitation (this is a feature for learning, not a bug): the whole row
# must fit in one block. When rows are too long for that in Phase 2,
# the fix is *tiling* — which is exactly the FlashAttention idea.
# ======================================================================
# This function will get compiled into a lower-level GPU program when it's used
# It's not "normal Python", it's compiled and run on the gpu
@triton.jit
def softmax_kernel(
out_ptr, in_ptr,
in_row_stride, out_row_stride,
n_cols,
BLOCK_SIZE: tl.constexpr,
):
# out_ptr, in_ptr are pointers -> memory addresses, it tells gpu where tensor data lives
# *_stride -> tells the kernel how far it needs to jump from one row to the next row
# n_cols -> number of columns per input
# BLOCK_SIZE -> how many elements the kernel will handle per row, usually rounder by a power of 2
# tl.constexpr -> value known at compile time, used to generate optimized GPU code
# when a kernel is launched, many copies of it are ran in parallel. one program per row
row = tl.program_id(0) # which row am I responsible for?
cols = tl.arange(0, BLOCK_SIZE) # 0..BLOCK_SIZE-1, create a vector of column indices
#### mask example:
# n_cols = 6
# BLOCK_SIZE = 8
# cols = [0, 1, 2, 3, 4, 5, 6, 7]
# mask = [True, True, True, True, True, True, False, False]
mask = cols < n_cols
# This computes the memory addresses for the current row.
in_ptrs = in_ptr + row * in_row_stride + cols
# Read row values from GPU memory
# mask -> only the valid columns
# other -> for invalid cols, use negative infinity (important, because next step takes the max()
x = tl.load(in_ptrs, mask=mask, other=-float('inf'))
x = x - tl.max(x, axis=0)
num = tl.exp(x) # since exp(-inf) = 0, fake columns don't affect
denom = tl.sum(num, axis=0)
y = num / denom
out_ptrs = out_ptr + row * out_row_stride + cols
# write the result back to the gpu, mask prevents fake values from being written
tl.store(out_ptrs, y, mask=mask)
The whole row has to fit in one block (BLOCK_SIZE = triton.next_power_of_2(n_cols)), which is the main limitation here — for very wide rows you’d need to tile, which is essentially the idea behind FlashAttention. But for the widths I’m sweeping, this one-pass-per-row kernel is already extremely close to the HBM roofline, matching torch.softmax almost exactly. That’s the bar the hand-written CUDA kernels below are trying to reach.
CUDA naive
The simplest possible CUDA version: one thread owns one entire row, and does the textbook three-pass softmax serially — a max pass, an exp+sum pass, and a normalize pass, each looping over every column in the row.
// ======================================================================
// SIMPLE VERSION: one thread handles one entire row, serially.
// No block-level cooperation, no reductions -- that's the next step.
// Use of __restrict__ helps the compiler optimize mem operations, which are often the bottleneck
// the compiler doesn't have to worry that a write to a restrict pointer will cause a value read
// from another restrict pointer to change
// ======================================================================
__global__ void softmax_kernel_naive(
const float* __restrict__ in,
float* __restrict__ out,
int n_rows,
int n_cols
) {
// one thread per row: figure out which row this thread owns
int row = blockIdx.x * blockDim.x + threadIdx.x;
if (row >= n_rows) return; // guard for the last, partially-filled block
const float* in_row = in + row * n_cols;
float* out_row = out + row * n_cols;
// pass 1: row max (mirrors x.max(dim=1) in naive_softmax)
float row_max = -INFINITY;
for (int col = 0; col < n_cols; col++) {
row_max = fmaxf(row_max, in_row[col]);
}
// pass 2: exp(x - max), accumulate sum. stash exp() in-place in out_row
// so we don't need a second read of `in_row`.
float row_sum = 0.0f;
for (int col = 0; col < n_cols; col++) {
float e = expf(in_row[col] - row_max);
out_row[col] = e;
row_sum += e;
}
// pass 3: normalize
for (int col = 0; col < n_cols; col++) {
out_row[col] /= row_sum;
}
}
This is 3n reads of global memory per row (n for max, n for exp, n again to re-read what was just written for normalize — well, actually normalize reads out_row which it just wrote, but the point stands: three full sweeps). One thread per row also means zero cooperation between threads on the same row, so there’s no way to hide memory latency within a row — everything is serialized. Unsurprisingly, this is the worst performer of the bunch (the green line, pinned near the bottom of the roofline chart below).
CUDA optimized memory passes (online)
The first real optimization: compute the running max and running sum in a single pass, instead of two. This works because of the online softmax trick — every time we see a new max, we rescale the sum we’ve accumulated so far rather than throwing it away and starting over:
// The big difference across this and the naive is that online computes max and sum
// in a single pass. The reason we can do this is because when we find a new max,
// we rescale our running sum.
// Keep in mind that one of the best ways to optimise kernels is cluster READS and WRITES;
// That's what really saves memory traffic.
__global__ void softmax_kernel_optimized(
const float* __restrict__ in,
float* __restrict__ out,
int n_rows,
int n_cols
) {
// one thread per row: figure out which row this thread owns
int row = blockIdx.x * blockDim.x + threadIdx.x;
if (row >= n_rows) return; // guard for the last, partially-filled block
const float* in_row = in + row * n_cols;
float* out_row = out + row * n_cols;
float row_max = -INFINITY; // running max
float row_sum = 0.0f; // running sum
// first pass
for (int col = 0; col < n_cols; col++) {
float xi = in_row[col];
float new_max = fmaxf(row_max, xi);
// rescale running sum: https://arxiv.org/pdf/1805.02867
row_sum = row_sum * expf(row_max - new_max) + expf(xi - new_max);
row_max = new_max;
}
// second pass: normalize and write output
for (int col = 0; col < n_cols; col++) {
float xi = in_row[col];
float result = expf(xi - row_max) / row_sum;
out_row[col] = result;
}
}
Down from 3n memory passes to 2n. It’s still one thread per row (still no intra-row parallelism), but cutting a third of the redundant global memory traffic shows up clearly: the red “raw cuda optimized” line sits well above naive, though both are still far below the roofline — one thread per row just can’t move enough bytes/cycle to saturate HBM bandwidth, no matter how few passes it makes.
CUDA optimized for block per row
The real fix isn’t fewer passes, it’s more parallelism per row. That’s the selling point of GPUs! Instead of one thread doing all the work for a row, give it a whole block of threads that split the row between them, then reduce their partial results in shared memory:
// TODO: This kernel re-reads in_row three times (max pass, sum pass, normalize pass)
// Worth folding the online rescale in later once this baseline is verified correct,
// but reducing a running (max, sum) pair across threads in one shared-memory
// reduction is meaningfully trickier than reducing two independent scalars
__global__ void softmax_kernel_optimized_block2row(
const float* __restrict__ in,
float* __restrict__ out,
int n_cols
) {
// one block per row
int row = blockIdx.x;
const float* in_row = in + row * n_cols;
float* out_row = out + row * n_cols;
// size = blockDim.x float
extern __shared__ float sdata[];
// pass 1: each thread's partial max over its stride slice.
// starting offset = threadIdx.x, stride = blockDim.x -> thread t visits
// columns t, t+blockDim.x, t+2*blockDim.x, ... so every column in the row
// is touched by exactly one thread, exactly once (a partition, not a
// re-scan). When n_cols == blockDim.x each thread's loop runs once and
// covers a single column; when n_cols > blockDim.x (the common case --
// threads_per_block is fixed at 128 below) each thread serially walks
// n_cols/blockDim.x columns before the block can even start reducing.
float local_max = -INFINITY; // running max
for (int col = threadIdx.x; col < n_cols; col += blockDim.x)
local_max = fmaxf(local_max, in_row[col]);
sdata[threadIdx.x] = local_max;
// tree reduction: fold blockDim.x partial maxes down to 1 in log2(blockDim.x) steps.
// stride halves each round (e.g. 128 -> 64 -> 32 -> ... -> 1). On each round, only
// the lower half of threads (threadIdx.x < stride) are still active: each one combines
// its own slot with the slot `stride` away, so the "live" range of sdata keeps
// shrinking in half. __syncthreads() after every round is mandatory -- without it,
// a thread could read sdata[threadIdx.x + stride] before that slot's write from the
// *previous* round has landed. After the loop, sdata[0] holds the row-wide max.
__syncthreads();
for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
if (threadIdx.x < stride)
sdata[threadIdx.x] = fmaxf(sdata[threadIdx.x], sdata[threadIdx.x + stride]);
__syncthreads();
}
float row_max = sdata[0];
__syncthreads(); // all threads must read row_max before sdata gets reused below
// pass 2: each thread partial sum of exp(x-row_max)
float local_sum = 0.0f;
for (int col = threadIdx.x; col < n_cols; col += blockDim.x)
local_sum += expf(in_row[col] - row_max);
sdata[threadIdx.x] = local_sum;
// same tree reduction as the max pass above, just summing instead of maxing.
__syncthreads();
for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
if (threadIdx.x < stride)
sdata[threadIdx.x] += sdata[threadIdx.x + stride];
__syncthreads();
}
float row_sum = sdata[0];
// perf optimization: Division is one of the more expensive ALU ops
// on GPUs (much higher latency/throughput cost than multiply).
// precompute inv_sum = 1.0f / row_sum once, multiply instead of dividing per element;
float inv_sum = 1.0f / row_sum;
for (int col = threadIdx.x; col < n_cols; col += blockDim.x)
out_row[col] = expf(in_row[col] - row_max) * inv_sum;
}
Each thread only handles n_cols / blockDim.x elements now — but it’s worth being precise about what that means, since it’s easy to misread for (int col = threadIdx.x; col < n_cols; col += blockDim.x) as every thread scanning the whole row. It isn’t: threadIdx.x sets each thread’s starting column and blockDim.x is the stride, so thread 0 visits columns 0, blockDim.x, 2*blockDim.x, …, thread 1 visits 1, blockDim.x+1, …, and so on. That’s a partition of the row across the block — every column gets touched exactly once, by exactly one thread, never recomputed. The total memory traffic per row is identical to the single-thread version; it’s just spread across blockDim.x threads instead of piled onto one.
The catch is that this only collapses to “one thread, one column” when n_cols == blockDim.x. With threads_per_block fixed at 128 and rows up to 16384 columns wide in the sweep, each thread still serially walks up to 128 columns before the block can even start the shared-memory reduction below — a smaller serial bottleneck than the single-thread kernels, but a real one, and exactly what the “CUDA scaling threads” section further down tries to shrink by matching threads_per_block to the row width instead of leaving it fixed.
The tree-reduction pattern (for stride = blockDim.x/2; stride > 0; stride /= 2) then folds all the per-thread partial maxes/sums down to a single value in log2(blockDim.x) steps instead of a linear scan. With threads_per_block = 128, this closes almost the entire gap to the roofline for narrow-to-mid row widths — the purple line jumps up to sit right alongside Triton and torch.softmax.
But notice the purple line still droops badly at the widest rows (16384 columns). With a fixed 128 threads per block, a row that wide means each thread strides through 128 columns per pass — that’s a lot of serialized memory traffic per thread before the block can even start its reduction, and it’s the next thing to fix.
CUDA scaling threads
The obvious fix for the wide-row droop is to give wide rows more threads. Derive threads_per_block on the fly based on the input instead of hardcoding it:
int threads_per_block = 128;
while (threads_per_block < n_cols && threads_per_block < 512)
threads_per_block *= 2;
Why is there a gap when we scale and why the fixed-128 degrades on wide rows
Why the autoscale version has a gap
Scaling threads per block, trades away occupancy (how many blocks fit resident on an SM at once) for fewer iterations per thread. Bigger blocks = fewer blocks per SM = less spare warps to hide the reductions __syncthreads. The dip is always where thread count first reaches its cap (ratio = 1). By moving the cap, we move the dip.
Why the fixed-128 degrades on wide rows
With threads pinned at 128, a wide row forces the thread loop to loop 128 times per pass (x3 passes = 384 global memory reads). Occupancy is great, but the sheer volume of redundant memory traffic is the bottleneck.
Autotuning
Deriving a scaling formula by hand is fragile — it bakes in assumptions (the < 512 cap, the doubling schedule) that only happen to work for the hardware and shapes I tested on. A more realistic approach: expose threads_per_block as a parameter on the kernel launch, and let Python-side code sweep a small set of candidates for the actual input shape, keeping whichever one benchmarks fastest.
// Same kernel as above, but threads_per_block is passed in explicitly instead of
// derived from n_cols -- lets Python-side code sweep candidate values per shape
// (a tiny hand-rolled version of what autotuning frameworks do) instead of trusting
// one hardcoded scaling heuristic for every width.
torch::Tensor softmax_cuda_optimized_block2row_tuned(torch::Tensor x, int64_t threads_per_block) {
TORCH_CHECK(x.is_cuda(), "input must be a CUDA tensor");
TORCH_CHECK(x.dim() == 2, "input must be 2D (n_rows, n_cols)");
TORCH_CHECK(x.scalar_type() == torch::kFloat32, "input must be float32");
TORCH_CHECK(threads_per_block > 0 && (threads_per_block & (threads_per_block - 1)) == 0,
"threads_per_block must be a power of 2 (reduction assumes it)");
x = x.contiguous();
int n_rows = x.size(0);
int n_cols = x.size(1);
auto out = torch::empty_like(x);
const int num_blocks = n_rows; // one block per row
softmax_kernel_optimized_block2row<<<num_blocks, (int)threads_per_block, threads_per_block * sizeof(float)>>>(
x.data_ptr<float>(),
out.data_ptr<float>(),
n_cols
);
return out;
}
BLOCK2ROW_THREAD_CANDIDATES = [128, 256, 512, 1024]
def autotune_block2row(x):
"""Times each candidate threads_per_block on this exact input, returns a
ready-to-call fn bound to whichever candidate was fastest."""
best_ms, best_threads = float("inf"), None
for t in BLOCK2ROW_THREAD_CANDIDATES:
ms = triton.testing.do_bench(lambda t=t: cuda_block2row_with_threads(x, t))
if ms < best_ms:
best_ms, best_threads = ms, t
print(f" [autotune] block2row picked threads_per_block={best_threads}")
return lambda t: cuda_block2row_with_threads(t, best_threads)
By timing every candidate config and keeping the fastest, we sidestep the whole “why does the gap move when I change the cap” question — the benchmark just tells us directly. In production you obviously wouldn’t re-run this sweep on every call: benchmark once per input shape, cache the winning config, and reuse it for every subsequent call with that shape. Re-benchmarking on every forward pass would burn far more compute than it saves.
Plotting the autotuned version (brown) against everything so far — including the plain fixed-128 block-per-row kernel (purple) for reference — the autotuned line stays right on the roofline across the entire width sweep, including exactly the wide rows where the fixed-128 version was drooping:
Next steps
We will profile all CUDA kernels, and deep dive on what we can conclude from the profiler when cross checking against our kernel. Will there be any optimizations that we can still make? TODO: Making the CUDA block per row, autotuned version, also online.