Metric<\/th>\n	mlx-mamba (ours)<\/th>\n	mlx-lm (neighborhood)<\/th>\n	mamba.py<\/th>\n<\/tr>\n<\/thead>\n
Decode<\/td>\n	152 tok\/s<\/td>\n	116 tok\/s<\/td>\n	40 tok\/s<\/td>\n<\/tr>\n
Prefill L=512<\/td>\n	5,751 tok\/s<\/td>\n	329 tok\/s<\/td>\n	1,089 tok\/s<\/td>\n<\/tr>\n
Prefill L=1024<\/td>\n	6,010 tok\/s<\/td>\n	327 tok\/s<\/td>\n	1,127 tok\/s<\/td>\n<\/tr>\n
Prefill L=2048<\/td>\n	6,612 tok\/s<\/td>\n	326 tok\/s<\/td>\n	1,092 tok\/s<\/td>\n<\/tr>\n
Prefill L=4096<\/td>\n	6,743 tok\/s<\/td>\n	339 tok\/s<\/td>\n	1,042 tok\/s<\/td>\n<\/tr>\n<\/tbody>\n<\/table>\n<\/div>\n Desk 1:<\/strong> Prefill and decode throughput on mamba-370m (f16, M1 Max 64GB). mlx-mamba (ours) reaches ~20\u00d7 larger prefill throughput than the neighborhood mlx-lm baseline, whereas decode stays comparable.<\/p>\n The ~20\u00d7 prefill speedup over mlx-lm comes down to 1 distinction: mlx-lm doesn’t implement a parallel scan for the SSM. The state recurrence<\/p>\n [h_t = bar{a}_t h_{t-1} + bar{b}_t]<\/p>\n seems inherently sequential, however every step will be written as a pair $(bar{a}_t, bar{b}_t)$ beneath the associative mix<\/p>\n [(a_2, b_2) circ (a_1, b_1) = left(a_2 a_1, a_2 b_1 + b_2right),]<\/p>\n which reproduces the recurrence precisely. As a result of the operator is associative, the entire sequence will be evaluated with a parallel (prefix) scan in $O(log N)$ dependent steps as an alternative of $O(N)$. mlx-lm skips this and processes tokens one after the other, leaving most of Apple Silicon\u2019s compute idle; our advanced Metallic kernel applies the scan and makes a lot fuller use of GPU throughput. The achieve exhibits up in prefill, the place the total sequence is accessible to scan in parallel, and never in single-token decode, the place there is just one new token per step and no scan to parallelize \u2014 which is why the decode row is roughly flat whereas prefill is ~20\u00d7.<\/p>\n mamba.py is sluggish on each prefill and decode as a result of it’s a PyTorch reference implementation that falls again to CPU or MPS on Apple Silicon, forgoing the hardware-specific optimizations that MLX\u2019s Metallic backend makes doable.<\/p>\n What\u2019s subsequent?<\/h2>\n On the 2 kernels we studied, AI-driven evolutionary kernel search grounded in structured cross-platform translation data reached near-expert efficiency on Apple Silicon with no group of GPU specialists ranging from scratch. We don’t but know the way far this generalizes, however the result’s encouraging.<\/p>\n For us the primary takeaway is that the bottleneck was not the LLM\u2019s capacity to write down Metallic code, however the high quality of the context and constraints we gave it. Our CUDA translation layer converts present NVIDIA kernel experience into actionable steering for Apple Silicon, and lets Ok-Search\u2019s evolutionary search do the remaining.<\/p>\n We’re actively extending this work in a number of instructions: supporting new architectures, with present efforts targeted on creating new kernels for the IBM Spyre AIU and broader {hardware} targets; including extra kernels reminiscent of paged consideration and fused MoE routing; and bettering integration with the Ok-Search evolution loop to make translation context much more computerized.<\/p>\n Acknowledgements<\/h2>\n This work was carried out by IBM Analysis and builds on Ok-Search from the UC Berkeley Sky Lab (Cao et al., 2026<\/a>). We welcome collaboration and suggestions from the MLX and broader AI programs communities. If you’re engaged on kernel optimization for non-CUDA {hardware}, we’d love to listen to from you.<\/p>\n\n Quotation<\/h2>\n \n \n @article<\/span>{<\/span>cao2026k<\/span>,<\/span>\n title<\/span>=<\/span>{Ok-Search: LLM Kernel Technology through Co-Evolving Intrinsic World Mannequin}<\/span>,<\/span>\n creator<\/span>=<\/span>{Cao, Shiyi and Mao, Ziming and Gonzalez, Joseph E and Stoica, Ion}<\/span>,<\/span>\n journal<\/span>=<\/span>{arXiv preprint arXiv:2602.19128}<\/span>,<\/span>\n 12 months<\/span>=<\/span>{2026}<\/span>\n}<\/span>\n<\/code><\/pre>\n<\/div>\n<\/div>\n\nAppendix: Attempt it your self<\/h2>\nThe MLX backend is constructed on prime of the open-source Ok-Search repo, so the outcomes right here will be reproduced instantly. The steps are:<\/p>\n 1. Clone and set up<\/p>\n \n\ngit clone https:\/\/github.com\/caoshiyi\/Ok-Search.git\ncd <\/span>Ok-Search\n\nuv pip set up <\/span>openai wandb\nuv pip set up <\/span>git+https:\/\/github.com\/caoshiyi\/flashinfer-bench-ksearch.git\n<\/code><\/pre>\n<\/div><\/div>\n2. Set your credentials<\/p>\n Open the related script beneath scripts\/ and set three variables on the prime:<\/p>\n \n\nKSEARCH_ROOT<\/span>=<\/span>\/path\/to\/Ok-Search\nAPI_KEY<\/span>=<\/span>your-llm-api-key\n<\/code><\/pre>\n<\/div><\/div>\n3. Run kernel search<\/p>\n \n\n# Optimize Flash Consideration on Apple Silicon (world-model mode)<\/span>\nbash scripts\/mac_flash_attention_wm.sh\n\n# Or a Mamba SSM kernel, e.g. the selective scan<\/span>\nbash scripts\/mamba_selective_scan_fwd_wm.sh\n<\/code><\/pre>\n<\/div><\/div>\nFull CLI reference and documentation are within the README.<\/p>\n<\/div>\n<\/div>\n\n","protected":false},"excerpt":{"rendered":" Determine 1: CUDA-to-MLX optimization translation map. CUDA optimization data will be translated into architecture-native MLX methods moderately than copied instruction-for-instruction. We face a brand new epoch in computing. {Hardware} is altering quickly \u2014 not simply sooner GPUs, however a rising vary of chips from totally different distributors, every with its personal structure and sometimes tailor-made […]<\/p>\n","protected":false},"author":2,"featured_media":17284,"comment_status":"open","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"footnotes":""},"categories":[55],"tags":[1395,311,310,110,414,7229,9992,312,4943,9991,193,3391],"class_list":["post-17282","post","type-post","status-publish","format-standard","has-post-thumbnail","hentry","category-machine-learning","tag-apple","tag-artificial","tag-berkeley","tag-blog","tag-brings","tag-decades","tag-expertise","tag-intelligence","tag-kernel","tag-ksearch","tag-research","tag-silicon"],"_links":{"self":[{"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/posts\/17282","targetHints":{"allow":["GET"]}}],"collection":[{"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/posts"}],"about":[{"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/types\/post"}],"author":[{"embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/users\/2"}],"replies":[{"embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=%2Fwp%2Fv2%2Fcomments&post=17282"}],"version-history":[{"count":1,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/posts\/17282\/revisions"}],"predecessor-version":[{"id":17283,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/posts\/17282\/revisions\/17283"}],"wp:featuredmedia":[{"embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/media\/17284"}],"wp:attachment":[{"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=%2Fwp%2Fv2%2Fmedia&parent=17282"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=%2Fwp%2Fv2%2Fcategories&post=17282"},{"taxonomy":"post_tag","embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=%2Fwp%2Fv2%2Ftags&post=17282"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}

\n
\n<\/p>\n

<\/p>\n

\n $\"Kernel$ <\/p>\n

Determine 1: CUDA-to-MLX optimization translation map.<\/strong> CUDA optimization data will be translated into architecture-native MLX methods moderately than copied instruction-for-instruction.<\/p>\n

We face a brand new epoch in computing. {Hardware} is altering quickly \u2014 not simply sooner GPUs, however a rising vary of chips from totally different distributors, every with its personal structure and sometimes tailor-made to particular AI workloads. Software program is altering simply as quick, and AI coding instruments now generate in minutes what took months of effort a couple of years in the past.<\/p>\n

<\/p>\n

With a lot of computing now centered on AI, GPU kernels are an important element of its success. These are the low-level applications that run contained in the GPU, and writing environment friendly ones is way from apparent \u2014 it takes years of experience to get proper. Transferring a kernel from one vendor\u2019s {hardware} to a different is tougher nonetheless, and sometimes means rediscovering the identical optimizations from scratch. The CUDA ecosystem, for instance, has collected a long time of hard-won kernel experience: hand-tuned implementations of consideration, state house fashions, and different vital operations representing 1000’s of engineering hours. Newer {hardware} ecosystems (Apple Silicon, customized AI accelerators, and others) are rising quick however lack this depth.<\/p>\n

On this work we ask whether or not that experience will be transferred mechanically. We constructed on Ok-Search<\/a>, an evolutionary kernel search framework launched by Cao et al. at Berkeley Sky Lab that makes use of AI to optimize GPU kernels, and prolonged it with a backend for MLX \u2014 Apple\u2019s machine-learning framework for its personal Apple Silicon chips. We developed a novel structured CUDA-to-MLX translation layer that lets Ok-Search take present CUDA kernels as a data base and adapt them into high-quality GPU kernels for Apple Silicon, moderately than rebuilding from scratch.<\/p>\n

We present that our method reaches near-expert stage efficiency on Apple Silicon with 0.97x speedup in comparison with the native MLX Consideration kernel, and as much as a 20x prefill speedup over the neighborhood mlx-lm implementation on the Mamba SSM kernel; we report the numbers, and the way a lot of the achieve comes from the interpretation layer, within the sections under. Though we concentrate on MLX kernels for Apple Silicon, the tactic just isn’t particular to MLX and applies to any ecosystem the place CUDA experience is transferable.<\/p>\n

Why MLX?<\/h2>\n
Apple\u2019s MLX framework has seen outstanding adoption since late 2023. With Apple Silicon in a whole bunch of thousands and thousands of MacBooks and Mac Studios, MLX permits native AI inference with out cloud prices. The unified reminiscence structure makes it particularly enticing for mid-sized fashions (7B\u201370B parameters on M sequence chips).<\/p>\n
But beneath this momentum lies a big hole: many performance-critical kernels that the NVIDIA ecosystem takes as a right: paged consideration, optimized SSM scan kernels, fused MoE routing are both absent or naive with out hardware-specific tuning. MLX runs fashions accurately however usually leaves important efficiency on the desk.<\/p>\n
This hole is what motivates the remainder of this put up.<\/p>\n
What’s Ok-Search?<\/h2>\n
Ok-Search is an evolutionary kernel optimization framework initially developed by our first creator Shiyi Cao at UC Berkeley Sky Lab. Given a naive kernel and a {hardware} specification, it runs an iterative optimization loop: an LLM causes about which optimizations to attempt subsequent, a code-writing mannequin generates candidate kernels, and people candidates are compiled and benchmarked on actual {hardware}.<\/p>\n
Measurements feed again into the search, which retains refining, pursuing promising instructions and dropping useless ends till efficiency converges.<\/p>\n
\n $\"Pseudocode$ <\/p>\n
Algorithm 1: Ok-Search through co-evolving world fashions.<\/strong> The search alternates between choosing probably the most promising motion, instantiating and evaluating code till enchancment stagnates, and evolving the world mannequin by means of insert, replace, and prune operations. Tailored from Cao et al. (2026)<\/a>.<\/p>\n
Search is grounded by a Spec: a domain-specific doc encoding {hardware} guidelines, optimization patterns, and mathematical constraints which retains generated code from hallucinating invalid primitives and ensures candidates will really compile and run effectively.<\/p>\n
In our runs, a single mannequin (Gemini 3.5 Professional Preview) performs each roles: it maintains the reasoning state and writes the kernels. The reasoning half is prompted as a \u201cGPU kernel efficiency engineer\u201d and requested to work by means of a set evaluation earlier than proposing something: classify the kernel (discount, scan, consideration\/softmax, \u2026), rewrite the reference computation in canonical kind, map out knowledge format and entry patterns, and hypothesize the seemingly bottleneck (bandwidth, latency, compute, or synchronization) in every runtime regime. Solely then does it emit candidate optimizations, every as a single change implementable in a single iteration.<\/p>\n
We name the persistent reasoning state a world mannequin<\/em>. Moderately than a flat record of issues to attempt, it’s a determination (prefix) tree: every root\u2192leaf path composes a full optimization plan, and sibling branches are competing options. Each node is scored \u2014 an overall_rating<\/code> in [0, 10], a confidence<\/code> in [0, 1], and per-node impacts<\/code> on reminiscence bandwidth, register stress, and compute\/{hardware} match \u2014 so the search can rank partial plans and increase probably the most promising ones. The tree persists and grows throughout rounds: refining an concept provides a baby node moderately than overwriting its guardian, and if the perfect rating fails to enhance for a couple of rounds (a stagnation window) the search backs off to discover an alternate department. A single node, because it seems mid-run on the eye kernel, seems like this:<\/p>\n
`\n`\n\n{<\/span>\n <\/span>\"motion\"<\/span>:<\/span> <\/span>\"Exchange the threadgroup-memory softmax discount\n with a register-only discount: every SIMD group\n owns 8 question rows and reduces throughout lanes with\n simd_shuffle_xor, eradicating a threadgroup_barrier.\"<\/span>,<\/span>\n <\/span>\"difficulty_1_to_5\"<\/span>:<\/span> <\/span>4<\/span>,<\/span>\n <\/span>\"impacts\"<\/span>:<\/span> <\/span>{<\/span>\n <\/span>\"memory_bandwidth\"<\/span>:<\/span> <\/span>8<\/span>,<\/span>\n <\/span>\"register_pressure\"<\/span>:<\/span> <\/span>4<\/span>,<\/span> <\/span>\/\/<\/span> <\/span>threat:<\/span> <\/span>spill<\/span> <\/span>if<\/span> <\/span>Br<\/span> <\/span>><\/span> <\/span>8<\/span>\n <\/span>\"compute_hw_fit\"<\/span>:<\/span> <\/span>9<\/span> <\/span>\/\/<\/span> <\/span>SIMD<\/span> <\/span>width<\/span> <\/span>32<\/span>;<\/span> <\/span>preserve<\/span> <\/span>tile<\/span> <\/span>8<\/span>x<\/span>8<\/span>\n <\/span>},<\/span>\n <\/span>\"overall_rating_0_to_10\"<\/span>:<\/span> <\/span>8<\/span>,<\/span>\n <\/span>\"confidence_0_to_1\"<\/span>:<\/span> <\/span>0.7<\/span>\n<\/span>}<\/span>\n<\/span><\/code><\/pre>\n<\/div><\/div>\n<\/div>\nItemizing 1: Instance Ok-Search world-model node.<\/strong> Every candidate optimization data a concrete motion, estimated {hardware} impacts, an general precedence ranking, and the mannequin’s confidence.<\/p>\n \n<\/p>\n Determine 2: Overview of Ok-Search.<\/strong> The framework operates on a Search State<\/strong> $S_t$ structured as a search tree. The tree consists of Closed nodes (blue, visited states with connected program like $x_{12}$) and a Frontier of Open nodes (orange, pending hypotheses like $u_{13}$). The workflow iterates by means of three phases: (1) Motion Choice<\/strong>, the place probably the most promising motion node is retrieved from the frontier based mostly on world mannequin estimated precedence rating $V$; (2) Native Refinement<\/strong>, the place a stochastic coverage $pi_{mathrm{code}}$ samples concrete implementations till stagnation; and (3) World Mannequin Replace<\/strong>, the place the LLM causes over the trajectory to replace the search tree through Insert<\/em> (including new actions), Replace<\/em> (adjusting $V$, e.g., $u_{11}$ dropping from 0.9 to 0.6), and Prune<\/em> (eradicating much less promising nodes like $u_{10}$).<\/p>\n The unique Ok-Search paper evaluated this search technique on CUDA kernels from FlashInfer. Throughout GQA decode, MLA decode, MLA prefill, and MoE, Ok-Search improved extra persistently than OpenEvolve and ShinkaEvolve over the identical 120-iteration finances. These outcomes set up the search framework we construct on right here; the rest of this put up asks whether or not its optimization data can switch past CUDA.<\/p>\n \n<\/p>\nDetermine 3: Essential outcomes from the unique Ok-Search paper.<\/strong> Throughout three runs, Ok-Search achieves stronger best-so-far search scores, per-workload kernel efficiency, and speedup distributions than OpenEvolve and ShinkaEvolve on 4 FlashInfer CUDA kernels. Reproduced precisely from Cao et al. (2026)<\/a>.<\/p>\n Constructing an MLX backend<\/h2>\nTo convey Ok-Search to Apple Silicon, we first constructed a local MLX backend. We carried out a full MLX-specific process adapter for Ok-Search, together with:<\/p>\n \nAn MLX process backend in k_search\/duties\/<\/code> dealing with kernel compilation and execution on Apple Silicon through MLX\u2019s Metallic\/C++ APIs.<\/li>\n Up to date kernel generator prompts for writing and modifying Metallic\/MLX kernels.<\/li>\n MLX-specific benchmarking integration utilizing mlx.core<\/code> measurement utilities.<\/li>\n<\/ul>\nTranslating CUDA experience to MLX<\/h2>\nNonetheless, the extra attention-grabbing problem was not merely working Ok-Search on MLX. The important thing perception is that knowledgeable CUDA kernels encode a long time of optimization data that’s transferable to Apple GPU should you can bridge the conceptual hole. Merely handing an LLM a CUDA kernel and asking it to port it isn’t sufficient: with out deep {hardware} context, it produces code that’s syntactically legitimate however architecturally unsuitable (unsuitable tile sizes, invalid primitives, mismatched reminiscence assumptions).<\/p>\n Our translation layer consists of:<\/p>\n \nIdea mapping tables:<\/strong> A structured glossary of CUDA primitives and their MLX\/Metallic equivalents with arduous constraints. For instance:\n\nshared<\/code> maps to Metallic threadgroup<\/code> reminiscence however with a tough 32 KB restrict (vs. NVIDIA\u2019s 48 KB)<\/li>\n warp_reduce<\/code> maps to MMA (most well-liked)<\/li>\n __syncthreads()<\/code> turns into threadgroup_barrier(mem_flags::mem_tg)<\/code><\/li>\n H100\u2019s ~3.35 TB\/s HBM3 maps to M3 Max\u2019s ~400 GB\/s unified DRAM a bandwidth distinction that reshapes which optimizations are price pursuing.<\/li>\n<\/ul>\n<\/li>\n MLX-specific hints and patterns:<\/strong> Concrete code-level patterns for operations with no direct CUDA equal, reminiscent of register-based row reductions utilizing simd_shuffle_xor<\/code> in an 8\u00d78 MMA tile format, or the \u201cexp2 trick\u201d (changing $exp(x)$ with $exp_2(x log_2 e)$) for sooner softmax on Apple\u2019s quick $exp_2$ {hardware} instruction.<\/li>\n Reusable assertions:<\/strong> Professional kernel behaviors reframed as properties the evolutionary search should protect, moderately than code to repeat.<\/li>\n<\/ul>\nMatching knowledgeable kernel efficiency: the Consideration kernel<\/h2>\nWe consider three configurations of an MLX consideration kernel for Apple Silicon: (1) a naive baseline, (2) pure evolution with no further supplied context, and (3) a full context translation layer, which provides the optimizer with architecture-specific implementation data extracted from high-performance kernels (e.g., FlashAttention-2), letting the evolutionary search cause about implementation methods moderately than ranging from a naive kernel. Collectively, these three configurations allow us to isolate the precise affect of the interpretation layer.<\/p>\n \n<\/p>\n Determine 4:<\/strong> Efficiency scaling of the Consideration Kernel by means of stacked optimizations. The “Full Context” configuration efficiently discovers and implements superior methods like double buffering and loop unrolling, reaching near-expert efficiency.<\/p>\n The leap from 0.26\u00d7 to 0.97\u00d7 the velocity of Apple\u2019s state-of-the-art consideration kernel \u2014 illustrates how a lot the interpretation layer issues. With full context, the advanced kernel independently discovers the important thing optimizations in FlashAttention 2: threadgroup reminiscence tiling, on-line softmax, Ok-transposition for reminiscence entry, and the exp2 trick. The final of those replaces each softmax exponential with a base-2 exponential,<\/p>\n [e^x = 2^{x log_2 e},]<\/p>\n which is actual and lets the kernel use Apple\u2019s quick quick::exp2()<\/code> {hardware} instruction instantly as an alternative of paying for a base conversion at runtime.<\/p>\n A 20\u00d7 sooner prefill: the Mamba SSM kernel<\/h2>\nTo judge whether or not Ok-Search generalizes past consideration kernels, we utilized it to the state-space mannequin (SSM) kernel utilized by Mamba. In contrast to consideration, the computational bottleneck is a recurrent state replace moderately than a softmax, offering a considerably totally different optimization problem. We evaluate the advanced implementation towards the neighborhood MLX implementation (mlx-lm) and the PyTorch reference implementation (mamba.py) on an M1 Max.<\/p>\n Evaluated on mamba-370m f16, M1 Max 64GB:<\/p>\n \n\n\n\n\n\n\n\n\n\nMetric<\/th>\n mlx-mamba (ours)<\/th>\n mlx-lm (neighborhood)<\/th>\n mamba.py<\/th>\n<\/tr>\n<\/thead>\n Decode<\/td>\n 152 tok\/s<\/td>\n 116 tok\/s<\/td>\n 40 tok\/s<\/td>\n<\/tr>\n Prefill L=512<\/td>\n 5,751 tok\/s<\/td>\n 329 tok\/s<\/td>\n 1,089 tok\/s<\/td>\n<\/tr>\n Prefill L=1024<\/td>\n 6,010 tok\/s<\/td>\n 327 tok\/s<\/td>\n 1,127 tok\/s<\/td>\n<\/tr>\n Prefill L=2048<\/td>\n 6,612 tok\/s<\/td>\n 326 tok\/s<\/td>\n 1,092 tok\/s<\/td>\n<\/tr>\n Prefill L=4096<\/td>\n 6,743 tok\/s<\/td>\n 339 tok\/s<\/td>\n 1,042 tok\/s<\/td>\n<\/tr>\n<\/tbody>\n<\/table>\n<\/div>\nDesk 1:<\/strong> Prefill and decode throughput on mamba-370m (f16, M1 Max 64GB). mlx-mamba (ours) reaches ~20\u00d7 larger prefill throughput than the neighborhood mlx-lm baseline, whereas decode stays comparable.<\/p>\n The ~20\u00d7 prefill speedup over mlx-lm comes down to 1 distinction: mlx-lm doesn’t implement a parallel scan for the SSM. The state recurrence<\/p>\n [h_t = bar{a}_t h_{t-1} + bar{b}_t]<\/p>\n seems inherently sequential, however every step will be written as a pair $(bar{a}_t, bar{b}_t)$ beneath the associative mix<\/p>\n [(a_2, b_2) circ (a_1, b_1) = left(a_2 a_1, a_2 b_1 + b_2right),]<\/p>\n which reproduces the recurrence precisely. As a result of the operator is associative, the entire sequence will be evaluated with a parallel (prefix) scan in $O(log N)$ dependent steps as an alternative of $O(N)$. mlx-lm skips this and processes tokens one after the other, leaving most of Apple Silicon\u2019s compute idle; our advanced Metallic kernel applies the scan and makes a lot fuller use of GPU throughput. The achieve exhibits up in prefill, the place the total sequence is accessible to scan in parallel, and never in single-token decode, the place there is just one new token per step and no scan to parallelize \u2014 which is why the decode row is roughly flat whereas prefill is ~20\u00d7.<\/p>\n mamba.py is sluggish on each prefill and decode as a result of it’s a PyTorch reference implementation that falls again to CPU or MPS on Apple Silicon, forgoing the hardware-specific optimizations that MLX\u2019s Metallic backend makes doable.<\/p>\n What\u2019s subsequent?<\/h2>\nOn the 2 kernels we studied, AI-driven evolutionary kernel search grounded in structured cross-platform translation data reached near-expert efficiency on Apple Silicon with no group of GPU specialists ranging from scratch. We don’t but know the way far this generalizes, however the result’s encouraging.<\/p>\n For us the primary takeaway is that the bottleneck was not the LLM\u2019s capacity to write down Metallic code, however the high quality of the context and constraints we gave it. Our CUDA translation layer converts present NVIDIA kernel experience into actionable steering for Apple Silicon, and lets Ok-Search\u2019s evolutionary search do the remaining.<\/p>\n We’re actively extending this work in a number of instructions: supporting new architectures, with present efforts targeted on creating new kernels for the IBM Spyre AIU and broader {hardware} targets; including extra kernels reminiscent of paged consideration and fused MoE routing; and bettering integration with the Ok-Search evolution loop to make translation context much more computerized.<\/p>\n Acknowledgements<\/h2>\nThis work was carried out by IBM Analysis and builds on Ok-Search from the UC Berkeley Sky Lab (Cao et al., 2026<\/a>). We welcome collaboration and suggestions from the MLX and broader AI programs communities. If you’re engaged on kernel optimization for non-CUDA {hardware}, we’d love to listen to from you.<\/p>\n\n Quotation<\/h2>\n\n\n@article<\/span>{<\/span>cao2026k<\/span>,<\/span>\n title<\/span>=<\/span>{Ok-Search: LLM Kernel Technology through Co-Evolving Intrinsic World Mannequin}<\/span>,<\/span>\n creator<\/span>=<\/span>{Cao, Shiyi and Mao, Ziming and Gonzalez, Joseph E and Stoica, Ion}<\/span>,<\/span>\n journal<\/span>=<\/span>{arXiv preprint arXiv:2602.19128}<\/span>,<\/span>\n 12 months<\/span>=<\/span>{2026}<\/span>\n}<\/span>\n<\/code><\/pre>\n<\/div>\n<\/div>\n\nAppendix: Attempt it your self<\/h2>\nThe MLX backend is constructed on prime of the open-source Ok-Search repo, so the outcomes right here will be reproduced instantly. The steps are:<\/p>\n 1. Clone and set up<\/p>\n \n\ngit clone https:\/\/github.com\/caoshiyi\/Ok-Search.git\ncd <\/span>Ok-Search\n\nuv pip set up <\/span>openai wandb\nuv pip set up <\/span>git+https:\/\/github.com\/caoshiyi\/flashinfer-bench-ksearch.git\n<\/code><\/pre>\n<\/div><\/div>\n2. Set your credentials<\/p>\n Open the related script beneath scripts\/ and set three variables on the prime:<\/p>\n \n\nKSEARCH_ROOT<\/span>=<\/span>\/path\/to\/Ok-Search\nAPI_KEY<\/span>=<\/span>your-llm-api-key\n<\/code><\/pre>\n<\/div><\/div>\n3. Run kernel search<\/p>\n \n\n# Optimize Flash Consideration on Apple Silicon (world-model mode)<\/span>\nbash scripts\/mac_flash_attention_wm.sh\n\n# Or a Mamba SSM kernel, e.g. the selective scan<\/span>\nbash scripts\/mamba_selective_scan_fwd_wm.sh\n<\/code><\/pre>\n<\/div><\/div>\nFull CLI reference and documentation are within the README.<\/p>\n<\/div>\n<\/div>\n\n","protected":false},"excerpt":{"rendered":" Determine 1: CUDA-to-MLX optimization translation map. CUDA optimization data will be translated into architecture-native MLX methods moderately than copied instruction-for-instruction. We face a brand new epoch in computing. {Hardware} is altering quickly \u2014 not simply sooner GPUs, however a rising vary of chips from totally different distributors, every with its personal structure and sometimes tailor-made […]<\/p>\n","protected":false},"author":2,"featured_media":17284,"comment_status":"open","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"footnotes":""},"categories":[55],"tags":[1395,311,310,110,414,7229,9992,312,4943,9991,193,3391],"class_list":["post-17282","post","type-post","status-publish","format-standard","has-post-thumbnail","hentry","category-machine-learning","tag-apple","tag-artificial","tag-berkeley","tag-blog","tag-brings","tag-decades","tag-expertise","tag-intelligence","tag-kernel","tag-ksearch","tag-research","tag-silicon"],"_links":{"self":[{"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/posts\/17282","targetHints":{"allow":["GET"]}}],"collection":[{"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/posts"}],"about":[{"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/types\/post"}],"author":[{"embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/users\/2"}],"replies":[{"embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=%2Fwp%2Fv2%2Fcomments&post=17282"}],"version-history":[{"count":1,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/posts\/17282\/revisions"}],"predecessor-version":[{"id":17283,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/posts\/17282\/revisions\/17283"}],"wp:featuredmedia":[{"embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=\/wp\/v2\/media\/17284"}],"wp:attachment":[{"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=%2Fwp%2Fv2%2Fmedia&parent=17282"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=%2Fwp%2Fv2%2Fcategories&post=17282"},{"taxonomy":"post_tag","embeddable":true,"href":"https:\/\/techtrendfeed.com\/index.php?rest_route=%2Fwp%2Fv2%2Ftags&post=17282"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}