ハードウェア認識動的投機的復号（DSD）| Cohere

投機的復号（Speculative Decoding, SD）は、品質を損なうことなく大規模言語モデル（LLM）の推論を高速化する一般的な手法です。標準的なLLM推論は一度に1トークンを生成しますが、SDはより小さなドラフトモデルを使用して複数のトークンを提案し、それらを大きなターゲットモデルが1タイムステップで検証します。各ステップで複数のトークンを受け入れられるため、生成が高速化されます。

SDは、GPUの計算とメモリ帯域幅のトレードオフを利用して高速化を実現します。複数のリクエストをバッチ処理することでGPU利用効率は向上しますが、動作環境はバッチサイズ（BS）に依存します。小規模BSでは推論はメモリ帯域幅制限となり、大規模BSでは計算制限となります。SDの時間は主に、ドラフト時間（小規模ドラフトモデル）と検証時間（大規模ターゲットモデル）で構成されます。

しかし、SDには課題があります。本番環境ではバッチサイズが動的に変化し、バッチサイズが大きい場合、通常のLLM推論は計算制限となり、SDが利用できるアイドル計算がなくなります。その結果、SDは通常の推論よりも遅くなる可能性があります。強化学習（RL）では、ロールアウトフェーズが主要なボトルネックであり、SDはロングテール生成に役立ちますが、高バッチサイズではスループットが低下するため、全体的な有用性が制限されます。

そこでCohereは、ハードウェアを考慮した動的投機的復号（DSD）を提案します。DSDは、ドラフトトークン数を適応的に調整することで標準SDを改善します。モデルとハードウェアの相互作用に基づいて最適なKを選択します。推論がメモリ帯域幅制限の場合はKを増やし、計算制限の場合はKを減らします。密モデルではKはバッチサイズとともに単調減少しますが、MoEモデルでは非単調になります。低BSでKは低く、中BSで高くなり、高BSで再び低下します。

最適なドラフトトークン数を見つけるために、DSDはgoodputメトリックを使用します：goodput = 受入長（AL）/ トークン間レイテンシ（ITL）。このメトリックは、ドラフトトークン追加の限界貢献と壁時計コストのトレードオフを捉えます。オフラインプロファイリングでALとITLを測定し、最適なKを見つけてルックアップテーブルに保存し、実行時に使用します。これによりコールドスタート問題が解決され、エンジンのランタイムメトリックを取り込むことでワークロードの変化に適応できます。

実験結果は、DSDがさまざまなバッチサイズで固定KのSDおよび通常推論よりも優れていることを示しています。MT-Benchデータセットにおいて、CohereのCommand A（密モデル）では、DSDは低BSでSDの高速化に匹敵し、高BS（64/128）ではSDおよび通常モデルよりも高速で、非常に高いBS（256）では通常モデルと同等の性能を示しました。具体的には、DSDはBS 128および256でSDよりも約23%高速であり、通常モデルと比較してBS 128で7.5%、BS 256で1.82%高速でした。Command A+（MoEモデル）では、DSDは固定K（K=3）とほぼ同じK値を選択したため、同様の高速化が得られましたが、将来の手法ではより明確な利得が期待されます。

Cohereはこの最適化をvLLMに貢献しました。DSDの実装には、非同期スケジューリングとの互換性とフルCUDAグラフとの互換性という2つの大きな課題がありました。非同期スケジューリングでは、スケジューラとモデルランナーが並行して動作します。DSDは、各タイムステップで同じ数のドラフトトークンが生成および検証されるという前提を破り、スケジューラとランナーの簿記を変更する必要があります。フルCUDAグラフでは、DSDはキャプチャされたタプルを拡張して異なるK値を含め、実行時にKを変更してもキャプチャされたCUDAグラフがヒットするようにします。これらの貢献により、DSDはvLLMの高度に最適化された効率と互換性を持ちます。