CUDA and ROCm optimizations

LICENSE

@@ -1,5 +1,5 @@
 The code in this repository currently relies on several other libraries, parts of libraries,
-or external files: clblast, composable_kernel_fmha, cudnn-frontend, filesystem-1.5.8, half-2.2.0,
+or external files: clblast, composable_kernel_fmha, cudnn-frontend, cutlass, filesystem-1.5.8, half-2.2.0,
 httplib, katagocoreml (which itself vendors components from Apple's coremltools and the FP16
 library), macos (Swift CMake modules), mozilla-cacerts, nlohmann_json, sgfmill, onnx, and
 tclap-1.2.5. For the licenses for those libraries and/or files, see the individual readmes and/or

cpp/CMakeLists.txt

@@ -563,6 +563,14 @@ set(USE_BIGGER_BOARDS_EXPENSIVE 0 CACHE BOOL "Allow boards up to size 50. Compil
 set(USE_CACHE_TENSORRT_PLAN 0 CACHE BOOL "Use TENSORRT plan cache. May use a lot of disk space. Only applies when USE_BACKEND is TENSORRT.")
 mark_as_advanced(USE_CACHE_TENSORRT_PLAN)
 
+# CUDA-backend opt-outs of optional compiled-in paths, replicating what a build gets when the
+# dependency is unavailable (no vendored CUTLASS / MSVC, cuDNN older than 8.9.3). Mainly for
+# testing those fallback configurations on a machine whose toolchain would otherwise enable them.
+set(NO_CUTLASS_FUSED_FFN 0 CACHE BOOL "CUDA backend: build without the CUTLASS fused FFN kernel even if CUTLASS is available.")
+set(NO_CUDNN_SDPA 0 CACHE BOOL "CUDA backend: build without the cudnn frontend graph SDPA attention path.")
+mark_as_advanced(NO_CUTLASS_FUSED_FFN)
+mark_as_advanced(NO_CUDNN_SDPA)
+
 # ---------- Auto-fetch missing third-party deps (e.g. zlib on a Windows ROCm/TheRock install that
 # doesn't ship a linkable zlib) via a local vcpkg clone in the build tree.
 # Defaults ON only on Windows ROCm builds, where the ROCm/TheRock toolchain genuinely may have no
@@ -682,7 +690,6 @@ if(USE_BACKEND STREQUAL "CUDA")
 
   enable_language(CUDA)
 
-  set(CUDA_STANDARD 11)
   # cudabackend.cpp, cudautils.cpp, and cudahelpers.cu are thin wrappers around
   # cudaandrocmbackend.inc, cudaandrocmutils.inc, and cudaandrocmhelpers.inc, which are shared
   # with the ROCm backend (tracked via compiler depfiles, so they do not appear here).
@@ -691,6 +698,26 @@ if(USE_BACKEND STREQUAL "CUDA")
     neuralnet/cudautils.cpp
     neuralnet/cudahelpers.cu
   )
+  # Optional vendored CUTLASS enables the fused transformer FFN kernel. The fused kernel is
+  # built on the DualGemm extension shipped in CUTLASS's example 45, so both the include tree
+  # and that example directory must be present. CUTLASS 4.x requires CUDA Toolkit 11.4+ and
+  # does not build under MSVC, so those builds skip the kernel rather than break.
+  if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/external/cutlass/include/cutlass/cutlass.h
+     AND EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/external/cutlass/examples/45_dual_gemm/device/dual_gemm.h
+     AND CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 11.4
+     AND NOT MSVC
+     AND NOT NO_CUTLASS_FUSED_FFN)
+    message(STATUS "Found vendored CUTLASS at external/cutlass, enabling the fused FFN kernel.")
+    set(KATAGO_USE_CUTLASS_FUSED_FFN 1)
+    set(NEURALNET_BACKEND_SOURCES ${NEURALNET_BACKEND_SOURCES} neuralnet/cudafusedffn.cu)
+    # CUTLASS requires C++17. CUDA_STANDARD is a target property that is silently ignored when
+    # set per-source, and nvcc before CUDA 12 defaults to C++14, so pass -std explicitly (no
+    # conflict: CMAKE_CUDA_STANDARD is unset, so CMake adds no -std flag of its own).
+    set_source_files_properties(neuralnet/cudafusedffn.cu PROPERTIES
+      COMPILE_OPTIONS "-std=c++17;--expt-relaxed-constexpr")
+  else()
+    message(STATUS "Vendored CUTLASS not found, not usable with this toolchain, or disabled via NO_CUTLASS_FUSED_FFN; building without the fused FFN kernel.")
+  endif()
   # The vendored header-only cudnn_frontend library trips GCC's -Wnull-dereference from its -O2
   # interprocedural analysis (no source location, so an in-file #pragma can't suppress it). It's
   # benign third-party code; silence it just for the one source that includes cudnn_frontend.h.
@@ -1042,6 +1069,7 @@ add_executable(katago
   command/commandline.cpp
   command/analysis.cpp
   command/benchmark.cpp
+  command/benchmarknn.cpp
   command/contribute.cpp
   command/dumponnx.cpp
   command/evalsgf.cpp
@@ -1071,6 +1099,16 @@ if(USE_BACKEND STREQUAL "CUDA")
   find_library(CUDNN_LIBRARY cudnn HINTS ${CUDNN_ROOT_DIR} ${CUDAToolkit_LIBRARY_DIR} PATH_SUFFIXES lib64)
   include_directories(SYSTEM ${CUDAToolkit_INCLUDE_DIRS} ${CUDNN_INCLUDE_DIR}) #SYSTEM is for suppressing some compiler warnings in thrust libraries
   include_directories(SYSTEM external/cudnn-frontend/include) #SYSTEM suppresses warnings inside the vendored header-only cudnn_frontend library
+  if(KATAGO_USE_CUTLASS_FUSED_FFN)
+    target_compile_definitions(katago PRIVATE USE_CUTLASS_FUSED_FFN)
+    target_include_directories(katago SYSTEM PRIVATE
+      external/cutlass/include
+      external/cutlass/examples/45_dual_gemm)
+  endif()
+  if(NO_CUDNN_SDPA)
+    message(STATUS "-DNO_CUDNN_SDPA=1 is set, building without the cudnn graph SDPA attention path.")
+    target_compile_definitions(katago PRIVATE NO_CUDNN_SDPA)
+  endif()
   # cudnn_frontend's experimental OSS engine headers reference NVRTC symbols at link time
   # (we don't actually use those engines, but the references are pulled in by the header-only includes).
   target_link_libraries(katago CUDA::cublas ${CUDNN_LIBRARY} CUDA::nvrtc)

cpp/command/benchmarknn.cpp

@@ -0,0 +1,228 @@
+#include "../core/global.h"
+#include "../core/config_parser.h"
+#include "../core/logger.h"
+#include "../core/rand.h"
+#include "../game/board.h"
+#include "../neuralnet/nneval.h"
+#include "../program/setup.h"
+#include "../command/commandline.h"
+#include "../main.h"
+
+#include <iomanip>
+#include <sstream>
+
+using namespace std;
+
+static string jsonEscape(const string& s) {
+  ostringstream out;
+  for(char c : s) {
+    if(c == '"' || c == '\\')
+      out << '\\' << c;
+    else if(c == '\n')
+      out << "\\n";
+    else if(c == '\t')
+      out << "\\t";
+    else if(c == '\r')
+      out << "\\r";
+    else
+      out << c;
+  }
+  return out.str();
+}
+
+int MainCmds::benchmarknn(const vector<string>& args) {
+  Board::initHash();
+  ScoreValue::initTables();
+  Rand seedRand;
+
+  ConfigParser cfg;
+  string modelFile;
+  int numIterations;
+  int numWarmups;
+  int batchSizeOverride;
+  string boardSizesStr;
+  bool requireExactNNLen;
+  bool jsonOut;
+  vector<int> boardSizes;
+
+  try {
+    KataGoCommandLine cmd(
+      "Benchmark raw neural net forward throughput, without search. Uses numNNServerThreadsPerModel "
+      "and per-thread GPU assignment from the config. Timing includes host-device transfers and "
+      "output postprocessing, excludes input feature generation and search."
+    );
+    cmd.addConfigFileArg(KataGoCommandLine::defaultGtpConfigFileName(),"gtp_example.cfg");
+    cmd.addModelFileArg();
+    TCLAP::ValueArg<int> iterationsArg(
+      "","iterations","Number of timed forward passes per NN server thread (default 200)",
+      false,200,"N"
+    );
+    TCLAP::ValueArg<int> warmupArg(
+      "","warmup","Untimed forward passes per NN server thread before timing (default 20)",
+      false,20,"N"
+    );
+    TCLAP::ValueArg<int> batchSizeArg(
+      "","batch-size","Batch size per NN server thread (default: nnMaxBatchSize from config, or 16)",
+      false,-1,"N"
+    );
+    TCLAP::ValueArg<string> boardSizesArg(
+      "","boardsize",
+      "Board size, or comma-separated sizes cycled across the rows of each batch, e.g. 19 or 9,13,19 "
+      "(default 19). The NN buffer size is the largest listed size.",
+      false,"19","SIZES"
+    );
+    TCLAP::SwitchArg requireExactNNLenArg(
+      "","require-exact-nnlen",
+      "Run with requireExactNNLen (backend may skip mask handling). Needs a single board size."
+    );
+    TCLAP::SwitchArg jsonArg("","json","Print results as JSON",false);
+    cmd.add(iterationsArg);
+    cmd.add(warmupArg);
+    cmd.add(batchSizeArg);
+    cmd.add(boardSizesArg);
+    cmd.add(requireExactNNLenArg);
+    cmd.add(jsonArg);
+    cmd.setShortUsageArgLimit();
+    cmd.addOverrideConfigArg();
+
+    cmd.parseArgs(args);
+
+    modelFile = cmd.getModelFile();
+    numIterations = iterationsArg.getValue();
+    numWarmups = warmupArg.getValue();
+    batchSizeOverride = batchSizeArg.getValue();
+    boardSizesStr = boardSizesArg.getValue();
+    requireExactNNLen = requireExactNNLenArg.getValue();
+    jsonOut = jsonArg.getValue();
+    cmd.getConfig(cfg);
+
+    if(numIterations <= 0)
+      throw StringError("benchmarknn: iterations must be positive");
+    if(numWarmups < 0)
+      throw StringError("benchmarknn: warmup must be nonnegative");
+    for(const string& piece : Global::split(boardSizesStr,',')) {
+      int bSize = Global::stringToInt(Global::trim(piece));
+      if(bSize < 2 || bSize > Board::MAX_LEN)
+        throw StringError("benchmarknn: invalid board size " + piece);
+      boardSizes.push_back(bSize);
+    }
+    if(boardSizes.size() <= 0)
+      throw StringError("benchmarknn: no board sizes specified");
+    if(requireExactNNLen && boardSizes.size() > 1)
+      throw StringError("benchmarknn: require-exact-nnlen needs a single board size");
+  }
+  catch(TCLAP::ArgException& e) {
+    cerr << "Error: " << e.error() << " for argument " << e.argId() << endl;
+    return 1;
+  }
+
+  // With -json, stdout must stay machine-readable, so route logs (including backend fallback
+  // warnings that would otherwise be lost) to stderr instead.
+  const bool logToStdout = !jsonOut;
+  const bool logToStderr = jsonOut;
+  Logger logger(NULL, logToStdout, logToStderr, false);
+  logger.write("Version " + Version::getGitRevisionWithBackend());
+
+  int nnLen = 0;
+  for(int bSize : boardSizes)
+    nnLen = std::max(nnLen, bSize);
+
+  const string expectedSha256 = "";
+  const int maxBatchSize =
+    batchSizeOverride > 0 ? batchSizeOverride :
+    cfg.contains("nnMaxBatchSize") ? cfg.getInt("nnMaxBatchSize",1,65536) :
+    16;
+  const int expectedConcurrentEvals = maxBatchSize;
+  const bool disableFP16 = false;
+
+  NNEvaluator* nnEval = NULL;
+  try {
+    nnEval = Setup::initializeNNEvaluator(
+      modelFile,modelFile,expectedSha256,cfg,logger,seedRand,expectedConcurrentEvals,
+      nnLen,nnLen,maxBatchSize,requireExactNNLen,disableFP16,
+      Setup::SETUP_FOR_BENCHMARKNN
+    );
+
+    NNEvalBenchmarkResult result = nnEval->benchmarkPureForward(numWarmups,numIterations,boardSizes);
+
+    if(jsonOut) {
+      ostringstream out;
+      out << "{";
+      out << "\"modelFile\":\"" << jsonEscape(nnEval->getModelFileName()) << "\",";
+      out << "\"modelName\":\"" << jsonEscape(nnEval->getInternalModelName()) << "\",";
+      out << "\"revision\":\"" << jsonEscape(Version::getGitRevisionWithBackend()) << "\",";
+      out << "\"boardSizes\":[";
+      for(size_t i = 0; i < boardSizes.size(); i++)
+        out << (i > 0 ? "," : "") << boardSizes[i];
+      out << "],";
+      out << "\"requireExactNNLen\":" << (nnEval->getRequireExactNNLen() ? "true" : "false") << ",";
+      out << "\"usingFP16\":" << (nnEval->isAnyThreadUsingFP16() ? "true" : "false") << ",";
+      out << "\"batchSize\":" << result.batchSize << ",";
+      out << "\"numThreads\":" << result.numThreads << ",";
+      out << "\"numIterations\":" << result.numIterations << ",";
+      out << "\"gpuIdxs\":[";
+      bool first = true;
+      for(int g : nnEval->getGpuIdxs()) {
+        out << (first ? "" : ",") << g;
+        first = false;
+      }
+      out << "],";
+      out << setprecision(10);
+      out << "\"perThreadMedianMs\":[";
+      for(int i = 0; i < result.numThreads; i++)
+        out << (i > 0 ? "," : "") << result.perThreadMedianSeconds[i] * 1000.0;
+      out << "],";
+      out << "\"perThreadNNEvalsPerSec\":[";
+      for(int i = 0; i < result.numThreads; i++)
+        out << (i > 0 ? "," : "") << result.perThreadNNEvalsPerSec[i];
+      out << "],";
+      out << "\"sumMedianNNEvalsPerSec\":" << result.sumMedianNNEvalsPerSec << ",";
+      out << "\"actualWallSeconds\":" << result.actualWallSeconds << ",";
+      out << "\"actualWallNNEvalsPerSec\":" << result.actualWallNNEvalsPerSec;
+      out << "}";
+      cout << out.str() << endl;
+    }
+    else {
+      cout << "=== benchmarknn ===" << endl;
+      cout << "model: " << nnEval->getModelFileName() << endl;
+      cout << "internal model: " << nnEval->getInternalModelName() << endl;
+      cout << "revision/backend: " << Version::getGitRevisionWithBackend() << endl;
+      cout << "board sizes:";
+      for(int bSize : boardSizes)
+        cout << " " << bSize;
+      cout << " (NN buffer " << nnLen << "x" << nnLen
+           << (nnEval->getRequireExactNNLen() ? ", requireExactNNLen" : "") << ")" << endl;
+      cout << "FP16: " << (nnEval->isAnyThreadUsingFP16() ? "true" : "false") << endl;
+      cout << "batch size per thread: " << result.batchSize << endl;
+      cout << "NN server threads: " << result.numThreads << endl;
+      cout << "GPU indices:";
+      for(int g : nnEval->getGpuIdxs())
+        cout << " " << g;
+      cout << endl;
+      cout << "timed iterations per thread: " << result.numIterations << endl;
+      for(int i = 0; i < result.numThreads; i++) {
+        cout << "thread " << i << ": "
+             << setprecision(6) << result.perThreadMedianSeconds[i] * 1000.0
+             << " ms/batch median, " << setprecision(8) << result.perThreadNNEvalsPerSec[i]
+             << " nnEval/s" << endl;
+      }
+      cout << "sum of per-thread median rates: "
+           << setprecision(8) << result.sumMedianNNEvalsPerSec << " nnEval/s" << endl;
+      cout << "wall time of timed region: "
+           << setprecision(6) << result.actualWallSeconds << " s" << endl;
+      cout << "overall throughput (wall): "
+           << setprecision(8) << result.actualWallNNEvalsPerSec << " nnEval/s" << endl;
+    }
+  }
+  catch(...) {
+    delete nnEval;
+    NeuralNet::globalCleanup();
+    ScoreValue::freeTables();
+    throw;
+  }
+
+  delete nnEval;
+  NeuralNet::globalCleanup();
+  ScoreValue::freeTables();
+  return 0;
+}

cpp/configs/analysis_example.cfg

@@ -282,6 +282,7 @@ nnRandomize = true
 # e.g. if FP16 is giving an error or too much numerical inaccuracy on your card.
 # rocmUseFP16 = auto
 # rocmUseNHWC = auto  # Uses NHWC tensor layout. Default: auto - NHWC with FP16 on GPUs where NHWC convolutions are faster (CDNA), else NCHW; transformers always NHWC.
+# rocmUse1x1Matmul = auto  # Whether 1x1 NHWC convs run as a hipBLAS GEMM instead of a MIOpen conv. auto = GEMM when possible (MIOpen's NHWC conv are slow on some GPUs)
 
 
 # OpenCL-specific GPU settings--------------------------------------

cpp/configs/gtp_example.cfg

@@ -617,6 +617,9 @@ searchFactorWhenWinningThreshold = 0.95
 # instructions make NHWC convolutions faster (CDNA), otherwise NCHW; transformer models
 # always use NHWC.
 # rocmUseNHWC = auto
+# Whether 1x1 NHWC convs run as a hipBLAS GEMM instead of a MIOpen conv.
+# auto = GEMM when possible (MIOpen's NHWC conv are slow on some GPUs)
+# rocmUse1x1Matmul = auto
 
 # ------------------------------
 # OpenCL GPU settings

cpp/external/cutlass/.github/ISSUE_TEMPLATE/bug_report.yml

@@ -0,0 +1,38 @@
+name: Bug Report
+description: Create a bug report to help us improve CUTLASS
+title: "[BUG] "
+labels: ["? - Needs Triage", "bug"]
+assignees: []
+
+body:
+  - type: dropdown
+    id: component
+    attributes:
+      label: Which component has the problem?
+      options:
+        - CuTe DSL
+        - CUTLASS C++
+    validations:
+      required: true
+  - type: textarea
+    id: bug-report
+    attributes:
+      label: Bug Report
+      description: Please fill out all sections below
+      value: |
+        **Describe the bug**
+        A clear and concise description of what the bug is.
+
+        **Steps/Code to reproduce bug**
+        Follow this guide http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports to craft a minimal bug report. This helps us reproduce the issue you're having and resolve the issue more quickly.
+
+        **Expected behavior**
+        A clear and concise description of what you expected to happen.
+
+        **Environment details (please complete the following information):**
+         - Environment location: [Bare-metal, Docker, Cloud(specify cloud provider)] 
+
+        **Additional context**
+        Add any other context about the problem here.
+    validations:
+      required: true 

cpp/external/cutlass/.github/ISSUE_TEMPLATE/config.yml

@@ -0,0 +1,5 @@
+blank_issues_enabled: true
+contact_links:
+  - name: CUTLASS Discord
+    url: https://discord.gg/nvidiadeveloper
+    about: Come chat about using and contributing to CUTLASS!