AIの推論速度は、サービスのコストと応答性を左右します。これまでspeculative decoding（投機的デコーディング）が高速化の主流でしたが、その「ドラフト生成」部分がボトルネックになっていました。
DFlash（Block Diffusion for Flash Speculative Decoding）は、そのボトルネックを根本から解消する手法です。Google Cloud TPU上での実装で平均3.13倍、数学タスクでは最大約6倍の高速化を記録しました。
この記事でわかること
- 既存speculative decodingが抱える逐次ドラフトの問題
- ブロック拡散によってO(K)をO(1)に変える仕組み
- Google TPU向け実装で必要だった3つの技術的工夫
- EAGLE-3との比較ベンチマーク結果
- vLLMへのオープンソース統合の状況
LLM推論の「逐次ドラフト」という壁
LLMは通常、トークンを1つずつ逐次生成します。speculative decodingはこの問題に対し、軽量な「ドラフトモデル」が複数トークンを先読みし、大きな「ターゲットモデル」が一括検証する仕組みで高速化を実現してきました。
しかし既存手法の多くは、ドラフト生成自体が逐次処理です。K個のトークンを先読みするには、K回の順次フォワードパスが必要になります。ドラフト段階の時間がかさみ、検証段階で節約した時間を食いつぶす構図になっていました。
EAGLE-3はその代表的な手法で、TPU v5p上で1.30倍の高速化を達成しています。ただしデフォルトのK値は2程度が限界で、それ以上に増やすと逐次コストが膨らみます。
DFlashの核心—「ブロック全体を1パスで描く」
https://arxiv.org/abs/2602.06036
DFlashは、UCSDs Z Lab（Hao Zhang教授ら）が開発したブロック拡散ベースのspeculative decoding手法です。paged attentionとprefill/decode分散サービングの共同発明者でもあるHao Zhang氏がリードしています。
鍵となるのはブロック拡散（block diffusion）の考え方です。ターゲットモデルの中間表現（隠れ状態）を条件として取り込みながら、ドラフトモデルがK個のトークンを1回のフォワードパスで一括生成します。逐次的なO(K)処理から、定数時間O(1)への転換です。
K=10でも16でも、ドラフトにかかるコストはほぼ変わりません。これはTPUのMXU（行列乗算ユニット）の特性と相性が良く、TPU上で特に効果が発揮されます。
Google TPU移植で解決した3つの技術課題
UCSDのチームがDFlashをGPU/PyTorchからTPU/JAX環境へ移植する際、以下の3つの課題が生じました。
デュアルキャッシュアーキテクチャ
DFlashのブロック拡散は非因果的（non-causal）な仕組みのため、TPUの高性能pagedアテンション（Pallasカーネル）と根本的に相性が悪いです。そこでチームは2つのキャッシュを分けて運用する設計を採用しました。ターゲットモデルは従来通りpaged KVキャッシュを使い、ドラフトモデルはオンデバイスのJAX静的配列で専用パスを確保します。
コンテキスト管理の最適化
DFlashのドラフトモデルはターゲットモデルの隠れ状態を参照し続けます。この「コンテキストバッファ」の転送効率を高めるため、2の累乗でパディングするストラテジーを実装しました。バッファを正確に追跡することで、重複処理やデータ欠損を防いでいます。
メタデータの同期
DFlashは、コンテキストバッファやKVキャッシュ位置、RoPEオフセットといった状態を複数のイテレーション間で持ち続けます。vLLMのTPUパイプラインでは、検証中のドラフトトークンをメタデータとして渡す仕様が「シーケンス長の膨張」を引き起こし、内部状態がずれていました。提案部分を実際に承認されたトークン数と同期させることで、この問題を解消しています。
K-Flat現象—検証コストはほぼ無料
移植作業中にチームが発見した知見が「K-Flat」です。TPU v5p上では、16トークンの検証と1024トークンの検証がほぼ同じコストになります。モデルの重みをロードする時間がボトルネックになっており、アテンションの計算量増加が実質的に無視できるためです。
これは設計上の意味が大きい発見です。投機ブロックを大きくしても検証コストが増えないなら、「いかに多く先読みするか」より「いかに正確なドラフトを生成するか」に研究の焦点を移すべきだということです。チームの分析では、承認確率（a）を上げることが、ブロックサイズKを増やすより2〜3倍高い効果があります。
ベンチマーク結果—EAGLE-3と直接比較
TPU v5p上でLlama-3.1-8B-Instructをターゲットモデルとし、DFlash（K=10）とEAGLE-3（K=2）を同条件で比較した結果は以下の通りです。
DFlashが達成した主な数値を整理します。エンドツーエンドのサービング速度でDFlashは2.29倍の高速化を達成し、EAGLE-3の1.30倍を大幅に上回りました。コーディングタスク（mbpp）では生成時間が9.81ms/トークンから3.48ms/トークンへ短縮され、2.83倍の改善になっています。
Qwen3-4BをターゲットモデルとするJAX単体ベンチマーク（サービング層を除いた測定）では、全データセット平均で3.13倍のスループット向上を確認。数学推論タスク（math500）では8.02ms/トークンが1.40ms/トークンに改善され、最大約5.7倍のスピードアップを記録しています。
この差が生まれる理由は構造的なものです。EAGLE-3が2トークンを逐次ドラフトする間、DFlashは10トークンを1パスで生成します。TPUの並列計算能力を有効に使える分、スループットの伸びが大きくなります。
vLLMへのオープンソース統合
https://github.com/z-lab/dflash
実装はすでに3本のPRとしてvLLM TPU inferenceリポジトリに提出されています。PR #1868でモデルとプロポーザーのアーキテクチャ、PR #1869でエンドツーエンドパイプライン統合、PR #1870でCIとテストフレームワークがそれぞれ追加されています。チームはさらに、torchaxプロポーザーの追加によってPyTorchサービングパスへの対応も進めています。
今後の展開として、投機キャッシュを活用した「Speculative Speculative Decoding（SSD）」や、TPU RLスタック（TunixとMaxText）を使った広いドラフトブロックへの拡張、拡散ベースのターゲットモデルへの対応なども視野に入っています。
まとめ
DFlashはLLM推論の「逐次ドラフト」というボトルネックを、ブロック拡散による並列化で解消します。Google TPU上での実装で平均3.13倍、数学・コーディングタスクで最大6倍近い高速化を達成し、vLLMのオープンソース実装として公開されています。
K-Flatの発見が示すように、今後の推論最適化の鍵は「より大きな投機ウィンドウ」ではなく「より正確なドラフト品質」にあります。DFlashはその方向性を技術的に具体化した最初の大きな成果と言えます。