NVIDIA Dynamoについて調べてみた

こんにちは。NTTコミュニケーションズの露崎です。本ブログでは2025年3月のGTCで紹介されたNVIDIA社のOSS Dynamoについて紹介します。

はじめに

こんにちは。NTTコミュニケーションズの露崎です。本ブログでは2025年3月のGTCで紹介されたNVIDIA社のOSS Dynamoについて紹介します。 NVIDIA Dynamoは発表されて間がなく、開発/変更が盛んに行われています。本ブログでは2025年5月の時点での最新版である0.2.0について紹介しますが、最新情報については公式をご参照ください。

NVIDIA DynamoはNVIDIA社が開発しOSSで公開している推論フレームワークです。LLMの推論時の処理をプロンプトの解釈を実施するPrefill、単語を逐次的に生成するDecodeなどのフェーズに分解し、フェーズ毎の処理を別々のGPUで実行させることにより、推論時の負荷を分散可能な分散推論フレームワークです。 各フェーズ間の通信をNVLINK経由のRDMAなどを用いることによって、より効率的でスケーラブルなAI推論が可能とする設計になっています。 PrefillやDecodeなどLLMの仕組みについての詳細はNVIDIA社のブログをご参照ください。

NVIDIA Dynamoのアーキテクチャ(出典: NVIDIA Dynamo)

特徴

NVIDIA Dynamoは、以下のような特徴を持っています。

  • 推論グラフの構築: 自由度の高い推論グラフを構築でき、パイプライン向けのコンポーネントをSDKで提供します。
  • 分散処理: Prefill、Decodeなどのフェーズに分解し、各フェーズ間の通信をRDMAを始めとした通信の抽象化を提供するNIXL、プロセス間通信向けのnatsで実現しています。
  • 基本言語と依存関係: 処理系はRust言語で書かれており、推論グラフの構築や分散処理の開発用SDKとしてPython言語のBindingsも提供しています。依存関係として状態管理のためのetcd、推論処理エンジンであるvllmSGLangTensorRT-LLMなどに依存しています。

このように、NVIDIA Dynamoは推論高速化のためのソフトウェアですが、独自の処理系を提供するわけではなく、さまざまなコンポーネントを組み合わせた最適化を実施するためのオーケストレータのような機能を提供するフレームワークです。

インストールと基本動作

NVIDIA Dynamoはpip packageで配布されており、以下のコマンドでインストールが可能です。

pip install ai-dynamo[all]

Python Package Index (PyPI)にはNVIDIA社のpackageを参照するスタブパッケージが提供されており、スタブ経由でNVIDIA社のリポジトリからインストールできます。NIXLなどの通信ライブラリについてもビルド済みのバイナリを梱包したものが配布されています。 分散処理に必要なetcd、natsについてはDynamo Serveの節でより詳細に解説します。

Note v0.2.0では、vllmエンジンとしてai_dynamo_vllm=0.8.4 が必要です。NIXLについてはバージョンの依存関係が定義されていませんが、筆者の環境では公式に従って検証時の最新のHEADをビルドし、動作確認を実施しました。

NVIDIA Dynamoでは基本的な動作確認のためのコマンドとして、dynamo run を提供するほか、分散処理のための dynamo serveKubernetes向けの dynamo deploy などのコマンドを提供しています。

NVIDIA Dynamoの動作環境については公式をご確認ください。

本ブログでは主に dynamo rundynamo serve について紹介します。

Dynamo Run

dynamo run は対話形式のshellを起動するコマンドです。以下のようにモデルを読み込み、対話的なシェルを起動します。モデルはHuggingFaceのモデルIDに対応している他、ローカルのファイルシステムからも読み込むことができます。

dynamo run out=vllm tokyotech-llm/Llama-3.1-Swallow-8B-Instruct-v0.3

Note dynamo run コマンドの out にはvllm、sglangなど推論を実行するエンジンを指定し、推論エンジン経由でGPUを利用します。outを指定しない場合、CPUが利用されるため、推論処理に時間がかかっている場合には、推論エンジンを正しく設定できているか確認してみてください。

上記のコマンドを実行すると以下のようなシェルが起動し、対話的な推論を実施できます。

$ dynamo run out=vllm tokyotech-llm/Llama-3.1-Swallow-8B-Instruct-v0.3
2025-05-01T07:52:54.065Z  WARN __init__.vllm_version_matches_substr: Using ai_dynamo_vllm
2025-05-01T07:52:54.074Z  INFO __init__.resolve_current_platform_cls_qualname: Automatically detected platform cuda.
2025-05-01T07:52:54.722Z  INFO nixl: NIXL is available
Loading safetensors checkpoint shards:   0% Completed | 0/4 [00:00<?, ?it/s]
Loading safetensors checkpoint shards:  25% Completed | 1/4 [00:00<00:01,  1.55it/s]
Loading safetensors checkpoint shards:  50% Completed | 2/4 [00:01<00:01,  1.49it/s]
Loading safetensors checkpoint shards:  75% Completed | 3/4 [00:01<00:00,  2.19it/s]
Loading safetensors checkpoint shards: 100% Completed | 4/4 [00:02<00:00,  1.85it/s]
Loading safetensors checkpoint shards: 100% Completed | 4/4 [00:02<00:00,  1.82it/s]
2025-05-01T07:53:11.006Z  INFO loader.load_model: Loading weights took 2.25 seconds   
2025-05-01T07:53:11.207Z  INFO model_runner.load_model: Model loading took 14.9596 GiB and 2.389781 seconds   
2025-05-01T07:53:11.913Z  INFO worker.determine_num_available_blocks: Memory profiling takes 0.57 seconds
the current vLLM instance can use total_gpu_memory (79.19GiB) x gpu_memory_utilization (0.90) = 71.27GiB
model weights take 14.96GiB; non_torch_memory takes 0.16GiB; PyTorch activation peak memory takes 1.26GiB; the rest of the memory reserved for KV Cache is 54.90GiB.   
2025-05-01T07:53:12.045Z  INFO executor_base.initialize_cache: # cuda blocks: 28108, # CPU blocks: 2048   
2025-05-01T07:53:12.046Z  INFO executor_base.initialize_cache: Maximum concurrency for 8192 tokens per request: 54.90x   
2025-05-01T07:53:13.427Z  INFO model_runner.capture_model: Capturing cudagraphs for decoding. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI. If out-of-memory error occurs during cudagraph capture, consider decreasing `gpu_memory_utilization` or switching to eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage.   
Capturing CUDA graph shapes: 100%|███████████████████████████████████████████████████████████████████| 35/35 [00:10<00:00,  3.40it/s]
2025-05-01T07:53:23.712Z  INFO model_runner.capture_model: Graph capturing finished in 10 secs, took 0.32 GiB   
2025-05-01T07:53:23.712Z  INFO llm_engine._initialize_kv_caches: init engine (profile, create kv cache, warmup model) took 12.51 seconds   
2025-05-01T07:53:25.279Z  INFO dynamo_run::input::text: Ctrl-c to exit
? User ›

終了時には Ctrl+C でプロセスを終了します。

この他、オプションなどについては公式を参照してください。

Dynamo Serve

dynamo serve は事前に定義した推論グラフに従って各コンポーネントの起動、オーケストレーションを提供するコマンドです。

各コンポーネント間の通信にはnatsプロトコルが利用され、状態管理をetcdで実施するため、 dynamo serve の実行前にnatsサービス、etcdサービスを起動しておく必要があります。ポートやコンフィグがデフォルト値でよければ、公式が配布するdocker compose用のyamlを利用するのが最も簡単で、ワーキングディレクトリをdynamoのリポジトリルートとした時に以下のコマンドで各サービスを起動できます。

docker compose -f deploy/docker-compose.yml up -d

推論グラフとコンポーネント

前述のとおり、 dynamo serve を利用するためには事前に分散処理を定義した推論グラフが必要です。ここでは dynamo serve を理解するために必要な推論グラフとコンポーネントについて解説します。

推論グラフは分散推論を行う際のコンポーネントの組み合わせや依存関係を示した定義です。コンポーネントは推論グラフを構成するためのパーツです。 具体的なコンポーネントにはFrontend、Processor、Router、Workerなどがあります。これらのコンポーネントの依存関係を整理し、Frontendで受け付けたリクエストをProcessorがRouterの情報をベースにスケジュール、実際の推論をWorkerで処理する、といった流れを定義するのが推論グラフです。

NVIDIA DynamoではLLMに関するコンポーネントと推論グラフのサンプル実装が公開されています。

推論グラフのサンプルは以下のとおりです。

# From examples/llm/graphs/agg.py
from components.frontend import Frontend
from components.processor import Processor
from components.worker import VllmWorker

Frontend.link(Processor).link(VllmWorker)

この例ではNVIDIA Dynamo SDKのPython Bindingを利用し依存関係を表現しています。 実際には components に含まれるFrontendやProcessorといったクラスがコンポーネントにあたり、これを実装する必要があります。

コンポーネントのサンプルは以下のとおりです。依存関係はコンポーネント内でも個別に定義できます。

# components/processor.py
class Processor(ProcessMixIn):
    worker = depends(VllmWorker)
    router = depends(Router)
    ...

依存関係の定義の仕方については、基本的にはclass側で設定した依存関係に応じて処理系を実装し、実際の起動時の関係を .link で表現するといった使い方になります。

ここからは各コンポーネントでこの依存関係を使ってどのように処理を実装するかを簡単に紹介します。 処理の実装では、コンポーネントで定義した依存関係を以下のようなコードで動的に解決し、依存先のコンポーネントを変数のように扱うことができます。ここでは上記のprocessorの起動時に依存するVllmWorkerを self.worker_client に初期化するコードを掲載します。

    comp_ns, comp_name = VllmWorker.dynamo_address()  # type: ignore
    self.worker_client = (
        await runtime.namespace(comp_ns)
        .component(comp_name)
        .endpoint("generate")
        .client()
    )

このように、分散処理におけるコミュニケーションをカプセル化し逐次的に処理を記述できます。ここで初期化された self.worker_client は以下のような形で別のプロセスに処理を引き継がせることができます。

    engine_generator = await self.worker_client.generate(
        vLLMGenerateRequest(
            engine_prompt=engine_prompt,
            sampling_params=sampling_params,
            request_id=request_id,
            prefix_hit_rate=prefix_hit_rate,
        ).model_dump_json()
    )

初期化の方法や、依存関係の解決の仕方には複数のパターンが存在するため、より詳細にコンポーネントを実装したい方は公式を参照してください。

dynamo serveの起動の流れ

ここでは前節で紹介したコンポーネントと推論グラフを用いて dynamo serve を使った推論APIサーバを起動する流れを紹介します。

1. nats/etcdの起動

dynamo serve を実行する前に通信ライブラリのnatsと状態管理用のetcdのサービスを起動します。NVIDIA Dynamoのデフォルト値を利用するのであれば前述した通りdocker composeを用いて以下のコマンドで起動します。

docker compose -f deploy/docker-compose.yml up -d

2. dynamo serveの起動

次に dynamo serve コマンドを使って各コンポーネントを起動します。以下はワーキングディレクトリをNVIDIA Dynamoのexamples/llmとした時のサンプルです。

dynamo serve graphs.agg:Frontend -f configs/agg.yaml

このコマンドでは examples/llm/graphs/agg.py 内に定義されているFrontendを起点とする推論グラフを -f で指定した設定ファイルに基づいて起動するということを実行します。

このサンプルではaggregated modeの依存関係にあるFrontend、Processor、VllmWokerが起動し VllmWorker has been initialized が画面に出力されれば準備完了です。

3. 動作確認

NVIDIA DynamoはOpenAI互換のAPIを提供します。このため、以下のようなOpenAIクライアントを使った推論で動作確認が可能です。

from openai import OpenAI

ENDPOINT = "http://localhost:8000/v1"

def main():
    client = OpenAI(
        base_url=ENDPOINT,
        api_key="needn't"
    )
    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": "Hello, how are you?",
            }
        ],
        model="tokyotech-llm/Llama-3.1-Swallow-8B-Instruct-v0.3",
        temperature=0.9,
    )
    print(chat_completion.choices[0].message.content)


if __name__ == "__main__":
    main()

4. 終了

dynamo serve を終了する際は Ctl+C など、親プロセスを終了させることで関連サービスを終了できます。

分散処理の仕組み

最後に dynamo serve を利用して分散処理を実施する場合、どのような分散処理が実行されるかをdisaggregatedのサンプル実装をベースに紹介します。

disaggregatedの実装はFrontend、Processor、VllmWorker、PrefillWorkerの4つのコンポーネントで構成されています。 LLMの推論は、入力された文章からKey/Valueの値を計算するPrefill処理とPrefillされた値から1トークンずつ回答を生成するDecode処理に分けることができます。NVIDIA Dynamoのdisaggregatedの実装では、このPrefill処理をPrefillWorker、Decode処理をVllmWorkerというそれぞれ別のコンポーネントで処理する機能を提供します。

より具体的な流れを解説します。disaggregatedの実装では Processor.link(VllmWorker).link(PrefillWorker) という依存関係を持っています。このため dynamo serve コマンドにより、それぞれの依存関係が解決され、初期化されます。初期化時にVllmWorkerは最大コンテキスト長に応じた計算用のメモリブロックをGPU上に確保します。

実際の推論リクエストがユーザから入力されるとFrontend経由でProcessorがVllmWorkerに対して処理をスケジュールします。この時にPrefillWorkerが利用可能であることを通知します。スケジュールされたVllmWorkerは起動時に確保したメモリブロックで利用可能なものを予約し、nats経由でPrefillWorkerへ通知します。Prefill Workerは通知された内容を元にPrefillを実行し、計算結果をRDMAで依頼元のVllmWorkerのメモリブロックに書き込みます。PrefillWorkerはメモリブロックを書き込んだことをRDMAのnotify機能を使ってメモリを確保しているプロセスに通知します。メモリの書き込み通知を受けたVllmWorkerはPrefillされたメモリブロックの結果を用いてDecode処理を実施します。

以下は、起動から推論リクエストをPrefill、Decodeし処理する流れを図にしたものです。

NVIDIA Dynamo Disaggregated Serving

Note 図にあるように、実際にはVllmWorkerとPrefillWorkerにはDynamoのコンポーネントであるWorkerと実際に推論を処理するVllmWorkerが存在し、コンポーネントとzmq経由でコミュニケーションを取っています。また、これらのPrefill、RDMAの通知を実現するためにOSSのvllmから変更されたソフトウェアを利用しており、変更点についてはこちらで確認することができます。

まとめ

本ブログではNVIDIA Dynamoについて紹介しました。NVIDIA Dynamoは、効率的でスケーラブルな分散推論を実現するためのフレームワークです。サンプルの実装のままでもLLMの推論を効率化するための分散処理が実現でき、さらに推論グラフやコンポーネントを実装することで複雑な処理形を実装することも可能です。今後、LLMの需要拡大に向けてこうしたGPUの処理効率を向上させるソフトウェア、仕組みを活用することはますます重要になると考えられます。