Google ML Drift公開：オンデバイスAI推論を変えるGPUエンジンの全貌

オンデバイスAIにおけるGPUアクセラレーションの課題

データセンターでのAI推論は、均質で予測可能なアクセラレータクラスタ上で動作します。しかしエッジデバイスは事情が異なります。多種多様なGPUアーキテクチャ、ドライババージョン、そしてOpenGL ES / OpenCL / Metal / WebGPUといった低レベルAPIが混在しており、開発者はアプリがどのハードウェア上で動作するかを事前に知ることができません。

従来のTensorFlow Lite GPU delegateは第一世代のGPUアクセラレーションを実現しましたが、現在のワークロードはそれを超えています。リアルタイムのコンピュータビジョンからパラメータ数の多い生成AIまで、コンシューマー向けシリコンを限界まで押し上げており、レガシーランタイムでは対応しきれない計算・メモリのボトルネックが生じています。

Google AI Edgeチームがこの課題に対して公開したのがML Driftです。Apache 2.0ライセンスの下でオープンソース化された、高性能クロスプラットフォームのオンデバイスGPUコンピュートエンジンであり、LiteRTの中核GPUアクセラレーションエンジンとして機能しつつ、スタンドアロンライブラリとしても利用できます。

💡 ポイント: ML Driftは「ハードウェア抽象化」と「性能」を両立させる統合ランタイムです。従来のTFLite GPU delegateは新機能の更新が停止され、ML Driftへの移行が推奨されています。

ML Driftの4つの主要アーキテクチャ変更

従来のTFLite GPU delegateに対するML Driftの構造的アップグレードは以下の通りです。

1. テンソル仮想化による統合シェーダー

従来はOpenGL、OpenCL、Metalのバックエンドごとに最適化シェーダーを個別に管理する必要がありました。ML Driftはテンソルの論理表現とGPU上の物理割り当てを分離するテンソル仮想化を導入しました。コンパイラの初期化フェーズで動的シェーダーテンプレートが座標を解決・変換するため、バックエンド固有のシェーダーコードベースを維持する必要がなくなります。

// ML Drift流のテンソル仮想化の概念例（擬似コード）
// 論理テンソル宣言 - バックエンド非依存
auto logical_tensor = env.CreateTensor(
    /*shape=*/{1, 3, 224, 224},
    /*dtype=*/DataType::FLOAT32
);

// コンパイラがバックエンドに応じて物理割り当てを自動決定
// OpenGL: texture, OpenCL: buffer, Metal: MTLBuffer など
auto shader = env.CompileShader(
    "conv2d",
    /*inputs=*/{logical_tensor},
    /*dynamic_layout=*/true  // 実行時にレイアウトを解決
);

// バックエンド分岐なしで実行
env.Execute(shader);

2. 拡張可能なカスタムOpフレームワーク

ML Driftは低レベルシェーディング言語へ直接アクセスできる登録APIを提供します。特にエージェント型SKILL.mdガイドが含まれており、コーディングエージェントが数分でカスタムシェーダーを作成・登録・検証できます。独自アーキテクチャを実行グラフに直接統合できる細やかな制御が可能です。

3. 5Dテンソル対応

TFLite GPU delegateは構造的に4Dテンソルにハードコードされており、5Dテンソルが必要な複雑なモデルではレイアウトハックを強いられていました。ML Driftは5Dテンソルを標準サポートし、3D畳み込みネットワーク（ボリューム/空間AI）や時空間モデル（YOLO 11n、MobileViT v2、Swin Transformer v2など）をエッジGPU上で直接実行できます。

4. エッジLLM向けステージ認識最適化

自己回帰LLM推論は2つの局面に分かれます。KVキャッシュのプリフィル段階は計算律速、デコード段階はメモリ帯域律速です。ML Driftは実行フェーズに応じてカーネルとレイアウト構成を動的に切り替えます。デコード段階では畳み込み整列KVキャッシュレイアウトを使用し、カーネル内アクティベーション量子化を積極的に適用して冗長なメモリ往復を排除します。

🔗 関連インサイト: GPU推論のボトルネックに関するより深い議論はNVIDIA Blackwell Ultra、AI推論の隠れたボトルネックSoftmaxを解決するをご参照ください。

実性能と導入事例

ML DriftはすでにGoogleエコシステム全体で数百万台のデバイス上で本番稼働しています。

またGemmaモデルのベンチマークでは、ML Driftは他フレームワーク比で最大12%低いメモリオーバーヘッドを達成しました。エッジ環境ではスループット以上にメモリフットプリントが重要となるため、この数値は実用的な意味が大きいです。

プラットフォーム拡張：デスクトッププレビュー

興味深い点として、ML DriftのWebGPUバックエンドは元々ブラウザ向けに設計されましたが、ChromiumのWebGPU実装であるDawnにより、同一コードベースをブラウザ外でもネイティブコンパイルできるようになりました。これによりWindowsとLinuxでもDirectX/Vulkanの断片化を回避しつつML Driftが動作します。macOSのネイティブMetalバックエンドと相互補完的な関係です。

⚠️ 注意事項と限界

• **レガシーTFLite GPU delegateは新機能の更新が停止されます。**既存モデルはML Driftへの移行が推奨され、後方互換性は保証されています。
• Androidでアンバンドルランタイムを使用する場合、ML Driftアクセラレーションは現在スタンドアロンのLiteRTパッケージで提供され、Google Play Services内のLiteRTには近日追加予定です。
• デスクトップ対応はまだ「プレビュー」段階であり、本番適用前に十分な検証が必要です。

まとめ：ML Driftが意味するもの

ML Driftは単なるGPU delegateのアップグレードではなく、オンデバイスAI推論の共通基盤を目指す統合ランタイムです。テンソル仮想化でバックエンドの断片化を減らし、5Dテンソルとステージ認識LLM最適化で生成AIまでカバーし、WebGPUベースのデスクトップ拡張で開発ワークフローまで包含します。

実務適用のアドバイス

1. 既存のTFLite GPU delegateユーザー: ML Drift GPU acceleratorへの移行を計画してください。後方互換が保証されているため既存モデルはそのまま動作し、即座に性能向上が得られます。
2. カスタムシェーダーが必要な場合: SKILL.mdガイドを活用し、コーディングエージェントでカスタムOpを素早くプロトタイピングしましょう。
3. エッジLLMを検討中の場合: ステージ認識最適化（prefill vs decode）が実際のワークロードでどの程度の差を生むか、ベンチマークで確認することをお勧めします。

次のステップ学習方向

• LiteRT公式ドキュメントでML Drift GPU acceleratorの移行ガイドを確認
• WebGPU/Dawnベースのネイティブコンパイルフローを学習（デスクトップ推論に関心がある場合）
• 量子化およびKVキャッシュレイアウト最適化に関するLLM推論技法を併せて学習

🔗 併せて読みたい記事: クラウド環境での大規模不確実性分析ワークフローに関心がある方は"SLB ARM + Azure 貯留層不確実性分析、クラウドでどう解決するか？"をご参照ください。

根拠資料: Google Developers Blog - ML Drift