The GPU kernel optimization problem has always been a bottleneck wrapped in expertise. NVIDIA’s CUDA ecosystem represents thousands of engineering hours distilled into finely-tuned implementations of attention, state space models, and other critical operations. But those years of accumulated knowledge sit trapped on CUDA hardware. When Apple Silicon, custom AI accelerators, or other emerging chips need the same optimizations, teams either rediscover them from scratch or accept significant performance penalties.
I think we’re watching the beginning of the end for that cycle.
A new research direction from UC Berkeley’s Sky Lab, extended with an Apple Silicon backend, suggests that GPU kernel expertise can transfer automatically across architectures. The work builds on K-Search, an evolutionary framework that uses LLMs to reason about and iteratively optimize GPU kernels. The novel contribution is a structured CUDA-to-MLX translation layer that converts hard-won NVIDIA optimization strategies into actionable guidance for Apple’s Metal framework, without requiring manual porting.
Why This Matters Now
Apple Silicon adoption is accelerating. Hundreds of millions of MacBooks and Mac Studios now run Apple’s chips, and MLX has emerged as a serious framework for local AI inference. But beneath the momentum lies a painful gap: critical kernels like paged attention, optimized SSM scans, and fused MoE routing either don’t exist or run naively on Apple hardware.
MLX works. It runs models correctly. But it often leaves substantial performance on the table.
This isn’t just an Apple problem. We’re entering an era where hardware diversity will be the norm, not the exception. Custom AI accelerators from startups, domain-specific processors from established vendors, and specialized chips from cloud providers will proliferate. Each one will need optimized kernels. We cannot afford to rediscover the same optimizations repeatedly across every new platform.
The translation layer approach offers a different path: instead of copying CUDA code line-for-line (which fails because architectures differ fundamentally), K-Search extracts the conceptual optimizations and adapts them to target hardware constraints.
How the Translation Works
The key insight is that an LLM prompted as a GPU performance engineer can reason about optimization strategies at the right level of abstraction. Before proposing any changes, the model:
- Classifies the kernel type (reduction, scan, attention)
- Rewrites the reference computation in canonical form
- Maps data layout and access patterns
- Hypothesizes the likely bottleneck (bandwidth, latency, compute, synchronization)
Only after this analysis does it emit candidate optimizations. Each is a single implementable change, tracked in a decision tree called the “world model.” The tree persists across rounds, allowing K-Search to expand promising branches and prune dead ends.
On the attention kernel, this approach achieved 0.97x speedup relative to Apple’s native MLX implementation. The evolved kernel independently discovered and implemented the key strategies from FlashAttention-2: threadgroup memory tiling, online softmax, K-transposition for memory patterns, and the exp2 trick (replacing base-e exponentials with faster base-2 hardware instructions).
But the attention result, impressive as it is, tells only part of the story.
The Mamba Test
Attention kernels are well-studied. The real test was whether the method generalizes. Researchers applied K-Search to Mamba’s SSM kernel, a fundamentally different computational pattern where the bottleneck is recurrent state updates rather than softmax operations.
The evolved MLX implementation achieved approximately 20x higher prefill throughput compared to the community mlx-lm baseline.
That difference traces back to a single missing optimization: the community implementation processes tokens sequentially, leaving GPU compute largely idle. The evolved kernel applies a parallel prefix scan, an associative operation that can evaluate the entire sequence in O(log N) dependent steps instead of O(N), making full use of Apple Silicon’s parallelism.
This is instructive because it shows the method working on a genuinely different problem. The framework didn’t just memorize attention optimizations and apply them elsewhere. It extracted more general reasoning about GPU optimization strategies and adapted them to a new kernel type and architectural context.
What This Enables
For developers, this approach suggests several immediate implications. First, performance gaps between CUDA and newer platforms may shrink faster than we expect. The expertise transfer becomes semi-automatic, removing one barrier to hardware adoption.
Second, the translation layer concept applies beyond MLX to any ecosystem where CUDA knowledge is transferable. IBM Spyre, emerging NPUs, and future custom accelerators could all benefit from the same approach with new backends.
Third, and perhaps most important, this shifts where human expertise matters. Rather than hand-tuning kernels for each new hardware platform, experts focus on building better translation layers and stronger domain specs that keep the LLM reasoning grounded in hardware reality.
Read more about AI-driven optimization and the broader challenges of cross-platform development in the evolving GPU ecosystem.
The question is no longer whether AI can optimize kernels, but how quickly we can move that optimization knowledge across the fragmented landscape of AI hardware.