MegaBee

Inspiration

Track 3 hands you a working PyTorch Transformer layer and asks a deceptively simple question: make it faster, without changing a single number it outputs.

What drew us in was the second half. Speed is negotiable; correctness is not. Every element has to land within 0.002 absolute or 2% relative of the reference, across 14 published input shapes, or the speed isn't even measured.

That constraint quietly disqualifies most of the famous "fast attention" literature. Sparse attention, linear attention, low-rank approximations — they all compute a different function. Excellent research, and useless to us. One paper we evaluated in detail replaces softmax's exponential with f(x) = 1 + x, a first-order Taylor truncation. Fast, well-engineered, and mathematically not the thing we were asked to reproduce.

So the real problem became: how much speed can you extract while changing nothing about the answer?

What it does

A drop-in replacement for the reference Transformer block that runs 5.40x faster by median, 6.54x by geometric mean, with all 13 verifiable shapes passing the correctness gate on an NVIDIA A100-SXM4-80GB.

It isn't one kernel. The rules permit shape-dependent implementations, so it's a dispatcher over six validated code paths, each chosen per shape by measurement rather than intuition:

route	technique	example
`compile`	`torch.compile` max-autotune	shape 1 → 4.92x
`reduce`	reduce-overhead + CUDA graphs	shape 4 → 9.54x
`fused`	fused QKV projection + SDPA	shape 12 → 11.24x
`fusedcg`	fused + Triton kernel + manual CUDA-graph capture	shape 2 → 13.98x
`amp`	fp16 autocast + Triton AddNorm	shape 13 → 14.11x
`chunked14amp`	exact batch chunking	the 100k-token shape

Individual shapes span 1.81x to 14.11x. Worst-case numerical error across all 13 is 0.00176 against a 0.002 budget.

The result we're proudest of is shape 14. Batch 32, sequence length 100,000. The organizer's own reference materializes the full attention score matrix — at that size, 18.6 terabytes. It cannot run its own largest test case on any single GPU, which also means no ground truth exists for it. Ours completes a full forward pass, by recognizing the 32 sequences are independent and processing them in exact, bit-identical chunks. We validated the chunking mechanism on shapes that do have references, then reported shape 14 with that caveat stated plainly.

To evaluate it, run the organizer's own script unmodified. torch_transformer_benchmark.py loads our implementation directly and prints [info] using candidates/v_router2_autotuned.py to confirm — no file swapping required, and it falls back safely to the original stub if our code is absent.

How we built it

Four techniques, layered, each kept only after beating the previous champion on a measured benchmark:

Memory-efficient attention replacing the explicit [B,H,S,S] score matrix — tiling with an online softmax, exact rather than approximate. This is what makes 100k tokens tractable at all.
torch.compile for operator fusion and CUDA-graph capture. The smallest shapes needed this most: one spends ~97% of its wall clock on dispatch rather than arithmetic.
A custom Triton kernel fusing residual-add + LayerNorm, at both boundaries per block — written after a profiler trace showed one boundary was unfused and costing 19% of a shape's CUDA time.
fp16 autocast on compute-bound shapes, with LayerNorm and softmax reductions held in fp32 to protect the error budget.

The benchmark harness reuses the organizer's own comparison and timing functions rather than reimplementing them, so our numbers come from the same code path the task defines. Every experiment appended a row to an append-only ledger; nothing counted until it had one. Sixty-one rows later, that ledger is the project.

We used AI tooling as a research process rather than an autopilot: a planning pass required to cite evidence for every claim, an adversarial review from a different model family, and a fact-checking pass whose only job was to open citations and prove the previous round wrong. It repeatedly did, including to us.

Challenges we ran into

An optimization that never ran. Our first attention rewrite had a fast path guarded by if mask is None. The mask is never None — the harness always passes an all-true one. The branch was unreachable, and it passed every CPU test because the slow path is correct, merely slow.

A correctness bug invisible off-GPU. torch.compile's autotuner silently selected TF32 kernels for fp32 inputs, drifting ~0.005 against a 0.002 tolerance on 9 of 12 shapes. Neither CPU nor Apple Silicon has TF32, so nothing we could run locally would ever have found it.

A headline number that was 24% wrong. Our first GPU result claimed 2.71x. The timing protocol didn't match the organizer's — wrong warmup, wrong repeat count, and the candidate was reloaded once per shape, discarding its compile cache. Fixed, re-measured: 2.18x. We retracted it in our own leaderboard.

Noise large enough to fake a win. Two back-to-back sweeps with zero code changes moved the aggregate geometric mean by 3.3%, and one shape swung 35.7%. Several improvements we were about to claim sat inside that band. We measured the noise floor before trusting anything near it.

A documentation bug that would have cost us more than any of the above. Our own README told a reader to reproduce results by running the original seed implementation rather than the champion — and the organizer's script still held the untouched stub. Anyone following our instructions verbatim would have measured roughly 1x, seeing none of the work. Caught and fixed in the final hours.

Accomplishments that we're proud of

13/13 shapes correct, worst error 0.00176 against a 0.002 budget, verified across three random seeds.
Running the test case the reference cannot. Shape 14 went from "confirmed infeasible" to a completed forward pass.
Six documented negative results with numbers attached — a fused-block kernel that missed target by 4.7x, its persistent variant at 0.55x, stream pipelining 6.6% slower, and three more. We kept every one, because knowing which ideas don't work on this hardware is a finding.
A measured error bar on our own headline, so we know which claims survive re-measurement and which don't.

What we learned

The correctness gate is a research filter, not a formality. Once you accept that the computed function is fixed, most of the fast-transformer literature falls away and the problem becomes removing overhead, not removing work.

Where attention matters is not where you'd assume. On twelve of fourteen shapes attention is only ~14% of the FLOPs, so even an infinitely fast attention buys at most 1.16x. We measured which backend actually fired: FlashAttention was eligible on 0 of 14 shapes at fp32 and 14 of 14 at fp16. Our gains came from launch overhead, precision and fusion — not from the technique everyone assumes is the answer.

Measure the noise before claiming a win. We nearly reported a 3.9% improvement that sat inside a 3.3% noise band.

Profile before writing kernels. Every kernel we wrote from a profiler trace landed. The one we wrote from a plausible theory failed by 4.7x.

Ship the thing a judge will actually run. Our fastest kernel was worth nothing while the entry point still pointed at the baseline.

What's next for MegaBee

Shape 8 is the unfinished business — 1.81x, roughly 30% of the card's fp16 ceiling, the largest identified headroom we didn't close.
Cross-device confirmation. Our canonical measurement device changed mid-project; every claim deserves re-running on both.
Automate the ledger's per-shape capture end to end. We fixed it late; had it existed from the start, several hand-transcribed tables — and one stale headline — would never have happened.