Team Fable's Transformer Kernel Optimisation

Inspiration

Transformer inference on GPU is rarely limited by arithmetic alone. On the graded shapes, the reference spends measurable time launching small kernels and materialising attention masks that cannot change the answer. The extreme shape needs an 18.6 TB score tensor per layer—the baseline cannot run on any hardware.

We asked: what can you win by changing how the same maths runs, without touching weights or tolerance?

What it does

We submit OptimizedTransformer—a drop-in replacement for the organisers' BaselineTransformer with an identical state_dict.

The reference computes attention as:

$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$$

and materialises the full ([B, H, S, S]) score matrix every layer. We instead use scaled_dot_product_attention, drop redundant padding masks for causal right-padded input (exact, not approximate), fuse Q/K/V into one GEMM, chunk extreme batches, and on CUDA run fp16 GEMMs with an fp32 residual stream plus CUDA-graph capture.

Results (Apple M4 CPU, fp32):

Correctness criterion (per element): (\lvert y_{\mathrm{opt}}-y_{\mathrm{ref}}\rvert \le 0.002) OR relative error (\le 0.02). Zero elements failed on any passing shape.

How we built it

1. Transcribed all fourteen appendix test shapes (causal=True, ffn_dim = d_model)—not the TensorFlow script's defaults.
2. Built a static cost model + three CPU probes (precision_probe, launch_count, roofline) to rank optimisations before GPU time.
3. Implemented six changes behind a pure dispatch policy (src/tko/dispatch.py), with 82 automated tests including the mask correctness proof.
4. Never edited the graded benchmark—model swapped at import time only (run_torch_benchmark.py). SHA-256 unchanged: ```text 5529c96a80799b51f68092e1444a30b17994554dffdf52da98ba701489a7f36e

Built With

• artificial-intelligence
• cuda
• deep-learning
• gpu
• machine-learning
• optimization
• python
• pytorch
• transformers