OpenDLSS-NR、DLSS 5のNRネットワークをVulkanでビット一致再実装

maanHimselfは、NVIDIAのDLSS 5 Neural RenderingネットワークをVulkanで再実装したリポジトリ「OpenDLSS-NR」をGitHubで公開しました。DLSS-NR build 310.8.0と同じ71ブロックのネットワークをテンソルコア上のFP8で実行し、最終画像だけでなく75のブロック境界すべてがバイト単位で一致するとしています。

NVIDIAはGPUを含むチップ、コンピューターシステム、ソフトウェアを開発する企業で、RTXによるグラフィックス技術を提供しています。

ネットワークの構成

OpenDLSS-NRが実装するのは、エンジンが描画したフレームを再レンダリングする生成型ニューラルレンダリングネットワークです。入力は1枚のレンダリングフレーム(LDRプロキシ)、3レーンのガウスノイズ、前フレームの再投影出力、5個の条件スカラーで、ピクセルごとにRGB残差と時間ブレンド用ロジット1つの計4つのf32チャンネルを出力します。入力と出力は同じ解像度で、アップスケーラーではありません。DLSS-SRは別ネットワークのため対象外です。

項目	内容
構造	U-Net型シフトウィンドウTransformer+最下部のグローバルViT
ブロック数	71(6段階のプーリング)
数値形式	FP8(E4M3)活性化、FP16累積
重み容量	141 MiB
ディスパッチ数	全解像度で241

ホスト側はC++20、参照ルートはcooperative-matrix FP8 GEMMなどのGLSLカーネル、高速ルートはPythonで生成するPTXカーネル(mma.sync E4M3、split-K GEMM、ストリーム化グローバルアテンションなど)で構成されます。どのチューニングスイッチでも出力はバイト単位で同一に保たれます。

性能

RTX 4070 SUPERで1フレームあたりネットワーク全体を実行し、40フレーム中の最小値を計測した結果です。

解像度	時間
768x768	2.8 ms
1920x1080	7.8 ms
2560x1440	12.6 ms
3840x2160	29.3 ms

ports/browser-webgpu/には独立したWebGPU版があり、テンソルコアもFP8もブロック間の融合も使わずに同じキャプチャとビット一致します。リポジトリはWebGPU版について「512x512で72 ms、こちらでは2.7 ms」と記しています。

要件とビルド手順

• Windows、NVIDIA Ada以降のGPU
• VK_KHR_cooperative_matrix、VK_NV_cooperative_matrix2、VK_EXT_shader_float8、VK_NV_cuda_kernel_launchを公開するドライバー
• Visual Studio 2022以降のC++ x64ツールセット、git、Python 3、Node.js+npm、Pillow
• Vulkan SDKのインストールは不要

powershell -File scripts\fetch_tools.ps1 [-Npm]
powershell -File scripts\build.ps1
powershell -File scripts\fetch_filament.ps1
powershell -File scripts\build_filament.ps1
powershell -File scripts\build_demo.ps1
build\dlss5vk.exe bench --model --width 768 --height 768
build\dlss5vk.exe profile --model --width 768 --height 768
build\dlss5vk.exe parity --model --fixture
build\dlss5vk.exe verify --model --fixture

デモにはFilament v1.77.0が必要で、ビルドに約15分、ビルドツリーは約6 GBです。verifyにはinputFeaturesとブロック0の参照が必要です。

利用条件

重みは利用者がmanifest.jsonを含むモデルディレクトリとして用意します。リポジトリはそのディレクトリを生成する手段も、元実装のキャプチャ(フィクスチャ)も含みません。対応するのは310.8.0の71ブロック構成のみで、ブロック数が異なるモデルはロード時に拒否されます。プロジェクトはNVIDIAと提携・承認・支援関係になく、ライセンスはMITです。GitHubでは577スター・58フォーク、Hacker Newsでは82ポイント・52コメントを集めています。