AI 基礎設施
Megatron-LM、プレフィックス認識ルーティングと動的バッチ処理を拡張――ハイブリッドモデル推論でKV・Mamba状態の共有が可能に
NVIDIAは9月第1週のMegatron-LMの変更点をまとめ、データ並列推論フロントエンド、プレフィックスキャッシュ、可変長シーケンス処理を、より完全なサービング経路へと前進させた。新機能により、Transformer/SSMハイブリッドモデルの状態再利用が改善されるが、スループット、Time to First Token(TTFT)、クロスノードスケーリングに関するデータはまだ公式に公開されていない。

NVIDIAは9月7日、[Megatron-LMの月次進捗](https://github.com/NVIDIA/Megatron-LM/issues/7100)を公開し、6日までに行われた58件の実質的なマージをまとめた。推論面で最も注目される変更は、HTTPフロントエンドがデータ並列レプリカに応じてスケールできるようになり、tokenizationなどのCPU処理がコーディネーターのイベントループをブロックしなくなったことだ。ルーターはレプリカの負荷とプレフィックスキャッシュのヒットを同時に考慮し、共通のシステムプロンプトや長い文書プレフィックスを持つリクエストを、関連する状態がすでに保存されているGPUへ送る可能性を高める。
基盤となる`DynamicInferenceContext`は、処理中のバッチとブロック方式のKV cacheを管理し、各ステップで長さの異なるリクエストを追加、一時停止、削除できる。新版では、アイドル時に使用するdummy forwardによってプレフィックスインデックスが消去されることも防いでいる。さらに、MambaやGatedDeltaNetなどのハイブリッドアーキテクチャ向けに状態スロットを追加した。これにより、反復するプレフィックスではAttention層のKVだけでなく、SSMと畳み込みの状態も再利用できる。公式の[APIドキュメント](https://docs.nvidia.com/megatron-core/developer-guide/latest/apidocs/core/core.inference.contexts.dynamic_context.html)によると、これらのキャッシュには引き続き、事前設定されたGPUメモリ予算とブロック境界による制約がある。
トレーニング経路では、可変長シーケンス、packed sequence、runtime context parallelismへの対応が拡大され、microbatchごとに異なるコンテキスト並列サブグループを使用できるようになった。MFSDPでも、通信順序、プリフェッチ、永続的なパラメーター/勾配ストレージが改善された。こうした取り組みは長いコンテキストやハイブリッドAttentionモデルにとって重要だ。固定トポロジーでは、シーケンス長のばらつきが大きい場合にアイドル時間やメモリの無駄が生じやすいためである。
ただし、これは安定版のリリースではなく、mainブランチの活動レポートだ。DeepSeek V4、prefill/decode分離、部分的なCUDA Graph機能の一部は、依然として進行中とされている。導入担当者は各項目のマージ状況を個別に確認したうえで、キャッシュヒット率、TTFT、リクエストのキャンセル後におけるメモリ回収、ルーティング係数の違いによる負荷の偏りを自ら測定すべきだ。現時点では、機能数だけから実際の性能向上を推測することはできない。