jetson-speculative-decoding

Jetson 投機的デコーディング（vLLM）

投機的デコーディングでは、小さな「ドラフト」モデルがトークンを提案し、ターゲットモデルが1回のフォワードパスで検証することで、トークンごとのレイテンシを削減します。Jetson では、効果の有無を左右するのはドラフトの品質ではなく、VRAM の余裕です。このスキルでは、LLM がすでに知っているはずの内容を除いた、重要な情報をまとめています。

目的

既存の Jetson vLLM デプロイに適切な --speculative-config を追加して調整し、単一ストリームのデコード速度が向上するかを検証します。

使用する場面

• ボトルネックが TPOT/ITL である場合（TTFT は問題なく、出力だけが遅い）。
• ワークロードが単一ストリーム、または低並行性（≤2）の場合。投機的デコーディングは通常、高並行性では効果がありません。
• Jetson のシリーズが Thor または AGX Orin の場合。Orin Nano/NX では EAGLE-3 を提案しないでください。ターゲットモデルとドラフトモデルの両方を動かすための VRAM の余裕がほとんどなく、起動時に OOM が発生するためです。

使用しない場面

• 高並行性のサービング（≥8）の場合：通常、バッチデコードのほうが投機的デコーディングより優れています。ドラフトモデルが VRAM を消費してしまいます。
• 公開済みの EAGLE-3 ヘッドがないモデルの場合。「修正」として場当たり的にトレーニングしないでください。
• jetson-inference-mem-tune のフラグを適用し、すでに --gpu-memory-utilization を上限近くまで引き上げている場合。まず少なくとも約2 GBの空きを確保してください。

前提条件

• jetson-llm-serve の、動作する vLLM サーバーレシピ。
• ターゲットモデルに加えて、ドラフトモデルまたは EAGLE-3 ヘッドを動かせるだけのメモリの余裕。
• 投機的デコーディングを有効にする前に、jetson-llm-benchmark で取得したベンチマークのベースライン。
• 互換性のある EAGLE-3 ヘッドを備えたターゲットモデル、またはフォールバック用の同じファミリーの小型ドラフトモデル。

手順

jetson-llm-serve に示されている vllm serve コマンドに --speculative-config を追加します。

EAGLE-3（ターゲットモデル用のヘッドが公開されている場合に推奨）：

--speculative-config '{
  "method": "eagle3",
  "model": "<eagle3-head-repo-id>",
  "num_speculative_tokens": 5,
  "draft_tensor_parallel_size": 1
}'

ドラフトモデル（フォールバック。同じファミリーの小型モデルを組み合わせます）：

--speculative-config '{
  "method": "draft_model",
  "model": "<small-draft-model-repo-id>",
  "num_speculative_tokens": 4,
  "draft_tensor_parallel_size": 1
}'

Jetson 固有の調整ルール

• num_speculative_tokens：Thor では 5、AGX Orin では 3 から始めます。値を上げる効果があるのは、ドラフトの受け入れ率が >0.6 の場合に限られます。
• jetson-llm-serve と同じ vLLM ランタイムパスを必ず使用してください。Thor では上流版 vLLM 0.20+（vllm/vllm-openai:latest）または検証済みのネイティブ vLLM 0.20+、Orin JetPack 7.2 / L4T r39+ では上流版 vLLM 0.20+、それ以前の Orin では NVIDIA-AI-IOT vLLM イメージを使用します。Orin 用 NVIDIA-AI-IOT vLLM イメージを Thor で使用しないでください。古いランタイムでは EAGLE-3 や現在の --speculative-config の形式に対応していない場合があります。
• ドラフトモデル用のメモリを確保するため、投機的デコーディングを使用しないベースラインと比べて --gpu-memory-utilization を約0.05下げます。

効果を確認する方法

1. 投機的デコーディングを有効にする前と後に、--concurrency 1 で jetson-llm-benchmark（vLLM パス）を実行します。
2. 合格基準：並行性1で throughput_tok_s が30%以上向上し、tpot_ms_p50 が20%以上低下すること。
3. 改善が10%未満の場合、または並行性8で throughput_tok_s が低下する場合は、投機的デコーディングを無効にしてください。ドラフトモデルのコストが効果を上回っています。

制限事項

• 投機的デコーディングはデコード主体のワークロードを改善しますが、TTFT が支配的なレイテンシは短縮しません。
• 継続的バッチ処理ですでに GPU が高稼働しているため、高並行性では効果がなくなることがあります。
• Orin Nano/NX では通常、ターゲットモデルとドラフトモデルの両方を動かすだけのメモリの余裕がありません。
• 受け入れ率とドラフトのオーバーヘッドはモデルごとに異なるため、速度向上を前提とせず、適用前後にベンチマークを実施してください。

エラー処理

• vLLM が --speculative-config を拒否する場合、Thor と Orin JetPack 7.2 / L4T r39+ では vLLM 0.20+ を使用していること、それ以前の Orin では JetPack に適合した NVIDIA-AI-IOT vLLM イメージを使用していることを確認してください。それでもランタイムが拒否する場合は、投機的デコーディングを使用しないサービングコマンドに戻してください。
• 起動時に OOM が発生した場合は、--gpu-memory-utilization を下げる、小型のドラフトモデルを使用する、または投機的デコーディングを無効にして jetson-inference-mem-tune に引き継ぐ、のいずれかを行ってください。
• ベンチマークのスループットが低下した場合は、--speculative-config を削除してください。不適切なドラフトパスは、投機的デコーディングを使わない場合より悪い結果になります。

引き継ぎ先

• 変更の効果を定量化するには、jetson-llm-benchmarkを使用してください。
• speculative decodingを有効にした後、起動時にOOMが発生する場合は、jetson-inference-mem-tuneを使用してください。

ソース

vLLM speculative decodingのドキュメントと、Jetson AI Lab GenAIチュートリアル。