From 680d37eedfc5027f5df4d379e39c2653977b2487 Mon Sep 17 00:00:00 2001 From: zsqdx Date: Tue, 28 Jul 2026 11:38:52 -0700 Subject: [PATCH] Default TensorRT transformer trunks to NHWC --- cpp/configs/gtp_example.cfg | 4 ++++ cpp/neuralnet/onnxmodelbuilder.cpp | 4 ++-- cpp/neuralnet/trtbackend.cpp | 12 ++++++------ 3 files changed, 12 insertions(+), 8 deletions(-) diff --git a/cpp/configs/gtp_example.cfg b/cpp/configs/gtp_example.cfg index 4ae1f0f46e..b1af90b473 100644 --- a/cpp/configs/gtp_example.cfg +++ b/cpp/configs/gtp_example.cfg @@ -410,6 +410,10 @@ searchFactorWhenWinningThreshold = 0.95 # ------------------------------ # These only apply when using the TENSORRT version of KataGo. +# Transformer trunks use NHWC by default. Set this to false only after benchmarking +# the particular GPU and TensorRT version. +# trtTransformerNHWC = true + # For one GPU: optionally uncomment this option and change if the GPU to # use is not device 0. # trtDeviceToUse = 0 diff --git a/cpp/neuralnet/onnxmodelbuilder.cpp b/cpp/neuralnet/onnxmodelbuilder.cpp index b172057bba..6469237bdb 100644 --- a/cpp/neuralnet/onnxmodelbuilder.cpp +++ b/cpp/neuralnet/onnxmodelbuilder.cpp @@ -15,8 +15,8 @@ using namespace std; namespace { // Builder that accumulates ONNX nodes and initializers into a single GraphProto, handing back -// tensor names as it goes. Tensors are float32; the trunk runs NCHW by -// default, or channel-last NHWC when transformerNHWC is set (see buildBlockStack / buildConv). +// tensor names as it goes. Tensors are float32; transformerNHWC selects a channel-last NHWC trunk +// instead of NCHW (see buildBlockStack / buildConv). struct Builder { onnx::GraphProto* graph; int nnXLen; diff --git a/cpp/neuralnet/trtbackend.cpp b/cpp/neuralnet/trtbackend.cpp index 8af6e623a8..775eb5fc82 100644 --- a/cpp/neuralnet/trtbackend.cpp +++ b/cpp/neuralnet/trtbackend.cpp @@ -94,7 +94,6 @@ ComputeContext* NeuralNet::createComputeContext( ConfigParser& cfg) { (void)gpuIdxs; (void)logger; - (void)loadedModel; ComputeContext* context = new ComputeContext(); context->nnXLen = nnXLen; @@ -105,11 +104,12 @@ ComputeContext* NeuralNet::createComputeContext( // nvonnxparser (the default). trtDisableOnnx=true falls back to the hand-built ModelParser, which // supports convnets only (transformer models will error in createComputeHandle). context->useOnnx = !(cfg.contains("trtDisableOnnx") ? cfg.getBool("trtDisableOnnx") : false); - // ONNX transformer emitter layout. Default is NCHW (genuine channel-major attention/FFN): benchmarks - // show it matches or slightly beats the NHWC bubble path at the saturated throughput operating point - // KataGo runs at, and it's the simpler graph. trtTransformerNHWC=true opts into the NHWC path (whole - // trunk channel-last with NCHW<->NHWC conversions around it), which wins on single-stream latency. - context->transformerNHWC = cfg.contains("trtTransformerNHWC") ? cfg.getBool("trtTransformerNHWC") : false; + // Transformer trunks default to NHWC, which benchmarks at least as fast as NCHW across tested GPUs + // and TensorRT versions. Explicit false remains available for hardware-specific tuning. Normalize + // convnets to false since the ONNX builder ignores this setting for models without transformers. + context->transformerNHWC = + (cfg.contains("trtTransformerNHWC") ? cfg.getBool("trtTransformerNHWC") : true) && + NeuralNet::getModelDesc(loadedModel).hasAnyTransformerBlocks(); // Debugging: if set, the ONNX-emitter path dumps the emitted ONNX model and the built engine's // per-layer info (precision/format/tactic, via a detailed-profiling build + IEngineInspector) into // this directory. Files are disambiguated by board size, FP16/FP32, and exact/max NN-length so the