Apple Neural Engine の Private API を叩いて LLM 推論を高速化しようとした話

概要

Apple SiliconのANE Private APIを直接制御してLLM推論の高速化を検証し、小規模モデルでの有効性と大規模モデルにおけるMetal GPUの優位性を明らかにした詳細な技術検証。

詳細内容

Apple Siliconに搭載された行列演算特化のアクセラレータ「ANE (Apple Neural Engine)」を、非公開APIを通じて直接制御し、LLM推論の高速化を目指した挑戦の記録です。単一演算ではデータ転送のオーバーヘッドによりMetal GPUに劣るものの、複数の演算を統合する「fused kernel」化により、隠れ層次元(D)が1024以下のモデルで最大12倍の高速化を達成しました。また、D=4096などの2のべき乗サイズで発生する「SRAM bank conflict」というハードウェア特有のボトルネックを特定し、ゼロパディングによる解決策を提示しています。最終的に、MLXのlazy evaluationによる最適化が非常に強力であるため、4B以上の大規模モデルではANEオフロードに伴うパイプライン分断コストが利得を上回るという構造的限界を結論付けています。ANEは巨大LLMよりも、音声認識や画像解析といった小〜中規模かつ固定形状の入力をGPUを占有せず処理する用途に最適であるとしています。

元記事を読む他のサマリーを見る