I Wrote My First GPU Kernel — Here’s What Changed

1.49x faster softmax. But only 1.074x faster forward pass. Amdahl’s Law explains the gap — and understanding it matters more than the speedup itself.

In Week 1 of this journey, I profiled attention on my RTX 3050 and found softmax consuming 21% of total GPU time at seq=2048. That was the cost of building a 2048×2048 scores matrix in VRAM three times over. I said I’d learn to fix it. This week, I wrote a fused Triton kernel that did — and what the profiler showed surprised me.

Why Naive Softmax Is Slow

Softmax looks simple: a max reduction, subtraction, exponentiation, and division by sum. Four steps. But on a GPU, each step means a round trip to VRAM — the slow global memory that all streaming multiprocessors share.

Here’s what naive softmax actually does for every row of the attention matrix:

Round trip 1: read X from VRAM → compute max        → write max back to VRAM
Round trip 2: read X from VRAM → compute exp(x-max)  → write exp back to VRAM
Round trip 3: read exp from VRAM → compute sum       → write sum back to VRAM
Round trip 4: read exp + sum from VRAM → divide      → write output to VRAM

At seq=2048, that’s a 2048×2048 matrix making four passes through a 192 GB/s memory bus. The GPU cores finish the arithmetic almost instantly — they’re waiting for data to arrive and leave. This is what “memory-bound” means in practice.

What Fusion Actually Does

A fused kernel does the same math but never writes intermediate results back to VRAM. Everything stays in registers — the fastest memory on the chip, private to each thread:

Fused (1 VRAM round trip):
  VRAM → registers: load row of X          ← only read
  registers: compute max                    ← stays in registers
  registers: compute exp(x - max)           ← stays in registers
  registers: compute sum                    ← stays in registers
  registers: divide exp by sum              ← stays in registers
  registers → VRAM: write final output      ← only write

One read, one write. The memory bus handles 2 trips instead of 8. The arithmetic is identical — same max subtraction to prevent overflow (exp(100) hits float32 limits), same normalization. The only difference is where intermediate values live.

The Kernel

I wrote this in Triton, a Python-like language for GPU programming. The key section:

@triton.jit
def fused_softmax_kernel(output_ptr, input_ptr, n_cols, BLOCK_SIZE: tl.constexpr):
    row_idx = tl.program_id(0)
    col_offsets = tl.arange(0, BLOCK_SIZE)
    mask = col_offsets < n_cols

    # Load entire row into registers — single VRAM read
    x = tl.load(input_ptr + row_idx * n_cols + col_offsets, mask=mask, other=-float('inf'))

    # All computation in registers — no VRAM writes
    max_val = tl.max(x, axis=0)
    exp_x = tl.exp(x - max_val)
    exp_x = tl.where(mask, exp_x, 0.0)  # zero masked positions
    sum_exp = tl.sum(exp_x, axis=0)
    result = exp_x / sum_exp

    # Single VRAM write
    tl.store(output_ptr + row_idx * n_cols + col_offsets, result, mask=mask)

One detail I learned the hard way: without tl.where(mask, exp_x, 0.0), the masked positions kept their -inf values, exp(-inf) returned 0 but accumulation noise crept in. The first version returned inf in the output. Small line, critical fix.

Results

I benchmarked naive Triton, fused Triton, and PyTorch’s built-in softmax. seq=2048 corresponds to a mid-length document or multi-turn conversation — the upper range of most production deployments:

| Seq Len | Naive Triton | Fused Triton | PyTorch ref | Fused Speedup | Bandwidth |
|---------|-------------|-------------|-------------|---------------|-----------|
| 256     | 0.0527ms    | 0.0381ms    | —           | 1.38x         | 14 GB/s (7%) |
| 512     | 0.0507ms    | 0.0386ms    | —           | 1.31x         | 54 GB/s (28%) |
| 1024    | 0.0714ms    | 0.0478ms    | 0.0481ms    | 1.50x         | 176 GB/s (91%) |
| 2048    | 0.2755ms    | 0.1851ms    | 0.1875ms    | 1.49x         | 181 GB/s (94%) |
| 4096    | 1.0923ms    | 0.7339ms    | —           | 1.49x         | 183 GB/s (95%) |

The low bandwidth numbers at seq=256 (7%) and seq=512 (28%) aren’t a kernel problem — they reflect kernel launch overhead dominating at small sizes. The actual computation is so fast that setup costs distort bandwidth measurements. The numbers become meaningful above seq=1024, where computation time exceeds launch overhead. At seq=1024 and above, the kernel runs at 91–95% of peak memory bandwidth.

At seq=1024, the fused kernel matched PyTorch’s built-in exactly (0.0478ms vs 0.0481ms). At seq=2048, it beat PyTorch by 1% (0.1851ms vs 0.1875ms). Bandwidth utilization climbed to 94% of my GPU’s 192 GB/s peak — nearly saturating the memory bus.

The Honest Part: Where The Speedup Actually Came From

I expected the speedup to come from eliminating VRAM round trips. The profiler told a different story:

Naive kernel (seq=2048):
  naive_softmax_kernel:           184μs   67%
  vectorized_elementwise_kernel:   90μs   33%  ← zeros_like overhead
  Total:                          274μs

Fused kernel (seq=2048):
  fused_softmax_kernel:           183μs  100%
  Total:                          183μs

The actual softmax computation took nearly identical time (184μs vs 183μs). The 1.49x speedup came almost entirely from eliminating zeros_like — a separate CUDA kernel that pre-initialized the output tensor with zeros. The naive implementation needed zeros_like because it writes output in multiple passes — zero-initialization prevents garbage values in positions not yet written. The fused kernel writes every position in a single pass, so empty_like is safe and the initialization step disappears entirely. Switching to empty_like removed that 90μs overhead.

The VRAM round-trip elimination is real and matters at larger scale. But at this matrix size, the measurable win was removing a wasted kernel launch.

This is worth being explicit about: the architectural improvement (fewer VRAM round trips) is real and will matter at larger scale. The measured win at this matrix size was the initialization overhead. Both are true simultaneously.

Amdahl’s Law: Why 1.49x Softmax ≠ 1.49x Model

Week 1 showed softmax is 21% of the forward pass at seq=2048. Amdahl’s Law sets the ceiling:

Maximum possible system speedup = 1 / (1 - 0.21) = 1.266x

My softmax improvement:    32.7% faster
Time saved:                0.21 × 0.327 = 6.87% of total forward pass
Actual system speedup:     1.074x
Fraction of ceiling:       (1.074 - 1) / (1.266 - 1) = 27.9%

I captured 27.9% of the theoretical maximum. That’s modest — the remaining gap is likely from register pressure and the kernel launch overhead visible at small sequence lengths. Closing it fully would require tiling across the sequence dimension, which is exactly what FlashAttention does. Even if softmax were infinitely fast, the forward pass can’t speed up more than 1.266x — the other 79% of computation doesn’t change. This connects directly to Week 2: below seq=6,144, MLP is 46% of GPU time. Optimizing a 21% component while ignoring a 46% component is working on the wrong bottleneck.

Amdahl’s Law is why profiling order matters. Optimize the largest fraction first.

What This Means for FlashAttention

Fused softmax is one building block of FlashAttention. FlashAttention fuses the entire attention computation — Q×K, softmax, and ×V — into a single kernel that tiles across the sequence dimension. It eliminates not just the softmax VRAM round trips, but the materialization of the full S×S scores matrix entirely.

My kernel fuses within softmax. FlashAttention fuses across the whole attention operation. The principle is the same: keep intermediate values in fast memory, minimize VRAM traffic.

What’s Next

Phase 3 begins: tearing down an MLSys paper and implementing its core optimization from scratch. I’ve profiled bottlenecks (Phase 1) and written kernels to fix them (Phase 2). Now I need to read how researchers think about these problems — and reproduce their results.

Code & profiler output: https://github.com/falakshair01/mlsys-journey

Previous post: [“I Profiled LLM Inference From First Principles — Here’s What I Found”](LINK)

Papers referenced:

• Dao et al., “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness,” NeurIPS 2022
• He, Horace, “Making Deep Learning Go Brrrr From First Principles,” 2024