Image: bair.berkeley.edu · rights & removal

From CUDA to MLX: How K-Search Brings Decades of Kernel Expertise to Apple Silicon

Reporting by BAIR BlogRead the original at bair.berkeley.edu

Executive Summary

The work investigates the transferability of deep kernel optimization expertise from the CUDA ecosystem to Apple Silicon's MLX framework using an AI-driven evolutionary search method called K-Search. The research proposes a novel structured translation layer to bridge this gap, allowing existing CUDA kernel knowledge to inform the optimization process for new hardware. The study demonstrates that by providing architectural context via a translation layer—including concept mapping tables and hardware-specific hints—K-Search can achieve near-expert performance on Apple Silicon with significant speedups over native implementations. Specifically, the approach yielded a 0.97x speedup compared to the MLX Attention kernel and a 20x prefill speedup over community MLX implementations for Mamba SSM kernels. The success is attributed to how the translation layer guides an evolutionary search through architecture-specific constraints, enabling the discovery of optimizations like using Apple's specific hardware instructions.

Facts Only

* The research built on K-Search, an evolutionary kernel search framework developed by Cao et al. at Berkeley Sky Lab.
* A novel structured CUDA-to-MLX translation layer was developed to adapt existing CUDA kernels for Apple Silicon.
* The approach achieved a 0.97x speedup compared to the native MLX Attention kernel on Apple Silicon.
* It achieved up to a 20x prefill speedup over the community mlx-lm implementation on the Mamba SSM kernel.
* The optimization gain resulted from identifying optimizations such as threadgroup memory tiling and using the exp2 trick for softmax based on high-performance CUDA knowledge.
* A second evaluation applied K-Search to the Mamba State Space Model (SSM) kernel, showing a 20x prefill speedup over mlx-lm.
* The speedup in SSM was attributed to applying parallel scan optimization due to the associativity of the recurrence relation.
* The work involved evaluating three configurations: naive baseline, pure evolution, and full context translation layer.

Full Take

The central pattern observed is that transferring expertise requires not just translating syntax but encoding implicit architectural knowledge into constraints for an adaptive search process. The transition from a flat optimization list to a structured decision tree (the world model) demonstrates a crucial cognitive shift: moving the burden of architectural understanding from the LLM's generative capacity to explicitly encoded rules. This suggests that deep domain expertise is not merely a repository of code, but a formalizable set of hardware invariants and performance constraints. The success hinges on establishing reliable mappings between disparate instruction sets (CUDA primitives vs. Metal/MLX idioms) and measuring these mismatches (e.g., memory bandwidth differences). The finding regarding the Mamba kernel suggests that optimization challenges differ based on the operation: attention focuses on fused algebraic operations, while SSM focuses on recurrence and parallelization potential. The system's effectiveness implies a pattern where abstract, context-rich knowledge can serve as an effective scaffold for low-level physical realization across different computational paradigms. What is the underlying mechanism by which this structured translation layer generalizes beyond the specific examples provided?

From the original · BAIR Blog

We face a new epoch in computing. Hardware is changing rapidly — not just faster GPUs, but a growing range of chips from different vendors, each with its own architecture and often tailored to specific AI workloads.

Read the full story at bair.berkeley.edu

Sentinel — provisional

Raw score, 0 to 1 (uncalibrated)0.30

No strong signs of machine writing were found in the source article. Provisional estimate, not a finding that a person wrote it.

The model's own note (unverified)

The text reads like a technical report or research paper detailing novel methods for transferring AI kernel optimization knowledge between different hardware architectures, possessing high internal consistency and specific empirical detail.

This looks only at the wording of the original source article, not at this page's AI-written sections. A small local AI model made this estimate. It has not been checked against known human and machine texts, so treat it as provisional. It cannot show who wrote an article.

What the model noticed in the source article

Varied sentence length and complex rhetorical structure; use of domain-specific jargon integrated fluidly.

Strong, focused argument built around a specific technical methodology (K-Search) and its application across different hardware paradigms.

Logical flow connecting the high-level problem (kernel transfer) to the technical solution (K-Search, translation layer) and empirical results (speedups).

Presence of highly specific, novel technical details (e.g., K-Search structure, MLX/Metal mappings, specific speedup metrics) suggesting deep domain expertise or direct research output.

Signs the model read as against machine writing

Use of very dense, specific, and highly nuanced technical terminology (e.g., 'SSM scan kernels', 'threadgroup-memory', 'exp2 trick') indicative of deep domain knowledge.

The narrative structure moves from general observation to a novel method proposal, execution details, comparative results, and future directions—a hallmark of research communication.

Inclusion of specific citations (e.g., Cao et al., 2026) and Appendix instructions points toward an academic or research paper derived style.