NVIDIA DGX SparkでMoE推論を高速化するための3つの最適化

大規模言語モデル（LLM）を最適化するエンジニアは、計算能力を増やしても生成速度が上がらないという壁に直面することがよくあります。NVIDIA DGX Sparkハードウェア上でLing-3.0-flashモデルに対する最近の技術評価では、デコーディング性能が生の計算能力ではなく、メモリ帯域幅によって制限されていることが特定されました。重み読み取りとトークン検証における特定のボトルネックに対処することで、研究者たちは精度を犠牲にすることなく、顕著なレイテンシ削減を達成しました。

何が起きたか

この評価は、Inclusion AI製の混合専門家（MoE）モデルであるLing-3.0-flashに焦点を当てました。このモデルは総パラメータ数が1,240億ですが、トークンあたりのアクティブパラメータ数は51億です。バッチサイズ1で実行した場合、プロファイリングツールは、トークン生成にかかった時間の大部分が2つの操作、すなわちMoE計算そのものと、LMヘッドとして知られる出力層射影によって占められていることを示しました。これらのステップではメモリから大量の重みデータを読み取る必要があり、標準的な最適化手法では完全には対処できないボトルネックを生み出していました。

これを解決するため、チームは3つの異なる最適化をテストしました。第一に、デフォルトのFlashInfer MXFP4 MoEバックエンドをHummingカーネルライブラリに置き換えました。第二に、LMヘッドに特化してオンラインFP8量子化を適用し、メモリフットプリントを削減しました。第三に、モデルに組み込まれたNEXTN層を使用して、マルチトークン予測による投機的デコーディングを実装しました。各変更は、カーネル実行効率からデータ量、ステップ数の削減まで、推論パイプラインの異なる側面を対象としています。

結果は、合成ワークロードと完全なGSM8Kテストセットの両方で測定されました。合成テストでは反復的な出力パターンにより大幅な向上が見られましたが、1,319問のGSM8K質問に対する実世界の評価では、デコードスループットが42.5%増加し、平均レイテンシが24.5%減少しました。重要なのは、これらの改善がモデルの精度に統計的に検出可能な変化をもたらさず、ベースラインの96.89%に対して96.82%のスコアを維持したことです。

仕組み

最初の最適化は、NVIDIA CUDAランタイムコンパイルを使用して、モデルの特定の次元に合わせてコードを生成する独立したカーネルライブラリであるHummingを活用しています。多くの汎用的ターゲットをコンパイルするベースラインのFlashInferパスとは異なり、HummingはLing-3.0-flashで使用される実際のマトリクス形状に特化したバイナリを作成します。この専門化により、GPUがモデルの狭い中間サイズに必要な命令のみを実行するように保証することで、出力トークンあたりの時間が短縮されます。

2番目のテクニックは、157,000以上のトークンのボキャブラリーへ隠れ状態を射影するLMヘッドに対処しています。ベースライン構成では、この層はBF16精度を使用しており、生成された各トークンに対して相当量のメモリ帯域幅を必要とします。この層のみをオンラインFP8量子化に切り替えることで、システムは各ステップ中に読み取られる重みデータの量を半分に減らします。この変更は、ベースラインのFP32 LMヘッドオプションと排他的であり、モデルの他の部分との互換性問題を防ぐために慎重な設定が必要です。

3番目の方法は、単一の検証ラウンドで複数のトークンを生成するために投機的デコーディングを使用します。Ling-3.0-flashには、ドラフトモデルとして機能するネイティブのNEXTN層が含まれており、後続のトークンを予測し、メインモデルがそれを検証します。システムは一度に1つのトークンを生成する代わりに、3つのトークンをドラフトし、4つの位置を同時に検証します。このアプローチは必要なデコードステップの総数を減らしますが、専門家アクティベーション数が増加するため、各検証ステップの複雑さを増します。

主要な詳細

• Ling-3.0-flashは、42層にわたってトークンあたり51億のアクティブパラメータを持つ、総計1,240億のパラメータを持っています。
• FlashInferバックエンドをHummingに置き換えたことにより、出力トークンあたりの時間は27.71 msから25.35 msに短縮され、9.3%の向上が見られました。
• LMヘッドのオンラインFP8量子化により、その層の重み読み取り量が半分になり、累積で21.7%のスループット向上に貢献しました。
• NEXTNを用いたマルチトークン予測は、合成データにおいてラウンドあたり3.90個の承認済みトークンを達成し、その特定のテストでスループットを185.1%向上させました。
• GSM8Kベンチマークでは、組み合わせられた最適化によりデコードスループットが42.5%増加し、平均レイテンシが24.5%減少しました。
• 精度は安定しており、96.89%から96.82%へとわずか0.08パーセンテージポイント低下しただけで、これは統計的に有意ではありません。

なぜ重要なのか

エージェント型ワークフローやリアルタイムアプリケーションを開発するエンジニアにとって、レイテンシはしばしば主要な制約となります。この分析は、低バッチサイズでの多くのMoEモデルにおいて、計算能力だけでなくメモリ帯域幅が制限要因であることを確認しています。重みトラフィックの削減とステップあたりのトークン数の増加に注力することで、エンジニアはより大きなハードウェアクラスターを必要とせずに、 substantialなパフォーマンス向上を達成できます。これは、すべてのミリ秒が重要なエッジ展開やコストに敏感なクラウド環境で特に重要です。

しかし、この研究は、深い統合知識なしに最適化を組み合わせた場合のリスクも浮き彫りにしています。研究者たちは、FP8 LMヘッド量子化とマルチトークン予測を有効にするには、ドラフトモデルがLMヘッドモジュール全体を共有する必要があることを発見しました。この特定のインターフェースがない場合、システムは静かに遅い構成にフォールバックし、FP8重みをBF16として解釈していました。これは、個々のコンポーネントテストに基づいて互換性を仮定するのではなく、起動時にすべての最適化を検証することの重要性を強調しています。

あなたができること

• 最適化を適用する前に、現在のMoE推論パイプラインをプロファイリングし、デコーディングがメモリバウンドなのか計算バウンドなのかを特定してください。
• モデルが特殊なマトリクス次元を持っている場合、FlashInferのドロップイン代替品としてHummingカーネルライブラリをテストしてください。
• SGLANG_ENABLE_FP8_LM_HEAD環境変数を使用してLMヘッドのオンラインFP8量子化を有効にし、メモリ帯域幅の使用量を削減してください。
• NEXTNなどのネイティブドラフト層を使用してマルチトークン予測を実装し、連続的なデコードステップの数を減らしてください。
• 量子化と投機的デコーディングを組み合わせる際に、ドラフトモデルとターゲットモデルが正しいモジュールを共有していることを確認し、静かなパフォーマンス低下を避けてください。
• 合成テストだけに頼るのではなく、GSM8Kのような現実的なワークロードで変更をベンチマークしてください。合成テストは投機的デコーディングの利点を過大評価する可能性があります。