Amazon SageMaker AI で P-EAGLE を使用した投機的復号の並列化
この記事では、Amazon SageMaker AI 内で P-EAGLE を直接使用する方法を説明します。SageMaker JumpStart カタログから互換性のあるモデルを選択し、並列ドラフト仕様を構成し、高度に最適化されたリアルタイム SageMaker AI エンドポイントをデプロイして生成 AI アプリケーションを高速化する方法を示します。
大規模言語モデル（LLM）のサイズと複雑性が増す中、推論スループットを最大化しつつレイテンシを最小化することは、エンタープライズ本番環境での重要な課題です。投機的復号（Speculative decoding）は、軽量なドラフトモデルを使用して将来のトークンを推測し、ターゲットLLMが1回のフォワードパスで検証する戦略です。EAGLEのような最先端フレームワークは大幅な高速化を実現しましたが、ドラフトトークンが自己回帰的に生成されるため、推測深度に比例してレイテンシが増加するというアーキテクチャ上の限界がありました。
AWSが発明したP-EAGLE（Parallel-EAGLE）は、このボトルネックを打破します。P-EAGLEは、学習可能なプレースホルダー表現を用いて、すべてのドラフトトークンを1回のフォワードパスで並列に予測します。これにより、ドラフトトークン数と逐次フォワードパス数の関係を切り離し、レイテンシオーバーヘッドを増やすことなくより深い推測を可能にします。実際のベンチマークでは、P-EAGLEは標準のEAGLEフレームワークと比較して最大1.69倍のスループット向上を達成しています。
本日より、Amazon SageMaker JumpStart は人気の基盤モデル向けに P-EAGLE をネイティブサポートしています。SageMaker JumpStart は、ワンクリックまたは数行のコードで展開可能な最新のオープンウェイトモデルのキュレーションされたハブを提供します。P-EAGLEのモデル最適化とAmazon SageMaker AIのマネージド環境を組み合わせることで、開発者は複雑なCUDAカーネルや分散サービング設定を管理することなく、EAGLE-3よりも最大1.69倍高速な推論エンドポイントをデプロイできます。
本記事では、Qwen3-Coder-30B-A3B-Instruct を例に、SageMaker JumpStart からモデルを選択し、並列ドラフト仕様を設定し、エンドポイントをデプロイする手順を説明します。ベンチマークでは、HumanEval および SPEED-Bench Code において、P-EAGLE（K=11）が同時実行数1でそれぞれ1,167および873出力トークン/秒を記録し、ベースラインの3.97倍および2.97倍の性能を示しました。
P-EAGLEの中核は、自己回帰チェーンを断ち切る2つの学習可能パラメータにあります。マスクトークン埋め込み（emb_mask）は未知の前トークン埋め込みを代替し、共有隠れ状態（h_shared）は前位置のドラフター隠れ状態を代替します。理論的分析により、アテンションだけで十分な位置情報が得られるため、位置固有の隠れ状態は不要であることが示されています。これにより、すべてのKドラフト位置を並列に構築できます。
現在、SageMaker JumpStart では GPT-OSS-120B、GPT-OSS-20B、Qwen3-Coder-30B-A3B-Instruct、Gemma-4-31B-IT の4モデルが事前学習済みP-EAGLEヘッドとともに利用可能です。環境変数 SM_VLLM_SPECULATIVE_CONFIG により、ドラフター設定が自動構成されます（例: {"model": "/opt/ml/additional-model-data-sources/eagle", "method": "eagle3", "num_speculative_tokens": 3, "parallel_drafting": true}）。ユーザーは num_speculative_tokens を調整するだけで、並列ドラフトの深さを制御できます。
P-EAGLEは投機的復号のレイテンシ上限を並列化によって打破し、SageMaker AIのマネージド環境と組み合わせることで、生成AI推論のための即時利用可能な高速化ソリューションを提供します。不要になったエンドポイントは、コスト発生を避けるために必ず削除してください。