I’ve been watching the GPU hardware landscape fracture for years now. NVIDIA dominates, sure, but Apple Silicon, custom AI accelerators, and specialized chips are multiplying faster than anyone can keep up with. The problem isn’t new, but it’s getting urgent: there’s a mountain of optimized CUDA kernel knowledge that simply doesn’t transfer to these new platforms.
That expertise cost decades to accumulate. Hand-tuned implementations of attention mechanisms, state-space models, and other critical operations represent thousands of engineering hours. When a new chip arrives, we essentially start over. A team rebuilds the same optimizations from scratch, often rediscovering the same tricks independently. It’s wasteful and it’s slowing down the adoption of promising new hardware.
Recently I read about K-Search, an evolutionary framework from Berkeley that uses AI to optimize GPU kernels. What grabbed my attention wasn’t the framework itself, but the extension someone built for MLX on Apple Silicon. They didn’t just run the optimizer on a new platform. They created a structured translation layer that takes CUDA kernel expertise and converts it into actionable guidance for Apple’s Metal backend.
The Translation Problem
Here’s what matters: you can’t just hand an LLM a CUDA kernel and ask it to port it to Metal. Without deep hardware context, the resulting code is syntactically valid but architecturally wrong. Tile sizes don’t match the GPU’s memory hierarchy. Primitives don’t exist. Memory assumptions fail. The gap between “this compiles” and “this is actually fast” is where the real work lives.
The researchers solved this by building a translation layer that extracts high-performance patterns from CUDA kernels and reformulates them as architecture-native strategies. Instead of copying instructions one-to-one, it captures the why behind each optimization and maps it to equivalent techniques on the target hardware.
For attention kernels specifically, they showed that this approach could reach 0.97x the speed of Apple’s state-of-the-art implementation starting from a naive baseline. More impressively, the evolved kernel independently discovered the key optimizations from FlashAttention-2: threadgroup memory tiling, online softmax, K-transposition for memory access patterns, and the exp2 trick that replaces general exponentials with base-2 variants to hit Apple’s fast hardware instructions.
That last one is particularly telling. It’s a small optimization that means almost nothing in isolation but compounds into real speedup. An AI system copying CUDA code line-by-line would never discover it. But an evolutionary search guided by architectural knowledge finds it because it understands the GPU’s constraints.
Beyond Attention
I was convinced about attention kernels until I saw what happened with Mamba. The researchers applied K-Search to state-space models, which have a completely different computational bottleneck. Instead of a softmax, you’re dealing with a recurrent state update. It’s a fundamentally different optimization problem.
The result was striking. Their evolved Metal kernel achieved ~20x prefill speedup over the community MLX implementation. The difference came down to one thing: parallel scan. The original implementation processed tokens sequentially, leaving most of Apple Silicon’s compute idle. The evolved kernel applied a parallel prefix scan that reduces O(N) dependent steps to O(log N). Same algorithm, but structured to match the hardware.
This matters because it suggests the approach generalizes. Attention and SSM kernels have different bottlenecks, different data access patterns, different computational structures. If the translation layer helps with both, it might help across a broader range of operations.
What This Means for Developers
I think this points to a major shift in how we’ll approach kernel optimization. Instead of assembling teams of GPU experts for each new architecture, we’re outsourcing the search itself to AI. But that only works if we can guide the search effectively. The translation layer is the guidance system. It encodes what we know about one architecture in a way that’s transferable to another.
For companies adopting Apple Silicon or custom accelerators, this is significant. You no longer need to choose between adopting new hardware and accepting performance penalties. The translation-guided search means you can port critical operations to new platforms in weeks or months instead of years.
For the broader AI infrastructure stack, this suggests we’re moving toward a world where kernel optimization becomes more systematic and less heroic. Less “we need a GPU expert” and more “we need to encode architectural knowledge so search systems can find optimizations automatically.” That’s more scalable, though it requires thinking carefully about what knowledge to transfer and how to formalize it.
The real question isn’t whether this works for attention and SSM kernels. The question is whether we can scale this approach to the dozens or hundreds of operations that make up a full ML stack, and whether the translation patterns generalize to hardware platforms we haven’t even designed yet.