クアルコムのAI処理ユニット「Hexagon」向けに、大規模言語モデルの推論を高速化するFlash Attentionの大規模な再実装が行われた。タスクの統合やメモリ管理の最適化、数値精度の改善により、エッジデバイスでのAI処理能力が底上げされる。

行列演算のタスク統合が生む効率性

今回のアップデートの中核には、行列積和演算を担当するスレッドへのタスク集約がある。従来は独立して処理されていた量子化タスクを主要な行列演算スレッドに統合したことで、タスク間の同期待ちや余分なデータ転送が削減された。これは、工場のラインで別々の工程を一つのセルにまとめ、仕掛かり品の滞留をなくすのに似ている。さらに、行列積に加算処理を融合させることで、メモリへの書き戻しと再読み込みの無駄を省き、演算器の稼働率を高める設計に変更されている。エッジデバイスの限られた計算資源を、より有効に使うための工夫だ。

Flash Attentionの演算精度とキャッシュ効率

Flash Attentionのカーネルでは、マスク処理の高速化とキャッシュヒット率の向上が図られた。ループ順序の変更により、一度ロードしたマスクデータを最大限再利用できるようになり、高速なSRAMへの依存度を高めつつ、低速な外部メモリへのアクセスを抑制している。また、数値の安定性にも踏み込み、ソフトマックス計算の累積値を32ビット浮動小数点で保持するよう変更。内部処理には半精度浮動小数点を活用しながら、最終的な精度低下を防ぐバランスを取っている。これにより、推論結果の品質を落とさずに計算速度を稼ぐことが可能になる。

多様な命令セットとプラットフォームへの対応

このリビジョンでは、演算の共通化とホスト側へのパラメータ計算の移行が進められた。カーネルを個別のヘッダーに分離し、命令セットに依存しない共通部分を統一したことで、コードの保守性が向上している。動作確認対象は、Android端末やWindows on Armパソコン、Linuxプラットフォームに加え、VulkanやROCmといった様々なGPU・AIアクセラレータAPI環境にまで及ぶ。特定のハードウェアに閉じない最適化戦略は、Snapdragonプラットフォームを開発するクアルコムが、OEMやクラウド事業者にシームレスな開発体験を提供する上で欠かせない布石と言える。

エッジAI推論の「待ち」を削減する非同期化

推論全体のレイテンシを左右するのは、しばしば計算そのものよりデータの準備待ちである。この再実装では、行列のドット積計算やマスクデータのDMA転送を従来よりも早期に開始する非同期化が随所に盛り込まれた。例えば、後段の演算に必要なデータのフェッチを前段の処理と並行して走らせることで、処理パイプラインの「隙間」を埋める。1推論あたりの絶対時間が短いエッジデバイスでは、こうした細かな制御の積み重ねが、ユーザー体験における応答性の違いとなって現れる。