Flash-dLLM、I/Oを意識したKVキャッシュで拡散型LLMを高速化

拡散型大規模言語モデル（dLLM）は、自己回帰型モデルのようにトークンを一つずつ確定するのではなく、複数の位置を並列に更新できる点が特徴です。しかし、並列性がそのまま実運用の高速化につながるとは限りません。推論中にKVキャッシュを何度も読み書きすると、演算性能よりもGPUメモリーI/Oが支配的なボトルネックになるためです。

Flash-dLLMは、この問題をKVキャッシュと並列デコードを一体として設計することで解決しようとします。論文では、従来の加速手法がキャッシュ再利用と並列検証を別々に扱うことが多く、両者を組み合わせた際のデータ移動コストを見落としやすいと説明しています。提案手法は追加学習を必要とせず、主に二つの構成要素から成ります。

主なポイント

• **Flash-Cache：**QKV投影、RoPE処理、キャッシュ書き込みを一つのTritonカーネルに融合します。中間データの移動を減らし、バッチ内でクエリ長が異なる場合にはブロック単位でスケジューリングします。
• **選択的なキャッシュ更新：**毎回キャッシュ全体を書き換えるのではなく、新しくデコードされたトークンと、注目度の高い既デコードトークンの固定集合を中心に更新します。これにより、キャッシュの有用性を保ちながら書き込み量を抑えます。
• **Flash-Verify：**補助モデルを用いず、dLLM自身をドラフターと検証器の両方として利用します。二つの視点を持つ因果アテンションマスクにより、候補生成と検証を同じ仕組みで扱い、素材では1ステップあたりの受理トークン数が概ね2倍になるとされています。

LLaDA-1.5を使った結果として、素材には毎秒148～211トークンという速度が示されています。キャッシュなしの貪欲デコードに対しては22.3～148.2倍、GSM8KとHumanEvalではElastic-Cacheに対してそれぞれ5.1倍と11.0倍高速だったと報告されています。また、Fast-dLLMよりGPUメモリー使用量を約48％削減し、バッチサイズ32まで拡張できたとされています。

ただし、これらは提示された実験条件における結果です。モデルの規模、系列長、GPU、バッチ構成、カーネル実装によって効果は変わる可能性があります。それでも本研究は、dLLMのサービングではアルゴリズム上の並列性だけでなく、キャッシュの読み書きを含むメモリーシステムとの協調設計が重要だと示しています。

Flash-dLLMの意義は、キャッシュ最適化と自己検証型デコードを別々の工夫としてではなく、同じ推論経路に統合した点にあります。長い系列や大きなバッチを扱う場合、不要なデータ移動を削減する設計は、単純に演算並列度を高めるより有効な場面があります。今後は、より多様なモデルとハードウェアで報告された効果が再現されるかを確認する必要があります。

出典：Hugging Face Daily Papers