Roofline

Inspiration

This problem statement delved deeper into the inner workings of how transformers are being processed which is a fundamental concept of deep learning. This is a problem that felt refreshing to me as it doesn't focus on trying to build something new but rather how to optimise what we already have.

What it does

This project takes a fixed Transformer inference layer and replaces how it runs on the GPU, without changing what it computes, to make it substantially faster while staying numerically correct. Starting from a profiler trace, it fuses the attention computation and Q/K/V projections into fewer, larger GPU kernels. This eliminates a masking operation that was silently wasting up to 20% of runtime, and builds a custom chunked-attention kernel for sequences up to 100,000 tokens long, a scale where the naive reference implementation itself is mathematically impossible to run. The result is a 1.1×–4.5× speedup across 13 tested input shapes with zero accuracy loss, verified end-to-end against the original implementation at every step.

How we built it

We built it by profiling before touching any code: using torch.profiler on a range of input shapes to see exactly where time was going, rather than guessing. That surfaced the attention computation's naive kernel chain as the first target, which we replaced with F.scaled_dot_product_attention and a fused Q/K/V projection. A second profiling pass caught a masking call quietly firing on every layer for a mask that was always empty. We also tried a hand-written Triton kernel to fuse LayerNorm with its residual add, measured it end-to-end, found it was actually slower despite looking correct on paper, and reverted it, treating the benchmark as the final word over intuition at every step. For the 100,000-token case, where the reference implementation can't even run due to its memory requirements, we built a chunked attention kernel processing the sequence in bounded-memory windows against a growing key/value cache, validated against the baseline at every scale we could directly test.

Challenges we ran into

A GPU with hard limits. Getting from zero compute access to a working Colab T4 loop, then working around its 16GB ceiling for the largest test shapes.

Telling a real win from noise. A shared, free-tier GPU has real run-to-run variance. More than once, a result that looked like a regression (or an improvement) turned out to be measurement noise, or in one case, leftover GPU memory from an earlier profiling run silently causing an unrelated test to fail. Learning to distinguish "this changed" from "this is just noise" was its own skill.

A shape that breaks the reference, not just the target. Row 14 wasn't a normal performance problem. The correctness ground truth itself was infeasible to compute naively, which meant the usual "compare against baseline" workflow didn't apply and needed a different validation strategy (testing at smaller, provably-correct scales instead).

What we learned

GPU performance isn't just about FLOPs. Kernel launch overhead and memory bandwidth dominate at different scales, and the right fix genuinely depends on the shape, not just "make it faster."

A profiler trace is worth more than intuition. The one optimization I was most confident about on paper (the Triton LayerNorm fusion) was the one that measured worse.

Debugging isn't only about code. Some of the hardest problems in this project were about environment state, not logic, and the fix was verifying independently rather than assuming.

Long-context attention isn't just "the same algorithm, but bigger", at extreme scale, even linear-memory operations become the bottleneck, and the real fix is architectural (chunking/streaming), not just a faster kernel for the same computation.