概要
Google CloudとUCSDの研究チームが、TPU上で拡散型投機的デコード技術「DFlash」を実装し、LLM推論を従来の逐次的な手法より大幅に高速化することに成功した。
詳細内容
Google CloudとUCSDの研究チームは、Google TPU v5pにおいて拡散型投機的デコード（Diffusion-style speculative decoding）技術「DFlash」を実装し、LLM推論の劇的な高速化を達成しました。従来の投機的デコードは、ドラフトモデルが逐次的にトークンを生成するため O(K) のステップを要していましたが、DFlashは「ブロック拡散」を用いることで、単一のフォワードパス（O(1)）で候補トークンのブロック全体を生成します。
主な成果と技術的ハイライト：
- **パフォーマンス**: TPU v5pにおいて平均3.13倍（複雑な数学タスクでは最大約6倍）の速度向上。既存手法のEAGLE-3を大きく上回る実効性能を実現。
- **K-Flat検証**: 高性能ハードウェアでは、少数のトークンでも大量のトークンでも検証コストがほぼ変わらないという「K-Flat」特性を解明。推論のボトルネックが検証コストではなく、ドラフトの「質」にあることを示しました。
- **エコシステムへの貢献**: この実装はvLLMのTPU推論リポジトリに統合され、オープンソースとして提供されます。
- **技術的工夫**: Pallasカーネルを用いたページド・アテンションとDFlashの非因果的ブロック拡散を両立させる「デュアルキャッシュ」アーキテクチャを採用し、TPUの性能を最大限に引き出しています。