Inspiration

The problem statement includes a test case with sequence length 100,000. Its attention matrix, computed directly, requires about 20 terabytes of memory. The task had to run on our own hardware: a Windows laptop with a 6 GB NVIDIA GPU and a MacBook. The project examines how much of this workload consumer hardware can handle when the implementation is adapted to the device.

What it does

The project is a replacement implementation of the organiser's transformer layer. It produces outputs within the required tolerance on all fourteen official test cases, on both devices. It runs 2 to 15 times faster than the reference implementation, depending on the case. The 100,000-token case completes in 147 seconds on the 6 GB GPU. Its output was verified against float64 arithmetic.

How we built it

Every candidate change followed the same procedure. It first had to pass the correctness check. It was then timed against the current version in an interleaved comparison. Changes were kept only if they measured faster. Rejected changes were logged together with their measurements. In total, 25 experiments were recorded and 11 were rejected.

The implementation has four main components. A custom kernel merges the normalization step, which previously required three GPU launches and five memory passes, into a single pass; this kernel was written twice, in Metal and in CUDA. A custom Metal attention kernel completes in 66 microseconds, compared to 182 for Apple's library kernel. A mixed-precision scheme runs the heavy arithmetic in fp16 and keeps error-sensitive steps in fp32. Inputs larger than GPU memory are processed one sequence at a time, with attention computed in blocks.

Challenges we ran into

The timing method had to be corrected before any optimization. The original schedule reported a 21% speed difference between two identical programs, because the GPU was cooler during the first measurement block.

PyTorch on Windows ships without a kernel compiler, which normally prevents custom CUDA kernels. NVRTC, a runtime compiler included in the PyTorch installation, made compilation at startup possible.

At sequence length 16,384, Apple's attention kernel returns values about 700 times outside the tolerance without raising an error. On the NVIDIA GPU, the equivalent comparison gave the opposite outcome: the library kernel was 7 times faster than our implementation. The final code therefore stores a separate set of measured choices for each device.

Accomplishments that we're proud of

All fourteen cases pass on both devices, including the case that cannot be computed in its direct form. Hardware utilization on the NVIDIA GPU increased from about 10% for the reference to 40% counting executed work, or 78% under the convention that credits skipped causal work. Both figures are reported because the organisers have not specified a convention. A reproducible numerical defect in Apple's GPU library was identified, documented, and worked around. The experiment log records the 11 rejected methods with the measurements that rejected them. What we learned Several changes we expected to help made the code slower. This was only visible because every change was timed under controlled conditions. Tuning decisions did not transfer between devices: a threshold measured on one platform, or under one precision, was sometimes wrong on the other. Consumer laptops also introduced failure modes that required specific handling, such as thermal drift and silent memory paging.

What's next for "Squeezing 20 Terabytes into a Gaming Laptop"

The remaining measured inefficiency is kernel fusion around the matrix multiplications, which the runtime compiler cannot reach. A working Triton installation on Windows would make it accessible; earlier experiments indicate about a factor of two remains on the long-sequence path. A further step is extending the measurement harness so the code re-tunes itself on new hardware. The same procedure could then be applied to a full model rather than a single layer.

Built With

Share this project:

Updates

Submission history