何が起きたか

Berkeley Sky LabのShiyi Cao氏らが開発したK-Search（Cao et al., 2026）は、素朴なカーネルとハードウェア仕様を与えると、LLMが最適化方針を推論し、コード生成モデルが候補を書き、実機でコンパイル・ベンチマークするループを回す仕組みです。今回はここにAppleのMLX向けバックエンド（Metal/C++ API経由のコンパイル・実行、mlx.coreによる計測統合、Metal/MLX用に書き換えた生成プロンプト）と、CUDAカーネルを知識ベースとして再利用する「構造化CUDA→MLX変換層」が加わりました。著者らはこの手法がMLX固有ではなく、CUDAの知見が転用可能なあらゆるエコシステムに適用できると述べています。

なぜ重要か

GPUカーネルはGPU内部で動く低レベルのプログラムで、効率的に書けるようになるまで年単位の熟練を要します。しかもベンダーが変われば、同じ最適化を一から発見し直すのが常でした。CUDA側にはattentionや状態空間モデルなど数十年ぶんの手書き実装が蓄積している一方、Apple SiliconやカスタムAIアクセラレータにはその厚みがありません。実際MLXでは、NVIDIAエコシステムでは当たり前のpaged attention、最適化されたSSMスキャン、fused MoEルーティングが欠落しているか素朴な実装のままです。「MLXはモデルを正しく動かすが、性能をかなり取りこぼしている」という状態です。

ここで重要なのは、単にCUDAカーネルをLLMに渡して「移植して」と頼んでも失敗するという点です。構文的には妥当でも、タイルサイズが不適切で、存在しないプリミティブを呼び、メモリ前提が噛み合わないコードが出てきます。

変換層が効かせているもの

変換層は3層構造です。第1にCUDAプリミティブとMLX/Metal相当物を制約付きで対応させる概念マッピング表。__shared__はMetalのthreadgroupメモリに対応しますが上限は32KBで、NVIDIAの48KBより狭い。__syncthreads()はthreadgroup_barrier(mem_flags::mem_tg)に、warp_reduceはMMAに、H100のHBM3（約3.35TB/s）はM3 Maxのユニファイドメモリ帯域（約400GB/s）に対応づけられます。第2にMLX固有のパターン——8x8のMMAタイル上でsimd_shuffle_xorを使うレジスタ内行リダクション、そしてexp(x)をexp2(x·log2 e)に置き換える「exp2トリック」。この恒等式は厳密で、Appleの高速なfast::exp2()命令を変換コストなしに使えます。第3が、熟練者のカーネル挙動を「探索が壊してはならない性質」として再定義したアサーション群です。

探索の土台にはSpec（ハードウェア規則・最適化パターン・数学的制約を書いた文書）があり、無効なプリミティブの幻覚を抑えます。推論状態は「world model」と呼ばれる決定木で保持され、根から葉までの経路が一つの最適化プランを、兄弟ノードが競合案を表します。各ノードはoverall_rating（0〜10）、confidence（0〜1）、メモリ帯域・レジスタ圧・計算/ハード適合度への影響で採点され、改善が数ラウンド止まると別枝に退避します。実行中の一例では、threadgroupメモリを使うsoftmaxリダクションを、各SIMDグループが8クエリ行を持ちレーン間でsimd_shuffle_xorする レジスタ限定リダクションに置き換え、barrierを一つ消す案が提示されています。

結果

attentionでは、素朴なベースライン／文脈なしの純粋進化／FlashAttention-2などから抽出した知識を与えるフル文脈、の3条件を比較。フル文脈の探索はthreadgroupメモリのタイリング、オンラインsoftmax、メモリアクセス向けのK転置、exp2トリックを自力で再発見し、さらにダブルバッファリングとループアンローリングにも到達しました。性能はApple最先端カーネル比0.26倍から0.97倍へ。ボトルネックがsoftmaxではなく再帰的な状態更新にあるMamba向けSSMカーネルでは、M1 Max 64GB・mamba-370m f16の条件で、デコードが116トークン/秒（mlx-lmコミュニティ実装）に対し152トークン/秒、prefillは最大20倍という結果になりました。

💼 事業会社視点：これは自社にどう効くか

経営側の含意は「ローカル推論のコスト計算をやり直す時期が来た」ことです。Apple Siliconは数億台のMacBook/Mac Studioに載り、ユニファイドメモリのおかげで7B〜70B級を手元で動かせますが、これまでMLXの実効性能はカーネル最適化の薄さに縛られていました。prefillが最大20倍変わるなら、社内文書要約や議事録処理のように長文入力が支配的な用途では、クラウドAPI課金と自社Mac群での実行の損益分岐が明確に動きます。まず自社の推論ワークロードをprefill偏重かデコード偏重かで分類し、前者から実測し直すべきです。

SaaS事業者にとっては、オンプレ/端末内推論を求める金融・医療・自治体案件への回答が変わります。「データを外に出さない」構成の性能上の言い訳が減るからです。受託開発では、Metalカーネルを書ける人材がいなくても、CUDA資産と探索フレームワークで穴を埋められる可能性が出てきた点が効きます——ただし0.97倍という数字は「専用実装に匹敵」であって超えたわけではなく、探索には120イテレーション級の実機ベンチ予算がかかります。ECのようにレイテンシ要求が厳しくリク��ストが分散する領域では、まだクラウド推論が現実解でしょう。国内ハードウェアやAIアクセラレータに関わる企業にとっては、著者らが「MLX固有ではない」と明言している点が最も重い——自社チップのソフトウェア資産不足を、人海戦術以外で埋める筋道が示されたということです。

関連リンク