TensorRT 10 × Blackwell 移行ガイド【後編】5D grid_sample をカスタムプラグインで通す

こんにちは！

「TensorRT 10 × Blackwell 移行ガイド」シリーズもいよいよ最終回です。前編では TensorRT 化の基本と「5D（volumetric）の grid_sample はビルドが拒否される」という壁を、中編では「ビルドが通っても正しいとは限らない」沈黙劣化の数々を扱いました。

後編は、前編で保留にした最大の宿題に挑みます。TensorRT がネイティブに対応していない 5D grid_sample を、カスタムプラグインで通す話です。ゴールは明確で、前編でこのエラーを出したのと同じモデル・同じ入力から export した ONNX を、5D の GridSample ノードだけプラグイン用ノードへ置き換えて、最後までビルド・実行することです。

INetworkDefinition::addGridSample: Error Code 3: API Usage Error
(Parameter check failed, condition: input.getDimensions().nbDims == 4. ...)

先に結論を言うと、通りました。PyTorch の F.grid_sample（5D）と最大絶対誤差 7.2e-07 で一致するエンジンができています。ただしそこまでの道には、Blackwell 世代ならではの壁がいくつも立っています。本記事ではその壁と越え方を、実機ログつきで順に共有します。

この記事の対象環境 WSL2 Ubuntu 24.04 ／ NVIDIA RTX PRO 4000 Blackwell・GeForce RTX 5060 Ti（ともに Compute Capability 12.0 = sm_120）／ PyTorch 2.11.0 + cu128 ／ TensorRT 10.16.1.11 ／ システム nvcc は CUDA 12.0（あえて古いまま使います。理由は本文で）。 題材のプラグインは OSS の grid-sample3d-trt-plugin（Apache License 2.0）です。リポジトリに正式リリースは無いため、本記事で使用したコミット（f964750）を明記しておきます。改変版のソースや .so を再配布する場合は、同ライセンスの条件（ライセンス文・著作権表示の同梱など）に従ってください。本記事のコード・ログ・数値はすべて実機の実測です。

第1部: プラグインという仕組み — TensorRT は「知らない op」をどう扱うか

前編で見たとおり、TensorRT の ONNX parser は、GridSample ノードの入力が 5 次元だとその場でビルドを拒否します。ネイティブ実装が存在しないからです。

こういうときのための正式な拡張機構がカスタムプラグインです。仕組みはこうなっています。

1. 演算の CUDA カーネルと、TensorRT へのアダプタ（プラグインクラス）を C++ で書き、共有ライブラリ（.so）にする
2. プロセス起動時にその .so をロードすると、プラグインが TensorRT の plugin registry に登録される
3. ONNX parser は、知らない op_type に出会うと plugin registry を検索し、同名のプラグインが見つかればそれをレイヤとして組み込む

つまりやることは 2 つです。(a) 5D grid_sample のプラグインを Blackwell で動く形にビルドする。(b) ONNX の GridSample ノードを、プラグインの op 名に張り替える。

幸い、(a) のカーネルをゼロから書く必要はありません。TensorRT 8 の時代から使われてきた OSS 実装 grid-sample3d-trt-plugin があり、5D grid_sample の CUDA カーネルとプラグインクラスが揃っています。問題は、これを TensorRT 10 × Blackwell の環境でビルドし直すところに集中しています。

第2部: ビルドの壁 3 連発

まず、プラグイン本体を本記事と同じソースで取得しておきます。

git clone https://github.com/SeanWangJS/grid-sample3d-trt-plugin.git
cd grid-sample3d-trt-plugin
git checkout f964750

壁その 1: pip で入れた TensorRT には、開発用ヘッダが入っていない

プラグインのビルドには TensorRT の C++ ヘッダ（NvInfer.h など）とリンク用の libnvinfer.so が必要です。ところが、pip の tensorrt-cu12 パッケージに入っているのは実行用の .so だけで、ヘッダは一切含まれていません。

NVIDIA が案内する正規の方法は、tar / Debian / RPM などヘッダを含む配布形態を使うことです。以下は「pip 環境を二重化したくない」場合に当社環境で実証した回避策、という位置づけで読んでください。NVIDIA/TensorRT の OSS リポジトリから、pip 版とバージョンの合う公開ヘッダだけを取ってくる方法です。

# ヘッダ: OSS リポの対応タグ（v10.16）を blob:none で軽量 clone
git clone --depth 1 --branch v10.16 --filter=blob:none \
    https://github.com/NVIDIA/TensorRT.git ~/trt_headers
# → 使うのは ~/trt_headers/include/ だけ

# リンク用 lib: pip が入れた実行用 .so へ symlink を張る
#（リンカは -lnvinfer で libnvinfer.so を探すので、バージョン付きの実体へ繋ぐ）
# .so の場所は環境ごとに違うので、Python に聞くのが確実
TRT_LIB_DIR="$(python -c "import pathlib, tensorrt_libs; print(pathlib.Path(tensorrt_libs.__file__).parent)")"
mkdir -p ~/trt_link
ln -sf "$TRT_LIB_DIR/libnvinfer.so.10" ~/trt_link/libnvinfer.so

TensorRT の公開 API・ABI はセマンティックバージョニングに従うため、v10.16 タグのヘッダと pip の 10.16.1.11 という近接した組み合わせは理にかなっており、実機でビルド・動作を確認しました（異なる minor 系列を混ぜる場合、特に新しいヘッダの API を古いライブラリへリンクする向きは成立しない可能性があるので、実際の組み合わせで確認してください）。

もうひとつ、元リポジトリの CMakeLists.txt は、この 2 つの場所をそのままでは参照してくれません。次の変更を入れて、外から渡せるようにします（CUDA_ARCHITECTURES の行は次の壁で説明します）。

# 追加: 外部から渡した TensorRT のヘッダ・lib を使う
target_include_directories(${PROJECT_NAME} PRIVATE
    "./src" ${CUDAToolkit_INCLUDE_DIRS} ${TensorRT_INCLUDE_DIR})
target_link_directories(${PROJECT_NAME} PRIVATE ${TensorRT_LIB_DIR})
target_link_libraries(${PROJECT_NAME} PRIVATE nvinfer CUDA::cudart)

ビルドコマンドはこうです。

mkdir build && cd build
cmake .. -DCMAKE_BUILD_TYPE=Release \
  -DTensorRT_INCLUDE_DIR="$HOME/trt_headers/include" \
  -DTensorRT_LIB_DIR="$HOME/trt_link"
cmake --build . --parallel

壁その 2: 手元の nvcc が Blackwell を知らない

次の壁が本記事の山場です。プラグインの CUDA カーネルをコンパイルしようとすると、こういう状況に直面します。

• Blackwell（sm_120）向けのネイティブコードを吐けるのは CUDA 12.8 以降の nvcc
• ところが手元のシステム nvcc は CUDA 12.0（nvcc --list-gpu-arch の上限が compute_90）
• CUDA toolkit を入れ替えれば解決するが、既存環境への影響が大きい

ここで効くのが、PTX の forward compatibility（前方互換） です。CUDA のコンパイルには 2 段階あります。nvcc がソースを PTX（仮想アーキテクチャ向けの中間表現）へ落とす段階と、PTX を実 GPU 向けのネイティブコード（SASS）へ落とす段階です。後者は GPU ドライバが実行時に JIT コンパイルすることもできます。

つまり、「nvcc では compute_90 の PTX までを作り、sm_120 への最終変換は Blackwell 対応済みのドライバに任せる」という分担にすれば、古い nvcc のままでもビルドできます。CMake ではこう書きます。

# compute_90 の PTX のみを埋め込む（"90" だと SASS も作ろうとして失敗する）。
# 実行時に driver が PTX → sm_120 へ JIT コンパイルする
set_target_properties(${PROJECT_NAME} PROPERTIES CUDA_ARCHITECTURES "90-virtual")

ビルドした .so に本当に PTX しか入っていないことは、cuobjdump で確認できます。

$ cuobjdump libgrid_sample_3d_plugin.so | grep -E "arch|Fatbin" | sort | uniq -c
      1 Fatbin ptx code:
      1 arch = sm_90

ネイティブコード（Fatbin elf code）の行がなく、PTX が 1 本だけ。この状態で sm_120 の GPU にロードすると、ドライバが JIT コンパイルして動きます（後述の数値照合が通っていることが、JIT が正しく機能した証明になります）。初回実行時には PTX の JIT コンパイル時間が発生します（本記事ではその時間は計測していません）。

余談ですが、JIT の検証手段として CUDA_FORCE_PTX_JIT=1（全カーネルを PTX からの JIT に強制する環境変数）も知られています。ただし当環境で試したところ、プロセス全体に効くため、SASS しか持たないカーネルを含む PyTorch 側が no kernel image is available で先に落ちました。PyTorch と同居するスクリプトでは使えないので、プラグイン単体の確認は上記の cuobjdump と数値照合で行うのが実務的です。

壁その 3: deprecated 警告の山

ビルドログには、IPluginV2DynamicExt 系 API への deprecated 警告が並びます。この API は TensorRT 10.0 で deprecated になっており、新規実装には V3 系が推奨されています。本記事の TensorRT 10.16.1.11 では警告止まりで、実際にビルド・実行できました。ただし TensorRT 11 では V2 系プラグイン API 自体が削除されているため、このプラグインをそのまま 11 系へ持っていくことはできません。今回は OSS 実装を無改変で活かす方針なので、警告は許容して先へ進みます（新規にプラグインを書き起こすなら V3 系 API を使ってください）。

ここまでの 3 つの壁を整理すると、次のとおりです。

第3部: ONNX ノードを張り替えて、ビルドを通す

プラグインができたら、次は (b) の張り替えです。前編で使った 5D のダミーモデル VolumetricWarp を、まったく同じように export します。

class VolumetricWarp(nn.Module):
    """5D（volumetric）grid_sample を含む最小モジュール（前編と同じ）"""
    def forward(self, vol, grid):
        return F.grid_sample(vol, grid, align_corners=False)

torch.manual_seed(0)
vol = torch.randn(1, 32, 16, 64, 64, device="cuda")
grid = torch.rand(1, 16, 64, 64, 3, device="cuda") * 2 - 1
model = VolumetricWarp().eval().cuda()

torch.onnx.export(model, (vol, grid), "vol_warp.onnx",
                  input_names=["vol", "grid"], output_names=["y"],
                  opset_version=20, dynamo=True)

この ONNX の GridSample ノードを、プラグインの op 名 GridSample3D へ張り替えます。属性もプラグインが期待する形式（int 型 3 つ）に置き換えます。

import onnx
from onnx import helper

def swap_gridsample(src_onnx, dst_onnx):
    """ONNX の GridSample(5D) ノードをプラグイン op『GridSample3D』へ張り替える"""
    m = onnx.load(src_onnx, load_external_data=False)
    n = 0
    for node in m.graph.node:
        if node.op_type != "GridSample":
            continue
        ins, outs, name = list(node.input), list(node.output), node.name
        del node.attribute[:]
        node.op_type = "GridSample3D"
        node.domain = ""      # 非標準 op → parser が plugin registry を検索しにいく
        node.attribute.extend([
            helper.make_attribute("interpolation_mode", 0),   # 0 = linear
            helper.make_attribute("padding_mode", 0),         # 0 = zeros
            helper.make_attribute("align_corners", 0),        # 0 = False
        ])
        del node.input[:]
        node.input.extend(ins)
        del node.output[:]
        node.output.extend(outs)
        node.name = name
        n += 1
    onnx.save(m, dst_onnx)
    return n

なお、この swap_gridsample はグラフ内のすべての GridSample を無条件に置き換えます。本記事のダミーは「5D の GridSample が 1 個だけ」という前提なので、それを検査してから使います（4D と 5D が混在する実モデルで使う場合は、shape 推論で入力の次元数を確認するか、対象ノード名を指定できるようにしてください）。

count = swap_gridsample("vol_warp.onnx", "vol_warp_plugin.onnx")
if count != 1:
    raise RuntimeError(f"expected exactly one GridSample node, found {count}")

もうひとつ、プラグインのロードには重要な注意があります。本記事の構成では、プラグインの共有ライブラリを plan へシリアライズしていないため（TensorRT には version-compatible engine へプラグインを同梱する仕組みもありますが、今回は使っていません）、.so のロードはビルド時だけでなく、保存済みエンジンを別プロセスで deserialize する前にも必要です。共通のロード関数を作っておきます。

import ctypes
from pathlib import Path

_PLUGIN_HANDLES = []   # GC でアンロードされないよう、プロセス中は参照を保持する

def load_grid_sample_plugin(plugin_path):
    # プラグイン .so を RTLD_GLOBAL でロードする。ロードの副作用（ライブラリ内の
    # 登録処理）で、プラグインが TensorRT の plugin registry に登録される。
    # ビルド前と deserialize 前の両方で呼ぶこと
    handle = ctypes.CDLL(str(Path(plugin_path).resolve()), mode=ctypes.RTLD_GLOBAL)
    _PLUGIN_HANDLES.append(handle)

ビルド側はこうなります。GridSample3D を registry に登録しているのは上のロード処理で、init_libnvinfer_plugins は NVIDIA 標準プラグイン群の初期化 API です（今回のカスタムプラグインの登録主体ではありません）。

import tensorrt as trt

def build_engine_with_plugin(onnx_path, engine_path, fp16=False):
    load_grid_sample_plugin("./libgrid_sample_3d_plugin.so")
    trt.init_libnvinfer_plugins(TRT_LOGGER, "")   # NVIDIA 標準プラグインの初期化（念のため）

    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, TRT_LOGGER)
    if not parser.parse_from_file(onnx_path):
        errs = [str(parser.get_error(i)) for i in range(parser.num_errors)]
        raise RuntimeError("ONNX parse failed:\n" + "\n".join(errs))
    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 2 << 30)
    if fp16:
        config.set_flag(trt.BuilderFlag.FP16)
    plan = builder.build_serialized_network(network, config)
    if plan is None:
        raise RuntimeError("build failed")
    with open(engine_path, "wb") as f:
        f.write(plan)
    return engine_path

関数を定義したら、張り替え済みの ONNX からエンジンを作ります。

build_engine_with_plugin("vol_warp_plugin.onnx", "vol_warp_plugin.plan", fp16=False)

実行側（保存済みエンジンを使うプロセス）では、deserialize の前にロードします。

load_grid_sample_plugin("./libgrid_sample_3d_plugin.so")
runner = TRTRunner("vol_warp_plugin.plan")   # deserialize はロードの後

そして——前編で拒否されたのと同じ計算グラフ（プラグインノードへの置き換えのみ）のビルドが通ります。実行して、PyTorch の F.grid_sample（5D）と突き合わせた結果がこちらです（照合の作法は前編・中編のとおり、export に使った同じモデル・同じ入力・fp32 基準です）。

計測条件: GeForce RTX 5060 Ti ／ 固定形状 [1,32,16,64,64]・batch=1 ／ 両者とも fp32 ／ ウォームアップ 30 回・計測 200 回・CUDA Event の median ／ TensorRT 側は入出力バッファを事前確保し enqueue 区間のみ ／ 1 回のビルドでの結果です。約 15% の差はモデルや GPU で容易に変わるので、「当環境のこの単体演算ではこうだった」として読んでください。

fp32 基準に対して非常に小さい誤差に収まりました。compute_90 の PTX を Blackwell のドライバが JIT した結果としてこの数値が出ている、というのが本記事の要点です。速度は単体ではわずかに速い程度ですが、本当の価値は「これまで PyTorch に残すしかなかった 5D grid_sample 入りのモジュールを、まるごと 1 個の TensorRT エンジンにできる」ことにあります。torch と TensorRT の境界そのものが消えるからです。

第4部: 中編の教訓をここでも — fp16 を受理する構成を検証する

さて、ここで終わると中編の教訓が泣きます。「ビルドが通っても、正しいとは限らない」 —— プラグインにも同じ検証を通します。

この OSS プラグインは fp16 カーネルも実装しており、supportsFormatCombination（TensorRT に「この精度・形式を受け入れるか」を答えるメソッド）は fp32 と fp16 の両方を受理する実装になっています。そこで、fp16 を受理する元実装のままの .so を別にビルドし、FP16 tactic を許可してエンジンを作って実測しました（プラグイン内部で実際に fp16 経路が選ばれたことのログ確認まではしていないため、正確には「fp16 を受理する構成のエンジン」の結果です）。

fp16 を受理する構成では、fp32 比で誤差が 5 桁悪化しました。nRMSE 2.1% は単一演算の fp16 丸め単位より大幅に大きく、少なくとも本用途では許容できない水準です（補間では座標の量子化・重み計算・積和などで誤差が増幅し得ますが、どの処理で増幅しているかまでは特定していません）。さらに私たちが実務モデルの内部で fp16 を受理する構成を使った際には、出力が数百のオーダーでずれる完全な数値破綻も経験しています（入力の分布次第で被害は大きく変わります。これも「乱数入力の結果を信用するな」の一例です）。

対処 — プラグイン側で fp16 を「受理しない」

対処はシンプルで、プラグインの supportsFormatCombination から fp16 の受理を外し、fp32 専用にすることです。

bool GridSample3DPlugin::supportsFormatCombination(int32_t pos,
                                                   PluginTensorDesc const* inOut,
                                                   int32_t nbInputs,
                                                   int32_t nbOutputs) noexcept {
    assert(nbInputs == 2 && nbOutputs == 1 && pos < (nbInputs + nbOutputs));
    bool condition = inOut[pos].format == TensorFormat::kLINEAR;
    // fp32-only: fp16 を受理する構成で許容できない誤差を観測したため kHALF を受理しない。
    // FP16 tactic を許可したエンジン内でも、TensorRT が本レイヤの前後に
    // reformat（型変換）を自動挿入し、この演算だけ fp32 で実行してくれる
    condition &= inOut[pos].type == DataType::kFLOAT;
    condition &= inOut[pos].type == inOut[0].type;
    return condition;
}

ここが面白いところで、こうしてもエンジン全体を fp32 にする必要はありません。FP16 tactic を許可してビルドすると、TensorRT はプラグインの申告を見て、このレイヤの前後にだけ型変換（reformat）を自動で挿入できます。結果として「grid_sample は fp32、周辺のレイヤは fp16 候補から選択」という構成が 1 個のエンジンの中で成立し得ます。ただし reformat のコストや、実際に選ばれた精度・エンジン全体の数値は、完成したエンジンで改めて検証してください（前編・中編で書いたとおり、「許可した」と「実際に使われた」は別物です）。

まだ確かめていないこと

• fp16 を受理する構成での誤差の根本原因 — プラグイン内部で実際に fp16 経路が選ばれたことのログ確認、および誤差がどの処理で増幅しているかの特定はしていません（本記事は「検証して、危なければ受理しない」という運用側の対処に留めています）
• JIT コンパイルの初回コスト — 初回ロード時の JIT 時間は計測していません。デプロイ形態によっては起動時間に効く可能性があります
• 新世代プラグイン API（V3 系）への移行 — 本記事は旧 API（IPluginV2DynamicExt）のまま動かしています。TensorRT の将来のメジャーバージョンでは旧 API が削除される可能性があります
• 可変形状（dynamic shape）での動作 — 本記事の検証はすべて固定形状です

まとめ — シリーズを締めて

後編の要点です。

1. TensorRT が知らない演算は、カスタムプラグイン + ONNX ノードの張り替えで通せる
2. pip の TensorRT にはヘッダが無い → OSS リポの対応タグから include を取り、lib は pip の .so へ symlink
3. nvcc が新しい GPU を知らなくても、PTX（仮想アーキテクチャ）でビルドして driver の JIT に任せる手がある（CUDA_ARCHITECTURES "90-virtual"）
4. プラグインにも中編の検証をそのまま適用する。許容できない誤差が出る型は supportsFormatCombination で受理対象から外す。TensorRT が reformat を自動挿入できるので、周辺は fp16 候補を残したまま該当演算だけ守れる（実際に選ばれた精度と数値は完成エンジンで確認する）

シリーズ 3 本を通して伝えたかったことを一言でまとめると、「TensorRT 化は『変換して速くなったら終わり』ではなく、export・ビルド・実行・数値のすべてを実測で確かめて初めて完成する」ということです。Blackwell という新しい世代はその検証の重要性を一段引き上げましたが、本シリーズの手順とチェックリストがそのまま皆さまの移行作業の地図になれば幸いです。

それでは、また次回、お会いしましょう！

参考資料（一次情報・主要ソース）

• SeanWangJS/grid-sample3d-trt-plugin（5D grid_sample プラグインの OSS 実装・Apache-2.0。使用コミット f964750）
• NVIDIA/TensorRT（OSS リポジトリ・公開ヘッダの取得元）
• NVIDIA TensorRT ドキュメント（プラグイン / Developer Guide）
• CUDA C++ Programming Guide: PTX と JIT コンパイル（forward compatibility）
• NVIDIA/TensorRT issue #3890: 5D GridSample のネイティブ対応要望
• 当社関連記事: 前編 / 中編 / 2026年 NVIDIA GPU 一発検索ツール