AI 推論系統
WIDE、token ごとにモデル幅を動的にプルーニングし、デコードをエンドツーエンドで1.55倍高速化
WIDEでは、各tokenが使用するattention headとFFN channelのグループを自ら選択し、動的プルーニングの粒度をレイヤー単位からニューロンブロック単位へと細分化する。論文では、50%のスパース性においてprefillとデコードをエンドツーエンドでそれぞれ1.68倍、1.55倍高速化したと報告しているが、依然としてカスタムカーネルと追加学習が必要となる。
大規模言語モデルの「幅」は通常、デプロイ前に固定されている。すべてのtokenが同じ数のattention headとFeed-Forward Network（FFN）channelを通過する。7月30日に公開されたWIDEは、token単位で計算量を割り当て、tokenごとにattention-head groupとFFN-channel groupを動的に選択できるようにする。処理が容易なtokenはより多くのブロックをスキップし、難しいtokenにはより完全なモデル幅を維持する。これはレイヤー全体のみをスキップするdynamic depthよりも粒度が細かく、prefillとtoken単位のdecodeの両方に対応する。
学習は2段階で行われる。まず、各tokenのスパースな選択を学習し、その後、動的な実行経路に適応するようモデルを調整する。真の難題はハードウェアにある。散在するmaskではGPU時間を必ずしも短縮できないため、チームはmask reorderingを導入し、選択されたheadとchannelをより連続したブロックに並べ替える。さらに、ハードウェア非依存のblock skippingと、デバイスごとに実装するintra-block skippingを組み合わせた。これにより、理論上のFLOPsを削減するだけでなく、カーネル性能とエンドツーエンドのレイテンシーも改善できる。
論文によると、50%のスパース性かつキャリブレーションデータのみを使用する設定で、WIDEは既存のdynamic depth pruningよりも高い品質を維持し、性能指標を55.1%向上させた。カスタムカーネルでは最大でprefillを1.98倍、decodeを4.95倍高速化したが、その他の推論オーバーヘッドを含めると、エンドツーエンドの高速化はそれぞれ1.68倍と1.55倍に低下する。この差こそ、デプロイ時に注目すべき点だ。計算量を半分に削減しても、サービスのスループットが2倍になるとは限らない。
コードはEIT-NLPのLLM-Pruningリポジトリで公開されており、学習方法やカーネル設計を確認できる。ただし、現時点の結果は、著者が選定したモデル、GPU、バッチサイズ、シーケンス長に基づくものだ。動的ルーティングには、mask生成、並べ替え、カーネル保守の追加コストも伴う。次の焦点は、vLLMやSGLangなどのcontinuous batching executorに統合できるか、そして異なるGPU、量子化モデル、高並行性のサービス環境でも効果を維持できるかどうかだ。