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 into how kernels are actually implemented and optimized. To get a better sense of this “world”, I treated it as an experiment: pick one op, implement it several different ways, hold the inputs fixed, and directly observe what each optimization actually buys you. This is mostly a personal learning exercise — writing up how I reasoned through each optimization, and the mental model I built along the way, more than a from-the-experts guide.
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.
Baseline & Triton Softmax kernel Link to heading
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: in simple terms, 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 ops, each round-tripping through HBM. 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.
# ======================================================================
# 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.
# ======================================================================
# 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. 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 Link to heading
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). 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) Link to heading
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 Link to heading
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:
__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
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: combine blockDim.x partial maxes down to 1 in log2(blockDim.x) steps
__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;
}
This kernel is definitely more complex than the others. But let’s go over all the details carefully.
Per-thread accumulation: local_max and local_sum
Link to heading
local_max and local_sum are each thread’s private accumulator — “running” in the sense that they update once per loop iteration, absorbing one more column each time, rather than looking at the whole slice at once. They follow the exact same shape, just with a different starting value and combining op: local_max starts at -INFINITY and combines with fmaxf, local_sum starts at 0.0f and combines with +=.
This is also where n_cols > blockDim.x shows up. With blockDim.x = 4 (again, this is the number of threads per block) and n_cols = 10 (the actual “values” the threads have to process), the loop for (int col = threadIdx.x; col < n_cols; col += blockDim.x) hands out columns like this:
thread 0: cols 0, 4, 8 (3 columns)
thread 1: cols 1, 5, 9 (3 columns)
thread 2: cols 2, 6 (2 columns)
thread 3: cols 3, 7 (2 columns)
Every column from 0-9 is covered exactly once, but since 10 doesn’t divide evenly by 4, threads 0-1 do one more loop iteration than threads 2-3. Take thread 1: its Pass 1 loop runs 3 times, once per column it owns, each time merging one more value into local_max:
local_max = -INFINITY
local_max = fmaxf(local_max, in_row[1]) // = in_row[1]
local_max = fmaxf(local_max, in_row[5])
local_max = fmaxf(local_max, in_row[9])
By the time the loop exits, local_max holds the max over just thread 1’s 3 columns — not the row’s max yet, only this thread’s slice of it. Pass 2 (the sum) runs the exact same pattern once row_max is known: same column assignment, same loop shape, just += instead of fmaxf and a different seed value.
Each thread ends its loop holding one number — its private partial max or partial sum over however many columns it was assigned (3 here, 2 for threads 2-3) — which is exactly what gets written into sdata[threadIdx.x] and fed into the tree reduction above. So there are two levels of work happening: the “running” accumulation is within a thread, serially walking its own columns one at a time; the tree reduction is across threads, combining those already-finished per-thread results into one row-wide value.
Partitioning the row across the block Link to heading
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 “how long is the step”, 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.
The tree reduction Link to heading
The tree-reduction pattern (for stride = blockDim.x/2; stride > 0; stride /= 2) then combines 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.
To see why it’s log2(blockDim.x) steps, walk it with blockDim.x = 4. Say Pass 1 leaves these 4 partial maxes in sdata, one per thread:
index: 0 1 2 3
value: 3.1 7.2 1.0 9.5
Round 1, stride = 2: threads 0-1 are active. Each compares its own slot against the slot 2 away and overwrites its own slot with the winner:
stride = 2, because, from the code:
stride = blockDim.x/2
sdata[0] = max(3.1, 1.0) = 3.1
sdata[1] = max(7.2, 9.5) = 9.5
The useful values are now packed into slots 0-1 ([3.1, 9.5]); slots 2-3 still hold stale numbers, but nothing reads them again.
Round 2, stride = 1: only thread 0 is active:
stride = 1, because, from the code:
stride /= 2. The previous stride from round 1 was2.
sdata[0] = max(3.1, 9.5) = 9.5
Loop ends — sdata[0] = 9.5, the max of all 4 original values, reached in 2 rounds (log2(4)) instead of sequential comparisons. Each round halves how many threads have work left, which is the “live range keeps shrinking” behavior the loop condition threadIdx.x < stride forces.
__syncthreads() after every round is what makes this safe: without it, a thread could read sdata[threadIdx.x + stride] before the thread that owns that slot has written this round’s value into it. The sum reduction a few lines down is the identical dance with += instead of fmaxf.
Results for this specific kernel Link to heading
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 walks 128 columns per pass serially. The traffic is the same as always — what’s missing is enough loads in flight at once to hide HBM latency, so achieved bandwidth drops. That’s the next thing to fix.
CUDA scaling threads Link to heading
The fix for the fixed-128 kernel’s wide-row droop is to stop hardcoding threads_per_block: give wider rows more threads, so each thread’s share of the row stays small no matter how wide it gets. Hand-deriving the exact scaling formula for that is fragile though — it bakes in assumptions about the specific hardware and shapes you tested, which won’t generalize. So instead of hardcoding a formula, the better move is autotuning: expose threads_per_block as a knob and let it get picked empirically per shape.
Why the fixed-128 kernel degrades on wide rows Link to heading
threads_per_block never moves — it’s stuck at 128 no matter how wide the row is. On a 16384-column row, each of those 128 threads has to serially walk 128 columns per pass (×3 passes = 384 reads per thread) before the block can even start reducing. The traffic per row is unchanged — the problem is that 128 threads can’t keep enough loads outstanding to hide latency. Giving wider rows more threads doesn’t move fewer bytes; it just splits the row into shorter slices so more loads are in flight at once, which is what pushes bandwidth back toward the roofline.
Autotuning Link to heading
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:
Conclusion Link to heading
The experiment paid off: each optimization’s effect on achieved bandwidth was directly observable in the sweep, rather than something I had to take on faith. Going from a naive, unfused PyTorch softmax to a hand-tuned CUDA kernel makes the cost of each optimization concrete: fusing ops into one pass cuts redundant HBM round-trips, but without intra-row parallelism (one thread per row) there’s a hard ceiling on how much bandwidth a kernel can actually achieve. Splitting a row across a block of cooperating threads, reducing their partial results in shared memory, is what closes most of the gap to the roofline — and autotuning threads_per_block per shape is what closes the rest, especially for the widest rows. By the end, the autotuned block-per-row kernel tracks torch.softmax and Triton almost exactly across the entire width sweep.
More than the specific kernel, this was my first real pass at building intuition for how GPU kernels are implemented and optimized — reasoning about memory traffic, parallelism, and reductions instead of just calling torch.softmax and moving on.