powered by TechFeed
表示モード
Deep Dive

NVIDIAのエッジAI推論がリファレンス実装の6.4倍速に — プリフィルトークンの96%をKVキャッシュから供給する仕組み

9月17日、NVIDIAが「TensorRT Edge-LLM Completes the MLPerf Edge Agentic Benchmark 6.4x Faster on Jetson AGX Thor」と題した記事を公開した。この記事では、NVIDIAのエッジ推論ソフトウェア「TensorRT Edge-LLM」がJetson AGX Thor上でMLPerfエッジエージェントベンチマークをリファレンス実装の6.4倍の速度で完了したことについて詳しく紹介されている。

9月17日、NVIDIAが「TensorRT Edge-LLM Completes the MLPerf Edge Agentic Benchmark 6.4x Faster on Jetson AGX Thor」と題した記事を公開した。この記事では、NVIDIAのエッジ推論ソフトウェア「TensorRT Edge-LLM」がJetson AGX Thor上でMLPerfエッジエージェントベンチマークをリファレンス実装の6.4倍の速度で完了したことについて詳しく紹介されている。


6.4倍という数字の中身

MLPerf Inference v6.1のEdge Agenticベンチマークにおいて、NVIDIAのTensorRT Edge-LLMはQwen3.6-27Bを1台のJetson AGX Thor Developer Kit上で動作させ、52.33トークン/秒を達成した。ベンチマーク全体の1,007ターンを24分36秒で完了しており、llama.cppを使ったリファレンス実装の2時間37分に対して6.4倍高速である。

この差を生み出しているのは主に3つの技術だ。

  • NVFP4量子化:重みと活性化を4ビット浮動小数点で表現
  • ツリーベースのMulti-Token Prediction(MTP):複数トークンを並列に予測・検証
  • KVキャッシュの再利用:エージェントの会話履歴を効率的に再利用

それぞれの技術が何をどう変えているのかを以下で詳しく見ていく。


KVキャッシュ再利用で「同じ計算を繰り返さない」

エージェント型のワークロードで特に効いているのがKVキャッシュの再利用だ。エージェントが会話を進めるたびに、入力トークン列には直前までの会話履歴がまるごと含まれる。何も工夫しなければ、毎ターン同じ長い履歴を再処理(プリフィル)することになる。

TensorRT Edge-LLMは再利用可能なプロンプトのプレフィックスを検出し、その注意機構のKVページをキャッシュから復元する。Qwen3.6-27Bはハイブリッドアーキテクチャ(TransformerのSelf-Attention層とMamba系の再帰層を組み合わせた構成)を採用しているため、ランタイムは再帰状態と部分的なKVページ状態も合わせて復元した上で、新しいサフィックス部分だけをプリフィルする。

なお、ここで「96%の計算をスキップ」と表現しているのは正確にはプリフィル計算の削減を指しており、デコード計算(各トークン生成時のフォワードパス)は含まれない点に注意されたい。

その効果は数字に直接表れている。全1,007ターンで発生するプロンプトトークン合計約1,360万トークンのうち、実際にプリフィル処理したのは約50万トークン(〜0.5M)のみ。約96%のトークンがホットキャッシュから供給された。入力長がターンを重ねるごとに最大約23,500トークンまで伸びていくこのベンチマークでは、この最適化が全体の処理時間に直結する。


ツリーベースMTPで「1ステップで複数トークンを確定」

通常の自己回帰デコードはモデルを1回呼ぶごとに1トークンを生成する。Multi-Token Prediction(MTP)はドラフトモデルが複数の未来トークンを予測し、ターゲットモデルが一括して検証することでスループットを上げる手法だ。

TensorRT Edge-LLMが実装しているのは「ツリーベース」のMTPである。線形なMTPが1本の候補列だけを保持するのに対し、ツリーベースでは高確率の候補を木構造に展開する。ターゲットモデルが1回のフォワードパスでこれらをまとめて検証し、一致するパスを採択する。複数の候補が同時に受理されれば、1ステップで複数トークン分の生成が進む。

MLPerfの提出設定ではドラフトステップ数8、各深さでトップ2候補、検証ツリーのノード数16という構成を使用している。ツリー構造はファンクションコールのベンチマークと特に相性が良いとされており、ツール名やJSON構文、共通の引数構造は予測しやすい一方、個々の引数値は複数の分岐で候補を保持できるためだ。線形MTP(ドラフト3ステップ)と比較して、このワークロードでは約40%のデコード性能向上を確認している。


NVFP4量子化でメモリ帯域幅の制約を緩和

エッジ環境での低バッチLLM推論は、DRAM帯域幅がボトルネックになりやすい。Jetson AGX Thorに搭載されているNVIDIA Blackwell GPUはNVFP4(4ビット浮動小数点)をネイティブサポートしており、重みと活性化の両方にNVFP4を適用することでカーネルのメモリフットプリントを削減し、デコード速度を向上させる。KVキャッシュにはFP8を使用している。

量子化済みチェックポイントはHuggingFaceで公開されており、自前でキャリブレーションを行わずにそのまま使用可能だ。また、128GBのユニファイドメモリのうち量子化で節約できた領域が、長いコンテキストや投機的デコードのステートに割り当てられる点もエッジ環境では重要な副次効果だ。


ベンチマーク結果の数値

指標 結果
出力スループット 52.33 トークン/秒
最初のトークン生成までの中央値 247.12 ms
1出力トークンあたりの時間(中央値) 14.68 ms
BFCL全体精度 87.94%

ベンチマークのAccuracyフェーズではBerkeley Function Calling Leaderboard(BFCL)v4のプロンプトを使用し、モデルが正しい関数を選択・引数を生成できているか、また不要なツール呼び出しを避けられているかを評価している。


再現手順の概要

今回の提出実装はTensorRT Edge-LLMのrelease/0.9.1-mlpinfブランチで公開されている。大まかな手順は以下の通りだ。

  1. TensorRT Edge-LLMをクローンし、サブモジュールを初期化:通常のgit操作で完結する。依存関係の取得も含め数分程度。
  2. キャリブレーション済みNVFP4チェックポイントをダウンロード:HuggingFace Hubからの取得となる。27Bモデルのため、ネットワーク環境によってはダウンロードに時間がかかる点に注意。
  3. mlperf/README.mdに従い、チェックポイントのエクスポートとTensorRTエンジンのビルドを実施:このステップが最も時間を要する。エンジンビルドはJetson AGX Thor上でGPUを使って行うため、初回は数十分単位を見込んでおくと良い。
  4. OpenAI互換のTensorRT Edge-LLMサーバーを起動:起動コマンド自体はシンプルで、OpenAI互換APIとして動作するため既存ツールとの接続も容易だ。
  5. MLCommonsのエンドポイントハーネスをクローンし、BFCLの依存関係をインストールしてベンチマークを実行:Pythonの依存関係インストール後、inference-endpoint benchmark from-configコマンド一発でパフォーマンスフェーズと精度フェーズの両方を実行できる。精度フェーズのみを実行する場合は--accuracy-onlyオプションを使う。

詳細はTensorRT Edge-LLM Completes the MLPerf Edge Agentic Benchmark 6.4x Faster on Jetson AGX Thorを参照していただきたい。