GPU Kernel for a Transformer Layer

Inspiration

We wanted to go beyond calling existing fused-attention libraries and actually understand the memory and compute characteristics of a transformer layer at the kernel level — motivated by roles (e.g. recommendation-system/inference infrastructure work) where this kind of low-level optimization directly matters.

What we learned

Why naive attention is memory-bound: the explicit [batch, heads, seq, seq] score matrix materializes fully in memory, causing thrashing on memory-constrained GPUs at large batch sizes.
How online softmax (FlashAttention-style tiling) avoids this by maintaining a running max/sum/accumulator across key/value blocks instead of computing the full score matrix at once.
Which operations are safe to run in reduced precision (elementwise ops, tensor-core-accumulated matmuls) versus which need FP32 (LayerNorm and softmax, since both are reduction-based and compound rounding error).
That a hand-written kernel's performance is highly shape-dependent: our custom Triton attention kernel beats PyTorch's SDPA by 2-12x on compute-heavy shapes (large batch, long sequences), but loses to SDPA on very small batches (fixed per-launch overhead dominates) and needs further work at very large batch sizes (grid-instance count explodes).

How we built it

We started from PyTorch's scaled_dot_product_attention as a fast, correct baseline, then hand-wrote a Triton kernel implementing online-softmax causal attention from scratch, validated it against SDPA on isolated test shapes, then integrated it into the full transformer with a safety fallback to SDPA for shapes it doesn't yet support (padding masks). We used torch.compile for automatic fusion of the surrounding LayerNorm/FFN/residual operations, and built our own same-process A/B testing methodology to reliably compare kernel variants under GPU thermal/timing noise.

In parallel, a teammate explored mixed FP16/FP32 precision on a second RTX 5060 GPU: FP16 Tensor Core operands for the linear projections, with FP32 kept for numerically sensitive reductions (LayerNorm statistics, softmax, exact GELU) and the residual stream — following the same principle we'd independently derived (reductions compound rounding error; elementwise ops and tensor-core-accumulated matmuls don't). This included shape-specific Triton projection kernels and a controlled, same-process tournament methodology to validate which optimization variant actually won per shape before retaining it.

Challenges we faced

GPU memory pressure and thermal throttling on a 4GB mobile GPU, requiring careful shape-conditional handling.
Getting torch.compile and a hand-written Triton kernel to coexist (had to exclude the custom kernel from Dynamo's tracing).
Case 14 (batch=32, seq_len=100,000) is provably intractable for the reference implementation on any real GPU — the Q tensor alone requires ~13GB, and the full attention score matrix would require ~20TB. We documented this precisely rather than reporting a number that couldn't exist.
Discovering, through rigorous A/B testing, that our custom kernel doesn't win everywhere — and building shape-aware dispatch logic to use the right implementation per case rather than claiming a uniform win.

How we used AI

We used AI deliberately as a learning and reasoning partner, not as a replacement for understanding the code we shipped.

One teammate worked conversationally with Claude throughout: discussing GPU architecture concepts, working line-by-line through the Triton fused-attention kernel (grid/block indexing, pointer arithmetic, online softmax's running max/sum/rescale mechanics), debugging real compilation and integration errors (shared-memory limits, tl.dot's minimum contraction size, torch.compile/Triton interaction), and reasoning through why the custom kernel wins or loses on specific shapes rather than just accepting benchmark numbers at face value.
The other teammate used an agentic coding tool (Codex) to iteratively explore and implement the mixed-precision optimization path across many experiment iterations, with results and rejected approaches logged throughout the process.
Before merging any AI-assisted work into the shared codebase, we held it to the same standard: could we explain what it did and why. Claims or results neither of us could independently verify were treated as unverified rather than included as fact.