Transferring GPU Expertise Across Hardware with AI-Driven Kernel Search

Transferring GPU Expertise Across Hardware with AI-Driven Kernel Search

The GPU optimization problem we face today is fundamentally different from five years ago. Back then, CUDA dominated, and if you wanted peak performance, you optimized for NVIDIA. Now? Apple Silicon is in hundreds of millions of MacBooks, custom AI accelerators are proliferating, and every chip vendor has their own instruction set, memory hierarchy, and preferred programming model.

The real bottleneck isn’t computing power anymore. It’s expertise. CUDA’s ecosystem represents decades of hard-won kernel knowledge: hand-tuned implementations of attention, state space models, fused operations. That accumulated wisdom is gold, but it’s locked to NVIDIA hardware. When a new chip arrives, we essentially start over, rediscovering the same optimizations from scratch on unfamiliar architectures.

I find this wasteful. And according to recent work extending K-Search to Apple Silicon, it doesn’t have to be.

The Translation Problem

The naive approach is obvious: take a high-performance CUDA kernel and ask an LLM to port it to Metal or any other target language. The problem is equally obvious when you try it. You get syntactically valid code that runs on the wrong hardware assumptions. Tile sizes that made sense for NVIDIA’s memory model perform terribly on Apple’s unified memory. Memory access patterns optimized for one cache hierarchy wreck another. The LLM has no concept of why the original optimizations existed.

What K-Search’s authors realized is that the gap between architectures is not insurmountable. It’s a translation problem. CUDA expertise encodes general principles about kernel optimization: how to hide latency through double buffering, when to tile computation, which memory operations are cheap vs. expensive on your target hardware. These principles are transferable. You just need to ground the LLM in architecture-specific context.

Their solution was elegant: build a structured translation layer that extracts implementation strategies from high-performance CUDA kernels (like FlashAttention-2) and feeds them to K-Search as actionable guidance for Apple Silicon. The LLM doesn’t start with “here’s Metal syntax, now port this.” It starts with “here’s what high-performance kernels do on CUDA, here’s how Apple Silicon’s memory model works differently, now reason about equivalent strategies.”

The results speak for themselves. On Apple’s MLX attention kernel, this approach achieved 0.97x the speed of the native baseline. On Mamba’s state space model kernel, the evolved implementation hit roughly 20x prefill speedup over the community implementation. The 20x gain came from recognizing that the SSM recurrence, which looks inherently sequential, can be rewritten as a parallel scan operation. The algorithm was there in the literature. The optimization required knowing both the math and the hardware well enough to recognize the opportunity.

Why This Matters Beyond Benchmarks

Let me be direct about what’s happening here: we’re decoupling kernel optimization expertise from hardware-specific knowledge. That’s a massive shift in how we should think about performance engineering going forward.

For years, optimizing kernels required holding two things simultaneously: deep understanding of the algorithm and deep familiarity with specific hardware. Those two concerns are intertwined in practice, but conceptually they’re separate. K-Search’s translation layer lets us separate them. You don’t need a team of GPU experts for every new chip that arrives. You need the expertise once, embedded as context and constraints that an AI optimizer can reason about.

The practical implication is that newer hardware ecosystems can catch up much faster. Apple Silicon, custom accelerators, and emerging chips no longer face the expertise tax that NVIDIA had to pay over decades. That accelerates adoption and shifts competitive advantages from “who has the best kernel library” to more interesting questions about architecture design and API usability.

For developers, this means MLX and similar frameworks for non-CUDA hardware become genuinely viable for performance-critical workloads. Local inference on MacBooks wasn’t just possible before; it was slow. That barrier is lower now, which changes the calculus for where inference actually runs.

The Harder Problem Still Ahead

The work focuses on attention and SSM kernels because they’re well-studied and the optimizations are somewhat understood. Real production systems need paged attention, fused MoE routing, custom operators that don’t exist in any reference implementation. The question of whether this translation approach generalizes to novel kernel patterns, where you don’t have a high-performance CUDA reference to learn from, remains open.

There’s also the question of whether the LLM is really doing the optimization work, or if the translation layer is. The 0.97x result on attention is impressive, but how much comes from the structured context versus the evolutionary search? The authors isolate this explicitly: going from 0.26x (pure evolution) to 0.97x (evolution plus context) shows the translation layer is doing most of the work. That’s honest, and it’s the right lesson: the bottleneck in kernel optimization isn’t code generation. It’s context quality.

If we can systematize that context well enough, the tedious work of manual kernel tuning becomes something algorithms handle for us. And if we can do that across multiple hardware targets simultaneously, we’ve solved not just a performance problem but an economic one.

Read Next