• 

• 

AutoMetal

Overview

AutoMetal is my TikTok TechJam 2026 Track 3 submission. I built it on an 8 GB Apple M2 MacBook Air for the official Transformer inference benchmark, which has 14 fixed workloads ranging from tiny launch-sensitive shapes to a batch-32, sequence-100,000 stress case.

The benchmark asks for a compatible UserOptimizedTransformer that preserves the organizer's model layout, forward interface, and numerical tolerance while improving useful throughput. One generic implementation left too much performance behind. AutoMetal measures the available paths and dispatches each shape to its fastest precision-passing implementation.

The final system combines packed QKV projections, native MPS scaled-dot-product attention, compiled MLX graphs, head-dimension padding, zero-copy PyTorch/MLX exchange, batch tiling, and mixed FP16/FP32 kernels. It keeps the fastest validated path for every shape instead of picking a single backend for the entire benchmark.

Results

Result	Measurement
Completion	All 14 official logical workloads have a completion path
Precision, Shapes 1–13	Organizer FP32 precision gate passes with zero failed elements
Overall speedup	4.98× geometric-mean speedup over the challenge reference
SDPA comparison	4.00× geometric-mean speedup over the strong SDPA baseline
Mean MFU	33.76% arithmetic mean across all 14 logical workloads
Peak throughput	2.005 useful TFLOP/s on Shape 8
Peak estimated FP32 MFU	70.12% on Shape 8
Shape 13	7.15× reference speedup, 1.259 useful TFLOP/s, 44.02% estimated FP32 MFU
Shape 14	Full logical streamed FP32 workload completed in 1,169.986 s at 1.189 useful TFLOP/s, with 4.326 GB peak active/cache memory

The speedup over SDPA came from more than replacing the organizer attention with a framework primitive. Packed QKV projection cut launches, compiled MLX graphs helped launch-bound shapes, padding opened faster attention paths, and the dispatcher kept different choices for different configurations. The 4.98× and 4.00× geometric means are unweighted internal diagnostics because the organizer does not publish the official shape weighting.

How it works

The submitted UserOptimizedTransformer inherits a measured DispatchTransformer. The dispatcher reads the model configuration and selects the retained implementation for that shape.

Small shapes often lose time in graph setup and kernel launches, so they use low-overhead compiled execution where it wins. Larger conventional shapes use packed QKV projection and native MPS attention. The extreme-batch case uses batch tiling. Long-sequence paths block or stream attention to reduce intermediate memory traffic. MLX remains a specialized backend rather than a universal replacement because it only stays in the dispatcher when its measured result wins.

Precision decides whether a path stays. Full-model FP16 produced very small average error in early experiments but still failed individual elements under the organizer check. AutoMetal therefore keeps FP32 residual and normalization state where required and only uses reduced precision inside validated linear and attention kernels.

Shape 14

Shape 14 is the benchmark's largest logical workload:

Parameter	Value
Batch size	32
Sequence length	100,000
Model dimension	1,024
Attention heads	16
Layers	2
FFN dimension	1,024

The untouched organizer FP32 harness allocates the whole input before it calls participant code. That one input needs 12.21 GiB. The M2/Metal environment rejects the allocation before forward() starts. Even the monolithic FP16 input is 6.10 GiB and fails in the same interface.

I preserved that result, then tested the actual compute problem separately. First, streaming attention with online softmax removed the quadratic attention matrix. A batch-2 implementation completed the full logical workload in 934.623 seconds at 1.489 useful TFLOP/s with 5.760 GB peak active memory.

I then changed the input and execution plumbing for an end-to-end run. The runner advances the organizer-style FP32 random stream one batch tile at a time, processes all 100,000 tokens, stores the output digest, releases the tile, and repeats until all 32 batch elements complete. Attention streams keys and values with online softmax, so it never materializes a 100,000 × 100,000 attention matrix. The active path uses FP32 input and state, FP16 linear and attention compute, and an FP32 output digest.

Shape 14 stage	Result
Untouched organizer FP32 interface	Fails before forward() because the 12.21 GiB input does not fit
First memory-efficient compute run	Completed batch 32 × sequence 100,000 in 934.623 s, 1.489 useful TFLOP/s, 5.760 GB peak
Current streamed official-style FP32 run	Completed in 1,169.986 s, 1.189 useful TFLOP/s, 4.326 GB peak active/cache memory

Development tools

I developed and benchmarked AutoMetal locally on a MacBook Air with an 8-core Apple M2 GPU and 8 GB of unified memory. The runtime used macOS 26.6.2, Python 3.12.14, PyTorch 2.13 with MPS, and MLX 0.32.2. I used VS Code, the terminal, Git, local Python profiling and validation scripts, plus Swift and Metal probes to inspect device memory limits.

I also used the OpenAI Codex desktop coding agent with GPT-5-Sol for repository analysis, experiment planning, implementation, benchmark automation, debugging, and technical documentation. All runs happened on the local M2 machine. I did not use a remote GPU or cloud notebook.

APIs, libraries, data, and assets

The project uses no external web, commercial, or third-party service APIs. Its local platform and framework APIs include PyTorch MPS device APIs with synchronized MPS timing, MLX tensors, compilation and fused attention, Metal device queries such as maxBufferLength and recommendedMaxWorkingSetSize, DLPack for PyTorch/MLX tensor exchange, and Python memory-mapped storage for streamed Shape 14 output.

The main libraries are PyTorch 2.13.0, Apple MLX 0.32.2, NumPy 2.5.2, psutil 7.2.2, the Python standard library, Swift, and Apple Metal system frameworks.

There is no external dataset. The benchmark provides synthetic random inputs, the 14 official Transformer configurations, and organizer-compatible model weights. I used no pretrained external model, scraped data, manually labelled data, or third-party visual assets.

Challenges and what I would improve

The difficult part was finding speed without creating a result that only looked correct. I rejected paths after exact evaluation found individual failures, even when their mean error was tiny. I also had to separate Shape 14's algorithmic memory cost from its interface allocation failure. That led to the streaming attention and incremental input runner.

With more time, I would write custom Metal kernels for the remaining launch-sensitive shapes and reduce framework transitions in the Shape 14 pipeline. I would also validate longer sequences on a higher-memory reference host while retaining the 8 GB M2 as the target machine.

Built With

• apple-mlx
• numpy
• python
• pytorch