👻
vLLMにおける低並列での推論性能向上の試行
はじめに
前回の記事ではAMD Ryzen AI Max+ 395でgpt-oss-20bをllama.cppとvLLMで推論させたときの性能と精度を評価した。
その結果、vLLMでは低並列リクエスト時の性能がllama.cppに対して著しく低いことが分かったが、納得がいかなかったのでvLLMで推論性能が低い原因を調べてみた。
PyTorch Profilerによる調査
詳細は割愛するが、PyTorch Profilerを使用してvLLMで1トークンを生成するときのタイムラインを取得した。水色と紫のパターンが24回繰り返しているのは、Transformerを24ブロック実行していることを意味している。
拡大するとmatmul(行列積)カーネルの実行に時間がかかっていることがわかる。
以下はカーネルごとの実行時間を集計してみた図。
最も時間のかかっているカーネルの名前は_matmul_ogs_NNT_bf16xbf16xmxfp4_32x256x128x1_swigluとなっていて、おそらく以下を意味していると考えられる。
- 入力と出力の型はbf16
- 重みの型はmxfp4
- block_m = 32
- block_n = 256
- block_k = 128
- split_k = 1
- SwiGLU活性化関数を融合
TransformerのFFNでは大きな行列の積を計算するが、それをそのまま計算するのではなくてタイル状に分割して計算する。block_{m,n,k}はそのタイルの大きさを表している。
行列積カーネルの最適化
上記の行列積の計算はTritonにより行われていて、行列の大きさに応じてどのようにタイル分割するかはヒューリスティックな手法で決定されているようだ。試しにブロックサイズを変更するとカーネルの実行時間が大きく変化し、短縮できる可能性があることがわかった。
最適なブロックサイズを求めるために、並列度（conc）とブロックサイズを変えながらカーネル実行時間を測定した。
$ python sweep_tiles.py
conc=1 bm=16 bn=32 bk=128 sk=1 median=10.101 ms samples=[10.100704000000007]
conc=1 bm=16 bn=64 bk=128 sk=1 median=11.037 ms samples=[11.036509000000002]
conc=1 bm=16 bn=128 bk=128 sk=1 median=16.656 ms samples=[16.65572499999999]
conc=1 bm=32 bn=32 bk=128 sk=1 median=11.430 ms samples=[11.430138999999992]
conc=1 bm=32 bn=64 bk=128 sk=1 median=24.284 ms samples=[24.283716000000002]
conc=1 bm=32 bn=128 bk=128 sk=1 median=28.853 ms samples=[28.853489000000042]
BEST 1 {'bm': 16, 'bn': 32, 'bk': 128, 'sk': 1, 'median_ms': 10.100704000000007}
conc=2 bm=16 bn=32 bk=128 sk=1 median=15.147 ms samples=[15.147296000000022]
conc=2 bm=16 bn=64 bk=128 sk=1 median=17.775 ms samples=[17.775126000000007]
conc=2 bm=16 bn=128 bk=128 sk=1 median=17.025 ms samples=[17.024792000000016]
conc=2 bm=32 bn=32 bk=128 sk=1 median=11.497 ms samples=[11.496612999999986]
conc=2 bm=32 bn=64 bk=128 sk=1 median=37.420 ms samples=[37.42023899999997]
conc=2 bm=32 bn=128 bk=128 sk=1 median=29.012 ms samples=[29.011639999999993]
BEST 2 {'bm': 32, 'bn': 32, 'bk': 128, 'sk': 1, 'median_ms': 11.496612999999986}
conc=3 bm=16 bn=32 bk=128 sk=1 median=15.170 ms samples=[15.170093999999994]
conc=3 bm=16 bn=64 bk=128 sk=1 median=17.639 ms samples=[17.638684999999995]
conc=3 bm=16 bn=128 bk=128 sk=1 median=24.430 ms samples=[24.430322000000032]
conc=3 bm=32 bn=32 bk=128 sk=1 median=11.533 ms samples=[11.532536999999977]
conc=3 bm=32 bn=64 bk=128 sk=1 median=37.682 ms samples=[37.68203199999997]
conc=3 bm=32 bn=128 bk=128 sk=1 median=46.550 ms samples=[46.54976800000002]
BEST 3 {'bm': 32, 'bn': 32, 'bk': 128, 'sk': 1, 'median_ms': 11.532536999999977}
conc=4 bm=16 bn=32 bk=128 sk=1 median=15.256 ms samples=[15.256412]
conc=4 bm=16 bn=64 bk=128 sk=1 median=17.777 ms samples=[17.77655900000001]
conc=4 bm=16 bn=128 bk=128 sk=1 median=24.520 ms samples=[24.520478999999973]
conc=4 bm=32 bn=32 bk=128 sk=1 median=11.505 ms samples=[11.505040000000015]
conc=4 bm=32 bn=64 bk=128 sk=1 median=37.579 ms samples=[37.578684999999986]
conc=4 bm=32 bn=128 bk=128 sk=1 median=46.486 ms samples=[46.486446]
BEST 4 {'bm': 32, 'bn': 32, 'bk': 128, 'sk': 1, 'median_ms': 11.505040000000015}
conc=5 bm=16 bn=32 bk=128 sk=1 median=80.182 ms samples=[80.18241600000006]
conc=5 bm=16 bn=64 bk=128 sk=1 median=94.947 ms samples=[94.94713600000009]
conc=5 bm=16 bn=128 bk=128 sk=1 median=15.922 ms samples=[15.922151999999999]
conc=5 bm=32 bn=32 bk=128 sk=1 median=43.256 ms samples=[43.25561900000002]
conc=5 bm=32 bn=64 bk=128 sk=1 median=23.605 ms samples=[23.605318999999987]
conc=5 bm=32 bn=128 bk=128 sk=1 median=19.858 ms samples=[19.85752999999998]
BEST 5 {'bm': 16, 'bn': 128, 'bk': 128, 'sk': 1, 'median_ms': 15.922151999999999}
conc=6 bm=16 bn=32 bk=128 sk=1 median=80.110 ms samples=[80.10993799999994]
conc=6 bm=16 bn=64 bk=128 sk=1 median=95.799 ms samples=[95.79853300000006]
conc=6 bm=16 bn=128 bk=128 sk=1 median=15.987 ms samples=[15.987370999999992]
conc=6 bm=32 bn=32 bk=128 sk=1 median=43.207 ms samples=[43.20698200000002]
conc=6 bm=32 bn=64 bk=128 sk=1 median=23.584 ms samples=[23.583591999999996]
conc=6 bm=32 bn=128 bk=128 sk=1 median=19.790 ms samples=[19.789950999999984]
BEST 6 {'bm': 16, 'bn': 128, 'bk': 128, 'sk': 1, 'median_ms': 15.987370999999992}
conc=7 bm=16 bn=32 bk=128 sk=1 median=80.658 ms samples=[80.657629]
conc=7 bm=16 bn=64 bk=128 sk=1 median=96.565 ms samples=[96.56531200000012]
conc=7 bm=16 bn=128 bk=128 sk=1 median=15.989 ms samples=[15.989277000000014]
conc=7 bm=32 bn=32 bk=128 sk=1 median=43.456 ms samples=[43.45628800000001]
conc=7 bm=32 bn=64 bk=128 sk=1 median=23.845 ms samples=[23.84533799999998]
conc=7 bm=32 bn=128 bk=128 sk=1 median=20.238 ms samples=[20.23824]
BEST 7 {'bm': 16, 'bn': 128, 'bk': 128, 'sk': 1, 'median_ms': 15.989277000000014}
conc=8 bm=16 bn=32 bk=128 sk=1 median=81.026 ms samples=[81.02575999999995]
conc=8 bm=16 bn=64 bk=128 sk=1 median=95.400 ms samples=[95.39969900000006]
conc=8 bm=16 bn=128 bk=128 sk=1 median=16.175 ms samples=[16.175329999999978]
conc=8 bm=32 bn=32 bk=128 sk=1 median=43.418 ms samples=[43.41767900000001]
conc=8 bm=32 bn=64 bk=128 sk=1 median=23.697 ms samples=[23.697203000000005]
conc=8 bm=32 bn=128 bk=128 sk=1 median=19.975 ms samples=[19.975342000000005]
BEST 8 {'bm': 16, 'bn': 128, 'bk': 128, 'sk': 1, 'median_ms': 16.175329999999978}
行列積のブロックサイズは以下の関数で決定されている。
測定した結果からパラメータを決定し、関数の入り口に次のコードを追加することで、並列数に応じてブロックサイズを強制するようにした。
constraints = dict(constraints)
if m <= 16:
constraints["block_m"] = 32
constraints["block_n"] = 32
constraints["block_k"] = 128
constraints["split_k"] = 1
else:
constraints["block_m"] = 16
constraints["block_n"] = 128
constraints["block_k"] = 128
constraints["split_k"] = 1
前回と同様に推論性能を測定した。
| Conc. | RPS | Avg Lat.(s) | P99 Lat.(s) | Gen. toks/s | Avg TTFT(s) | P99 TTFT(s) | Avg TPOT(s) | P99 TPOT(s) | Success Rate |
|---|---|---|---|---|---|---|---|---|---|
| 1 | 0.06 | 15.851 | 15.851 | 32.30 | 0.116 | 0.116 | 0.031 | 0.031 | 100.0% |
| 2 | 0.10 | 19.188 | 19.188 | 53.37 | 0.177 | 0.177 | 0.037 | 0.037 | 100.0% |
| 3 | 0.12 | 25.627 | 25.640 | 59.91 | 0.275 | 0.336 | 0.050 | 0.050 | 100.0% |
| 4 | 0.16 | 24.882 | 24.895 | 82.26 | 0.326 | 0.385 | 0.048 | 0.048 | 100.0% |
| 5 | 0.20 | 25.403 | 25.418 | 100.71 | 0.389 | 0.448 | 0.049 | 0.049 | 100.0% |
| 6 | 0.22 | 26.923 | 26.934 | 113.95 | 0.396 | 0.458 | 0.052 | 0.052 | 100.0% |
| 7 | 0.25 | 27.975 | 28.005 | 127.98 | 0.487 | 0.564 | 0.054 | 0.054 | 100.0% |
| 8 | 0.29 | 27.591 | 27.617 | 148.31 | 0.551 | 0.626 | 0.053 | 0.054 | 100.0% |
| 16 | 0.46 | 34.785 | 34.824 | 235.23 | 1.517 | 1.643 | 0.065 | 0.068 | 100.0% |
| 32 | 0.82 | 39.085 | 39.120 | 418.71 | 1.965 | 2.064 | 0.073 | 0.076 | 100.0% |
トークン生成速度をグラフ化すると以下のようになる（緑が最適化後）。
行列積のパラメータを変更することですべての並列数においてトークン生成性能が向上した。TTFTも改善し、RPSが大きく向上している。
まとめ
低並列のときvLLMの推論性能が低かったのはTritonによる行列積のパラメータが最適でなかったためであることがわかった。パラメータを調整することで2~4倍くらい性能が向上した。
llama.cppと比べると並列数が1のときのみ大きく劣っている。今回最適化したmatmul以外にも時間がかかっているカーネルがあるのでそちらも改善の余地があるかもしれない。
バックアップ（コード修正に問題がありうまく最適化できなかったときの結果）
測定した結果からパラメータを決め、簡易的に次のようなコードを追加した。
if tokens_per_expt <= 16:
block_m, block_n, block_k, split_k = 32, 32, 128, 1
elif tokens_per_expt <= 32:
block_m, block_n, block_k, split_k = 16, 128, 128, 1
最適化の結果
前回と同様に推論性能を測定した。
| Conc. | RPS | Avg Lat.(s) | P99 Lat.(s) | Gen. toks/s | Avg TTFT(s) | P99 TTFT(s) | Avg TPOT(s) | P99 TPOT(s) | Success Rate |
|---|---|---|---|---|---|---|---|---|---|
| 1 | 0.06 | 15.963 | 15.963 | 32.07 | 0.218 | 0.218 | 0.031 | 0.031 | 100.0% |
| 2 | 0.10 | 19.529 | 19.529 | 52.43 | 0.177 | 0.177 | 0.038 | 0.038 | 100.0% |
| 3 | 0.12 | 26.022 | 26.035 | 59.00 | 0.545 | 0.686 | 0.050 | 0.050 | 100.0% |
| 4 | 0.16 | 24.700 | 24.701 | 82.91 | 0.584 | 0.584 | 0.047 | 0.047 | 100.0% |
| 5 | 0.08 | 58.950 | 58.966 | 43.41 | 0.730 | 0.847 | 0.114 | 0.115 | 100.0% |
| 6 | 0.10 | 62.247 | 62.283 | 49.32 | 0.815 | 0.941 | 0.120 | 0.121 | 100.0% |
| 7 | 0.10 | 71.727 | 71.762 | 49.94 | 0.911 | 1.032 | 0.139 | 0.140 | 100.0% |
| 8 | 0.13 | 63.062 | 63.125 | 64.88 | 0.882 | 1.149 | 0.122 | 0.122 | 100.0% |
| 16 | 0.21 | 74.561 | 74.605 | 109.81 | 3.317 | 3.554 | 0.139 | 0.145 | 100.0% |
| 32 | 0.33 | 97.719 | 97.785 | 167.48 | 3.381 | 4.189 | 0.185 | 0.189 | 100.0% |
グラフ化すると以下のようになる（緑が最適化後）。
並列数が4以下の場合の推論性能が大きく向上した。それ以外の並列数ではなぜか性能が下がってしまったので、もう一度きれいな状態からやり直してみるつもり。
Tritonのmatmul以外にも時間がかかっているカーネルがあるのでそちらも最適化の余地があるかもしれない。
Tritonについて気がついたこと
vLLM v0.13.0ではTriton v3.5.0のtriton_kernelsを使用している。Triton v3.5.1でmatmulのブロックサイズの決定ロジックに変更が入っているので、Tritonのバージョンを変えるだけでまたパフォーマンスが変わるかもしれない。これも試してみるつもり。 → 追記：この変更はNvidia GPUのみが対象、かつv3.5.1で入った変更ではなさそう。
Discussion