Introduction

The problem statement given to us was twofold: firstly, optimise the Transformer layer from the BaselineTransformer benchmark, and secondly, explore how AI can help developers analyse Transformer workloads, identify bottlenecks, and generate more efficient implementations for specific GPUs.

After attending the information session, we also understood that part of the problem statement was to perform these optimisations on consumer-grade hardware.

With that in mind, we set out to optimise the Transformer layer on our personal computers. Throughout the optimisation process, we worked on a MacBook Pro with an M3 Pro and its integrated Apple GPU.

Methodology

We first identified the trade-offs we were dealing with when shifting from industrial/data-centre GPUs to consumer-grade GPUs.

We lose the higher FLOPs and memory bandwidth that industrial-grade GPUs provide, as well as their ability to scale distributed workloads across multiple GPUs in a data centre.
Consumer-grade GPUs are also a less common target for machine learning optimisation, meaning that libraries and open-source implementations may be less extensively optimised for them.
However, Apple Silicon provides a unified memory architecture, where the CPU and GPU share the same physical memory pool. While this does not eliminate all memory copies, it creates opportunities to reduce unnecessary data movement between different representations and stages of execution.

With these trade-offs in mind, we ported the code from PyTorch to MLX, Apple's array framework designed specifically for Apple Silicon. MLX exposes the unified-memory programming model and custom Metal kernels more directly, giving us greater flexibility to investigate platform-specific optimisations.

Subsequently, we followed two main approaches:

Identify possible mathematical and implementation-level optimisations for the Transformer layer that the benchmark missed.
Profile the implementation using the roofline model to determine whether different stages of the Transformer layer were bottlenecked by memory I/O or compute for each test-case shape.

The roofline analysis allowed us to distinguish compute-bound stages from memory-bound stages. This distinction is important because a memory-bound operation cannot necessarily be improved by simply reducing arithmetic. Instead, optimisation may require reducing memory traffic, eliminating unnecessary memory operations, or fusing multiple operations into a single kernel.

Our optimisation process therefore followed an iterative loop:

Profile
  ↓
Identify Bottleneck
  ↓
Form Hypothesis
  ↓
Implement / Microbenchmark
  ↓
Correctness Test
  ↓
Benchmark Across Shapes
  ↓
Keep / Revert
  │
  └──────────────→ Repeat

We followed this loop throughout the development process to arrive at our final solution.

Our Solution

Architecture

Our solution architecture involves a kernel planner, which selects the kernels for different parts of the Transformer layer based on the input shape and generates a kernel plan to be executed.

The planning flow can be seen in the images attached above.

The plan determines, among other decisions, which attention implementation to use, whether to use our custom LayerNorm and fused FFN kernels, whether residual biases can be deferred, and whether large batches should be chunked.

This kernel selection is necessary because optimisations may serve different input shapes differently. A kernel that is substantially faster for one shape can be slower than a plain MLX operation for another.

For example, this can be seen with shapes 6 and 8, which use different attention kernels.

Shape 6 has a head width of 32, for which MLX does not normally dispatch to its fused attention kernel. Our planner therefore selects a custom-compiled version of Apple's Steel attention kernel, resulting in a 1.32x improvement.

In contrast, shape 8 has a head width of 256. Candidate Steel configurations either exceed the GPU's threadgroup-memory limit or perform worse than MLX's existing implementation. The default block shape requires 68.5 KiB of threadgroup memory against the 32 KiB hardware limit, while a narrower configuration that fits achieves only 0.904x the performance of the existing path. The planner therefore deliberately falls back to the standard MLX implementation.

As an example, we have attached the plan generated for test case 6.

The kernel dispatch rules were derived empirically through profiling and controlled benchmarks. Different kernels have hardware constraints and performance trade-offs involving arithmetic intensity, memory traffic, kernel-launch overhead, threadgroup memory, correctness, and input shape.

Consequently, the fastest kernel for one shape may regress another, requiring the planner to select kernels conditionally rather than applying every optimisation globally.

Challenges Faced

One of the major challenges was distinguishing plausible optimisation hypotheses from mechanisms that were actually supported by profiling.

As this optimisation process was performed with a human in the loop, and the human was not initially particularly well versed in machine learning or the individual stages of a Transformer layer, there were occasions where we had to wrestle with unfamiliar terminology and initially rely on the AI's explanations of the underlying behaviour.

However, we found that AI-generated explanations, while often plausible, were not necessarily correct.

For example, an apparent performance penalty from strided Q, K and V tensors was initially attributed to inefficient memory access. Further profiling showed that the custom Steel attention kernel was actually silently materialising contiguous copies of these tensors before execution. The real optimisation was therefore not to improve the strided memory-access pattern, but to remove these hidden copies and allow the kernel to consume the strided views directly.

This reinforced an important principle throughout the project: AI-generated hypotheses should be treated as candidates rather than facts. Optimisations were accepted only after their proposed mechanisms were investigated through profiling, their outputs passed correctness tests, and their performance improvements survived controlled benchmarking.

This also influenced the design of our optimisation log. Successful, unsuccessful, and ruled-out approaches were recorded so that subsequent optimisation attempts could build upon previous findings rather than repeatedly exploring the same ideas.

Future Developments

Autonomous agentic kernel optimisation is within reach if a well-informed agent can orchestrate the optimisation process while specialised agents investigate different hypotheses.

A future system could automate the profiling-hypothesis-benchmark loop used in this project. A profiling agent could identify bottlenecks, while specialised agents investigate different optimisation strategies such as kernel fusion, memory-layout optimisation, algorithmic changes, and GPU-specific kernels. A deterministic correctness and benchmarking harness could then decide which candidate implementations should be retained.

Importantly, we believe agents should first build and share knowledge of platform-specific characteristics before beginning the optimisation process.

Our results demonstrate that seemingly minor implementation details—such as MLX's attention dispatch rules, GPU threadgroup-memory limits, tensor layouts, and hidden tensor materialisations—can determine whether an optimisation improves or degrades performance.

Rather than simply asking an agent to "make the Transformer faster", an autonomous optimisation system should therefore be capable of learning these platform-specific constraints, maintaining a shared record of successful and unsuccessful experiments, and using profiling evidence to decide which hypotheses should be investigated next.

Development Tools

We used Claude Code to perform the kernel implementations and tests, and occasionally used Xcode's native GPU profiling tools to investigate GPU behaviour.

APIs Used

No calls to external services are used in this project.

Libraries and Frameworks Used

We used MLX, Apple's framework for array operations designed specifically for Apple Silicon. The default benchmark implementation uses PyTorch.

MLX gave us a more direct route to investigating Apple-specific GPU behaviour, including custom Metal kernels and the unified-memory programming model, while allowing the original PyTorch-facing benchmark interface to remain intact.