ランタイム Draft Model 学習：Speculative Decoding を実環境ワークロードに適応させる

概要

大規模言語モデル（LLM）の規模と複雑さが増し続ける中、推論コストの削減はAIデータセンターおよびサービスプロバイダーにとって最重要課題となっています。これらのモデルを大規模に提供するための計算コストは、モデル品質を損なうことなく大幅なコスト削減を実現できる効率的な最適化手法の必要性を高めています。

推論最適化に対処するため、disaggregation、KV cache対応ルーティング、test-timeルーティング、quantization、speculative decodingなど、さまざまなアプローチが登場しています。これらの手法の中でも、speculative decodingは主要なクラウドサービスプロバイダー（CSP）から大きな注目を集めています。モデル品質の維持を保証しながら他の最適化手法と組み合わせ可能であり、ほとんどの実用的なシナリオでパフォーマンス向上を実現するためです。

Speculative decodingは、小型で高速なモデル（draft model）を使用してdraft tokenを生成し、それを元のモデル（target model）が並列で検証することで推論を高速化します。このアプローチは、target modelのみで逐次的にトークンを生成するよりも根本的に効率的です。

上の図は、4つのdraft tokenを使用したspeculative decodingと標準decodingで10個のトークンを生成するタイムラインを示しています。平均して2つのdraft tokenしか受理されなくても、speculative decodingは約2倍の高速化を達成します。

この効率向上が可能なのは、target modelが複数のトークンを並列で検証するのに必要な時間が、単一のトークンを生成するのに必要な時間とほぼ同じであるためです。これは、LLM推論が計算バウンドではなく主にメモリバウンドであることに起因します。ボトルネックは実際の計算ではなく、メモリからのモデル重みの読み込みにあります。1つのトークンを検証する場合でも複数のトークンを同時に検証する場合でも、target modelは同じ重みを読み込む必要があるため、レイテンシはほぼ同じになります。このメモリアクセスコストを複数のトークン検証に分散させることで、speculative decodingはシーケンス生成に必要なコストの高いtarget model forward passの総数を大幅に削減します。

しかし、従来のspeculative decodingアプローチには重要な制限があります：draft modelは通常、実際のプロダクショントラフィックの分布と一致しない可能性がある汎用ワークロードで事前学習されます。ワークロード分布はサービスごとに大きく異なり、同じサービス内でも時間とともに変化するため、汎用データで学習されたdraft modelは、専門的または変化するユースケースに対して最適でないパフォーマンスを示すことが多くあります。

自己改善LLM推論のためのTemporal Incremental Draft Engine (TIDE)

TIDEは、ランタイムdraft model学習によってこの課題に対処します。TIDEは推論エンジンとしてSGLangを使用し、draft model学習にSpecForgeを活用し、EAGLE3 speculative decoding技術を基盤としています。ライブサービスワークロードに基づいてdraft modelを継続的に適応させることで、TIDEは手動介入なしに自動的に推論パフォーマンスを向上させます。

EAGLE3: Hidden Stateからのdraft model学習

EAGLE3はdraft model学習に独自のアプローチを取ります。独立した言語モデルとして機能する従来のdraft modelとは異なり、EAGLE3のdraft modelはtarget modelの複数の中間レイヤーからのhidden stateを入力として受け取り、target modelの出力分布を予測するよう学習します。

TIDEシステム設計

TIDEのアーキテクチャは、オーバーヘッドと複雑さを最小限に抑えながら、ランタイム学習をプロダクション推論システムにシームレスに統合するよう設計されています。

1. Hidden Stateロギング付き推論エンジン

推論エンジン（SGLang）は、ユーザーリクエストを処理するために標準的なprefillおよびdecoding計算を実行します。重要なのは、これらの操作中にtarget modelの中間レイヤーからhidden stateをキャプチャし、ストレージにダンプすることです。これらのhidden stateは、実際のプロダクション推論中に計算された内部表現を表し、draft modelの学習データとして機能します。

2. 学習とモデル更新

推論エンジンがリクエストの処理を続ける中、hidden stateがストレージに蓄積されます。現在のワークロード分布の意味のあるサンプルを表す十分な量のデータが収集されると、学習プロセスが自動的にトリガーされます。学習エンジンは蓄積されたhidden stateを読み込み、最近のワークロードに基づいてtarget modelの出力分布をより正確に予測するようEAGLE3 draft modelを学習させます。学習が収束すると、更新されたdraft modelが以前のバージョンを置き換えて推論エンジンにデプロイされます。これで1つの適応サイクルが完了し、新しいhidden stateが蓄積されるにつれてプロセスが継続され、変化するワークロードパターンへの継続的な適応が保証されます。

このアーキテクチャは、ランタイム学習の主要な課題を優雅に解決します：

• ほぼゼロの推論オーバーヘッド： Hidden stateロギングは推論計算と非同期で実行されるため、オーバーヘッドはほぼ完全に隠蔽され、リクエスト処理にレイテンシをほとんど追加しません。
• 非同期学習： 学習は推論とは独立して行われるため、リクエスト処理をブロックしたり遅延させたりしません。
• 自動適応： サイクル全体が手動介入なしに自律的に実行され、ワークロードの変化に継続的に適応します。
• リソース効率： 異種GPU環境では、学習を別のハードウェアにオフロードし、高性能GPUを推論に集中させることができます。

このアーキテクチャのシンプルさと自動化により、手動チューニングや介入がコスト高で非実用的なプロダクション環境での展開に適しています。

評価

ランタイムdraft model学習の有効性を評価するため、TIDEのパフォーマンスを汎用データで事前学習された静的draft modelと比較しました。評価には韓国語の会話クエリを含むdbdu/ShareGPT-74k-koデータセットを使用し、ワークロード固有の適応のための現実的なテストベッドを提供します。実験ではlmsys/gpt-oss-120b-bf16をtarget modelとして、lmsys/EAGLE3-gpt-oss-120b-bf16をベースラインの事前学習済みdraft modelとして使用し、推論はAMD MI300XまたはNVIDIA H100 GPUで実行、draft model学習はAMD MI250 GPUで行い、TIDEが異種GPUリソースを効果的に活用できることを実証しています。

Figure 3に示すように、TIDEは並行処理レベルに応じて、事前学習済みdraft modelと比較して1.14倍から1.35倍の出力トークンスループット向上を達成しています。

Figure 4はランタイム学習の重要な利点を示しています：TIDEがワークロードから継続的に学習するにつれてスループットが向上します。事前学習済みdraft modelは、パフォーマンスが事前学習によって固定されているため、実験全体を通じて比較的一定のスループットを維持します。対照的に、TIDEは韓国語会話データセットの特定のパターンに適応するにつれて、スループットが段階的に向上する様子を示しています。

TIDEのパフォーマンス向上メカニズムをより深く理解するため、acceptance length——各検証ステップでtarget modelが受理するdraft tokenの数——が異なる並行処理レベルで時間とともにどのように変化するかを分析しました。

Figure 5はランタイム学習の重要な特性を明らかにしています：並行処理レベルが高いほど、acceptance lengthの向上が速くなります。その理由は以下の通りです：

• より速いデータ蓄積： 高い並行処理レベルでは、より多くのリクエストが同時に処理されるため、TIDEはより速い速度で学習データを収集できます。これにより学習プロセスが加速され、ワークロード分布へのより迅速な適応が可能になります。
• より頻繁なモデル更新： 単位時間あたりより多くの学習サンプルが利用可能なため、draft modelを統計的に有意なバッチでより頻繁に更新でき、より速い収束とより広範なパターンカバレッジにつながります。

結論

TIDEは、ランタイムdraft model学習がプロダクション環境におけるspeculative decodingのパフォーマンスを大幅に向上させ得ることを実証しています。韓国語会話データセットでの実験により、静的な事前学習済みdraft modelと比較して1.14倍から1.35倍の出力トークンスループット向上が示され、システムがライブワークロードに適応するにつれてパフォーマンスは継続的に改善されます。

パフォーマンス向上に加え、TIDEは異種GPUシステムにおいて説得力のある利点を提供します。旧世代のGPUや活用度の低いハードウェアなど、アイドル状態または学習に最適化されたリソースをdraft model学習に活用し、高性能GPUを推論に専念させることで、TIDEはリソース活用率と推論効率の両方を同時に向上させます。これにより、システムレベルでの全体的なコストパフォーマンスが向上します。

ランタイム学習をプロダクション推論システムにシームレスに統合するため、オープンソースプロジェクトに貢献し、TIDEの機能をより広いコミュニティに提供しています：

• SpecForge: PR#1, PR#2
• SGLang: PR#3

AIワークロードが多様化・進化し続ける中、変化するパターンに自動的に適応するTIDEのようなシステムは、大規模で効率的かつコスト効果の高い推論を維持するためにますます不可欠になるでしょう。