DGX Sparkでは、仕事に合わせて複数のローカルLLMをとっかえひっかえ使っている。モデルを切り替えるたびに待たされるので、切替時間を短くしたい。Super VLの記事に続き、今回はNemotron 3.5 Bolt 74B-A9B(24GB版)を扱う。
切替に393.8秒かかっていたのを91.8秒にした。同時に保持できる文脈の量(KVキャッシュ)は変えていない。
改善前:どこに時間がかかっていたか
切替は自作の管理ツールが行う。旧モデルの停止から新モデルが応答できるまでの各段階に時刻を記録し、vLLMのログと突き合わせて5つに分けた。
| 段階 | 改善前(秒) | 中身 |
|---|---|---|
| 重みの読込 | 162 | 本体と補助モデルの91ファイルを読み、GPUへ配置する |
| 画像の事前計測 | 27 | 最大サイズの画像をエンコーダに通し、必要なメモリを測る |
| コンパイル・試運転 | 75 | モデルの計算を準備し、Mamba層などを試しに動かす |
| Graph作成・カーネル準備 | 93 | CUDA Graphの記録と、それに使うカーネルのコンパイル |
| その他 | 36.8 | コンテナ起動、vLLM起動、API起動、応答確認 |
| 合計 | 393.8 |
重みの読込が41%、コンパイル・試運転とGraph作成・カーネル準備が合わせて43%だった。
改善の内容
1. 高速ローダーに替える(重みの読込 162→22秒)
vLLM標準のローダーは、モデルファイルをmmapで開き、テンソルを1つずつ取り出してGPUへ配置する。
--load-format fastsafetensors に替えると、ファイルをまとめて読んでからGPUへ渡す。重みの読込は22秒になった。
2. CUDA・FlashInfer・b12xのキャッシュを残す(Graph作成・カーネル準備 93→24秒)
CUDA Graphは、GPUで繰り返す計算の手順を記録する仕組みだ。記録の際に、CUDA・FlashInfer・量子化モデル用のb12xがカーネル(GPU用の計算部品)をその場でコンパイルする。保存先がコンテナ内だったため、起動のたびに作り直していた。
3つの保存先を、ホストへマウントした /root/.cache/vllm の下へ移した。
ENV CUDA_CACHE_PATH=/root/.cache/vllm/p052-bolt/cuda-compute-cache
ENV B12X_COMPILE_CACHE_DIR=/root/.cache/vllm/p052-bolt/b12x-compile
ENV FLASHINFER_WORKSPACE_BASE=/root/.cache/vllm/p052-bolt/flashinfer-workspace
3. TritonとInductorのキャッシュを残す(コンパイル・試運転 75→6秒、Graph作成・カーネル準備 24→12秒)
Boltは、Transformerと、長い文脈を少ないメモリで扱うMamba層を組み合わせたモデルだ。Mamba層の計算部品は、起動時にTriton(GPU用のカーネルをその場でコンパイルする仕組み)がコンパイルする。PyTorchのコンパイル機構Inductorも、同じように部品を作る。
この2つのキャッシュはコンテナ内に置かれ、再起動で消えていた。環境変数を2つ足し、残るボリュームへ移した。
ENV TRITON_CACHE_DIR=/root/.cache/vllm/p065-bolt/triton
ENV TORCHINDUCTOR_CACHE_DIR=/root/.cache/vllm/p065-bolt/inductor
Mamba層の準備運転がほぼ不要になり、コンパイル・試運転は6秒になった。Graph作成・カーネル準備でもTritonのコンパイルが減り、12秒になった。
画像の事前計測とその他には直接手を入れていない。
改善後も、生成速度(1要求あたり約146トークン/秒)、画像の読み取り、ツール呼出し、4要求の同時処理、長文入力は改善前と同じように動く。
環境
- NVIDIA DGX Spark(GB10、メモリ128GB)
- vLLM
0.32.0rc1.dev100+g77154ca61、PyTorch2.13.0+cu130、fastsafetensors0.4.0、FlashInfer0.7.0.post1、b12x1.5.1 - モデル:Nemotron 3.5 Bolt 74B-A9B(24GB版、IQ2_XXSとNVFP4の混合)、文脈長131,072、同時4要求、投機デコードDFlash(7トークン)

