Devpost submission — Problem Statement #3
Inspiration
The obvious way to attack this problem is to pick the strongest-sounding
optimization — flash attention, or FP16 tensor cores, or torch.compile — and
apply it everywhere. Profiling the baseline made it clear that this would leave
most of the win on the table, because the fourteen required shapes are not one
workload. They are three:
B=1, S=128spends ~97% of its wall time on kernel-launch overhead. It never gets close to saturating the GPU.B=64, S=128, d=1024is arithmetic-bound, sitting at the card's GEMM roofline.B=10000andS=100000are memory-bound — the baseline's[B,H,S,S]score tensor simply does not fit in 8 GB.
An optimization that is decisive for one of these is worthless or actively
harmful for another. Measured: CUDA graphs are worth 23.2× on batch 1 and are a
small regression at d=1024. FP16 is worth 2.7× on the long-sequence shape
and is slower than FP32 at batch 1, because the cast kernels cost more than
the tiny GEMMs save.
So the interesting problem is not "which optimization is best" — it is "how does the layer decide, at runtime, which one it needs".
What it does
A drop-in replacement for UserOptimizedTransformer in the official benchmark
harness. Parameter names are unchanged, so copy_model_weights(strict=True)
works untouched and the only edit to the harness is inside the designated
your codes here block.
At the first forward pass for a given shape it inspects the input and picks a strategy:
| Condition | Strategy |
|---|---|
| ≤ 4 Mi elements | Inductor-compiled forward, captured as a CUDA graph |
| > 4 Mi, fits in VRAM | Inductor-compiled, plain launches |
| exceeds free VRAM | batch-chunked streaming |
| no CUDA / no Triton / capture failed | eager fallback |
Underneath, all paths share: flash attention via
scaled_dot_product_attention, a fused [3d, d] QKV projection with a
zero-copy head split, FP16 tensor-core GEMMs with an FP32 residual stream, and a
padding mask that is proved away under causal attention rather than passed to
the kernel.
Result: 11.30× geometric-mean speedup across all 14 measurable shapes, with zero failing elements out of ~4.2 billion compared, on a 115 W laptop GPU (RTX 4060 Laptop, 8 GB). Best case 33.0×; worst case 2.60×, which is 73% of the card's measured GEMM roofline — i.e. near the physical limit, not a soft spot.
It also makes shapes run that previously could not. At B=1, S=100000 the
baseline needs a 640 GB score tensor per layer and dies; ours completes in
1.56 s, sustaining 27.8 TFLOP/s — 96% of the card's measured FP16 GEMM
ceiling on a workload that is 94% attention.
How we built it
Measure first. Before touching the transformer, a microbenchmark
(microbench.py) established what the card can actually do: FP16 tensor cores
run at 2.6× TF32, and DRAM tops out around 240 GB/s. Every later claim is judged
against those numbers, which is what made it possible to stop optimizing the
d=1024 shape once it hit 73% of roofline instead of grinding on it.
Every optimization then went in behind an environment flag (TTJ_MODE,
TTJ_DTYPE, TTJ_RESIDUAL) so ablation.py could measure each one
independently, per shape. The shipped runtime policy is read directly off that
table rather than guessed.
Two findings came from error analysis rather than profiling:
- The residual stream must stay FP32. The final LayerNorm divides by σ, so
an output element near the layer mean has its error amplified by
|residual|/σ. A 5×10⁻⁴ relative rounding error on a residual of magnitude ~4 lands as ~2×10⁻³ of absolute error — exactly theatolbudget. Predicted from the algebra, then confirmed: an FP16 residual produces 5.7× more out-of-tolerance elements. - A causal prefix mask is provably a no-op. A valid query at
i < lengthalready only attends toj ≤ i < length, so masking invalid keys removes nothing. That lets the mask collapse to a single output multiply and keeps the fused causal kernel, instead of falling back to a slower masked path.
Challenges we ran into
BF16 cannot meet the stated tolerance, and proving that took an experiment.
BF16 failed 56/56 correctness cases. Rather than assume a bug, we compared the
baseline against itself — BF16 versus the same model computed in FP32 — and
found the reference already misses the tolerance on 6% of elements. atol=0.002
is smaller than one BF16 ULP for any output above magnitude 0.5. No
implementation can pass without bit-reproducing the reference's operation order.
FP32 (the harness default, and what the appendix specifies) passes everything.
Appendix shape #14 is unreachable on 8 GB by anyone. B=32, S=100000,
d=1024 needs 13.1 GB just to hold the FP32 input; the harness fails before
either model is called. We report that plainly and anchor the regime with
reduced-batch variants that keep every other dimension identical.
A 2× phantom regression. Shape #6 ran at roughly double the time it should
have until profiling showed that mem_get_info was returning slightly different
free-VRAM readings on each call, changing the chunk size, and triggering an
Inductor recompile every time. Caching the decision halved that shape's latency.
A contaminated measurement, caught. An early sweep started immediately after a large install showed launch-bound baselines inflated by up to 2.1×. A clean re-run did not reproduce it, and the compute-bound shapes were unaffected in both — localizing it to transient host load rather than anything real. The headline number was corrected downward as a result.
What's next
The remaining ~1.4× on the arithmetic-bound shape is entirely in the elementwise
chain between GEMMs. A hand-written Triton kernel fusing
residual add + LayerNorm + FP16 cast with FP32 accumulation would recover most
of it without the accuracy cost an FP16 residual carries — the one place a
custom kernel still pays, given Inductor already fuses the rest and the GEMMs
are at the cuBLAS ceiling.
Built with
Development tools — Claude Code (Claude Opus 5) as the primary development
environment, driving a terminal on the target machine with direct GPU access;
WSL2 running Ubuntu 22.04; nvidia-smi; git.
Libraries and frameworks — PyTorch 2.6.0 (+cu118), Triton 3.2.0 (via
torch.compile / TorchInductor), PyTorch's scaled_dot_product_attention
(Flash Attention and memory-efficient backends), CUDA Graphs, CUDA 11.8,
cuDNN 9.1.0.
Profiling and measurement — torch.profiler, torch.cuda.Event,
torch.cuda.mem_get_info, plus a purpose-built roofline microbenchmark
(microbench.py) and ablation harness (ablation.py).
APIs — none. No external or hosted APIs are used at runtime.
Datasets and assets — none. All inputs are the synthetic random tensors
generated by the official benchmark harness (generate_random_case), seeded for
reproducibility. No third-party data, models, or pretrained weights.
Hardware — NVIDIA GeForce RTX 4060 Laptop GPU (Ada Lovelace, sm_89,
24 SMs, 8 GB), Intel Core i9-13900HX, 16 GB RAM.
Log in or sign up for Devpost to join the conversation.