【dlshogi】TensorRTの推論処理の最適化でNPSが1.3倍になった件

先日の世界コンピュータ将棋選手権の会場で「あすとら将棋」さんから、GPUへのデータ転送を推論と並列化すると1割くらい速くなるという話を伺って、さっそく実験してみた。

データ転送の並列化

これまでは、一つのGPUを複数スレッドで共有して、1つのスレッドが推論中はロックして排他的に利用していた。
しかし、推論中にもデータ転送は行うことができるため、次のスレッドは先にデータ転送を行って、純粋に推論部分だけを排他すればよい。

また、GPT-5.5で実装したところ、圧縮した特徴量をfloat32にするunpack処理をCUDAで実装しているが、その部分も並列化して問題ないことに気づいた。
初期のCUDAのプログラミングモデルでは、同時に実行できるカーネルは一つだけだったが、Compute Capability 2.0以上では、ストリームが分かれていれば、複数カーネルをSMに分配して実行できるようになっている。

測定結果

benchmark.pyを使用して、floodgateから抽出した100局面で、1秒探索した際のNPSは以下の通り。

5回測定した平均値の100局面の統計
H100 PCI 1枚、3スレッド

	変更前	変更後
平均	10591	11751
中央値	10625	11818
最大値	12468	14168
最小値	9076	10678

NPSが平均で、1.11倍になった。

実行コンテキストによる非同期化

上記を実装する際、GPT-5.5が実行コンテキスト(IExecutionContext)を使うと、推論も排他不要である指摘をしてくれた。

NVIDIAのドキュメントにも、1 つの重みセットを複数の重複する推論タスクに使用できると記載がある。
Developer Guide :: NVIDIA Deep Learning TensorRT Documentation

IExecutionContextを使うことで、各スレッドは非同期で実行できる。
それにより、ロック処理を外すことができた。

github.com

IExecutionContextを使う場合、推論プロファイルの作成方法が変わるため、シリアライズファイルもスレッド数に応じて別になっている。

測定結果

	変更前	変更後
平均	10591	13953
中央値	10625	13855
最大値	12468	16660
最小値	9076	13227

NPSが平均で、1.32倍になった。

強さ

NPSが強さに反映されるか連続対局で測定した。

中終盤互角局面集を使用
基準として氷彗8スレッド(hayabusa-8th)を加えている
H100 PCI x 2、3スレッド

   # PLAYER            :  RATING  ERROR  POINTS  PLAYED   (%)  CFS(%)    W    D    L  D(%)
   1 hayabusa-8th      :    34.3   31.7   125.5     221    57      85  120   11   90     5
   2 pre68_40x512_e    :     4.8   32.5   114.0     223    51      94  108   12  103     5
   3 pre68_40x512      :   -39.1   31.9    95.5     226    42     ---   91    9  126     4

pre68_40x512_eが実行コンテキストに対応した版で、pre68_40x512は変更前の版である。
測定数は少ないが速報値としては、R+33.9になっている。

2GPU、3スレッドで測定している。
変更前はスレッド数 3が最適だったが、実行コンテキスを使う場合、3よりも増やすことで、さらに強くできる可能性がある。

ただし、他のスレッドの結果を待たずに実行するため、ツリーの成長の仕方が変わりNPSが上がるだけで強くならない可能性もあるので、強さを計測してチューニングする必要がある。

まとめ

GPUへのデータ転送と特徴量unpack処理を推論と並列化し、H100環境でNPSが平均1.11倍向上した。
また、TensorRTのIExecutionContextを利用して推論自体も並列化し、ロック不要化によってNPSは平均1.32倍まで向上した。
連続対局の速報値では実行コンテキスト対応版が約+34 Elo強くなった。

この改良はすでにmasterブランチにマージしている。