From CUDA to MLX: How K-Search Brings Decades of Kernel Expertise to Apple Silicon
Source
bair.berkeley.edu
Date
Key takeaways · AI-distilled
Handing an LLMA large language model — the neural network behind tools like Claude and ChatGPT, trained on huge amounts of text to predict what comes next.Full definition → a CUDA kernel to port directly produces syntactically valid but architecturally wrong code; the translation layer instead supplies concept-mapping tables, e.g. __shared__'s 32KB Metal limit versus NVIDIA's 48KB memory.
On the attentionThe mechanism that lets a model weigh which earlier words matter for the word it's currently processing — the core operation of a transformer.Full definition → kernel, giving the evolutionary search full architecture-specific context windowThe maximum amount of text a model can consider at once — its working memory for the current conversation or task.Full definition → jumped performance from 0.26x to 0.97x of Apple's own state-of-the-art attention kernel, independently rediscovering FlashAttention-2 techniques like online softmax and the exp2 trick.
The roughly 20x Mamba SSM prefill speedup over mlx-lm comes from exploiting that the state recurrence step is associative, letting it run as an O(log N) parallel scan instead of mlx-lm's O(N) tokenThe chunk of text a model reads and writes in — roughly three-quarters of a word — and the unit AI usage is billed in.Full definition →-by-token processing; the gain shows only in prefill, not decode.
K-Search's search state is a scored decision tree, not a flat list: each candidate node carries an overall rating (0-10), a confidence (0-1), and separate impact scores for memory bandwidth, register pressure, and compute fit, with a stagnation window triggering a branch switch.
Terms in this piece · Glossary
LLM — A large language model — the neural network behind tools like Claude and ChatGPT, trained on huge amounts of text to predict what comes next.
attention — The mechanism that lets a model weigh which earlier words matter for the word it's currently processing — the core operation of a transformer.
context window — The maximum amount of text a model can consider at once — its working memory for the current conversation or task.
token — The chunk of text a model reads and writes in — roughly three-quarters of a word — and the unit AI usage is billed in.
Why it matters
It shows AI can automatically rediscover architecture-specific kernel optimizations instead of hand-porting CUDA code, relevant for anyone deploying models on non-NVIDIA silicon.
