diff --git a/Compiling.md b/Compiling.md index abe7de36fc..208412b577 100644 --- a/Compiling.md +++ b/Compiling.md @@ -33,6 +33,7 @@ As also mentioned in the instructions below but repeated here for visibility, if * If using the OpenCL backend, a modern GPU that supports OpenCL 1.2 or greater, or else something like [this](https://software.intel.com/en-us/opencl-sdk) for CPU. But if using CPU, Eigen should be better. * If using the CUDA backend, CUDA 11 or later and a compatible version of CUDNN based on your CUDA version (https://developer.nvidia.com/cuda-toolkit) (https://developer.nvidia.com/cudnn) and a GPU capable of supporting them. * If using the TensorRT backend, in addition to a compatible CUDA Toolkit (https://developer.nvidia.com/cuda-toolkit), you also need TensorRT (https://developer.nvidia.com/tensorrt) that is at least version 8.5. + * If using the ROCm backend, ROCm 6.4 or later and a GPU capable of supporting them. More information about installation(https://rocm.docs.amd.com/projects/install-on-linux/en/latest/) and please install all possible ROCm developer packages, instead of just ROCm runtime packages. * If using the Eigen backend, Eigen3. With Debian packages, (i.e. apt or apt-get), this should be `libeigen3-dev`. * zlib, libzip. With Debian packages (i.e. apt or apt-get), these should be `zlib1g-dev`, `libzip-dev`. * If you want to do self-play training and research, probably Google perftools `libgoogle-perftools-dev` for TCMalloc or some other better malloc implementation. For unknown reasons, the allocation pattern in self-play with large numbers of threads and parallel games causes a lot of memory fragmentation under glibc malloc that will eventually run your machine out of memory, but better mallocs handle it fine. @@ -41,7 +42,7 @@ As also mentioned in the instructions below but repeated here for visibility, if * `git clone https://github.com/lightvector/KataGo.git` * Compile using CMake and make in the cpp directory: * `cd KataGo/cpp` - * `cmake . -DUSE_BACKEND=OPENCL` or `cmake . -DUSE_BACKEND=CUDA` or `cmake . -DUSE_BACKEND=TENSORRT` or `cmake . -DUSE_BACKEND=EIGEN` depending on which backend you want. + * `cmake . -DUSE_BACKEND=OPENCL` or `cmake . -DUSE_BACKEND=CUDA` or `cmake . -DUSE_BACKEND=TENSORRT` or `cmake . -DUSE_BACKEND=EIGEN` or `cmake . -DUSE_BACKEND=ROCM` depending on which backend you want. * Specify also `-DUSE_TCMALLOC=1` if using TCMalloc. * Compiling will also call git commands to embed the git hash into the compiled executable, specify also `-DNO_GIT_REVISION=1` to disable it if this is causing issues for you. * Specify `-DUSE_AVX2=1` to also compile Eigen with AVX2 and FMA support, which will make it incompatible with old CPUs but much faster. (If you want to go further, you can also add `-DCMAKE_CXX_FLAGS='-march=native'` which will specialize to precisely your machine's CPU, but the exe might not run on other machines at all). @@ -54,6 +55,36 @@ As also mentioned in the instructions below but repeated here for visibility, if * You will probably want to edit `configs/gtp_example.cfg` (see "Tuning for Performance" above). * If using OpenCL, you will want to verify that KataGo is picking up the correct device when you run it (e.g. some systems may have both an Intel CPU OpenCL and GPU OpenCL, if KataGo appears to pick the wrong one, you can correct this by specifying `openclGpuToUse` in `configs/gtp_example.cfg`). + * **ROCm backend (Linux) — additional notes:** + * Install ROCm following the [official guide](https://rocm.docs.amd.com/en/7.12.0-preview/install/rocm.html). Install the full developer stack (not just runtime): `sudo apt install rocm-dev miopen-hip rocblas hipblas`. + * Build: + ``` + cd KataGo/cpp + mkdir build && cd build + cmake .. -DUSE_BACKEND=ROCM -DCMAKE_BUILD_TYPE=Release + make -j$(nproc) + ``` + No `-DCMAKE_PREFIX_PATH` is needed in the common case: the build auto-detects the ROCm + install location, preferring the newer `/opt/rocm/core-/` layout used since + ROCm ~7.9 (picking the highest version present) and falling back to the older flat + `/opt/rocm/` layout. Pass `-DCMAKE_PREFIX_PATH=...` explicitly to override. + * GPU architecture: by default the build targets a broad set of AMD GPU architectures + (CDNA + all RDNA generations) in a single "fat" binary, probing the installed compiler for + which ones it actually supports, so the resulting `katago` runs on more than just the + machine it was built on. Pass `-DCMAKE_HIP_ARCHITECTURES=gfx1100` (replace with your GPU's + gfx target) to build for only your own GPU instead, which is faster to compile. + * The build bakes the resolved ROCm lib directory into the binary's RPATH, so the built + `katago` doesn't depend on `/opt/rocm` still pointing at the same install later (e.g. after + installing a different ROCm version) or on `LD_LIBRARY_PATH` being set when run. + * On first run, MIOpen will search for optimal convolution algorithms for your specific GPU and network size. This may take up to a minute and results are cached in `~/.config/miopen/` for subsequent runs. + * **Transformer/attention models (model version 17+):** supported on all architectures via a + built-in kernel. If a matching version of AMD's Composable Kernel (`composablekernel-dev` / + `amdrocm-ck*`) is also installed, the build additionally enables a fused-attention fast path + (measured ~2x faster on a gfx1100/RX 7900 XTX) for CDNA and RDNA3/RDNA3.5/RDNA4 GPUs - + RDNA1/RDNA2 always use the built-in kernel, as does any GPU if the fused path isn't + available or is explicitly disabled with `rocmDisableFusedAttention = true` in the config. + See `cpp/external/composable_kernel_fmha/README.md` for details. + ## Windows * TLDR: * Building from source on Windows is actually a bit tricky, depending on what version you're building, there's not necessarily a super-fast way. @@ -117,6 +148,58 @@ As also mentioned in the instructions below but repeated here for visibility, if * You will probably want to edit `configs/gtp_example.cfg` (see "Tuning for Performance" above). * If using OpenCL, you will want to verify that KataGo is picking up the correct device (e.g. some systems may have both an Intel CPU OpenCL and GPU OpenCL, if KataGo appears to pick the wrong one, you can correct this by specifying `openclGpuToUse` in `configs/gtp_example.cfg`). + * **ROCm backend (Windows) — building via AMD TheRock:** + * The ROCm (MIOpen) backend supports Windows via [AMD TheRock](https://github.com/ROCm/TheRock) (tested with TheRock 7.13 / ROCm 7.13, RX 7900 XTX / gfx1100), including transformer/attention models (model version 17+) and the optional CK fused-attention fast path. + * **Prerequisites:** + * Install ROCm following the [official guide](https://rocm.docs.amd.com/en/7.13.0-preview/install/rocm.html?fam=all&os=windows). For Windows, download [AMD TheRock](https://github.com/ROCm/TheRock) and extract to e.g. `C:\TheRock\build`. + * Install **Visual Studio Build Tools or Community** with the "Desktop development with C++" workload, for the MSVC toolchain and Windows SDK the HIP compiler needs. Any MSVC toolset version is fine to install - if more than one ends up installed side by side, `CMakeLists.txt` automatically probes them at configure time and picks a compatible one itself (see "Fully automatic" below), no manual toolset selection needed. + * Install [Ninja](https://ninja-build.org) build tool: `winget install Ninja-build.Ninja`. + * Set the following **system environment variables** (via System Properties → Advanced → Environment Variables): + ``` + HIP_PATH=C:/TheRock/build + HIP_PLATFORM=amd + HIP_DEVICE_LIB_PATH=C:/TheRock/build/lib/llvm/amdgcn/bitcode + LLVM_PATH=C:/TheRock/build/lib/llvm + ``` + * Add to system `PATH`: + ``` + C:\TheRock\build\bin + C:\TheRock\build\lib\llvm\bin + ``` + * Reboot after setting environment variables so they take effect system-wide. + * **Build - fully automatic**, just like Linux: + ``` + cd KataGo/cpp + mkdir build + cd build + cmake .. -G Ninja -DUSE_BACKEND=ROCM -DCMAKE_BUILD_TYPE=Release + ninja -j $env:NUMBER_OF_PROCESSORS + ``` + No manual environment setup, no `-D` flags, no `vcvarsall`, and no external package manager + install are needed beyond the prerequisites above. `CMakeLists.txt` handles the rest of the + Windows-specific setup automatically at configure/build time: + * **MSVC toolset selection:** if more than one MSVC toolset is installed side by side, a + newer one can conflict with TheRock's bundled clang (a brand-new MSVC STL declaring math + functions in a way clang's CUDA/HIP compatibility headers don't yet handle, or a + different SSE2-intrinsics-resolution conflict). `CMakeLists.txt` finds all installed + toolsets via `vswhere` and probes each with a real compile until it finds one that works, + with no user action needed. + * **zlib:** TheRock's Windows package ships `zlib.h` but (as of 7.13) no longer ships a + linkable `.lib`. `CMakeLists.txt` automatically bootstraps a local + [vcpkg](https://github.com/microsoft/vcpkg) clone under `cpp/build/deps/vcpkg` (this + needs internet access and `git` on `PATH` the first time; subsequent reconfigures reuse + the same local install) and builds zlib through it - this is the same + `KATAGO_AUTO_FETCH_DEPS`/vcpkg-in-build-tree mechanism the ONNX backend uses, so it stays + consistent across backends rather than using a different fetch method here. + * **Runtime DLLs:** all the ROCm/HIP DLLs (`amdhip64_7.dll`, `MIOpen.dll`, `hipblas.dll`, + `rocblas.dll` + its `library/` subfolder, `libhipblaslt.dll` + its `library/` subfolder, + `amdocl64.dll`, `hiprtc*.dll`, `amd_comgr*.dll`) and the vcpkg zlib runtime DLL are + automatically copied next to `katago.exe` as a post-build step - nothing to copy by hand. + `ck_tile` headers for the optional fused-attention path are still auto-detected from + `HIP_PATH` the same way as on Linux. + * **First-run note:** MIOpen will search for optimal convolution algorithms on the first run. This may take 45+ seconds per network configuration and results are cached in `%USERPROFILE%\.miopen\` for subsequent runs. Do not terminate the process during this initial tuning. + * **Performance note:** GPU utilization on Windows may be somewhat lower than on Linux due to the Windows Driver Model (WDDM) adding overhead to GPU kernel submissions. This is a known limitation of ROCm on Windows. For example, the CK fused-attention path measured ~2x faster than the built-in kernel on Linux (gfx1100), but only ~1.3x faster on Windows on the same GPU — still a real win, just smaller due to WDDM overhead. + ## MacOS * TLDR (Metal backend - recommended for most users, hybrid CPU+GPU+Neural Engine for maximum throughput): ``` diff --git a/README.md b/README.md index 4ea665e7f5..c86a1a7cf0 100644 --- a/README.md +++ b/README.md @@ -1,27 +1,30 @@ # KataGo -* [Overview](#overview) -* [Training History and Research](#training-history-and-research) -* [Where To Download Stuff](#where-to-download-stuff) -* [Setting Up and Running KataGo](#setting-up-and-running-katago) - * [GUIs](#guis) - * [Windows and Linux](#windows-and-linux) - * [MacOS](#macos) - * [OpenCL vs CUDA vs TensorRT vs Eigen](#opencl-vs-cuda-vs-tensorrt-vs-eigen) - * [How To Use](#how-to-use) - * [Tuning for Performance](#tuning-for-performance) - * [Common Questions and Issues](#common-questions-and-issues) - * [Issues with specific GPUs or GPU drivers](#issues-with-specific-gpus-or-gpu-drivers) - * [Common Problems](#common-problems) - * [Other Questions](#other-questions) -* [Features for Developers](#features-for-developers) - * [GTP Extensions](#gtp-extensions) - * [Analysis Engine](#analysis-engine) -* [Compiling KataGo](#compiling-katago) -* [Source Code Overview](#source-code-overview) -* [Selfplay Training](#selfplay-training) -* [Contributors](#contributors) -* [License](#license) +- [KataGo](#katago) + - [Overview](#overview) + - [Training History and Research and Docs](#training-history-and-research-and-docs) + - [Where To Download Stuff](#where-to-download-stuff) + - [Setting Up and Running KataGo](#setting-up-and-running-katago) + - [GUIs](#guis) + - [Windows and Linux](#windows-and-linux) + - [MacOS](#macos) + - [OpenCL vs CUDA vs TensorRT vs ROCm vs Eigen](#opencl-vs-cuda-vs-tensorrt-vs-rocm-vs-eigen) + - [How To Use](#how-to-use) + - [Human-style Play and Analysis](#human-style-play-and-analysis) + - [Other Commands:](#other-commands) + - [Tuning for Performance](#tuning-for-performance) + - [Common Questions and Issues](#common-questions-and-issues) + - [Issues with specific GPUs or GPU drivers](#issues-with-specific-gpus-or-gpu-drivers) + - [Common Problems](#common-problems) + - [Other Questions](#other-questions) + - [Features for Developers](#features-for-developers) + - [GTP Extensions:](#gtp-extensions) + - [Analysis Engine:](#analysis-engine) + - [Compiling KataGo](#compiling-katago) + - [Source Code Overview:](#source-code-overview) + - [Selfplay Training:](#selfplay-training) + - [Contributors](#contributors) + - [License](#license) ## Overview @@ -85,8 +88,8 @@ The community also provides KataGo packages for [Homebrew](https://brew.sh) on M Use `brew install katago`. The latest config files and networks are installed in KataGo's `share` directory. Find them via `brew list --verbose katago`. A basic way to run katago will be `katago gtp -config $(brew list --verbose katago | grep 'gtp.*\.cfg') -model $(brew list --verbose katago | grep .gz | head -1)`. You should choose the Network according to the release notes here and customize the provided example config as with every other way of installing KataGo. -### OpenCL vs CUDA vs TensorRT vs Eigen -KataGo has four backends, OpenCL (GPU), CUDA (GPU), TensorRT (GPU), and Eigen (CPU). (On macOS there is also a Metal backend, most easily obtained via homebrew - see above.) +### OpenCL vs CUDA vs TensorRT vs ROCm vs Eigen +KataGo has five backends, OpenCL (GPU), CUDA (GPU), TensorRT (GPU), ROCm (GPU) and Eigen (CPU). (On macOS there is also a Metal backend, most easily obtained via homebrew - see above.) As of v1.17, KataGo supports transformer neural nets, which are generally much stronger for the same compute cost and which the main training run is switching to. Transformer models are more demanding on the GPU backend than the older convolutional nets, so the backend recommendations below matter more for them - in particular OpenCL is noticeably slower on transformers, and on NVIDIA the CUDNN and TensorRT versions make a large difference. @@ -96,11 +99,13 @@ The quick summary is: * Use Eigen with AVX2 if you don't have a GPU or if your GPU is too old/weak to work with OpenCL, and you just want a plain CPU KataGo. * Use Eigen without AVX2 if your CPU is old or on a low-end device that doesn't support AVX2. * The CUDA+CUDNN backend can also work well for NVIDIA GPUs. It has faster startup than TensorRT and is competitive on transformers if using CUDNN >= 9.8.0, though TensorRT 10.16 is often still slightly faster. + * The ROCm backend can work for AMD GPUs with ROCm+MIOpen installed. More in detail: * OpenCL is a general GPU backend should be able to run with any GPUs or accelerators that support [OpenCL](https://en.wikipedia.org/wiki/OpenCL), including NVIDIA GPUs, AMD GPUs, as well CPU-based OpenCL implementations or things like Intel Integrated Graphics. This is the most general GPU version of KataGo and doesn't require a complicated install like CUDA does, so is most likely to work out of the box as long as you have a fairly modern GPU. **However, it also need to take some time when run for the very first time to tune itself.** For many systems, this will take 5-30 seconds, but on a few older/slower systems, may take many minutes or longer. Also, the quality of OpenCL implementations is sometimes inconsistent, particularly for Intel Integrated Graphics and for AMD GPUs that are older than several years, so it might not work for very old machines, as well as specific buggy newer AMD GPUs, see also [Issues with specific GPUs or GPU drivers](#issues-with-specific-gpus-or-gpu-drivers). OpenCL is not as optimized as the NVIDIA-specific backends and will generally be slower, particularly for transformer models. * CUDA is a GPU backend specific to NVIDIA GPUs (it will not work with AMD or Intel or any other GPUs) and requires installing [CUDA](https://developer.nvidia.com/cuda-zone) and [CUDNN](https://developer.nvidia.com/cudnn) and a modern NVIDIA GPU. For older convolutional nets, on many GPUs the OpenCL implementation can match or beat NVIDIA's own CUDA/CUDNN, with the exception of top-end NVIDIA GPUs that support FP16 and tensor cores. For transformer nets, CUDA+CUDNN clearly outperforms OpenCL, but you should use CUDNN >= 9.8.0 if at all possible - the older CUDNN 8.9.7 is a LOT slower on transformer models. Compared to TensorRT, CUDA+CUDNN has faster startup times and is often only slightly slower (and occasionally faster) on transformers. * TensorRT is similar to CUDA, but uses NVIDIA's TensorRT framework to run the neural network with more optimized kernels. For modern NVIDIA GPUs it should work whenever CUDA does, and will usually be the fastest backend, though it has much longer startup times on every launch. As of v1.17.0, TensorRT versions older than 10 are no longer supported. For transformer models, recent versions like CUDA 13 + TensorRT 10.16 are best, while older TensorRT versions can be outperformed by CUDA+CUDNN. + * ROCm is a GPU backend specific to AMD GPUs (it will not work with NVIDIA or Intel or any other GPUs) and requires installing [ROCm](https://rocm.docs.amd.com) and [MIOpen](https://rocm.docs.amd.com/projects/MIOpen) and a modern AMD GPU. Supports both **Linux** (via official ROCm packages, ROCm 6.4+) and **Windows** (via [AMD TheRock](https://github.com/ROCm/TheRock) builds). On most GPUs, the OpenCL implementation will actually beat AMD's own ROCm/MIOpen at performance. The exception is for top-end AMD GPUs that support FP16 and stream processors, in which case sometimes one is better and sometimes the other is better. Transformer/attention-based neural nets (model version 17+) are supported on all AMD GPUs, and get an additional fused-attention speedup on CDNA and RDNA3/RDNA3.5/RDNA4 GPUs when AMD's Composable Kernel library is also installed (see [Compiling.md](Compiling.md)). * Eigen is a *CPU* backend that should work widely *without* needing a GPU or fancy drivers. Use this if you don't have a good GPU or really any GPU at all. It will be quite significantly slower than OpenCL or CUDA, but on a good CPU can still often get 10 to 20 playouts per second if using the smaller (15 or 20) block neural nets. Eigen can also be compiled with AVX2 and FMA support, which can provide a big performance boost for Intel and AMD CPUs from the last few years. However, it will not run at all on older CPUs (and possibly even some recent but low-power modern CPUs) that don't support these fancy vector instructions. For **any** implementation, it's recommended that you also tune the number of threads used if you care about optimal performance, as it can make a factor of 2-3 difference in the speed. See "Tuning for Performance" below. However, if you mostly just want to get it working, then the default untuned settings should also be still reasonable. @@ -178,6 +183,7 @@ This section summarizes a number of common questions and issues when running Kat #### Issues with specific GPUs or GPU drivers If you are observing any crashes in KataGo while attempting to run the benchmark or the program itself, and you have one of the below GPUs, then this is likely the reason. +* **AMD GPUs** - If you choose to use the ROCm backend, you need a GPU on the official [System requirements list](https://rocm.docs.amd.com/projects/install-on-linux/en/latest/reference/system-requirements.html) (at least AMD Radeon RX 7700 XT). ROCm backend supports both Linux (via official ROCm packages) and Windows (via [AMD TheRock](https://github.com/ROCm/TheRock) builds). On Linux, install the full ROCm developer stack. On Windows, see the ROCm Windows build instructions in [Compiling.md](Compiling.md). * **AMD Radeon RX 5700** - AMD's drivers for OpenCL for this GPU have been buggy ever since this GPU was released, and as of May 2020 AMD has still never released a fix. If you are using this GPU, you will just not be able to run KataGo (Leela Zero and other Go engines will probably fail too) and will probably also obtain incorrect calculations or crash if doing anything else scientific or mathematical that uses OpenCL. See for example these reddit threads: [[1]](https://www.reddit.com/r/Amd/comments/ebso1x/its_not_just_setihome_any_mathematic_or/) or [[2]](https://www.reddit.com/r/BOINC/comments/ebiz18/psa_please_remove_your_amd_rx5700xt_from_setihome/) or this [L19 thread](https://lifein19x19.com/viewtopic.php?f=18&t=17093). * **OpenCL Mesa** - These drivers for OpenCL are buggy. Particularly if on startup before crashing you see KataGo printing something like `Found OpenCL Platform 0: ... (Mesa) (OpenCL 1.1 Mesa ...) ...` diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index fb8bb130fa..9a73cdaa3a 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1,6 +1,478 @@ cmake_minimum_required(VERSION 3.18.2) + +# Helper: produce the default broad list of AMD GPU architectures the ROCm +# backend should target when the user has not passed -DCMAKE_HIP_ARCHITECTURES=... +# Required archs are always included (well-established across many ROCm releases); optional archs +# are probed via an actual compile test and only included if the installed toolchain accepts them - +# this covers newer/rarer archs (RDNA1, RDNA2 sub-variants, RDNA3.5, RDNA4, ...) without hard-failing +# the whole build on an older toolchain that doesn't know about them yet. Keep this list in sync with +# whatever `apt-cache search '^amdrocm-core-sdk-gfx'` (or the CK codegen's own supported archs, +# see cpp/external/composable_kernel_fmha/README.md) shows as AMD adds new chips. +function(katago_default_hip_archs out_var) + set(_required_archs + gfx906 gfx908 gfx90a gfx942 gfx950 + gfx1030 gfx1100 gfx1101) + set(_optional_archs + gfx1010 gfx1011 gfx1012 + gfx1031 gfx1032 gfx1033 gfx1034 gfx1035 gfx1036 + gfx1102 gfx1103 + gfx1150 gfx1151 gfx1152 gfx1153 + gfx1200 gfx1201) + set(_result "${_required_archs}") + + if(CMAKE_HIP_COMPILER) + set(_probe_dir "${CMAKE_BINARY_DIR}/_katago_hip_arch_probe") + file(MAKE_DIRECTORY "${_probe_dir}") + set(_probe_src "${_probe_dir}/probe.cpp") + if(NOT EXISTS "${_probe_src}") + file(WRITE "${_probe_src}" "int main(){return 0;}\n") + endif() + foreach(_arch IN LISTS _optional_archs) + execute_process( + COMMAND "${CMAKE_HIP_COMPILER}" --offload-arch=${_arch} -x hip -c + "${_probe_src}" -o "${_probe_dir}/probe_${_arch}.o" + RESULT_VARIABLE _rc + OUTPUT_QUIET ERROR_QUIET) + if(_rc EQUAL 0) + list(APPEND _result ${_arch}) + message(STATUS "HIP arch ${_arch} accepted by compiler — including") + else() + message(STATUS "HIP arch ${_arch} not accepted by compiler — skipping") + endif() + endforeach() + else() + message(STATUS "CMAKE_HIP_COMPILER not yet known; optional archs not probed") + endif() + + set(${out_var} "${_result}" PARENT_SCOPE) +endfunction() + +# Helper: resolve the actual ROCm install prefix on Linux. Since ROCm ~7.9, AMD's packaging puts +# real per-version installs under /opt/rocm/core-/ (e.g. /opt/rocm/core-7.13), with +# /opt/rocm itself just a distro-managed symlink/alternatives entry that can point anywhere (or get +# silently repointed by installing another package) - relying on it directly is what caused this +# project to intermittently pick up a stale/mismatched ROCm install. Prefer the highest-versioned +# /opt/rocm/core-* directory if any exist; otherwise fall back to the old flat /opt/rocm layout. +function(katago_find_rocm_prefix out_var) + set(_result "") + if(EXISTS "/opt/rocm") + file(GLOB _core_dirs "/opt/rocm/core-*") + set(_versioned "") + foreach(_dir ${_core_dirs}) + if(IS_DIRECTORY "${_dir}") + get_filename_component(_name "${_dir}" NAME) + string(REGEX REPLACE "^core-" "" _ver "${_name}") + list(APPEND _versioned "${_ver}|${_dir}") + endif() + endforeach() + if(_versioned) + # Compare the version fields with real version semantics ("7.*" > "7"), not the raw + # "ver|path" strings - a plain string sort would rank the major-only "core-7" symlink + # above the real "core-7.*" dir ('.' < '|'), silently re-selecting the repointable + # alternatives symlink this function exists to avoid. + set(_best_ver "") + foreach(_entry ${_versioned}) + string(REGEX REPLACE "\\|.*$" "" _ver "${_entry}") + string(REGEX REPLACE "^[^|]*\\|" "" _dir "${_entry}") + if(_best_ver STREQUAL "" OR _ver VERSION_GREATER _best_ver) + set(_best_ver "${_ver}") + set(_result "${_dir}") + endif() + endforeach() + endif() + endif() + if(_result STREQUAL "" AND EXISTS "/opt/rocm") + set(_result "/opt/rocm") + endif() + set(${out_var} "${_result}" PARENT_SCOPE) +endfunction() + +# Helper: locate vswhere.exe via the well-known Program Files env vars (Microsoft's own documented, +# stable install location for it: "/Microsoft Visual Studio/Installer/"), so we +# never hardcode a user- or edition-specific path. +function(katago_win_find_vswhere out_var) + foreach(_pf_var "ProgramFiles(x86)" "ProgramFiles") + if(DEFINED ENV{${_pf_var}}) + file(TO_CMAKE_PATH "$ENV{${_pf_var}}" _pf) + set(_candidate "${_pf}/Microsoft Visual Studio/Installer/vswhere.exe") + if(EXISTS "${_candidate}") + set(${out_var} "${_candidate}" PARENT_SCOPE) + return() + endif() + endif() + endforeach() + set(${out_var} "" PARENT_SCOPE) +endfunction() + +# Helper: fallback Windows Kits (SDK) install root candidates, used only when the registry lookup +# ("Installed Roots" key) doesn't resolve. Built from the ProgramFiles env vars rather than a +# hardcoded drive letter, so this still works on a system with Windows installed to a non-C: drive. +function(katago_win_winsdk_fallback_dirs out_var) + set(_dirs) + foreach(_pf_var "ProgramFiles(x86)" "ProgramFiles") + if(DEFINED ENV{${_pf_var}}) + file(TO_CMAKE_PATH "$ENV{${_pf_var}}" _pf) + list(APPEND _dirs "${_pf}/Windows Kits/10") + endif() + endforeach() + list(REMOVE_DUPLICATES _dirs) + set(${out_var} "${_dirs}" PARENT_SCOPE) +endfunction() + +# Helper: apply the environment vcvarsall.bat would set for a given MSVC toolset version into the +# CURRENT cmake process's environment (equivalent to the user running vcvarsall.bat before invoking +# cmake, but automatic). This is enough for every check that happens during THIS SAME configure run +# (probing below, and the actual project()-triggered HIP compiler ABI test), since execute_process +# children inherit it - but it does NOT persist into a later, separate `ninja` invocation (a sibling +# process, not a child of cmake.exe), which is why katago_win_autoselect_msvc_toolset (below) also +# bakes the resulting INCLUDE/LIB paths into cached compiler/linker flags for that to keep working. +function(katago_win_apply_vcvarsall vcvarsall_bat vcvars_ver) + # Generate a tiny .bat wrapper instead of trying to pass the whole "call ... && set" pipeline as + # one execute_process COMMAND argument - CMake's own argument quoting mangles a string that + # complex (embedded quotes plus '&&'), even though the identical command works fine typed + # directly into a shell. + set(_bat_file "${CMAKE_BINARY_DIR}/_katago_vcvars_dump.bat") + set(_out_file "${CMAKE_BINARY_DIR}/_katago_vcvars_env.txt") + file(WRITE "${_bat_file}" "@echo off\r\ncall \"${vcvarsall_bat}\" x64 -vcvars_ver=${vcvars_ver} >NUL 2>NUL\r\nset\r\n") + execute_process( + COMMAND cmd /c "${_bat_file}" + OUTPUT_FILE "${_out_file}" + RESULT_VARIABLE _rc) + if(NOT _rc EQUAL 0 OR NOT EXISTS "${_out_file}") + return() + endif() + # Read line-by-line via file(STRINGS) rather than splitting the raw output ourselves - INCLUDE/ + # LIB/PATH are themselves semicolon-separated, which would corrupt a naive string(REPLACE "\n" ";" + # ...) split; file(STRINGS) keeps each real line (with its embedded semicolons) as one list entry. + file(STRINGS "${_out_file}" _env_lines) + foreach(_line IN LISTS _env_lines) + if(_line MATCHES "^([A-Za-z_][A-Za-z0-9_().]*)=(.*)$") + set(ENV{${CMAKE_MATCH_1}} "${CMAKE_MATCH_2}") + endif() + endforeach() +endfunction() + +# Helper: with the given HIP compiler and whatever environment is currently active, try compiling a +# tiny HIP program that includes , to check whether the active MSVC toolset's STL conflicts +# with clang's own CUDA/HIP math forward declarations (see katago_win_autoselect_msvc_toolset below). +function(katago_win_probe_hip_cmath hip_compiler out_var) + set(_probe_dir "${CMAKE_BINARY_DIR}/_katago_msvc_toolset_probe") + file(MAKE_DIRECTORY "${_probe_dir}") + set(_probe_src "${_probe_dir}/probe.cpp") + if(NOT EXISTS "${_probe_src}") + file(WRITE "${_probe_src}" "#include \nint main(){ return 0; }\n") + endif() + # gfx900 (the oldest/most universally-recognized ROCm target ID, not tied to any GPU actually + # installed on this machine) is used here only to give --offload-arch a valid value so the HIP + # compile path (and its host-side resolution, which is what's actually being tested) runs + # at all - this probe is unrelated to katago_default_hip_archs()'s real per-arch compatibility + # probing for the final build's architecture list. + execute_process( + COMMAND "${hip_compiler}" --offload-arch=gfx900 -x hip -c "${_probe_src}" -o "${_probe_dir}/probe.o" + RESULT_VARIABLE _rc + OUTPUT_QUIET ERROR_QUIET) + set(${out_var} ${_rc} PARENT_SCOPE) +endfunction() + +# Helper: turn a semicolon-separated path list (like the env var INCLUDE or LIB) into a single +# string of "-Iflagprefix" (or "-Lflagprefix") tokens, skipping empty entries. +function(katago_win_pathlist_to_flags pathlist flagprefix out_var) + set(_flags "") + # CMake's own ';' list-splitting would fight with the semicolons already in the pathlist - split + # it as a plain string via regex instead of list(...). + string(REGEX MATCHALL "[^;]+" _dirs "${pathlist}") + foreach(_dir IN LISTS _dirs) + if(NOT _dir STREQUAL "") + set(_flags "${_flags} ${flagprefix}\"${_dir}\"") + endif() + endforeach() + set(${out_var} "${_flags}" PARENT_SCOPE) +endfunction() + +# Helper: find an MSVC "v143" toolset (possibly one of several 14.3x/14.4x versions installed +# side-by-side under one or more VS installs) whose STL doesn't conflict with the HIP compiler's +# own CUDA/HIP math forward-declare headers, apply the vcvarsall.bat environment for it (so the +# rest of THIS configure run - the HIP arch probing below, and project()'s own HIP compiler ABI +# test - see it), and also bake the resulting INCLUDE/LIB paths into cached compiler/linker flags +# so a later, separate `ninja` invocation (which does NOT inherit this configure run's environment +# - see katago_win_apply_vcvarsall's comment) keeps working too. +# +# Deliberately pinned to the v143 toolset family (MSVC 14.3x/14.4x) rather than probing every +# installed toolset regardless of version: a newer MSVC STL (seen with a VS "18" preview-channel +# toolset, MSVC 14.51+, i.e. the next "v144" family) declares math functions like isgreater/isless +# in a way that conflicts with clang's forward declarations for CUDA/HIP device overloads ("device +# function cannot overload host device function"), because that clang version predates the STL's +# newly-added declarations. v143 is the toolset actually verified to work with TheRock's clang; +# still probing within the v143 family (oldest first) rather than hardcoding one exact version, +# since multiple v143 point releases can be installed side by side and this needs one that's +# actually present. +function(katago_win_autoselect_msvc_toolset hip_compiler) + katago_win_find_vswhere(_vswhere) + if(NOT _vswhere) + message(FATAL_ERROR "vswhere.exe not found (expected under \"/Microsoft Visual Studio/Installer/\") - install Visual Studio Build Tools or Community with the \"Desktop development with C++\" workload.") + endif() + execute_process( + COMMAND "${_vswhere}" -all -products * -property installationPath + OUTPUT_VARIABLE _install_paths_raw + RESULT_VARIABLE _rc + OUTPUT_STRIP_TRAILING_WHITESPACE) + if(NOT _rc EQUAL 0 OR _install_paths_raw STREQUAL "") + message(FATAL_ERROR "vswhere found no Visual Studio installation - install Visual Studio Build Tools or Community with the \"Desktop development with C++\" workload.") + endif() + string(REPLACE "\r\n" "\n" _install_paths_raw "${_install_paths_raw}") + string(REPLACE "\n" ";" _install_paths "${_install_paths_raw}") + + set(_toolsets) + foreach(_install_path IN LISTS _install_paths) + if(_install_path STREQUAL "") + continue() + endif() + file(TO_CMAKE_PATH "${_install_path}" _install_path) + file(GLOB _tool_dirs "${_install_path}/VC/Tools/MSVC/*") + foreach(_t IN LISTS _tool_dirs) + if(IS_DIRECTORY "${_t}" AND EXISTS "${_t}/include/cmath") + get_filename_component(_ver "${_t}" NAME) + # v143 toolset versions are 14.3x/14.4x (the next-generation "v144" toolset, seen with + # VS 18 preview builds, starts at 14.5x and is the one known to conflict - see above). + if(NOT _ver MATCHES "^14\\.[34][0-9]\\.") + continue() + endif() + set(_vcvarsall "${_install_path}/VC/Auxiliary/Build/vcvarsall.bat") + if(EXISTS "${_vcvarsall}") + list(APPEND _toolsets "${_ver}|${_vcvarsall}") + endif() + endif() + endforeach() + endforeach() + + if(NOT _toolsets) + message(FATAL_ERROR "No installed MSVC v143 toolset (14.3x/14.4x) found. Open the Visual Studio Installer, choose \"Modify\" on your VS installation, go to \"Individual Components\", and install \"MSVC v143 - VS 2022 C++ x64/x86 build tools\" (a newer toolset alone, e.g. a VS 18 preview's v144, is not compatible with TheRock's bundled clang).") + endif() + list(REMOVE_DUPLICATES _toolsets) + list(SORT _toolsets COMPARE NATURAL ORDER ASCENDING) + + foreach(_entry IN LISTS _toolsets) + string(REPLACE "|" ";" _entry_list "${_entry}") + list(GET _entry_list 0 _ver) + list(GET _entry_list 1 _vcvarsall) + message(STATUS "Trying MSVC toolset ${_ver} for HIP compiler compatibility...") + katago_win_apply_vcvarsall("${_vcvarsall}" "${_ver}") + katago_win_probe_hip_cmath("${hip_compiler}" _probe_rc) + if(_probe_rc EQUAL 0) + message(STATUS "MSVC toolset ${_ver} is compatible with the HIP compiler; using it") + # Bake this toolset's resolved INCLUDE/LIB (now sitting in ENV{} courtesy of + # katago_win_apply_vcvarsall) into cached flags, so a later separate `ninja` invocation - + # which won't have this environment - still finds the same headers/libs. + # + # Clang's own resource-dir/include (containing its clang-compatible emmintrin.h/immintrin.h + # etc) must come FIRST, ahead of the MSVC toolset's own include dir: pulls in + # , and MSVC's own / declare SSE/AVX intrinsics using + # __declspec(intrin_type), a cl.exe-only mechanism clang doesn't understand - it just sees an + # ordinary extern function declaration with no body, so any code that ends up calling one + # (e.g. the MSVC STL's wmemcmp) links with an "undefined symbol: _mm_loadu_si128"-style error. + # Explicitly preferring clang's own compatible headers for those specific filenames avoids + # this, while still falling through to the MSVC/SDK -I entries below for headers clang has no + # equivalent for (the actual MSVC STL, ucrt, etc). This needs a plain -I (not -isystem/ + # -idirafter): the earlier __clang_hip_runtime_wrapper.h force-included in every HIP + # compilation uses #include_next to reach the "real" starting from wherever clang's + # own cuda_wrappers/cmath shim was found, and that continuation only walks through the same + # -I search list, not -isystem/-idirafter entries. + execute_process(COMMAND "${hip_compiler}" -print-resource-dir + OUTPUT_VARIABLE _clang_resource_dir OUTPUT_STRIP_TRAILING_WHITESPACE) + file(TO_CMAKE_PATH "${_clang_resource_dir}" _clang_resource_dir) + set(_inc_flags "-I\"${_clang_resource_dir}/include\"") + katago_win_pathlist_to_flags("$ENV{INCLUDE}" "-I" _msvc_inc_flags) + katago_win_pathlist_to_flags("$ENV{LIB}" "-L" _lib_flags) + set(_inc_flags "${_inc_flags} ${_msvc_inc_flags}") + set(CMAKE_HIP_FLAGS "${CMAKE_HIP_FLAGS} ${_inc_flags}" CACHE STRING "" FORCE) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${_inc_flags}" CACHE STRING "" FORCE) + set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${_inc_flags}" CACHE STRING "" FORCE) + set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} ${_lib_flags}" CACHE STRING "" FORCE) + return() + else() + message(STATUS "MSVC toolset ${_ver} conflicts with the HIP compiler's CUDA/HIP math headers; trying another") + endif() + endforeach() + message(FATAL_ERROR "None of the installed MSVC v143 toolset(s) (14.3x/14.4x) compiled cleanly against TheRock's HIP compiler. This indicates a problem beyond the usual v143-vs-newer-toolset conflict this check is meant to catch - please report this along with the MSVC/TheRock versions involved.") +endfunction() + if(USE_BACKEND STREQUAL "METAL") project(katago LANGUAGES CXX Swift) +elseif(USE_BACKEND STREQUAL "ROCM") + if(WIN32) + # Normalize HIP_PATH to forward slashes. + # cmake's HIP detection calls `hipconfig --rocmpath` which returns + # backslash paths on Windows; it writes the result into the generated + # CMakeHIPCompiler.cmake without normalizing, causing "Invalid character + # escape '\T'" errors. Pre-setting CMAKE_HIP_COMPILER_ROCM_ROOT (and + # CMAKE_PREFIX_PATH) before project() bypasses that detection entirely. + if(DEFINED ENV{HIP_PATH}) + file(TO_CMAKE_PATH "$ENV{HIP_PATH}" _hip_path_fwd) + set(ENV{HIP_PATH} "${_hip_path_fwd}") + list(APPEND CMAKE_PREFIX_PATH "${_hip_path_fwd}") + if(NOT CMAKE_HIP_COMPILER_ROCM_ROOT) + set(CMAKE_HIP_COMPILER_ROCM_ROOT "${_hip_path_fwd}" CACHE PATH "" FORCE) + endif() + endif() + # ---------- C/C++ compiler (clang++ from HIP SDK) ---------- + # hipcc.exe was tried here (to mirror the Linux code path, where CMAKE_CXX_COMPILER=hipcc + # transparently treats every .cpp as HIP source) but it mis-tokenizes quoted paths containing + # spaces (e.g. "C:/Program Files (x86)/Windows Kits/..."), breaking the compiler-detection step + # entirely - not usable as CMAKE_CXX_COMPILER on Windows. Use clang++ directly instead; the + # rocmbackend.cpp / CK generated sources that need real HIP compilation get it via an explicit + # LANGUAGE HIP source property instead (see below and the CK FMHA section further down). + if(NOT CMAKE_CXX_COMPILER) + if(DEFINED ENV{HIP_PATH}) + if(EXISTS "$ENV{HIP_PATH}/lib/llvm/bin/clang++.exe") + set(CMAKE_CXX_COMPILER "$ENV{HIP_PATH}/lib/llvm/bin/clang++.exe" CACHE FILEPATH "" FORCE) + set(CMAKE_C_COMPILER "$ENV{HIP_PATH}/lib/llvm/bin/clang.exe" CACHE FILEPATH "" FORCE) + elseif(EXISTS "$ENV{HIP_PATH}/bin/clang++.exe") + set(CMAKE_CXX_COMPILER "$ENV{HIP_PATH}/bin/clang++.exe" CACHE FILEPATH "" FORCE) + set(CMAKE_C_COMPILER "$ENV{HIP_PATH}/bin/clang.exe" CACHE FILEPATH "" FORCE) + endif() + endif() + endif() + # HIP compiler = same binary as C++ compiler + if(NOT CMAKE_HIP_COMPILER AND CMAKE_CXX_COMPILER) + set(CMAKE_HIP_COMPILER "${CMAKE_CXX_COMPILER}" CACHE FILEPATH "" FORCE) + endif() + # ---------- MSVC toolset auto-selection (avoid newer-STL / clang HIP header conflicts) ---------- + # Automatically finds and applies (via vcvarsall.bat) a locally-installed MSVC toolset that's + # compatible with this HIP compiler - see katago_win_autoselect_msvc_toolset()'s comment above + # for why this is needed. Runs once per fresh CMakeCache (cached so a plain re-run of cmake + # doesn't redo the probing every time). + if(CMAKE_HIP_COMPILER AND NOT KATAGO_WIN_MSVC_TOOLSET_CHECKED) + katago_win_autoselect_msvc_toolset("${CMAKE_HIP_COMPILER}") + set(KATAGO_WIN_MSVC_TOOLSET_CHECKED TRUE CACHE INTERNAL "") + endif() + # ---------- HIP architectures (must be set before project() / enable_language(HIP)) ---------- + # Default to the broad set of AMD GPU architectures KataGo supports unless + # the user explicitly passed -DCMAKE_HIP_ARCHITECTURES=... on the command line. + if(NOT DEFINED CMAKE_HIP_ARCHITECTURES) + katago_default_hip_archs(_default_archs) + set(CMAKE_HIP_ARCHITECTURES "${_default_archs}" CACHE STRING "Default broad set of AMD GPU targets") + message(STATUS "Pre-project default CMAKE_HIP_ARCHITECTURES=${CMAKE_HIP_ARCHITECTURES}") + else() + message(STATUS "Pre-project user-specified CMAKE_HIP_ARCHITECTURES=${CMAKE_HIP_ARCHITECTURES}") + endif() + # ---------- Windows SDK includes (needed by HIP compiler test during project()) ---------- + # The HIP runtime wrapper includes MSVC headers that require Windows SDK ucrt/shared/um. + # These flags must be set before project() so cmake's HIP compiler test can compile. + if(NOT KATAGO_WINSDK_ROOT) + get_filename_component(_pre_winsdk_root + "[HKEY_LOCAL_MACHINE\\SOFTWARE\\Microsoft\\Windows Kits\\Installed Roots;KitsRoot10]" + ABSOLUTE) + if(NOT EXISTS "${_pre_winsdk_root}") + katago_win_winsdk_fallback_dirs(_pre_winsdk_candidates) + foreach(_p IN LISTS _pre_winsdk_candidates) + if(EXISTS "${_p}") + set(_pre_winsdk_root "${_p}") + break() + endif() + endforeach() + endif() + if(EXISTS "${_pre_winsdk_root}") + set(KATAGO_WINSDK_ROOT "${_pre_winsdk_root}" CACHE INTERNAL "") + endif() + endif() + if(KATAGO_WINSDK_ROOT) + file(GLOB _pre_sdk_ver_dirs "${KATAGO_WINSDK_ROOT}/Include/*/ucrt") + if(_pre_sdk_ver_dirs) + list(SORT _pre_sdk_ver_dirs ORDER DESCENDING) + list(GET _pre_sdk_ver_dirs 0 _pre_ucrt_dir) + get_filename_component(_pre_sdk_ver_dir "${_pre_ucrt_dir}" DIRECTORY) + set(_winsdk_cflags + "-I\"${_pre_sdk_ver_dir}/ucrt\" -I\"${_pre_sdk_ver_dir}/shared\" -I\"${_pre_sdk_ver_dir}/um\"") + set(CMAKE_HIP_FLAGS "${CMAKE_HIP_FLAGS} ${_winsdk_cflags}" CACHE STRING "" FORCE) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${_winsdk_cflags}" CACHE STRING "" FORCE) + set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${_winsdk_cflags}" CACHE STRING "" FORCE) + message(STATUS "Pre-project: injected Windows SDK includes from ${_pre_sdk_ver_dir}") + endif() + endif() + # ---------- RC compiler (Windows resource compiler, highest SDK version) ---------- + if(NOT CMAKE_RC_COMPILER) + # Try registry first + get_filename_component(_winsdk_root + "[HKEY_LOCAL_MACHINE\\SOFTWARE\\Microsoft\\Windows Kits\\Installed Roots;KitsRoot10]" + ABSOLUTE) + if(NOT EXISTS "${_winsdk_root}") + katago_win_winsdk_fallback_dirs(_winsdk_candidates) + foreach(_p IN LISTS _winsdk_candidates) + if(EXISTS "${_p}") + set(_winsdk_root "${_p}") + break() + endif() + endforeach() + endif() + if(EXISTS "${_winsdk_root}") + file(GLOB _rc_candidates "${_winsdk_root}/bin/*/x64/rc.exe") + if(_rc_candidates) + list(SORT _rc_candidates ORDER DESCENDING) + list(GET _rc_candidates 0 _rc_exe) + set(CMAKE_RC_COMPILER "${_rc_exe}" CACHE FILEPATH "" FORCE) + endif() + # Persist for use in compiler flags section below + set(KATAGO_WINSDK_ROOT "${_winsdk_root}" CACHE INTERNAL "") + endif() + endif() + else() + # Linux: resolve the real ROCm install prefix first (see katago_find_rocm_prefix - handles the + # /opt/rocm/core- layout used since ROCm ~7.9, preferring the highest version found, + # falling back to the old flat /opt/rocm layout). Respect a user-provided CMAKE_PREFIX_PATH. + if(NOT DEFINED CMAKE_PREFIX_PATH OR CMAKE_PREFIX_PATH STREQUAL "") + katago_find_rocm_prefix(KATAGO_ROCM_PREFIX) + if(KATAGO_ROCM_PREFIX) + list(APPEND CMAKE_PREFIX_PATH "${KATAGO_ROCM_PREFIX}") + message(STATUS "Pre-project: resolved ROCm install prefix to ${KATAGO_ROCM_PREFIX}") + endif() + endif() + + # Use hipcc. Its location varies across ROCm packaging generations (a classic apt "hipcc" + # package at /usr/bin, one bundled under the ROCm prefix's bin/, or newer installs that only + # ship amdclang++ with no hipcc wrapper at all) - probe candidates instead of hardcoding one + # path, since a missing hardcoded path fails project() below with a confusing compiler-not- + # found error rather than falling back. + find_program(KATAGO_HIPCC_EXECUTABLE NAMES hipcc + HINTS ${CMAKE_PREFIX_PATH} + PATH_SUFFIXES bin) + if(NOT KATAGO_HIPCC_EXECUTABLE AND EXISTS "/usr/bin/hipcc") + set(KATAGO_HIPCC_EXECUTABLE "/usr/bin/hipcc") + endif() + # Search for clang++ unconditionally (not just as a hipcc-missing fallback): CMAKE_HIP_COMPILER + # always needs the clang++ binary directly below, even when hipcc is used for C/CXX, since CMake + # rejects the hipcc wrapper script for CMAKE_HIP_COMPILER ("...not supported"). + find_program(KATAGO_CLANGXX_EXECUTABLE NAMES clang++ + HINTS ${CMAKE_PREFIX_PATH} + PATH_SUFFIXES lib/llvm/bin bin) + if(KATAGO_HIPCC_EXECUTABLE) + set(CMAKE_C_COMPILER "${KATAGO_HIPCC_EXECUTABLE}" CACHE FILEPATH "" FORCE) + set(CMAKE_CXX_COMPILER "${KATAGO_HIPCC_EXECUTABLE}" CACHE FILEPATH "" FORCE) + elseif(KATAGO_CLANGXX_EXECUTABLE) + message(STATUS "hipcc not found; falling back to ${KATAGO_CLANGXX_EXECUTABLE}") + get_filename_component(_clang_dir "${KATAGO_CLANGXX_EXECUTABLE}" DIRECTORY) + set(CMAKE_CXX_COMPILER "${KATAGO_CLANGXX_EXECUTABLE}" CACHE FILEPATH "" FORCE) + set(CMAKE_C_COMPILER "${_clang_dir}/clang" CACHE FILEPATH "" FORCE) + endif() + # ---------- HIP architectures (must be set before project()/enable_language(HIP)) ---------- + # project(... HIP) below triggers CMake's own HIP language detection, which auto-populates + # CMAKE_HIP_ARCHITECTURES with just the native/current GPU's arch if it isn't already set by + # then. That makes the broad-default-arch logic further down (before enable_language(HIP)) + # a no-op, since by that point CMAKE_HIP_ARCHITECTURES already looks "user-specified". Set the + # broad default here instead, before project(), mirroring the Windows pre-project block above. + # Note: CMAKE_HIP_COMPILER must be the clang++ binary directly - CMake rejects the hipcc + # wrapper script here ("CMAKE_HIP_COMPILER is set to the hipcc wrapper... not supported"). + if(NOT CMAKE_HIP_COMPILER AND KATAGO_CLANGXX_EXECUTABLE) + set(CMAKE_HIP_COMPILER "${KATAGO_CLANGXX_EXECUTABLE}" CACHE FILEPATH "" FORCE) + endif() + if(NOT DEFINED CMAKE_HIP_ARCHITECTURES) + katago_default_hip_archs(_default_archs) + set(CMAKE_HIP_ARCHITECTURES "${_default_archs}" CACHE STRING "Default broad set of AMD GPU targets") + message(STATUS "Pre-project default CMAKE_HIP_ARCHITECTURES=${CMAKE_HIP_ARCHITECTURES}") + else() + message(STATUS "Pre-project user-specified CMAKE_HIP_ARCHITECTURES=${CMAKE_HIP_ARCHITECTURES}") + endif() + endif() + project(katago LANGUAGES C CXX HIP) else() project(katago) endif() @@ -44,7 +516,7 @@ endif() set(BUILD_DISTRIBUTED 0 CACHE BOOL "Build with http support for contributing to distributed training") set(USE_BACKEND CACHE STRING "Neural net backend") string(TOUPPER "${USE_BACKEND}" USE_BACKEND) -set_property(CACHE USE_BACKEND PROPERTY STRINGS "" CUDA TENSORRT OPENCL EIGEN METAL) +set_property(CACHE USE_BACKEND PROPERTY STRINGS "" CUDA TENSORRT OPENCL EIGEN METAL ROCM) set(USE_TCMALLOC 0 CACHE BOOL "Use TCMalloc") set(NO_GIT_REVISION 0 CACHE BOOL "Disable embedding the git revision into the compiled exe") @@ -54,6 +526,111 @@ 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) +# ---------- 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. Mirrors the same +# KATAGO_AUTO_FETCH_DEPS / vcpkg-in-build-tree pattern used on the Intel_NPU branch, so this stays +# consistent across branches instead of reinventing a different fetch mechanism here. +if(WIN32 OR (UNIX AND NOT APPLE)) + set(_katago_auto_fetch_default ON) +else() + set(_katago_auto_fetch_default OFF) +endif() +option(KATAGO_AUTO_FETCH_DEPS "Automatically fetch missing dependencies into build/deps (Windows/Linux use vcpkg)." ${_katago_auto_fetch_default}) +set(KATAGO_DEPS_DIR "${CMAKE_SOURCE_DIR}/build/deps" CACHE PATH "Directory for auto-fetched third-party dependencies") +if(WIN32) + set(_katago_vcpkg_triplet_default "x64-windows") +elseif(UNIX AND NOT APPLE) + if(CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64)$") + set(_katago_vcpkg_triplet_default "arm64-linux") + else() + set(_katago_vcpkg_triplet_default "x64-linux") + endif() +else() + set(_katago_vcpkg_triplet_default "x64-windows") +endif() +set(KATAGO_VCPKG_TRIPLET "${_katago_vcpkg_triplet_default}" CACHE STRING "vcpkg triplet used by KATAGO_AUTO_FETCH_DEPS") +set(KATAGO_VCPKG_ROOT "${KATAGO_DEPS_DIR}/vcpkg" CACHE PATH "Path to local vcpkg clone used by KATAGO_AUTO_FETCH_DEPS") +mark_as_advanced(KATAGO_VCPKG_TRIPLET KATAGO_VCPKG_ROOT) + +function(katago_vcpkg_bootstrap_if_needed) + if(NOT WIN32 AND NOT (UNIX AND NOT APPLE)) + message(FATAL_ERROR "katago_vcpkg_bootstrap_if_needed is only supported on Windows and Linux") + endif() + + if(NOT KATAGO_AUTO_FETCH_DEPS) + message(FATAL_ERROR "KATAGO_AUTO_FETCH_DEPS is OFF, cannot auto-fetch missing dependency") + endif() + + file(MAKE_DIRECTORY "${KATAGO_DEPS_DIR}") + + if(WIN32) + set(_katago_vcpkg_exe "${KATAGO_VCPKG_ROOT}/vcpkg.exe") + else() + set(_katago_vcpkg_exe "${KATAGO_VCPKG_ROOT}/vcpkg") + endif() + + if(NOT EXISTS "${_katago_vcpkg_exe}") + if(NOT EXISTS "${KATAGO_VCPKG_ROOT}/.git") + find_package(Git QUIET) + if(NOT GIT_FOUND) + message(FATAL_ERROR "KATAGO_AUTO_FETCH_DEPS requires git to clone vcpkg") + endif() + message(STATUS "Auto-fetch deps: cloning vcpkg into ${KATAGO_VCPKG_ROOT}") + execute_process( + COMMAND "${GIT_EXECUTABLE}" clone --depth=1 https://github.com/microsoft/vcpkg.git "${KATAGO_VCPKG_ROOT}" + RESULT_VARIABLE _clone_result + OUTPUT_VARIABLE _clone_out + ERROR_VARIABLE _clone_err + ) + if(NOT _clone_result EQUAL 0) + message(FATAL_ERROR "Failed to clone vcpkg.\n${_clone_out}\n${_clone_err}") + endif() + endif() + + message(STATUS "Auto-fetch deps: bootstrapping vcpkg") + if(WIN32) + execute_process( + COMMAND "${KATAGO_VCPKG_ROOT}/bootstrap-vcpkg.bat" -disableMetrics + WORKING_DIRECTORY "${KATAGO_VCPKG_ROOT}" + RESULT_VARIABLE _bootstrap_result + ) + else() + execute_process( + COMMAND sh "${KATAGO_VCPKG_ROOT}/bootstrap-vcpkg.sh" -disableMetrics + WORKING_DIRECTORY "${KATAGO_VCPKG_ROOT}" + RESULT_VARIABLE _bootstrap_result + ) + endif() + if(NOT _bootstrap_result EQUAL 0) + message(FATAL_ERROR "Failed to bootstrap vcpkg") + endif() + endif() +endfunction() + +function(katago_vcpkg_install_if_needed package_name) + katago_vcpkg_bootstrap_if_needed() + + if(WIN32) + set(_katago_vcpkg_exe "${KATAGO_VCPKG_ROOT}/vcpkg.exe") + else() + set(_katago_vcpkg_exe "${KATAGO_VCPKG_ROOT}/vcpkg") + endif() + if(NOT EXISTS "${_katago_vcpkg_exe}") + message(FATAL_ERROR "vcpkg executable not found after bootstrap: ${_katago_vcpkg_exe}") + endif() + + set(_spec "${package_name}:${KATAGO_VCPKG_TRIPLET}") + message(STATUS "Auto-fetch deps: ensuring ${_spec} via vcpkg") + execute_process( + COMMAND "${_katago_vcpkg_exe}" install "${_spec}" --disable-metrics + WORKING_DIRECTORY "${KATAGO_VCPKG_ROOT}" + RESULT_VARIABLE _install_result + ) + if(NOT _install_result EQUAL 0) + message(FATAL_ERROR "Failed to install ${_spec} via vcpkg") + endif() +endfunction() + #--------------------------- NEURAL NET BACKEND ------------------------------------------------------------------------ message(STATUS "Building 'katago' executable for GTP engine and other tools.") @@ -164,6 +741,72 @@ elseif(USE_BACKEND STREQUAL "EIGEN") set(NEURALNET_BACKEND_SOURCES neuralnet/eigenbackend.cpp ) +# --------------------------- ROCM backend（AMD GPU / HIP MIOpen） --------------------------- +elseif(USE_BACKEND STREQUAL "ROCM") + message(STATUS "-DUSE_BACKEND=ROCM, using AMD ROCm backend.") + + if(CMAKE_PREFIX_PATH STREQUAL "" OR NOT DEFINED CMAKE_PREFIX_PATH) + if(WIN32) + # Windows: HIP SDK installed via installer or manually + if(DEFINED ENV{HIP_PATH}) + list(APPEND CMAKE_PREFIX_PATH "$ENV{HIP_PATH}") + message(STATUS "Auto-detected HIP_PATH=$ENV{HIP_PATH} → CMAKE_PREFIX_PATH") + elseif(DEFINED ENV{ROCM_PATH}) + list(APPEND CMAKE_PREFIX_PATH "$ENV{ROCM_PATH}") + message(STATUS "Auto-detected ROCM_PATH=$ENV{ROCM_PATH} → CMAKE_PREFIX_PATH") + else() + message(FATAL_ERROR "ROCM backend on Windows requires HIP_PATH or ROCM_PATH environment variable (or -DCMAKE_PREFIX_PATH=). Please install the HIP SDK for Windows.") + endif() + else() + # Linux: resolve the real ROCm install prefix (same core- vs flat /opt/rocm + # resolution as the pre-project block above; this block re-runs it standalone in case + # USE_BACKEND wasn't ROCM early enough for that block to have fired). + katago_find_rocm_prefix(KATAGO_ROCM_PREFIX) + if(KATAGO_ROCM_PREFIX) + list(APPEND CMAKE_PREFIX_PATH "${KATAGO_ROCM_PREFIX}") + message(STATUS "CMAKE_PREFIX_PATH not given; defaulting to ${KATAGO_ROCM_PREFIX}") + endif() + endif() + endif() + + # Ensure CMAKE_HIP_COMPILER is set BEFORE we run the optional-arch probe and before + # enable_language(HIP). On Windows the pre-project block already sets it; on Linux it should + # already be set by the pre-project block too - this is just a safety net. Must be the clang++ + # binary directly, not the hipcc wrapper (CMake rejects hipcc for CMAKE_HIP_COMPILER). + if(NOT WIN32 AND NOT CMAKE_HIP_COMPILER) + if(NOT KATAGO_CLANGXX_EXECUTABLE) + find_program(KATAGO_CLANGXX_EXECUTABLE NAMES clang++ HINTS ${CMAKE_PREFIX_PATH} PATH_SUFFIXES lib/llvm/bin bin) + endif() + if(KATAGO_CLANGXX_EXECUTABLE) + set(CMAKE_HIP_COMPILER "${KATAGO_CLANGXX_EXECUTABLE}" CACHE FILEPATH "" FORCE) + endif() + endif() + + # Default GPU architectures must be set BEFORE enable_language(HIP). + # Users who pass -DCMAKE_HIP_ARCHITECTURES=... on the command line get + # exactly that list; otherwise we default to the broad set KataGo supports + # (plus gfx1031/gfx1032 if the compiler accepts them). + if(NOT DEFINED CMAKE_HIP_ARCHITECTURES) + katago_default_hip_archs(_default_archs) + set(CMAKE_HIP_ARCHITECTURES "${_default_archs}" CACHE STRING "Default broad set of AMD GPU targets") + message(STATUS "Default CMAKE_HIP_ARCHITECTURES=${CMAKE_HIP_ARCHITECTURES}") + else() + message(STATUS "User-specified CMAKE_HIP_ARCHITECTURES=${CMAKE_HIP_ARCHITECTURES}") + endif() + + enable_language(HIP) + set(CMAKE_HIP_STANDARD 17) + + # 2) Specify backend source code. rocmhelpers.hip contains GPU kernels, don't forget it + set(NEURALNET_BACKEND_SOURCES + neuralnet/rocmbackend.cpp + neuralnet/rocmutils.cpp + neuralnet/rocmhelpers.hip + ) + + # Optional: Enable model-size‑based autotuning and other macros + # add_compile_definitions(HIP_SUPPORTS_FP16) + elseif(USE_BACKEND STREQUAL "") message(WARNING "${ColorBoldRed}WARNING: Using dummy neural net backend, intended for non-neural-net testing only, will fail on any code path requiring a neural net. To use neural net, specify -DUSE_BACKEND=CUDA or -DUSE_BACKEND=TENSORRT or -DUSE_BACKEND=OPENCL or -DUSE_BACKEND=EIGEN to compile with the respective backend.${ColorReset}") set(NEURALNET_BACKEND_SOURCES neuralnet/dummybackend.cpp) @@ -504,6 +1147,217 @@ elseif(USE_BACKEND STREQUAL "OPENCL") link_directories(${OpenCL_LIBRARY}) target_link_libraries(katago ${OpenCL_LIBRARY}) endif() +# --------------------------- ROCM linking stage --------------------------- +elseif(USE_BACKEND STREQUAL "ROCM") + # Macro: used in source code with #ifdef USE_ROCM_BACKEND + target_compile_definitions(katago PRIVATE USE_ROCM_BACKEND) + target_compile_definitions(katago PRIVATE HIP_TARGET_VERSION=${CMAKE_HIP_COMPILER_VERSION}) + + string(TOLOWER "${CMAKE_HIP_ARCHITECTURES}" _gfxlist) # e.g. "90a;942" + # Every architecture in katago_default_hip_archs()'s full candidate list (CDNA gfx906/908/90a/942/ + # 950, RDNA1 gfx101x, RDNA2 gfx103x, RDNA3 gfx110x, RDNA3.5 gfx115x, RDNA4 gfx120x) supports packed + # FP16 ops. The regex below matches each of them so user-specified subset lists still work. + if(_gfxlist MATCHES "(gfx)?(90[06a]|908|94[02]|950|101[012]|103[0-6]|110[0-3]|115[0-3]|120[01])") + target_compile_definitions(katago PRIVATE HIP_SUPPORTS_FP16) + message(STATUS "Detected FP16‑capable GFX arch (${CMAKE_HIP_ARCHITECTURES}); defining HIP_SUPPORTS_FP16") + endif() + + # 3) Find ROCm runtime & libraries. Since ROCm 6.x, CMake config-mode packages are included. If not found, add -DCMAKE_PREFIX_PATH=/opt/rocm (Linux) or HIP SDK path (Windows) + find_package(hip QUIET CONFIG) # Export hip::device / hip::host + find_package(hipblas QUIET CONFIG) # Export roc::hipblas + find_package(miopen QUIET CONFIG) # Export roc::miopen or MIOpen + + # Explicitly (re-)add the resolved HIP include dir via a raw -isystem compile option (not + # target_include_directories(... SYSTEM ...)): hip::device's own INTERFACE_INCLUDE_DIRECTORIES + # already carries this same path as a plain -I, and CMake's include-directory de-duplication + # drops our SYSTEM entry as "already covered" rather than upgrading it - so the plain -I from + # hip::device wins, which matters here because some systems have a stale hip_runtime.h under + # /usr/include (e.g. from an old standalone "hipcc"/"libamdhip64-dev" apt package) that's + # incompatible with a newer compiler's builtins (__AMDGCN_WAVEFRONT_SIZE etc). A plain -I loses to + # that stale /usr/include; a raw -isystem (bypassing the dedup) wins. + find_path(KATAGO_HIP_SYSTEM_INCLUDE_DIR hip/hip_runtime.h + HINTS ${CMAKE_PREFIX_PATH} ENV HIP_PATH ENV ROCM_PATH + PATH_SUFFIXES include) + if(KATAGO_HIP_SYSTEM_INCLUDE_DIR) + target_compile_options(katago PRIVATE "-isystem${KATAGO_HIP_SYSTEM_INCLUDE_DIR}") + endif() + + # ---------- fallback：HIP Runtime ---------- + if(NOT hip_FOUND) + if(WIN32) + # Windows: Search in HIP SDK installation + find_path(HIP_INCLUDE_DIR hip/hip_runtime.h + HINTS ${CMAKE_PREFIX_PATH} ENV HIP_PATH ENV ROCM_PATH + PATH_SUFFIXES include) + find_library(HIP_RUNTIME_LIB + NAMES amdhip64 amdhip64_6 + HINTS ${CMAKE_PREFIX_PATH} ENV HIP_PATH ENV ROCM_PATH + PATH_SUFFIXES lib bin) + else() + # Linux: Search in /opt/rocm + find_path(HIP_INCLUDE_DIR hip/hip_runtime.h + HINTS ${CMAKE_PREFIX_PATH} /opt/rocm + PATH_SUFFIXES include) + find_library(HIP_RUNTIME_LIB amdhip64 + HINTS ${CMAKE_PREFIX_PATH} /opt/rocm + PATH_SUFFIXES lib lib64) + endif() + + if(NOT HIP_INCLUDE_DIR OR NOT HIP_RUNTIME_LIB) + if(WIN32) + message(FATAL_ERROR "HIP headers or runtime NOT found; install HIP SDK for Windows or set CMAKE_PREFIX_PATH to HIP SDK installation path.") + else() + message(FATAL_ERROR "HIP headers or runtime NOT found; install ROCm or set CMAKE_PREFIX_PATH.") + endif() + endif() + add_library(hip::device UNKNOWN IMPORTED) + set_target_properties(hip::device PROPERTIES + IMPORTED_LOCATION "${HIP_RUNTIME_LIB}" + INTERFACE_INCLUDE_DIRECTORIES "${HIP_INCLUDE_DIR}") + target_include_directories(katago SYSTEM PRIVATE ${HIP_INCLUDE_DIR}) + endif() + + # ---------- fallback：hipBLAS / MIOpen ---------- + foreach(_pkg hipblas miopen) + if(NOT ${_pkg}_FOUND) + if(WIN32) + # Windows naming conventions + if(_pkg STREQUAL "hipblas") + set(_lib_names hipblas) + else() + set(_lib_names MIOpen) + endif() + find_library(${_pkg}_LIB + NAMES ${_lib_names} + HINTS ${CMAKE_PREFIX_PATH} ENV HIP_PATH ENV ROCM_PATH + PATH_SUFFIXES lib bin) + else() + # Linux naming + find_library(${_pkg}_LIB ${_pkg} + HINTS ${CMAKE_PREFIX_PATH} /opt/rocm + PATH_SUFFIXES lib lib64) + endif() + + if(${_pkg}_LIB) + add_library(roc::${_pkg} UNKNOWN IMPORTED) + set_target_properties(roc::${_pkg} PROPERTIES + IMPORTED_LOCATION "${${_pkg}_LIB}") + target_include_directories(katago SYSTEM PRIVATE ${HIP_INCLUDE_DIR}) + message(STATUS "Found ${_pkg} at ${${_pkg}_LIB}") + else() + if(WIN32) + message(FATAL_ERROR "Required ROCm component ${_pkg} not found – install HIP SDK for Windows or set CMAKE_PREFIX_PATH.") + else() + message(FATAL_ERROR "Required ROCm component ${_pkg} not found – install it or set CMAKE_PREFIX_PATH.") + endif() + endif() + endif() + endforeach() + + # 4) Link libraries + # Note: On Windows, MIOpen might need to be linked as "MIOpen" directly if the target doesn't exist + if(TARGET MIOpen) + set(_miopen_target MIOpen) + elseif(TARGET roc::miopen) + set(_miopen_target roc::miopen) + else() + message(FATAL_ERROR "Neither 'MIOpen' nor 'roc::miopen' imported target exists after find_package + fallback. This is a bug in the ROCM backend detection logic in cpp/CMakeLists.txt.") + endif() + + target_link_libraries(katago + hip::device # HIP runtime & kernel offload + roc::hipblas # BLAS + ${_miopen_target} # DNN primitives + ) + + # See the comment in core/win_rocm_sse_shim.cpp: on Windows, ucrt's own ends up calling + # SSE2 intrinsics that resolve to bodyless declarations under this toolchain, causing "undefined + # symbol" link errors. This file is a no-op everywhere except Windows (guarded internally). + if(WIN32) + target_sources(katago PRIVATE core/win_rocm_sse_shim.cpp) + endif() + + # KATAGO_ROCM_PREFIX may not be set if the user passed an explicit -DCMAKE_PREFIX_PATH (which + # skips the auto-detect branches that populate it) - fall back to CMAKE_PREFIX_PATH's first entry + # in that case, so RPATH still gets set below regardless of how the prefix was determined. + if(NOT KATAGO_ROCM_PREFIX) + if(CMAKE_PREFIX_PATH) + list(GET CMAKE_PREFIX_PATH 0 KATAGO_ROCM_PREFIX) + else() + katago_find_rocm_prefix(KATAGO_ROCM_PREFIX) + endif() + endif() + + # Bake the resolved ROCm lib dir into the binary's RPATH. Without this, the built katago relies + # on /opt/rocm (a distro-managed symlink that can get silently repointed by installing another + # ROCm package, or a subsequent /opt/rocm/core- release) resolving to the same + # install this was actually built and tested against - which is not guaranteed, especially for a + # binary meant to be copied to another machine. CMAKE_SKIP_BUILD_RPATH/INSTALL_RPATH aren't set by + # this project, so this also takes effect for the build-tree binary, not just `make install`. + if(KATAGO_ROCM_PREFIX) + set_target_properties(katago PROPERTIES + BUILD_RPATH "${KATAGO_ROCM_PREFIX}/lib" + INSTALL_RPATH "${KATAGO_ROCM_PREFIX}/lib" + ) + message(STATUS "Set katago RPATH to ${KATAGO_ROCM_PREFIX}/lib") + endif() + + # ---------- Optional: Composable Kernel FMHA (fused attention) ---------- + # Mirrors the CUDA backend's optional cudnn-frontend SDPA path: KataGo vendors a small set of + # glue headers + pre-generated ck_tile FMHA kernel instantiations under + # external/composable_kernel_fmha, but these only compile against a matching version of CK's own + # core ck_tile headers (system dependency, not vendored - part of a "composablekernel-dev" / + # "amdrocm-ck*" package). If a compatible version isn't installed, we just skip this and the + # backend always uses its own plain (non-fused) attention kernel - this is a pure performance + # optimization, not required for correctness. + find_path(KATAGO_CK_TILE_INCLUDE_DIR + NAMES ck_tile/ops/fmha_fwd.hpp + HINTS ${CMAKE_PREFIX_PATH} ENV ROCM_PATH + PATH_SUFFIXES include + ) + if(KATAGO_CK_TILE_INCLUDE_DIR) + message(STATUS "Found ck_tile headers at ${KATAGO_CK_TILE_INCLUDE_DIR}; enabling optional CK FMHA fused attention path") + file(GLOB KATAGO_CK_FMHA_GENERATED_SOURCES external/composable_kernel_fmha/generated/*.cpp) + # ck_tile's own arch.hpp (get_compiler_target()) has no case at all for gfx906 or RDNA1 + # (gfx1010/1011/1012) - its FMHA kernels were never ported to these (no MFMA/WMMA), so its + # constexpr dispatch function silently falls through to nothing on a device-compile pass for + # them, and merely including ck_tile headers then hard-fails ("member reference base type + # 'void' is not a structure or union"). These generated files exist solely to provide CK kernel + # instantiations, so simply exclude them from those particular archs' compilation entirely via + # a dedicated object library with a narrowed HIP_ARCHITECTURES - the plain (non-fused) attention + # kernel already covers those archs regardless. + set(KATAGO_CK_FMHA_SUPPORTED_ARCHS ${CMAKE_HIP_ARCHITECTURES}) + list(REMOVE_ITEM KATAGO_CK_FMHA_SUPPORTED_ARCHS gfx906 gfx1010 gfx1011 gfx1012) + set_source_files_properties(${KATAGO_CK_FMHA_GENERATED_SOURCES} PROPERTIES LANGUAGE HIP) + add_library(katago_ck_fmha_kernels OBJECT ${KATAGO_CK_FMHA_GENERATED_SOURCES}) + set_target_properties(katago_ck_fmha_kernels PROPERTIES HIP_ARCHITECTURES "${KATAGO_CK_FMHA_SUPPORTED_ARCHS}") + # SYSTEM/-isystem is required, not just for warning suppression: ck_tile's own headers + # #include , and on systems with a stale libamdhip64-dev under /usr/include, + # a plain -I here still loses to /usr/include for angle-bracket resolution. -isystem does not. + target_include_directories(katago_ck_fmha_kernels SYSTEM PRIVATE ${KATAGO_CK_TILE_INCLUDE_DIR} external/composable_kernel_fmha) + # As above (see the KATAGO_HIP_SYSTEM_INCLUDE_DIR block for `katago` itself): the SYSTEM + # keyword alone isn't enough here either - CMake's include-directory de-duplication drops it as + # "already covered" against hip::device's own plain -I of the same path, so the stale + # /usr/include/hip still wins for this target too unless we force a raw -isystem. This is a + # separate target from `katago`, so it needs its own copy of the same compile option - compile + # options aren't inherited across targets, only via target_link_libraries propagating usage + # requirements, which compile options set via target_compile_options(PRIVATE) are not. + if(KATAGO_HIP_SYSTEM_INCLUDE_DIR) + target_compile_options(katago_ck_fmha_kernels PRIVATE "-isystem${KATAGO_HIP_SYSTEM_INCLUDE_DIR}") + endif() + target_link_libraries(katago katago_ck_fmha_kernels) + target_include_directories(katago SYSTEM PRIVATE ${KATAGO_CK_TILE_INCLUDE_DIR} external/composable_kernel_fmha) + target_compile_definitions(katago PRIVATE KATAGO_ROCM_HAS_CK_FMHA=1) + # rocmbackend.cpp itself (unlike the generated files above) still needs to compile for every + # arch, since it also holds the non-CK ROCm backend logic - it guards its own CK usage per-arch + # internally (see KATAGO_ROCM_CK_FMHA_ARCH_OK in rocmbackend.cpp) rather than being excluded here. + if(WIN32) + set_source_files_properties(neuralnet/rocmbackend.cpp PROPERTIES LANGUAGE HIP) + endif() + else() + message(STATUS "ck_tile headers not found; ROCm backend will only use its built-in (non-fused) attention kernel") + target_compile_definitions(katago PRIVATE KATAGO_ROCM_HAS_CK_FMHA=0) + endif() elseif(USE_BACKEND STREQUAL "EIGEN") target_compile_definitions(katago PRIVATE USE_EIGEN_BACKEND) # Allow EIGEN3_INCLUDE_DIRS as a manual override (for users who downloaded @@ -535,6 +1389,7 @@ elseif(USE_BACKEND STREQUAL "EIGEN") message(STATUS "Found Eigen3 at ${EIGEN3_INCLUDE_DIRS}") endif() endif() + endif() if(USE_BIGGER_BOARDS_EXPENSIVE) @@ -545,6 +1400,59 @@ if(NO_GIT_REVISION AND (NOT BUILD_DISTRIBUTED)) target_compile_definitions(katago PRIVATE NO_GIT_REVISION) endif() +# On Windows ROCm builds, zlib is bundled inside the HIP SDK (TheRock layout) - try that first. +if(WIN32 AND USE_BACKEND STREQUAL "ROCM" AND DEFINED ENV{HIP_PATH}) + if(NOT ZLIB_INCLUDE_DIR AND EXISTS "$ENV{HIP_PATH}/lib/rocm_sysdeps/include/zlib.h") + set(ZLIB_INCLUDE_DIR "$ENV{HIP_PATH}/lib/rocm_sysdeps/include" CACHE PATH "" FORCE) + endif() + if(NOT ZLIB_LIBRARY) + foreach(_zlib_name "zlibstatic.lib" "zlib.lib" "zlibstaticd.lib") + if(EXISTS "$ENV{HIP_PATH}/lib/rocm_sysdeps/lib/${_zlib_name}") + set(ZLIB_LIBRARY "$ENV{HIP_PATH}/lib/rocm_sysdeps/lib/${_zlib_name}" CACHE FILEPATH "" FORCE) + break() + endif() + endforeach() + endif() +endif() + +# TheRock's Windows package (as of ~7.13) ships zlib.h under rocm_sysdeps/include but no linkable +# .lib anymore (the block above then leaves ZLIB_LIBRARY unset) - auto-fetch zlib via vcpkg instead +# of requiring the user to separately install a package manager (see KATAGO_AUTO_FETCH_DEPS above). +# +# Probe for a normal system zlib first (quiet - the "real", error-emitting find_package(ZLIB) call +# is further below): on a fresh Linux configure, ZLIB_INCLUDE_DIR/ZLIB_LIBRARY are never set by +# anything above this point regardless of whether the system already has zlib1g-dev installed (as +# Compiling.md's Linux instructions require), so without this probe the check below would always +# be true and vcpkg would fire unconditionally on every Linux build - needlessly cloning vcpkg, +# downloading its own CMake, and building zlib from source even when the system package is right +# there. This probe costs nothing extra when it does find the system zlib (the later find_package +# call just reuses the cached result), and Windows is unaffected since a bundled TheRock/system +# zlib usually still won't satisfy this either way, falling through to vcpkg as before. +if(KATAGO_AUTO_FETCH_DEPS AND (NOT ZLIB_INCLUDE_DIR OR NOT ZLIB_LIBRARY)) + find_package(ZLIB QUIET) +endif() + +set(KATAGO_ZLIB_IS_VCPKG_DLL FALSE) +if(KATAGO_AUTO_FETCH_DEPS AND (NOT ZLIB_INCLUDE_DIR OR NOT ZLIB_LIBRARY)) + katago_vcpkg_install_if_needed("zlib") + set(_katago_vcpkg_installed_root "${KATAGO_VCPKG_ROOT}/installed/${KATAGO_VCPKG_TRIPLET}") + if(NOT ZLIB_INCLUDE_DIR AND EXISTS "${_katago_vcpkg_installed_root}/include/zlib.h") + set(ZLIB_INCLUDE_DIR "${_katago_vcpkg_installed_root}/include" CACHE PATH "Path to directory with zlib.h and other header files" FORCE) + endif() + if(NOT ZLIB_LIBRARY) + find_library(_katago_zlib_lib NAMES zlib z zlibstatic HINTS "${_katago_vcpkg_installed_root}/lib" NO_DEFAULT_PATH) + if(_katago_zlib_lib) + set(ZLIB_LIBRARY "${_katago_zlib_lib}" CACHE FILEPATH "Path to 'libz.so' on Linux or 'libz.lib' on Windows" FORCE) + # vcpkg's default triplets (e.g. x64-windows) build zlib as a DLL, not a static lib - the + # .lib found above is just an import lib, so the matching runtime DLL needs to be copied + # next to katago.exe at build time (see the POST_BUILD step further below). + if(WIN32) + set(KATAGO_ZLIB_IS_VCPKG_DLL TRUE) + endif() + endif() + endif() +endif() + find_package(ZLIB) if(ZLIB_FOUND) include_directories(${ZLIB_INCLUDE_DIRS}) @@ -556,6 +1464,50 @@ else() message(SEND_ERROR "${ColorBoldRed}zlib was not found, if zlib is actually installed but not being found you can set ZLIB_INCLUDE_DIR to the directory with zlib.h and other headers, and ZLIB_LIBRARY to the compiled library 'libz.so' on Linux or 'libz.lib' on Windows. On the command line, this is -DZLIB_INCLUDE_DIR=... and -DZLIB_LIBRARY=... ${ColorReset}") endif(ZLIB_FOUND) +if(KATAGO_ZLIB_IS_VCPKG_DLL) + get_filename_component(_katago_zlib_lib_dir "${ZLIB_LIBRARY}" DIRECTORY) + get_filename_component(_katago_vcpkg_pkg_root "${_katago_zlib_lib_dir}" DIRECTORY) + file(GLOB _katago_zlib_dlls "${_katago_vcpkg_pkg_root}/bin/*.dll") + foreach(_katago_zlib_dll IN LISTS _katago_zlib_dlls) + add_custom_command(TARGET katago POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy_if_different "${_katago_zlib_dll}" "$" + VERBATIM) + endforeach() +endif() + +# On Windows ROCm builds, copy the HIP/MIOpen/rocBLAS runtime DLLs next to katago.exe so it runs +# without the user needing to manually copy anything or add HIP_PATH/bin to PATH themselves. +# amdhip64_*.dll in particular must be copied (not just PATH-resolved): AMD's GPU driver installs +# its own, potentially older/incompatible copy into System32, which would otherwise take priority. +if(WIN32 AND USE_BACKEND STREQUAL "ROCM" AND DEFINED ENV{HIP_PATH}) + file(TO_CMAKE_PATH "$ENV{HIP_PATH}" _katago_hip_path) + file(GLOB _katago_rocm_runtime_dlls + "${_katago_hip_path}/bin/amdhip64_*.dll" + "${_katago_hip_path}/bin/hipblas.dll" + "${_katago_hip_path}/bin/MIOpen.dll" + "${_katago_hip_path}/bin/rocblas.dll" + "${_katago_hip_path}/bin/libhipblaslt.dll" + "${_katago_hip_path}/bin/amdocl64.dll" + "${_katago_hip_path}/bin/hiprtc*.dll" + "${_katago_hip_path}/bin/amd_comgr*.dll" + ) + foreach(_katago_rocm_dll IN LISTS _katago_rocm_runtime_dlls) + add_custom_command(TARGET katago POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy_if_different "${_katago_rocm_dll}" "$" + VERBATIM) + endforeach() + # rocBLAS/hipBLASLt look for their kernel library files relative to their own DLL's location. + foreach(_katago_rocm_lib_subdir IN ITEMS rocblas hipblaslt) + if(IS_DIRECTORY "${_katago_hip_path}/bin/${_katago_rocm_lib_subdir}/library") + add_custom_command(TARGET katago POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy_directory + "${_katago_hip_path}/bin/${_katago_rocm_lib_subdir}/library" + "$/${_katago_rocm_lib_subdir}/library" + VERBATIM) + endif() + endforeach() +endif() + find_library(LIBZIP_LIBRARY NAMES zip) find_path(LIBZIP_INCLUDE_DIR_ZIP NAMES zip.h) find_path(LIBZIP_INCLUDE_DIR_ZIPCONF NAMES zipconf.h) @@ -651,7 +1603,7 @@ if(MSVC) set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} /STACK:8388608") elseif(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" OR CMAKE_CXX_COMPILER_ID STREQUAL "Clang" OR CMAKE_CXX_COMPILER_ID STREQUAL "AppleClang") message(STATUS "Setting up build for GNU, Clang or MinGW.") - if(NOT (${CMAKE_SYSTEM_PROCESSOR} MATCHES "(arm|aarch32|aarch64)")) + if(NOT (${CMAKE_SYSTEM_PROCESSOR} MATCHES "(arm|aarch32|aarch64)") AND NOT USE_BACKEND STREQUAL "ROCM") set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -mfpmath=sse") else() # For ARM architecture, as a hack, ensure that char is signed @@ -688,7 +1640,10 @@ elseif(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" OR CMAKE_CXX_COMPILER_ID STREQUAL "C else() message(STATUS "Enabling Clang-specific build options.") set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wnull-dereference -Wdangling-else") - target_link_libraries(katago "atomic") + if(NOT WIN32) + # libatomic is a Linux GCC/Clang runtime; not needed (or available) on Windows + target_link_libraries(katago "atomic") + endif() endif() if(USE_TCMALLOC) @@ -700,3 +1655,41 @@ endif() target_include_directories(katago PUBLIC ${CMAKE_CURRENT_BINARY_DIR}) +# On Windows ROCm builds, clang compiles all files with -x hip which pulls in +# MSVC-compatibility headers that require Windows SDK ucrt/shared/um headers. +if(WIN32 AND USE_BACKEND STREQUAL "ROCM") + # Prefer the KATAGO_WINSDK_ROOT detected earlier in the pre-project() block + if(NOT KATAGO_WINSDK_ROOT) + get_filename_component(_winsdk_root2 + "[HKEY_LOCAL_MACHINE\\SOFTWARE\\Microsoft\\Windows Kits\\Installed Roots;KitsRoot10]" + ABSOLUTE) + if(EXISTS "${_winsdk_root2}") + set(KATAGO_WINSDK_ROOT "${_winsdk_root2}" CACHE INTERNAL "") + else() + katago_win_winsdk_fallback_dirs(_winsdk_candidates2) + foreach(_p IN LISTS _winsdk_candidates2) + if(EXISTS "${_p}") + set(KATAGO_WINSDK_ROOT "${_p}" CACHE INTERNAL "") + break() + endif() + endforeach() + endif() + endif() + if(KATAGO_WINSDK_ROOT) + # Pick highest version directory + file(GLOB _sdk_ver_dirs "${KATAGO_WINSDK_ROOT}/Include/*/ucrt") + if(_sdk_ver_dirs) + list(SORT _sdk_ver_dirs ORDER DESCENDING) + list(GET _sdk_ver_dirs 0 _ucrt_dir) + get_filename_component(_sdk_ver_dir "${_ucrt_dir}" DIRECTORY) + set(_winsdk_include "${_sdk_ver_dir}") + message(STATUS "Auto-detected Windows SDK include root: ${_winsdk_include}") + target_include_directories(katago PRIVATE + "${_winsdk_include}/ucrt" + "${_winsdk_include}/shared" + "${_winsdk_include}/um" + ) + endif() + endif() +endif() + diff --git a/cpp/README.md b/cpp/README.md index 1f5d8d21fc..5d1e97716a 100644 --- a/cpp/README.md +++ b/cpp/README.md @@ -9,13 +9,13 @@ Summary of source folders, in approximate dependency order, from lowest level to * `board.{cpp,h}` - Raw board implementation, without move history. Helper functions for Benson's algorithm and ladder search. * `boardhistory.{cpp,h}` - Datastructure that does include move history - handles superko, passing, game end, final scoring, komi, handicap detection, etc. * `graphhash.{cpp,h}` - History-sensitive hash used for [monte-carlo graph search](https://github.com/lightvector/KataGo/blob/master/docs/GraphSearch.md). -* `neuralnet` - Neural net GPU implementation and interface. Contains OpenCL, CUDA, Eigen, TensorRT backends along with common interfaces and model data structures. +* `neuralnet` - Neural net GPU implementation and interface. Contains OpenCL, CUDA, Eigen, TensorRT, ROCm, Metal backends along with common interfaces and model data structures. * `desc.{cpp,h}` - Data structure holding neural net structure and weights. * `modelversion.{cpp,h}` - Enumerates the various versions of neural net features and models. * `nninputs.{cpp,h}` - Implements the input features for the neural net. * `sgfmetadata.{cpp,h}` - Implements the input features for the [HumanSL neural net](https://github.com/lightvector/KataGo/blob/master/docs/Analysis_Engine.md#human-sl-analysis-guide), for conditioning on various SGF metadata about human players from training data. * `nninterface.h` - Common interface that is implemented by every low-level neural net backend. - * `{cuda,opencl,eigen,trt,dummy}backend.cpp` - Various backends. + * `{cuda,opencl,eigen,trt,rocm,mgx,metal,dummy}backend.cpp` - Various backends. * `nneval.{cpp,h}` - Top-level handle to the neural net used by the rest of the engine, implements thread-safe batching of queries. * `search` - The main search engine. * `timecontrols.cpp` - Basic handling of a few possible time controls. diff --git a/cpp/command/benchmark.cpp b/cpp/command/benchmark.cpp index 89ecd1bb19..d6d32aa12c 100644 --- a/cpp/command/benchmark.cpp +++ b/cpp/command/benchmark.cpp @@ -264,6 +264,9 @@ int MainCmds::benchmark(const vector& args) { cout << "If you have a strong GPU capable of FP16 tensor cores (e.g. RTX2080), " << "using the Cuda version of KataGo instead may give a mild performance boost." << endl; #endif +#ifdef USE_ROCM_BACKEND + cout << "You are currently using the ROCm version of KataGo." << endl; +#endif #ifdef USE_EIGEN_BACKEND cout << "You are currently using the Eigen (CPU) version of KataGo. Due to having no GPU, it may be slow." << endl; #endif diff --git a/cpp/configs/analysis_example.cfg b/cpp/configs/analysis_example.cfg index 4aed2f4f85..95a7a0894d 100644 --- a/cpp/configs/analysis_example.cfg +++ b/cpp/configs/analysis_example.cfg @@ -223,9 +223,7 @@ nnRandomize = true # cudaUseNHWC = auto -# ------------------------------ -# Metal GPU settings -# ------------------------------ +# Metal GPU settings-------------------------------------- # These only apply when using the METAL version of KataGo. # Metal backend dispatch is configured via numNNServerThreadsPerModel and metalDeviceToUseThread. @@ -261,6 +259,32 @@ nnRandomize = true # metalUseFP16 = true +# ROCm GPU settings-------------------------------------- +# These only apply when using the ROCm version of KataGo. + +# IF USING ONE GPU: optionally uncomment and change this if the GPU you want to use turns out to be not device 0 +# rocmDeviceToUse = 0 + +# IF USING TWO GPUS: Uncomment these two lines (AND set numNNServerThreadsPerModel above): +# rocmDeviceToUseThread0 = 0 # change this if the first GPU you want to use turns out to be not device 0 +# rocmDeviceToUseThread1 = 1 # change this if the second GPU you want to use turns out to be not device 1 + +# IF USING THREE GPUS: Uncomment these three lines (AND set numNNServerThreadsPerModel above): +# rocmDeviceToUseThread0 = 0 # change this if the first GPU you want to use turns out to be not device 0 +# rocmDeviceToUseThread1 = 1 # change this if the second GPU you want to use turns out to be not device 1 +# rocmDeviceToUseThread2 = 2 # change this if the third GPU you want to use turns out to be not device 2 + +# You can probably guess the pattern if you have four, five, etc. GPUs. + +# KataGo will automatically use FP16 or not based on the compute capability of your AMD GPU. If you +# want to try to force a particular behavior though you can uncomment these lines and change them +# to "true" or "false". E.g. it's using FP16 but on your card that's giving an error, or it's not using +# FP16 but you think it should. +# rocmUseFP16 = auto +# ROCm does not support NHWC, so this is always false. + + + # OpenCL-specific GPU settings-------------------------------------- # These only apply when using the OpenCL version of KataGo. diff --git a/cpp/configs/contribute_example.cfg b/cpp/configs/contribute_example.cfg index 6ca039f112..865433f83c 100644 --- a/cpp/configs/contribute_example.cfg +++ b/cpp/configs/contribute_example.cfg @@ -83,9 +83,8 @@ watchOngoingGameInFileName = watchgame.txt # cudaUseNHWC = auto -# ------------------------------ -# Metal GPU settings -# ------------------------------ +# Metal GPU settings-------------------------------------- + # These only apply when using the METAL version of KataGo. # For one Metal instance: KataGo will automatically use the default device. @@ -99,6 +98,32 @@ watchOngoingGameInFileName = watchgame.txt # The pattern continues for additional Metal instances. +# ROCm GPU settings-------------------------------------- +# These only apply when using the ROCm version of KataGo. + +# IF USING ONE GPU: optionally uncomment and change this if the GPU you want to use turns out to be not device 0 +# rocmDeviceToUse = 0 + +# IF USING TWO GPUS: Uncomment these two lines (AND set numNNServerThreadsPerModel above): +# rocmDeviceToUseThread0 = 0 # change this if the first GPU you want to use turns out to be not device 0 +# rocmDeviceToUseThread1 = 1 # change this if the second GPU you want to use turns out to be not device 1 + +# IF USING THREE GPUS: Uncomment these three lines (AND set numNNServerThreadsPerModel above): +# rocmDeviceToUseThread0 = 0 # change this if the first GPU you want to use turns out to be not device 0 +# rocmDeviceToUseThread1 = 1 # change this if the second GPU you want to use turns out to be not device 1 +# rocmDeviceToUseThread2 = 2 # change this if the third GPU you want to use turns out to be not device 2 + +# You can probably guess the pattern if you have four, five, etc. GPUs. + +# KataGo will automatically use FP16 or not based on the compute capability of your AMD GPU. If you +# want to try to force a particular behavior though you can uncomment these lines and change them +# to "true" or "false". E.g. it's using FP16 but on your card that's giving an error, or it's not using +# FP16 but you think it should. +# rocmUseFP16 = auto +# ROCm does not support NHWC, so this is always false. + + + # OpenCL GPU settings-------------------------------------- # These only apply when using the OpenCL version of KataGo. diff --git a/cpp/configs/gtp_example.cfg b/cpp/configs/gtp_example.cfg index 7c3a8f8341..eeea22c30b 100644 --- a/cpp/configs/gtp_example.cfg +++ b/cpp/configs/gtp_example.cfg @@ -462,9 +462,9 @@ searchFactorWhenWinningThreshold = 0.95 # "auto" (default) uses the GEMM only in FP16, where it is slightly faster. # cudaUse1x1Matmul = auto -# ------------------------------ -# Metal GPU settings -# ------------------------------ + +# Metal GPU settings-------------------------------------- + # These only apply when using the METAL version of KataGo. # Metal backend dispatch is configured via numNNServerThreadsPerModel and metalDeviceToUseThread. @@ -499,6 +499,33 @@ searchFactorWhenWinningThreshold = 0.95 # inference but much slower than the GPU path. # metalUseFP16 = true + +# ROCm GPU settings-------------------------------------- +# These only apply when using the ROCm version of KataGo. + +# IF USING ONE GPU: optionally uncomment and change this if the GPU you want to use turns out to be not device 0 +# rocmDeviceToUse = 0 + +# IF USING TWO GPUS: Uncomment these two lines (AND set numNNServerThreadsPerModel above): +# rocmDeviceToUseThread0 = 0 # change this if the first GPU you want to use turns out to be not device 0 +# rocmDeviceToUseThread1 = 1 # change this if the second GPU you want to use turns out to be not device 1 + +# IF USING THREE GPUS: Uncomment these three lines (AND set numNNServerThreadsPerModel above): +# rocmDeviceToUseThread0 = 0 # change this if the first GPU you want to use turns out to be not device 0 +# rocmDeviceToUseThread1 = 1 # change this if the second GPU you want to use turns out to be not device 1 +# rocmDeviceToUseThread2 = 2 # change this if the third GPU you want to use turns out to be not device 2 + +# You can probably guess the pattern if you have four, five, etc. GPUs. + +# KataGo will automatically use FP16 or not based on the compute capability of your AMD GPU. If you +# want to try to force a particular behavior though you can uncomment these lines and change them +# to "true" or "false". E.g. it's using FP16 but on your card that's giving an error, or it's not using +# FP16 but you think it should. +# rocmUseFP16 = auto +# ROCm does not support NHWC, so this is always false. + + + # ------------------------------ # OpenCL GPU settings # ------------------------------ diff --git a/cpp/configs/match_example.cfg b/cpp/configs/match_example.cfg index 7e5b4fc09f..f4d53a2af4 100644 --- a/cpp/configs/match_example.cfg +++ b/cpp/configs/match_example.cfg @@ -156,9 +156,8 @@ numNNServerThreadsPerModel = 1 # cudaUseNHWC = auto -# ------------------------------ -# Metal GPU settings -# ------------------------------ +# Metal GPU settings-------------------------------------- + # These only apply when using the METAL version of KataGo. # For one Metal instance: KataGo will automatically use the default device. @@ -172,6 +171,32 @@ numNNServerThreadsPerModel = 1 # The pattern continues for additional Metal instances. +# ROCm GPU settings-------------------------------------- +# These only apply when using the ROCm version of KataGo. + +# IF USING ONE GPU: optionally uncomment and change this if the GPU you want to use turns out to be not device 0 +# rocmDeviceToUse = 0 + +# IF USING TWO GPUS: Uncomment these two lines (AND set numNNServerThreadsPerModel above): +# rocmDeviceToUseThread0 = 0 # change this if the first GPU you want to use turns out to be not device 0 +# rocmDeviceToUseThread1 = 1 # change this if the second GPU you want to use turns out to be not device 1 + +# IF USING THREE GPUS: Uncomment these three lines (AND set numNNServerThreadsPerModel above): +# rocmDeviceToUseThread0 = 0 # change this if the first GPU you want to use turns out to be not device 0 +# rocmDeviceToUseThread1 = 1 # change this if the second GPU you want to use turns out to be not device 1 +# rocmDeviceToUseThread2 = 2 # change this if the third GPU you want to use turns out to be not device 2 + +# You can probably guess the pattern if you have four, five, etc. GPUs. + +# KataGo will automatically use FP16 or not based on the compute capability of your AMD GPU. If you +# want to try to force a particular behavior though you can uncomment these lines and change them +# to "true" or "false". E.g. it's using FP16 but on your card that's giving an error, or it's not using +# FP16 but you think it should. +# rocmUseFP16 = auto +# ROCm does not support NHWC, so this is always false. + + + # OpenCL GPU settings-------------------------------------- # These only apply when using OpenCL as the backend for inference. # (For GTP, we only ever have one model, when playing matches, we might have more than one, see match_example.cfg) diff --git a/cpp/core/win_rocm_sse_shim.cpp b/cpp/core/win_rocm_sse_shim.cpp new file mode 100644 index 0000000000..f8ab064b61 --- /dev/null +++ b/cpp/core/win_rocm_sse_shim.cpp @@ -0,0 +1,49 @@ +#ifdef USE_ROCM_BACKEND +#ifdef _WIN32 + +// On Windows ROCm builds, clang compiles every source file with -x hip (see the comment on this +// in CMakeLists.txt), which force-includes clang's own __clang_hip_runtime_wrapper.h before any +// user code. That wrapper's own includes (via /) transitively pull in ucrt's +// , which includes - and on this toolchain, that resolves to MSVC's own +// (chained in via clang's own intrin.h doing "#include_next "), not clang's +// compatible one. MSVC's declares SSE2 intrinsics like _mm_loadu_si128 as plain +// bodyless extern functions (relying on cl.exe's special-cased intrinsic recognition, which clang +// doesn't replicate for headers reached this way), rather than the "static __inline__" functions +// with actual bodies that clang's own provides. Any code that ends up calling one of +// these - e.g. ucrt's own SSE2-optimized wmemcmp/wmemchr in , used transitively by +// std::filesystem/std::wstring paths in fileutils.cpp/makedir.cpp/loadmodel.cpp - then links with +// "undefined symbol: _mm_loadu_si128"-style errors, since no definition exists anywhere in the +// link (HIP's --hip-link mode also passes -nostdlib, so nothing here falls back to a prebuilt +// ucrt.lib implementation either). +// +// Providing real definitions for just the 3 intrinsics ucrt's wmemcmp/wmemchr actually use fixes +// this: matching MSVC's own __m128i type (a union, already visible via the same transitive +// include chain above, avoiding a "conflicting types" redeclaration error) for the external +// signature, but implementing the body using GNU vector-extension types instead of any SSE +// header, which need no header at all and so sidestep the whole conflict. +extern "C" __m128i _mm_loadu_si128(__m128i const* p) { + __m128i r; + __builtin_memcpy(&r, p, sizeof(r)); + return r; +} + +extern "C" __m128i _mm_cmpeq_epi16(__m128i a, __m128i b) { + typedef short katago_v8hi __attribute__((__vector_size__(16))); + katago_v8hi va, vb; + __builtin_memcpy(&va, &a, sizeof(va)); + __builtin_memcpy(&vb, &b, sizeof(vb)); + katago_v8hi vr = (va == vb); + __m128i r; + __builtin_memcpy(&r, &vr, sizeof(r)); + return r; +} + +extern "C" int _mm_movemask_epi8(__m128i a) { + typedef char katago_v16qi __attribute__((__vector_size__(16))); + katago_v16qi va; + __builtin_memcpy(&va, &a, sizeof(va)); + return __builtin_ia32_pmovmskb128(va); +} + +#endif +#endif diff --git a/cpp/external/composable_kernel_fmha/LICENSE b/cpp/external/composable_kernel_fmha/LICENSE new file mode 100644 index 0000000000..68f6ae5746 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/LICENSE @@ -0,0 +1,28 @@ +Copyright (c) 2018- , Advanced Micro Devices, Inc. (Chao Liu, Jing Zhang) +Copyright (c) 2019- , Advanced Micro Devices, Inc. (Letao Qin, Qianfeng Zhang, Liang Huang, Shaojie Wang) +Copyright (c) 2022- , Advanced Micro Devices, Inc. (Anthony Chang, Chunyu Lai, Illia Silin, Adam Osewski, Poyen Chen, Jehandad Khan) +Copyright (c) 2019-2021, Advanced Micro Devices, Inc. (Hanwen Chang) +Copyright (c) 2019-2020, Advanced Micro Devices, Inc. (Tejash Shah) +Copyright (c) 2020 , Advanced Micro Devices, Inc. (Xiaoyan Zhou) +Copyright (c) 2021-2022, Advanced Micro Devices, Inc. (Jianfeng Yan) + +SPDX-License-Identifier: MIT +Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/cpp/external/composable_kernel_fmha/README.md b/cpp/external/composable_kernel_fmha/README.md new file mode 100644 index 0000000000..a6aa42462e --- /dev/null +++ b/cpp/external/composable_kernel_fmha/README.md @@ -0,0 +1,52 @@ +# Composable Kernel FMHA (fused multi-head attention) + +Vendored glue headers + pre-generated kernel instantiations from AMD's Composable Kernel (CK) +`ck_tile` FMHA implementation, used by the ROCm backend as an optional fused attention path +(mirrors the CUDA backend's optional cudnn-frontend SDPA graph path). Requires the ck_tile base +headers from a `composablekernel-dev`/`amdrocm-ck*` system package (not vendored here). If not +found at configure time (see `KATAGO_CK_TILE_INCLUDE_DIR` in `cpp/CMakeLists.txt`), the ROCm +backend just always uses its own plain (non-fused) attention kernel instead - this is a pure +performance optimization, not required for correctness. Measured ~2.2x nnEvals/s on a gfx1100 +(RX 7900 XTX) with a small transformer test model, FP16. + +Runtime opt-out: set `rocmDisableFusedAttention = true` in the KataGo config to force the plain +kernel even when the fused path is compiled in and available. + +Source: https://github.com/ROCm/rocm-libraries, tag `therock-7.13`, +`projects/composablekernel/example/ck_tile/01_fmha/`. Must match the ck_tile core headers version +from the installed `amdrocm-ck7.13` system package byte-for-byte (`fmha_fwd.hpp` etc. reference +internal ck_tile core APIs that change between releases) — verify with `diff` against the installed +`.../include/ck_tile/ops/fmha_fwd.hpp` before regenerating from a different tag. + +## What's here + +- `fmha_fwd.hpp`, `mask.hpp`, `bias.hpp`, `rotary.hpp`, `quant.hpp`: the example's glue headers + declaring `fmha_fwd()`/`fmha_fwd_traits`/`fmha_fwd_args` and friends. Copied unmodified. +- `generated/`: kernel instantiations produced by CK's `generate.py` codegen script, narrowed to + exactly what KataGo needs (see regeneration command below). `fmha_fwd_api.cpp` is the dispatcher + (`fmha_fwd()`); the rest are individual `fmha_fwd_` template instantiations it calls into. + +## Scope (matches what the CUDA backend's cudnn-frontend SDPA path actually uses) + +- fp16 only (CUDA's fused SDPA path is FP16-only too; FP32 always uses the plain kernel fallback) +- batch mode only (no group/variable-length mode) +- bias: no-bias or elementwise (matches the [B,1,S,S] additive mask-derived bias KataGo builds); + no alibi +- mask: none (KataGo has no causal masking; padding is handled via the elementwise bias instead) +- no LSE output, no dropout, no quantization scaling, no attention sink +- hdim buckets: 32, 64 (covers KataGo's (qHeadDim, vHeadDim) combos of 32/32, 32/16, 64/64, 64/32, + 32/64 — smaller actual head dims like 16 are handled via CK's own padding within the 32 bucket) +- targets: gfx9, gfx950, gfx11, gfx115, gfx12 (as of `therock-7.13`; CK has no FMHA codegen support + for gfx10/RDNA2 — gfx1030/1031/1032 always use the plain kernel fallback. gfx125 doesn't exist as + a target in this codegen version either.) + +## Regenerating + +From a checkout of `projects/composablekernel/example/ck_tile/01_fmha/` at tag `therock-7.13` in the +CK source repo: + +``` +python3 generate.py --output_dir --targets gfx9,gfx950,gfx11,gfx115,gfx12 -a fwd \ + -f "*_fp16_batch_*_nlogits_*bias_nmask_nlse_ndropout_nskip_nqscale_*nsink" \ + --optdim 32,64 --receipt 0 -m simplified +``` diff --git a/cpp/external/composable_kernel_fmha/bias.hpp b/cpp/external/composable_kernel_fmha/bias.hpp new file mode 100644 index 0000000000..b526204384 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/bias.hpp @@ -0,0 +1,114 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include +#include +#include "ck_tile/core.hpp" +#include "ck_tile/ops/fmha.hpp" + +// keep sync with BlockAttentionBiasEnum +enum class bias_enum +{ + no_bias = 0, + elementwise_bias = 1, + alibi = 2, +}; + +struct bias_info +{ + bias_enum type; + /* + * simple dispatch logic + * + * if type == elementwise_bias: + * if rank_info == 0: + * bias is 1*1*s*s + * elif rank_info == 1: + * bias is 1*h*s*s + * elif rank_info == 2: + * bias is b*h*s*s + * + * elif type == alibi: + * if rank_info == 0: + * alibi in 1*h + * elif rank_info == 1: + * alibi in b*h + */ + int rank_info; + + void serialize(std::ostream& os) const + { + if(type == bias_enum::no_bias) + os << "n"; + else if(type == bias_enum::elementwise_bias) + { + os << "e"; + if(rank_info != 0) + { + os << "[" << rank_info << "]"; + } + } + else if(type == bias_enum::alibi) + { + os << "alibi"; + if(rank_info != 0) + { + os << "[" << rank_info << "]"; + } + } + } + + static bias_info decode(std::string str) + { + bias_info info{bias_enum::no_bias, 0}; + auto found_0 = str.find(':'); + if(found_0 != std::string::npos) + { + std::string t = str.substr(0, found_0); + std::string v = str.substr(found_0 + 1); + if(t == "e" || t == "elementwise") + { + info.type = bias_enum::elementwise_bias; + info.rank_info = std::stoi(v); + if(info.rank_info < 0 || info.rank_info > 2) + throw std::invalid_argument("invalid bias rank: " + str); + } + else if(t == "a" || t == "alibi") + { + info.type = bias_enum::alibi; + info.rank_info = std::stoi(v); + if(info.rank_info < 0 || info.rank_info > 1) + throw std::invalid_argument("invalid bias rank: " + str); + } + else + { + throw std::invalid_argument("invalid bias value: " + str); + } + } + else if(str == "0" || str == "n") + { + info.type = bias_enum::no_bias; + } + else if(str == "1" || str == "e" || str == "elementwise") + { + info.type = bias_enum::elementwise_bias; + } + else if(str == "2" || str == "a" || str == "alibi") + { + info.type = bias_enum::alibi; + } + else + { + throw std::invalid_argument("invalid bias value: " + str); + } + return info; + } + + friend std::ostream& operator<<([[clang::lifetimebound]] std::ostream& os, const bias_info& bi) + { + bi.serialize(os); + return os; + } +}; diff --git a/cpp/external/composable_kernel_fmha/fmha_fwd.hpp b/cpp/external/composable_kernel_fmha/fmha_fwd.hpp new file mode 100644 index 0000000000..98e2df2e1e --- /dev/null +++ b/cpp/external/composable_kernel_fmha/fmha_fwd.hpp @@ -0,0 +1,1779 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/host/device_prop.hpp" +#include "ck_tile/host/kernel_launch.hpp" +#include "ck_tile/ops/epilogue.hpp" +#include "ck_tile/ops/fmha.hpp" + +#include "bias.hpp" +#include "mask.hpp" +#include "quant.hpp" +#include "rotary.hpp" + +#include +#include +#include + +struct FmhaFwdFp32 +{ +}; + +struct FmhaFwdFp16 +{ +}; + +struct FmhaFwdBf16 +{ +}; + +struct FmhaFwdFp8 +{ +}; + +struct FmhaFwdBf8 +{ +}; + +struct FmhaFwdFp8Fp16 +{ +}; + +struct FmhaFwdFp8Bf16 +{ +}; + +struct FmhaFwdFp8Fp32 +{ +}; + +struct FmhaFwdMxFp8 +{ +}; + +struct FmhaFwdMxFp4 +{ +}; + +template +struct FmhaFwdTypeConfig; + +template <> +struct FmhaFwdTypeConfig +{ + using QDataType = float; + using KDataType = float; + using VDataType = float; + using BiasDataType = float; + using RandValOutputDataType = uint8_t; + using LSEDataType = float; // data type for lse(logsumexp L_j = max_j + log(l_j)) + using SaccDataType = float; // data type for first gemm accumulation + using SMPLComputeDataType = float; // data type for reduction, softmax + using PDataType = float; // data type for A matrix of second gemm + using OaccDataType = float; // data type for second gemm accumulation + using ODataType = float; +}; + +template <> +struct FmhaFwdTypeConfig +{ + using QDataType = ck_tile::half_t; + using KDataType = ck_tile::half_t; + using VDataType = ck_tile::half_t; + using BiasDataType = ck_tile::half_t; + using RandValOutputDataType = uint8_t; + using LSEDataType = float; // data type for lse(logsumexp L_j = max_j + log(l_j)) + using SaccDataType = float; // data type for first gemm accumulation + using SMPLComputeDataType = float; // data type for reduction, softmax + using PDataType = ck_tile::half_t; // data type for A matrix of second gemm + using OaccDataType = float; // data type for second gemm accumulation + using ODataType = ck_tile::half_t; +}; + +template <> +struct FmhaFwdTypeConfig +{ + using QDataType = ck_tile::bf16_t; + using KDataType = ck_tile::bf16_t; + using VDataType = ck_tile::bf16_t; + using BiasDataType = ck_tile::bf16_t; + using RandValOutputDataType = uint8_t; + using LSEDataType = float; // data type for lse(logsumexp L_j = max_j + log(l_j)) + using SaccDataType = float; // data type for first gemm accumulation + using SMPLComputeDataType = float; // data type for reduction, softmax + using PDataType = ck_tile::bf16_t; // data type for A matrix of second gemm + using OaccDataType = float; // data type for second gemm accumulation + using ODataType = ck_tile::bf16_t; +}; + +template <> +struct FmhaFwdTypeConfig +{ + using QDataType = ck_tile::fp8_t; + using KDataType = ck_tile::fp8_t; + using VDataType = ck_tile::fp8_t; + using BiasDataType = float; + using RandValOutputDataType = uint8_t; + using LSEDataType = float; // data type for lse(logsumexp L_j = max_j + log(l_j)) + using SaccDataType = float; // data type for first gemm accumulation + using SMPLComputeDataType = float; // data type for reduction, softmax + using PDataType = ck_tile::fp8_t; // data type for A matrix of second gemm + using OaccDataType = float; // data type for second gemm accumulation + using ODataType = ck_tile::fp8_t; +}; + +template <> +struct FmhaFwdTypeConfig +{ + using QDataType = ck_tile::bf8_t; + using KDataType = ck_tile::bf8_t; + using VDataType = ck_tile::bf8_t; + using BiasDataType = ck_tile::bf8_t; + using RandValOutputDataType = uint8_t; + using LSEDataType = float; // data type for lse(logsumexp L_j = max_j + log(l_j)) + using SaccDataType = float; // data type for first gemm accumulation + using SMPLComputeDataType = float; // data type for reduction, softmax + using PDataType = ck_tile::bf8_t; // data type for A matrix of second gemm + using OaccDataType = float; // data type for second gemm accumulation + using ODataType = ck_tile::bf8_t; +}; + +template <> +struct FmhaFwdTypeConfig +{ + using QDataType = ck_tile::fp8_t; + using KDataType = ck_tile::fp8_t; + using VDataType = ck_tile::fp8_t; + using BiasDataType = float; + using RandValOutputDataType = uint8_t; + using LSEDataType = float; // data type for lse(logsumexp L_j = max_j + log(l_j)) + using SaccDataType = float; // data type for first gemm accumulation + using SMPLComputeDataType = float; // data type for reduction, softmax + using PDataType = ck_tile::fp8_t; // data type for A matrix of second gemm + using OaccDataType = float; // data type for second gemm accumulation + using ODataType = ck_tile::bf16_t; +}; + +template <> +struct FmhaFwdTypeConfig +{ + using QDataType = ck_tile::fp8_t; + using KDataType = ck_tile::fp8_t; + using VDataType = ck_tile::fp8_t; + using BiasDataType = float; + using RandValOutputDataType = uint8_t; + using LSEDataType = float; // data type for lse(logsumexp L_j = max_j + log(l_j)) + using SaccDataType = float; // data type for first gemm accumulation + using SMPLComputeDataType = float; // data type for reduction, softmax + using PDataType = ck_tile::fp8_t; // data type for A matrix of second gemm + using OaccDataType = float; // data type for second gemm accumulation + using ODataType = float; +}; + +template <> +struct FmhaFwdTypeConfig +{ + using QDataType = ck_tile::fp8_t; + using KDataType = ck_tile::fp8_t; + using VDataType = ck_tile::fp8_t; + using BiasDataType = float; + using RandValOutputDataType = uint8_t; + using LSEDataType = float; // data type for lse(logsumexp L_j = max_j + log(l_j)) + using SaccDataType = float; // data type for first gemm accumulation + using SMPLComputeDataType = float; // data type for reduction, softmax + using PDataType = ck_tile::fp8_t; // data type for A matrix of second gemm + using OaccDataType = float; // data type for second gemm accumulation + using ODataType = float; + + using QScaleDataType = ck_tile::e8m0_t; + using KScaleDataType = ck_tile::e8m0_t; + using VScaleDataType = ck_tile::e8m0_t; + using PScaleDataType = ck_tile::e8m0_t; + + static constexpr ck_tile::index_t kQKScaleGranularity = 32; + static constexpr ck_tile::index_t kVScaleGranularity = 32; +}; + +template <> +struct FmhaFwdTypeConfig +{ + using QDataType = ck_tile::pk_fp4_t; + using KDataType = ck_tile::pk_fp4_t; + using VDataType = ck_tile::pk_fp4_t; + using BiasDataType = float; + using RandValOutputDataType = uint8_t; + using LSEDataType = float; // data type for lse(logsumexp L_j = max_j + log(l_j)) + using SaccDataType = float; // data type for first gemm accumulation + using SMPLComputeDataType = float; // data type for reduction, softmax + using PDataType = ck_tile::pk_fp4_t; // data type for A matrix of second gemm + using OaccDataType = float; // data type for second gemm accumulation + using ODataType = float; + + using QScaleDataType = ck_tile::e8m0_t; + using KScaleDataType = ck_tile::e8m0_t; + using VScaleDataType = ck_tile::e8m0_t; + using PScaleDataType = ck_tile::e8m0_t; + + static constexpr ck_tile::index_t kQKScaleGranularity = 32; + static constexpr ck_tile::index_t kVScaleGranularity = 32; +}; + +struct FmhaMasks +{ + using NoMask = ck_tile::GenericAttentionMask; + using GenericMask = ck_tile::GenericAttentionMask; + using CausalMask = ck_tile::GenericAttentionMask; +}; + +// runtime args, some will passed to karg, some will used to compute grids/blocks +struct fmha_fwd_args +{ + const void* q_ptr; + const void* k_ptr; + const void* v_ptr; + const void* bias_ptr; // bias or alibi_slope pointer + const void* q_descale_ptr; + const void* k_descale_ptr; + const void* v_descale_ptr; + void* rand_val_ptr; + void* lse_ptr; + void* o_ptr; + + // Usage notes for sequence length pointer parameters: + // + // [Note: Define "Group mode" vs "Batch mode" here if possible, e.g., "Group mode handles + // MQA/GQA..."] + // + // With padding: + // Group mode: + // - seqstart_q_ptr, seqstart_k_ptr: Record cumulative physical (including padding) sequence + // lengths. [array size: batch + 1] + // - seqlen_q_ptr/seqlen_k_ptr: Records logical (excluding padding) length for each + // sequence. [array size: batch] + // - cu_seqlen_q_ptr/cu_seqlen_k_ptr: Records cumulative logical (excluding padding) + // sequence lengths. [array size: batch + 1] + // - seqlen_q_ptr (per-sequence) and cu_seqlen_q_ptr (cumulative logical) are mutually + // exclusive. Use one set, not both. + // + // Batch mode: + // - cu_seqlen_q_ptr/cu_seqlen_k_ptr: Records cumulative logical (excluding padding) + // sequence lengths. [array size: batch + 1] + // - seqstart_* and seqlen_* pointers must be nullptr. + // + // Without padding: + // (Note: Physical length equals logical length) + // + // Group mode: + // - seqstart_q_ptr, seqstart_k_ptr: Record cumulative physical sequence lengths. [array + // size: batch + 1] + // - seqlen_q_ptr/seqlen_k_ptr and cu_seqlen_q_ptr/cu_seqlen_k_ptr must be nullptr. + // + // Batch mode: + // - All sequence length pointers (seqstart_*, seqlen_*, cu_seqlen_*) must be nullptr. + // + const void* seqstart_q_ptr = + nullptr; // Cumulative physical sequence length array [batch + 1]. (Used in Group mode) + const void* seqstart_k_ptr = + nullptr; // Cumulative physical sequence length array [batch + 1]. (Used in Group mode) + const void* seqlen_q_ptr = nullptr; // Per-sequence logical (excluding padding) length array + // [batch]. (Used in Group mode with padding) + const void* seqlen_k_ptr = nullptr; // Per-sequence logical (excluding padding) length array + // [batch]. (Used in Group mode with padding) + const void* cu_seqlen_q_ptr = nullptr; // Cumulative logical (excluding padding) sequence length + // array [batch + 1]. (Used with padding) + const void* cu_seqlen_k_ptr = nullptr; // Cumulative logical (excluding padding) sequence length + // array [batch + 1]. (Used with padding) + const void* block_scale_seqstart_q_ptr; + const void* block_scale_seqstart_k_ptr; + const void* seqstart_v_scale_ptr; + const void* sink_ptr; + + ck_tile::index_t seqlen_q; + ck_tile::index_t seqlen_k; + ck_tile::index_t batch; + ck_tile::index_t max_seqlen_q; + ck_tile::index_t hdim_q; + ck_tile::index_t hdim_v; + ck_tile::index_t nhead_q; + ck_tile::index_t nhead_k; + ck_tile::index_t num_head_q_total = 0; + ck_tile::index_t head_start = 0; + + float scale_s; + float logits_soft_cap; + + ck_tile::index_t stride_q; + ck_tile::index_t stride_k; + ck_tile::index_t stride_v; + ck_tile::index_t stride_bias; // if alibi, b*h need set this to h, 1*h need set this to 0 + ck_tile::index_t stride_randval; + ck_tile::index_t stride_o; + ck_tile::index_t stride_q_descale; + ck_tile::index_t stride_k_descale; + ck_tile::index_t stride_v_descale; + ck_tile::index_t nhead_stride_q; + ck_tile::index_t nhead_stride_k; + ck_tile::index_t nhead_stride_v; + ck_tile::index_t nhead_stride_bias; + ck_tile::index_t nhead_stride_randval; + ck_tile::index_t nhead_stride_lse; + ck_tile::index_t nhead_stride_o; + ck_tile::index_t nhead_stride_q_descale; + ck_tile::index_t nhead_stride_k_descale; + ck_tile::index_t nhead_stride_v_descale; + ck_tile::index_t batch_stride_q; + ck_tile::index_t batch_stride_k; + ck_tile::index_t batch_stride_v; + ck_tile::index_t batch_stride_bias; + ck_tile::index_t batch_stride_randval; + ck_tile::index_t batch_stride_lse; + ck_tile::index_t batch_stride_o; + ck_tile::index_t batch_stride_q_descale; + ck_tile::index_t batch_stride_k_descale; + ck_tile::index_t batch_stride_v_descale; + + ck_tile::index_t window_size_left; + ck_tile::index_t window_size_right; + ck_tile::index_t sink_size; + ck_tile::index_t mask_type; + ck_tile::index_t min_seqlen_q; + + float p_drop; + bool s_randval; + + std::variant, std::pair> + drop_seed_offset; + + ck_tile::index_t block_scale_size_q; + ck_tile::index_t block_scale_size_kv; +}; + +struct fmha_fwd_pagedkv_args +{ + const void* q_ptr; + const void* k_ptr; + const void* v_ptr; + const void* bias_ptr; // bias or alibi_slope pointer + void* lse_ptr; + void* o_ptr; + + void* block_table_ptr; + ck_tile::index_t batch_stride_block_table; // only used if 'block_table_ptr' is not nullptr + ck_tile::index_t page_block_size; // only used if 'block_table_ptr' is not nullptr + bool is_gappy; // differentiate seqstart_k_ptr usage. only used if 'block_table_ptr' is not + // nullptr. + + const void* cache_batch_idx; + + // the real seqlen_q & seqlen_k are decided by following: + // batch mode: seqlen_q = kargs.seqlen_q + // seqlen_k = kargs.seqlen_k + // group mode: seqlen_q = kargs.seqstart_q_ptr[b + 1] - kargs.seqstart_q_ptr[b] + // seqlen_k = kargs.seqstart_k_ptr[b + 1] - kargs.seqstart_k_ptr[b] + // or kargs.seqlen_k_ptr[b] + // + // batch mode (kvcache): + // seqlen_q = kargs.seqlen_q + // seqlen_k = kargs.seqlen_k_ptr[b] + // group mode (kvcache): + // seqlen_q = kargs.seqstart_q_ptr[b + 1] - kargs.seqstart_q_ptr[b] + // + // when is_gappy=true: + // seqlen_k = kargs.seqlen_k_ptr[b] + // seqstart_k_ptr[b] now store local offset of each batch + // + // when is_gappy=false: + // seqlen_k = kargs.seqstart_k_ptr[b + 1] - kargs.seqstart_k_ptr[b] + // or kargs.seqlen_k_ptr[b] + const void* seqstart_q_ptr; + const void* seqstart_k_ptr; + const void* seqlen_k_ptr; + const void* sink_ptr; + + ck_tile::index_t seqlen_q; + ck_tile::index_t seqlen_k; + ck_tile::index_t batch; + ck_tile::index_t max_seqlen_q; + ck_tile::index_t hdim_q; + ck_tile::index_t hdim_v; + ck_tile::index_t nhead_q; + ck_tile::index_t nhead_k; + + float scale_s; + float scale_p; + float scale_o; + + float logits_soft_cap; + + ck_tile::index_t stride_q; + ck_tile::index_t stride_k; + ck_tile::index_t stride_v; + ck_tile::index_t stride_bias; // if alibi, b*h need set this to h, 1*h need set this to 0 + ck_tile::index_t stride_o; + ck_tile::index_t nhead_stride_q; + ck_tile::index_t nhead_stride_k; + ck_tile::index_t nhead_stride_v; + ck_tile::index_t nhead_stride_bias; + ck_tile::index_t nhead_stride_lse; + ck_tile::index_t nhead_stride_o; + ck_tile::index_t batch_stride_q; + ck_tile::index_t batch_stride_k; + ck_tile::index_t batch_stride_v; + ck_tile::index_t batch_stride_bias; + ck_tile::index_t batch_stride_lse; + ck_tile::index_t batch_stride_o; + + ck_tile::index_t window_size_left; + ck_tile::index_t window_size_right; + ck_tile::index_t sink_size; + ck_tile::index_t mask_type; + ck_tile::index_t min_seqlen_q; +}; + +struct fmha_fwd_splitkv_args +{ + const void* q_ptr; + const void* k_ptr; + const void* v_ptr; + const void* bias_ptr; // bias or alibi_slope pointer + void* lse_acc_ptr; + void* o_acc_ptr; + void* lse_ptr; + void* o_ptr; + + void* block_table_ptr; + ck_tile::index_t batch_stride_block_table; // only used if 'block_table_ptr' is not nullptr + ck_tile::index_t page_block_size; // only used if 'block_table_ptr' is not nullptr + bool is_gappy; // differentiate seqstart_k_ptr usage. only used if 'block_table_ptr' is not + // nullptr. + + const void* cache_batch_idx; + + // the real seqlen_q & seqlen_k are decided by following: + // batch mode: seqlen_q = kargs.seqlen_q + // seqlen_k = kargs.seqlen_k + // group mode: seqlen_q = kargs.seqstart_q_ptr[b + 1] - kargs.seqstart_q_ptr[b] + // seqlen_k = kargs.seqstart_k_ptr[b + 1] - kargs.seqstart_k_ptr[b] + // or kargs.seqlen_k_ptr[b] + // + // batch mode (kvcache): + // seqlen_q = kargs.seqlen_q + // seqlen_k = kargs.seqlen_k_ptr[b] + // group mode (kvcache): + // seqlen_q = kargs.seqstart_q_ptr[b + 1] - kargs.seqstart_q_ptr[b] + // + // when is_gappy=true: + // seqlen_k = kargs.seqlen_k_ptr[b] + // seqstart_k_ptr[b] now store local offset of each batch + // + // when is_gappy=false: + // seqlen_k = kargs.seqstart_k_ptr[b + 1] - kargs.seqstart_k_ptr[b] + // or kargs.seqlen_k_ptr[b] + const void* seqstart_q_ptr; + const void* seqstart_k_ptr; + const void* seqlen_k_ptr; + const void* sink_ptr; + + ck_tile::index_t seqlen_q; + ck_tile::index_t seqlen_k; + ck_tile::index_t batch; + ck_tile::index_t max_seqlen_q; + ck_tile::index_t hdim_q; + ck_tile::index_t hdim_v; + ck_tile::index_t nhead_q; + ck_tile::index_t nhead_k; + ck_tile::index_t num_splits; + + float scale_s; + float scale_p; + float scale_o; + + float logits_soft_cap; + + ck_tile::index_t stride_q; + ck_tile::index_t stride_k; + ck_tile::index_t stride_v; + ck_tile::index_t stride_bias; // if alibi, b*h need set this to h, 1*h need set this to 0 + ck_tile::index_t stride_o_acc; + ck_tile::index_t stride_o; + ck_tile::index_t nhead_stride_q; + ck_tile::index_t nhead_stride_k; + ck_tile::index_t nhead_stride_v; + ck_tile::index_t nhead_stride_bias; + ck_tile::index_t nhead_stride_lse; + ck_tile::index_t nhead_stride_lse_acc; + ck_tile::index_t nhead_stride_o_acc; + ck_tile::index_t nhead_stride_o; + ck_tile::index_t batch_stride_q; + ck_tile::index_t batch_stride_k; + ck_tile::index_t batch_stride_v; + ck_tile::index_t batch_stride_bias; + ck_tile::index_t batch_stride_lse; + ck_tile::index_t batch_stride_lse_acc; + ck_tile::index_t batch_stride_o_acc; + ck_tile::index_t batch_stride_o; + ck_tile::index_t split_stride_lse_acc; + ck_tile::index_t split_stride_o_acc; + + ck_tile::index_t window_size_left; + ck_tile::index_t window_size_right; + ck_tile::index_t sink_size; + ck_tile::index_t mask_type; +}; + +struct fmha_fwd_appendkv_args +{ + void* q_ptr; + void* k_ptr; + const void* knew_ptr; + void* v_ptr; + const void* vnew_ptr; + + const void* seqlen_k_ptr; + + ck_tile::index_t seqlen_q; + ck_tile::index_t seqlen_knew; + ck_tile::index_t batch; + ck_tile::index_t hdim_q; + ck_tile::index_t hdim_v; + ck_tile::index_t nhead_q; + ck_tile::index_t nhead_k; + + const void* rotary_cos_ptr; // only used if 'rotary_dim' > 0 + const void* rotary_sin_ptr; // only used if 'rotary_dim' > 0 + ck_tile::index_t rotary_dim; + bool has_mask; + + void* block_table_ptr; + ck_tile::index_t batch_stride_block_table; // only used if 'block_table_ptr' is not nullptr + ck_tile::index_t page_block_size; // only used if 'block_table_ptr' is not nullptr + + const void* cache_batch_idx; // only used if block_table_ptr is nullptr -> batch mode (kvcache) + const void* sink_ptr; + + ck_tile::index_t stride_q; + ck_tile::index_t stride_k; + ck_tile::index_t stride_knew; + ck_tile::index_t stride_v; + ck_tile::index_t stride_vnew; + ck_tile::index_t nhead_stride_q; + ck_tile::index_t nhead_stride_k; + ck_tile::index_t nhead_stride_knew; + ck_tile::index_t nhead_stride_v; + ck_tile::index_t nhead_stride_vnew; + ck_tile::index_t batch_stride_q; + ck_tile::index_t batch_stride_k; + ck_tile::index_t batch_stride_knew; + ck_tile::index_t batch_stride_v; + ck_tile::index_t batch_stride_vnew; +}; + +struct fmha_batch_prefill_args +{ + const void* q_ptr; + const void* k_ptr; + const void* v_ptr; + const void* bias_ptr; // bias or alibi_slope pointer + const void* q_descale_ptr; + const void* k_descale_ptr; + const void* v_descale_ptr; + void* rand_val_ptr; + void* lse_ptr; + void* o_ptr; + + // the real seqlen_q & seqlen_k are decided by following: + // batch mode (kvcache): + // seqlen_q = kargs.seqlen_q + // seqlen_k = kargs.page_block_size * (kargs.kv_indptr[b + 1] - kargs.kv_indptr[b] - + // 1) + + // kargs.kv_last_page_lens[b] + // group mode (kvcache): + // seqlen_q = kargs.seqstart_q_ptr[b + 1] - kargs.seqstart_q_ptr[b] + // seqlen_k = kargs.page_block_size * (kargs.kv_indptr[b + 1] - kargs.kv_indptr[b] - + // 1) + + // kargs.kv_last_page_lens[b] + const void* seqstart_q_ptr; + const void* sink_ptr; + + ck_tile::index_t seqlen_q; + ck_tile::index_t seqlen_k; + ck_tile::index_t batch; + ck_tile::index_t max_seqlen_q; + ck_tile::index_t hdim_q; + ck_tile::index_t hdim_v; + ck_tile::index_t nhead_q; + ck_tile::index_t nhead_k; + + // KV cache page table fields (kv_lookup_table selects interpretation): + // - SGLANG_PAGE_TABLE_1D: + // kv_indptr: prefix-sum [batch+1] into kv_page_indices + // kv_page_indices: 1D list of physical page ids, length = num_total_pages + // kv_last_page_lens: per-batch last page lengths [batch] + // - VLLM_BLOCK_TABLE_2D: + // kv_page_indices: block_table [batch, max_blocks_per_seq] (2D) + // batch_stride_block_table: row stride for block_table + // seqlen_k_ptr: per-batch seqlen_k [batch] + int32_t num_total_pages; // total physical pages in KV cache (SGLang/vLLM) + ck_tile::index_t page_block_size; // tokens per page (SGLang/vLLM) + ck_tile::BlockAttentionKVCacheMemoryLayoutEnum + kv_memory_layout; // KV memory layout (SGLang/vLLM) + ck_tile::BlockAttentionKVCacheLookupTableEnum kv_lookup_table; // lookup table layout selector + void* kv_indptr; // SGLang: prefix-sum; vLLM: unused + void* kv_page_indices; // SGLang: 1D page list; vLLM: block_table 2D + void* kv_last_page_lens; // SGLang: last page lengths; vLLM: unused + void* seqlen_k_ptr; // vLLM: per-batch seqlen_k; SGLang: unused + ck_tile::index_t batch_stride_block_table; // vLLM: row stride; SGLang: unused + + float scale_s; + float scale_p; + float scale_o; + + float logits_soft_cap; + + ck_tile::index_t stride_q; + ck_tile::index_t stride_k; + ck_tile::index_t stride_v; + ck_tile::index_t stride_bias; // if alibi, b*h need set this to h, 1*h need set this to 0 + ck_tile::index_t stride_randval; + ck_tile::index_t stride_o; + ck_tile::index_t nhead_stride_q; + ck_tile::index_t nhead_stride_k; + ck_tile::index_t nhead_stride_v; + ck_tile::index_t nhead_stride_bias; + ck_tile::index_t nhead_stride_randval; + ck_tile::index_t nhead_stride_lse; + ck_tile::index_t nhead_stride_o; + ck_tile::index_t batch_stride_q; + ck_tile::index_t batch_stride_k; + ck_tile::index_t batch_stride_v; + ck_tile::index_t batch_stride_bias; + ck_tile::index_t batch_stride_randval; + ck_tile::index_t batch_stride_lse; + ck_tile::index_t batch_stride_o; + + ck_tile::index_t window_size_left; + ck_tile::index_t window_size_right; + ck_tile::index_t sink_size; + ck_tile::index_t mask_type; + + float p_drop; + bool s_randval; + + std::variant, std::pair> + drop_seed_offset; + + // KV_BLOCKSCALE: per-page K/V descales (Q per-tensor, K/V per-page) + // k_descale_ptr/v_descale_ptr are reused for KV_BLOCKSCALE mode: + // k_descale_ptr: [num_block, num_kv_head] - points to k block descale + // v_descale_ptr: [num_block, num_kv_head] - points to v block descale + ck_tile::index_t nblock_stride_kv_block_descale = 0; // Stride along num_block dimension + ck_tile::index_t nhead_stride_kv_block_descale = 0; // Stride along num_kv_head dimension +}; + +// Selects the KV-cache load mode for a batch-prefill dispatch arm. +// GLOBAL_LOAD_LDS: required when (a) the page is smaller than one K/V tile +// so per-page SRD is impossible, AND (b) the total KV-pool byte size +// exceeds INT32_MAX so SRD's 32-bit byte offset cannot address it. +// BUFFER_LOAD: every other case — the SGPR-resident SRD path is fastest. +// Inputs are taken as plain integers so the helper has no template parameter +// and can be called from each codegen-emitted dispatcher arm with the arm's +// compile-time kN0 / element_bytes substituted as constants. +inline ck_tile::BlockAttentionKVCacheLoadModeEnum +fmha_batch_prefill_select_kv_load_mode(ck_tile::index_t page_block_size, + ck_tile::index_t kN0, + ck_tile::index_t num_total_pages, + ck_tile::index_t batch_stride_k, + ck_tile::index_t element_bytes) +{ + // Promote every operand to long_index_t so overflow is impossible regardless + // of multiplication order. A bare `static_cast(num_total_pages) + // * batch_stride_k * element_bytes` only works because of left-to-right + // associativity — a future reorder of the operands would silently truncate. + const auto kv_pool_bytes = static_cast(num_total_pages) * + static_cast(batch_stride_k) * + static_cast(element_bytes); + return (page_block_size < kN0 && kv_pool_bytes > INT32_MAX) + ? ck_tile::BlockAttentionKVCacheLoadModeEnum::GLOBAL_LOAD_LDS + : ck_tile::BlockAttentionKVCacheLoadModeEnum::BUFFER_LOAD; +} + +template +auto fmha_fwd_create_kargs_and_grids(fmha_fwd_args args) +{ + assert(args.nhead_q % args.nhead_k == 0); + auto kargs = [&] { + // create group mode kernel arguments + if constexpr(FmhaKernel::kIsGroupMode) + { + return FmhaKernel::MakeKargsImpl(args.q_ptr, + args.k_ptr, + args.v_ptr, + args.bias_ptr, + args.q_descale_ptr, + args.k_descale_ptr, + args.v_descale_ptr, + args.rand_val_ptr, + args.lse_ptr, + args.o_ptr, + args.seqstart_q_ptr, + args.seqstart_k_ptr, + args.seqlen_q_ptr, + args.seqlen_k_ptr, + args.block_scale_seqstart_q_ptr, + args.block_scale_seqstart_k_ptr, + args.seqstart_v_scale_ptr, + args.hdim_q, + args.hdim_v, + args.nhead_q, + args.nhead_q / args.nhead_k, + args.scale_s, + args.logits_soft_cap, + args.stride_q, + args.stride_k, + args.stride_v, + args.stride_bias, + args.stride_randval, + args.stride_o, + args.stride_q_descale, + args.stride_k_descale, + args.stride_v_descale, + args.nhead_stride_q, + args.nhead_stride_k, + args.nhead_stride_v, + args.nhead_stride_bias, + args.nhead_stride_randval, + args.nhead_stride_lse, + args.nhead_stride_o, + args.nhead_stride_q_descale, + args.nhead_stride_k_descale, + args.nhead_stride_v_descale, + args.window_size_left, + args.window_size_right, + args.sink_size, + args.mask_type, + args.min_seqlen_q, + args.p_drop, + args.s_randval, + args.drop_seed_offset, + args.block_scale_size_q, + args.block_scale_size_kv, + args.cu_seqlen_q_ptr, + args.cu_seqlen_k_ptr, + args.sink_ptr, + args.num_head_q_total, + args.head_start); + } + else + { // create batch mode kernel arguments + return FmhaKernel::MakeKargsImpl(args.q_ptr, + args.k_ptr, + args.v_ptr, + args.bias_ptr, + args.q_descale_ptr, + args.k_descale_ptr, + args.v_descale_ptr, + args.rand_val_ptr, + args.lse_ptr, + args.o_ptr, + args.seqlen_q, + args.seqlen_k, + args.hdim_q, + args.hdim_v, + args.nhead_q, + args.nhead_q / args.nhead_k, + args.scale_s, + args.logits_soft_cap, + args.stride_q, + args.stride_k, + args.stride_v, + args.stride_bias, + args.stride_randval, + args.stride_o, + args.stride_q_descale, + args.stride_k_descale, + args.stride_v_descale, + args.nhead_stride_q, + args.nhead_stride_k, + args.nhead_stride_v, + args.nhead_stride_bias, + args.nhead_stride_randval, + args.nhead_stride_lse, + args.nhead_stride_o, + args.nhead_stride_q_descale, + args.nhead_stride_k_descale, + args.nhead_stride_v_descale, + args.batch_stride_q, + args.batch_stride_k, + args.batch_stride_v, + args.batch_stride_bias, + args.batch_stride_randval, + args.batch_stride_lse, + args.batch_stride_o, + args.batch_stride_q_descale, + args.batch_stride_k_descale, + args.batch_stride_v_descale, + args.window_size_left, + args.window_size_right, + args.sink_size, + args.mask_type, + args.p_drop, + args.s_randval, + args.drop_seed_offset, + args.block_scale_size_q, + args.block_scale_size_kv, + args.cu_seqlen_q_ptr, + args.cu_seqlen_k_ptr, + args.sink_ptr, + args.num_head_q_total, + args.head_start); + } + }(); + + if constexpr(FmhaKernel::kIsGroupMode) + { + dim3 grids = FmhaKernel::GridSize( + args.batch, args.nhead_q, args.max_seqlen_q, args.hdim_v, args.seqlen_k_ptr != nullptr); + return ck_tile::make_tuple(kargs, grids); + } + else + { + dim3 grids = + FmhaKernel::GridSize(args.batch, args.nhead_q, args.max_seqlen_q, args.hdim_v, false); + return ck_tile::make_tuple(kargs, grids); + } +} + +template +auto fmha_fwd_v3_create_kargs_and_grids(fmha_fwd_args args) +{ + /// NOTICE: This was borrowed from Aiter. Make sure the selected remap_opt setting truly + /// maximizes the kernel's performance. + int remap_opt = 2; + if(args.mask_type != static_cast(mask_enum::no_mask) && + ((args.nhead_q % 8 != 0) || (16384 < args.seqlen_q))) + { + if(65536 <= args.seqlen_q) + { + remap_opt = 0; + } + else + { + remap_opt = 1; + } + } + + auto kargs = [&] { + if constexpr(FmhaKernel::kIsGroupMode) + { + return FmhaKernel::MakeKargs(args.q_ptr, + args.k_ptr, + args.v_ptr, + args.q_descale_ptr, + args.k_descale_ptr, + args.v_descale_ptr, + nullptr, // lse_ptr + args.o_ptr, + args.seqstart_q_ptr, + args.seqstart_k_ptr, + args.seqlen_q_ptr, + args.seqlen_k_ptr, + args.hdim_q, + args.hdim_v, + args.nhead_q, + args.nhead_q / args.nhead_k, + args.scale_s, + args.logits_soft_cap, + args.stride_q, + args.stride_k, + args.stride_v, + args.stride_o, + args.nhead_stride_q, + args.nhead_stride_k, + args.nhead_stride_v, + 0, // nhead_stride_lse + args.nhead_stride_o, + args.window_size_left, + args.window_size_right, + args.mask_type, + remap_opt, + args.cu_seqlen_q_ptr, + args.cu_seqlen_k_ptr); + } + else + { + return FmhaKernel::MakeKargs(args.q_ptr, + args.k_ptr, + args.v_ptr, + args.q_descale_ptr, + args.k_descale_ptr, + args.v_descale_ptr, + nullptr, // lse_ptr + args.o_ptr, + args.seqlen_q, + args.seqlen_k, + args.hdim_q, + args.hdim_v, + args.nhead_q, + args.nhead_q / args.nhead_k, + args.scale_s, + args.logits_soft_cap, + args.stride_q, + args.stride_k, + args.stride_v, + args.stride_o, + args.nhead_stride_q, + args.nhead_stride_k, + args.nhead_stride_v, + 0, // nhead_stride_lse + args.nhead_stride_o, + args.batch_stride_q, + args.batch_stride_k, + args.batch_stride_v, + 0, // batch_stride_lse + args.batch_stride_o, + args.window_size_left, + args.window_size_right, + args.mask_type, + remap_opt, + args.cu_seqlen_q_ptr, + args.cu_seqlen_k_ptr); + } + }(); + + dim3 grids = FmhaKernel::GridSize(args.batch, args.nhead_q, args.max_seqlen_q, args.hdim_v); + + return ck_tile::make_tuple(kargs, grids); +} + +template +auto fmha_fwd_pagedkv_create_kargs_and_grids(fmha_fwd_pagedkv_args args) +{ + assert(args.nhead_q % args.nhead_k == 0); + auto kargs = [&] { + // create group mode kernel arguments + if constexpr(FmhaKernel::kIsGroupMode) + { + return FmhaKernel::MakeKargs(args.q_ptr, + args.k_ptr, + args.v_ptr, + args.bias_ptr, + args.lse_ptr, + args.o_ptr, + args.seqstart_q_ptr, + args.seqstart_k_ptr, + args.seqlen_k_ptr, + args.hdim_q, + args.hdim_v, + args.nhead_q, + args.nhead_q / args.nhead_k, + args.block_table_ptr, + args.batch_stride_block_table, + args.page_block_size, + args.is_gappy, + args.scale_s, + args.scale_p, + args.scale_o, + args.logits_soft_cap, + args.stride_q, + args.stride_k, + args.stride_v, + args.stride_bias, + args.stride_o, + args.nhead_stride_q, + args.nhead_stride_k, + args.nhead_stride_v, + args.nhead_stride_bias, + args.nhead_stride_lse, + args.nhead_stride_o, + args.batch_stride_k, + args.batch_stride_v, + args.window_size_left, + args.window_size_right, + args.sink_size, + args.mask_type, + args.min_seqlen_q, + args.sink_ptr); + } + else + { // create batch mode kernel arguments + return FmhaKernel::MakeKargs(args.q_ptr, + args.k_ptr, + args.v_ptr, + args.bias_ptr, + args.lse_ptr, + args.o_ptr, + args.seqlen_q, + args.seqlen_k, + args.seqlen_k_ptr, + args.hdim_q, + args.hdim_v, + args.nhead_q, + args.nhead_q / args.nhead_k, + args.block_table_ptr, + args.batch_stride_block_table, + args.page_block_size, + args.cache_batch_idx, + args.scale_s, + args.scale_p, + args.scale_o, + args.logits_soft_cap, + args.stride_q, + args.stride_k, + args.stride_v, + args.stride_bias, + args.stride_o, + args.nhead_stride_q, + args.nhead_stride_k, + args.nhead_stride_v, + args.nhead_stride_bias, + args.nhead_stride_lse, + args.nhead_stride_o, + args.batch_stride_q, + args.batch_stride_k, + args.batch_stride_v, + args.batch_stride_bias, + args.batch_stride_lse, + args.batch_stride_o, + args.window_size_left, + args.window_size_right, + args.sink_size, + args.mask_type, + args.sink_ptr); + } + }(); + + // FmhaKernel::PrintParameters(kargs, args.batch); + if constexpr(FmhaKernel::kIsGroupMode) + { + dim3 grids = FmhaKernel::GridSize( + args.batch, args.nhead_q, args.max_seqlen_q, args.hdim_v, args.seqlen_k_ptr != nullptr); + return ck_tile::make_tuple(kargs, grids); + } + else + { + dim3 grids = + FmhaKernel::GridSize(args.batch, args.nhead_q, args.max_seqlen_q, args.hdim_v, false); + return ck_tile::make_tuple(kargs, grids); + } +} + +template +auto fmha_fwd_splitkv_create_kargs_and_grids(fmha_fwd_splitkv_args args) +{ + assert(args.nhead_q % args.nhead_k == 0); + auto kargs = [&] { + // create group mode kernel arguments + if constexpr(Kernel::kIsGroupMode) + { + return Kernel::MakeKargs(args.q_ptr, + args.k_ptr, + args.v_ptr, + args.bias_ptr, + args.lse_acc_ptr, + args.o_acc_ptr, + args.batch, + args.seqstart_q_ptr, + args.seqstart_k_ptr, + args.seqlen_k_ptr, + args.hdim_q, + args.hdim_v, + args.nhead_q, + args.nhead_q / args.nhead_k, + args.num_splits, + args.block_table_ptr, + args.batch_stride_block_table, + args.page_block_size, + args.is_gappy, + args.scale_s, + args.scale_p, + args.logits_soft_cap, + args.stride_q, + args.stride_k, + args.stride_v, + args.stride_bias, + args.stride_o_acc, + args.nhead_stride_q, + args.nhead_stride_k, + args.nhead_stride_v, + args.nhead_stride_bias, + args.nhead_stride_lse_acc, + args.nhead_stride_o_acc, + args.batch_stride_k, // only used for paged-kvcache + args.batch_stride_v, // only used for paged-kvcache + args.split_stride_lse_acc, + args.split_stride_o_acc, + args.window_size_left, + args.window_size_right, + args.sink_size, + args.mask_type, + args.sink_ptr); + } + else + { // create batch mode kernel arguments + return Kernel::MakeKargs(args.q_ptr, + args.k_ptr, + args.v_ptr, + args.bias_ptr, + args.lse_acc_ptr, + args.o_acc_ptr, + args.batch, + args.seqlen_q, + args.seqlen_k, + args.seqlen_k_ptr, + args.hdim_q, + args.hdim_v, + args.nhead_q, + args.nhead_q / args.nhead_k, + args.num_splits, + args.block_table_ptr, + args.batch_stride_block_table, + args.page_block_size, + args.cache_batch_idx, + args.scale_s, + args.scale_p, + args.logits_soft_cap, + args.stride_q, + args.stride_k, + args.stride_v, + args.stride_bias, + args.stride_o_acc, + args.nhead_stride_q, + args.nhead_stride_k, + args.nhead_stride_v, + args.nhead_stride_bias, + args.nhead_stride_lse_acc, + args.nhead_stride_o_acc, + args.batch_stride_q, + args.batch_stride_k, + args.batch_stride_v, + args.batch_stride_bias, + args.batch_stride_lse_acc, + args.batch_stride_o_acc, + args.split_stride_lse_acc, + args.split_stride_o_acc, + args.window_size_left, + args.window_size_right, + args.sink_size, + args.mask_type, + args.sink_ptr); + } + }(); + + dim3 grids = Kernel::GridSize( + args.batch, args.nhead_q, args.nhead_k, args.max_seqlen_q, args.hdim_v, args.num_splits); + + return ck_tile::make_tuple(kargs, grids); +} + +template +auto fmha_fwd_splitkv_combine_create_kargs_and_grids(fmha_fwd_splitkv_args args) +{ + assert(args.nhead_q % args.nhead_k == 0); + auto kargs = [&] { + // create group mode kernel argumentszs + if constexpr(Kernel::kIsGroupMode) + { + return Kernel::MakeKargs(args.lse_acc_ptr, + args.o_acc_ptr, + args.lse_ptr, + args.o_ptr, + args.batch, + args.seqstart_q_ptr, + args.hdim_v, + args.num_splits, + args.scale_o, + args.stride_o_acc, + args.stride_o, + args.nhead_stride_lse_acc, + args.nhead_stride_o_acc, + args.nhead_stride_lse, + args.nhead_stride_o, + args.split_stride_lse_acc, + args.split_stride_o_acc); + } + else + { // create batch mode kernel arguments + return Kernel::MakeKargs(args.lse_acc_ptr, + args.o_acc_ptr, + args.lse_ptr, + args.o_ptr, + args.batch, + args.seqlen_q, + args.hdim_v, + args.num_splits, + args.scale_o, + args.stride_o_acc, + args.stride_o, + args.nhead_stride_lse_acc, + args.nhead_stride_o_acc, + args.nhead_stride_lse, + args.nhead_stride_o, + args.batch_stride_lse_acc, + args.batch_stride_o_acc, + args.batch_stride_lse, + args.batch_stride_o, + args.split_stride_lse_acc, + args.split_stride_o_acc); + } + }(); + + dim3 grids = Kernel::GridSize(args.batch, args.nhead_q, args.max_seqlen_q, args.hdim_v); + + return ck_tile::make_tuple(kargs, grids); +} + +template +auto fmha_fwd_appendkv_create_kargs_and_grids(fmha_fwd_appendkv_args args) +{ + assert(args.nhead_q % args.nhead_k == 0); + auto kargs = Kernel::MakeKargs(args.q_ptr, + args.k_ptr, + args.knew_ptr, + args.v_ptr, + args.vnew_ptr, + args.seqlen_q, + args.seqlen_k_ptr, + args.seqlen_knew, + args.hdim_q, + args.hdim_v, + args.nhead_q, + args.nhead_q / args.nhead_k, + args.rotary_cos_ptr, + args.rotary_sin_ptr, + args.rotary_dim, + args.has_mask, + args.block_table_ptr, + args.batch_stride_block_table, + args.page_block_size, + args.cache_batch_idx, + args.stride_q, + args.stride_k, + args.stride_knew, + args.stride_v, + args.stride_vnew, + args.nhead_stride_q, + args.nhead_stride_k, + args.nhead_stride_knew, + args.nhead_stride_v, + args.nhead_stride_vnew, + args.batch_stride_q, + args.batch_stride_k, + args.batch_stride_knew, + args.batch_stride_v, + args.batch_stride_vnew); + + dim3 grids = Kernel::GridSize(args.batch, args.nhead_q, args.seqlen_q, args.seqlen_knew); + + return ck_tile::make_tuple(kargs, grids); +} + +template +auto fmha_batch_prefill_create_kargs_and_grids(fmha_batch_prefill_args args) +{ + assert(args.nhead_q % args.nhead_k == 0); + using PageTableKargs = typename FmhaKernel::PageBlockTableKargs; + const PageTableKargs page_table = [&]() { + if constexpr(FmhaKernel::kKVLookupTable == + ck_tile::BlockAttentionKVCacheLookupTableEnum::SGLANG_PAGE_TABLE_1D) + { + return PageTableKargs{reinterpret_cast(args.kv_indptr), + reinterpret_cast(args.kv_page_indices), + reinterpret_cast(args.kv_last_page_lens)}; + } + else + { + return PageTableKargs{reinterpret_cast(args.kv_page_indices), + args.batch_stride_block_table, + reinterpret_cast(args.seqlen_k_ptr)}; + } + }(); + auto kargs = [&] { + // create group mode kernel arguments + if constexpr(FmhaKernel::kIsGroupMode) + { + return FmhaKernel::MakeKargs(args.q_ptr, + args.k_ptr, + args.v_ptr, + args.bias_ptr, + args.q_descale_ptr, + args.k_descale_ptr, + args.v_descale_ptr, + args.rand_val_ptr, + args.lse_ptr, + args.o_ptr, + args.seqstart_q_ptr, + args.hdim_q, + args.hdim_v, + args.nhead_q, + args.nhead_q / args.nhead_k, + args.num_total_pages, + args.page_block_size, + page_table, + args.scale_s, + args.scale_p, + args.scale_o, + args.logits_soft_cap, + args.stride_q, + args.stride_k, + args.stride_v, + args.stride_bias, + args.stride_randval, + args.stride_o, + args.nhead_stride_q, + args.nhead_stride_k, + args.nhead_stride_v, + args.nhead_stride_bias, + args.nhead_stride_randval, + args.nhead_stride_lse, + args.nhead_stride_o, + args.batch_stride_k, + args.batch_stride_v, + args.window_size_left, + args.window_size_right, + args.sink_size, + args.mask_type, + args.p_drop, + args.s_randval, + args.drop_seed_offset, + args.sink_ptr, + args.nblock_stride_kv_block_descale, + args.nhead_stride_kv_block_descale); + } + else + { // create batch mode kernel arguments + return FmhaKernel::MakeKargs(args.q_ptr, + args.k_ptr, + args.v_ptr, + args.bias_ptr, + args.q_descale_ptr, + args.k_descale_ptr, + args.v_descale_ptr, + args.rand_val_ptr, + args.lse_ptr, + args.o_ptr, + args.seqlen_q, + args.hdim_q, + args.hdim_v, + args.nhead_q, + args.nhead_q / args.nhead_k, + args.num_total_pages, + args.page_block_size, + page_table, + args.scale_s, + args.scale_p, + args.scale_o, + args.logits_soft_cap, + args.stride_q, + args.stride_k, + args.stride_v, + args.stride_bias, + args.stride_randval, + args.stride_o, + args.nhead_stride_q, + args.nhead_stride_k, + args.nhead_stride_v, + args.nhead_stride_bias, + args.nhead_stride_randval, + args.nhead_stride_lse, + args.nhead_stride_o, + args.batch_stride_q, + args.batch_stride_k, + args.batch_stride_v, + args.batch_stride_bias, + args.batch_stride_randval, + args.batch_stride_lse, + args.batch_stride_o, + args.window_size_left, + args.window_size_right, + args.sink_size, + args.mask_type, + args.p_drop, + args.s_randval, + args.drop_seed_offset, + args.sink_ptr, + args.nblock_stride_kv_block_descale, + args.nhead_stride_kv_block_descale); + } + }(); + + dim3 grids = FmhaKernel::GridSize(args.batch, args.nhead_q, args.max_seqlen_q, args.hdim_v); + return ck_tile::make_tuple(kargs, grids); +} + +// this is used to pattern-match internl kernel implementation, not to instantiate kernel +template +struct fmha_fwd_traits_ +{ + static constexpr ck_tile::index_t HDim = HDim_; + using DataType = ck_tile::remove_cvref_t; + static constexpr bool kIsGroupMode = kIsGroupMode_; + static constexpr ck_tile::index_t kM0 = kM0_; + static constexpr ck_tile::index_t kN0 = kN0_; + static constexpr ck_tile::index_t kK0 = kK0_; + static constexpr ck_tile::index_t kN1 = kN1_; + static constexpr ck_tile::index_t kK1 = kK1_; + static constexpr ck_tile::index_t kK0BlockLength = kK0BlockLength_; + static constexpr bool kIsVLayoutRowMajor = kIsVLayoutRowMajor_; + static constexpr auto FmhaPipelineEnum = FmhaPipelineEnum_; + static constexpr bool kHasLogitsSoftCap = kHasLogitsSoftCap_; + using FmhaMask = ck_tile::remove_cvref_t; + static constexpr auto BiasEnum = BiasEnum_; + static constexpr bool kStoreLse = kStoreLse_; + static constexpr bool kHasDropout = kHasDropout_; + static constexpr auto QScaleEnum = QScaleEnum_; + static constexpr bool kPadS = kPadS_; + static constexpr bool kPadSK = kPadSK_; + static constexpr bool kPadD = kPadD_; + static constexpr bool kPadDv = kPadDv_; + static constexpr bool kUseTrLoad = kUseTrLoad_; + static constexpr bool kSkipMinSeqlenQ = kSkipMinSeqlenQ_; + static constexpr bool kHasSink = kHasSink_; +}; + +template +struct fmha_fwd_batch_prefill_traits_ : public fmha_fwd_traits_ +{ + static constexpr auto kKVMemoryLayout = kKVMemoryLayout_; + static constexpr auto kKVLookupTable = kKVLookupTable_; + static constexpr ck_tile::index_t kPageBlockSize = kPageBlockSize_; + static constexpr auto kKVLoadMode = kKVLoadMode_; + static_assert(kIsVLayoutRowMajor_, "Batch prefill only supports row-major V layout"); +}; + +template +float fmha_fwd_(const ck_tile::stream_config&, fmha_fwd_args); + +template +struct fmha_fwd_pagedkv_traits_ +{ + static constexpr ck_tile::index_t HDim = HDim_; + using DataType = ck_tile::remove_cvref_t; + static constexpr bool kIsGroupMode = kIsGroupMode_; + static constexpr ck_tile::index_t kM0 = kM0_; + static constexpr ck_tile::index_t kN0 = kN0_; + static constexpr ck_tile::index_t kK0 = kK0_; + static constexpr ck_tile::index_t kN1 = kN1_; + static constexpr ck_tile::index_t kK1 = kK1_; + static constexpr ck_tile::index_t kK0BlockLength = kK0BlockLength_; + static constexpr bool kIsVLayoutRowMajor = kIsVLayoutRowMajor_; + static constexpr auto FmhaPipelineEnum = FmhaPipelineEnum_; + static constexpr bool kHasLogitsSoftCap = kHasLogitsSoftCap_; + using FmhaMask = ck_tile::remove_cvref_t; + static constexpr auto BiasEnum = BiasEnum_; + static constexpr bool kStoreLse = kStoreLse_; + static constexpr bool kIsPagedKV = kIsPagedKV_; + static constexpr bool kDoFp8StaticQuant = kDoFp8StaticQuant_; + static constexpr bool kPadS = kPadS_; + static constexpr bool kPadSK = kPadSK_; + static constexpr bool kPadD = kPadD_; + static constexpr bool kPadDv = kPadDv_; + static constexpr bool kSkipMinSeqlenQ = kSkipMinSeqlenQ_; + static constexpr bool kHasSink = kHasSink_; +}; + +template +float fmha_fwd_pagedkv_(const ck_tile::stream_config&, fmha_fwd_pagedkv_args); + +template +struct fmha_fwd_splitkv_traits_ +{ + static constexpr ck_tile::index_t HDim = HDim_; + using DataType = ck_tile::remove_cvref_t; + static constexpr bool kIsGroupMode = kIsGroupMode_; + static constexpr ck_tile::index_t kM0 = kM0_; + static constexpr ck_tile::index_t kN0 = kN0_; + static constexpr ck_tile::index_t kK0 = kK0_; + static constexpr ck_tile::index_t kN1 = kN1_; + static constexpr ck_tile::index_t kK1 = kK1_; + static constexpr ck_tile::index_t kK0BlockLength = kK0BlockLength_; + static constexpr bool kIsVLayoutRowMajor = kIsVLayoutRowMajor_; + static constexpr auto FmhaPipelineEnum = FmhaPipelineEnum_; + static constexpr bool kHasLogitsSoftCap = kHasLogitsSoftCap_; + using FmhaMask = ck_tile::remove_cvref_t; + static constexpr auto BiasEnum = BiasEnum_; + static constexpr bool kStoreLse = kStoreLse_; + static constexpr bool kDoFp8StaticQuant = kDoFp8StaticQuant_; + static constexpr bool kPadS = kPadS_; + static constexpr bool kPadSK = kPadSK_; + static constexpr bool kPadD = kPadD_; + static constexpr bool kPadDv = kPadDv_; + static constexpr bool kIsPagedKV = kIsPagedKV_; + static constexpr bool kHasSink = kHasSink_; +}; + +template +void fmha_fwd_splitkv_oneshot_(const ck_tile::stream_config&, fmha_fwd_splitkv_args); + +template +std::string fmha_fwd_splitkv_get_name_(); + +template +struct fmha_fwd_splitkv_combine_traits_ +{ + static constexpr ck_tile::index_t HDim = HDim_; + using DataType = ck_tile::remove_cvref_t; + static constexpr bool kIsGroupMode = kIsGroupMode_; + static constexpr ck_tile::index_t kN1 = kN1_; + static constexpr bool kStoreLse = kStoreLse_; + static constexpr bool kDoFp8StaticQuant = kDoFp8StaticQuant_; + static constexpr bool kPadS = kPadS_; + static constexpr bool kPadDv = kPadDv_; +}; + +template +void fmha_fwd_splitkv_combine_oneshot_(const ck_tile::stream_config&, fmha_fwd_splitkv_args); + +template +std::string fmha_fwd_splitkv_combine_get_name_(); + +// this is used to pattern-match internl kernel implementation, not to instantiate kernel +template +struct fmha_fwd_appendkv_traits_ +{ + static constexpr ck_tile::index_t HDim = HDim_; + using DataType = ck_tile::remove_cvref_t; + static constexpr ck_tile::index_t kTileSizeS = kTileSizeS_; + static constexpr ck_tile::index_t kTileSizeSk = kTileSizeSk_; + static constexpr ck_tile::index_t kTileSizeD = kTileSizeD_; + static constexpr ck_tile::index_t kTileSizeDv = kTileSizeDv_; + static constexpr bool kIsVLayoutRowMajor = kIsVLayoutRowMajor_; + static constexpr bool kPadS = kPadS_; + static constexpr bool kPadSk = kPadSk_; + static constexpr bool kPadD = kPadD_; + static constexpr bool kPadDv = kPadDv_; + static constexpr auto RotaryEnum = RotaryEnum_; + static constexpr bool kIsPagedKV = kIsPagedKV_; +}; + +template +float fmha_fwd_appendkv_(const ck_tile::stream_config&, fmha_fwd_appendkv_args); + +template +float fmha_batch_prefill_(const ck_tile::stream_config&, fmha_batch_prefill_args); + +// This is the public API, will be generated by script +struct fmha_fwd_traits +{ + int hdim_q; + int hdim_v; + std::string data_type; + bool is_group_mode; + bool is_v_rowmajor; + bool has_logits_soft_cap; + mask_enum mask_type; + bias_enum bias_type; // 0:no bias, 1:elementwise bias, 2:alibi. sync with BlockAttentionBiasEnum + bool has_lse; + bool has_dropout; + quant_scale_enum qscale_type; + bool skip_min_seqlen_q = false; + bool has_sink = false; + // TODO: padding check is inside this api +}; +float fmha_fwd(fmha_fwd_traits, fmha_fwd_args, const ck_tile::stream_config&); + +struct fmha_fwd_pagedkv_traits +{ + int hdim_q; + int hdim_v; + std::string data_type; + bool is_group_mode; + bool is_v_rowmajor; + bool has_logits_soft_cap; + mask_enum mask_type; + bias_enum bias_type; // 0:no bias, 1:elementwise bias, 2:alibi. sync with BlockAttentionBiasEnum + bool has_lse = false; + bool use_pagedkv = true; + bool do_fp8_static_quant = false; + bool skip_min_seqlen_q = false; + bool has_sink = false; + // TODO: padding check is inside this api +}; + +float fmha_fwd_pagedkv(fmha_fwd_pagedkv_traits&, + fmha_fwd_pagedkv_args&, + const ck_tile::stream_config&); + +struct fmha_fwd_splitkv_traits +{ + int hdim_q; + int hdim_v; + std::string data_type; + bool is_group_mode; + bool is_v_rowmajor; + bool has_logits_soft_cap; + mask_enum mask_type; + bias_enum bias_type; // 0:no bias, 1:elementwise bias, 2:alibi. sync with BlockAttentionBiasEnum + bool has_lse; + bool do_fp8_static_quant = false; + bool has_sink = false; + // TODO: padding check is inside this api +}; +float fmha_fwd_splitkv(fmha_fwd_splitkv_traits, + fmha_fwd_splitkv_args, + const ck_tile::stream_config&); + +struct fmha_fwd_appendkv_traits +{ + int hdim_q; + int hdim_v; + std::string data_type; + bool is_v_rowmajor; + rope_enum rope_type; +}; +float fmha_fwd_appendkv(fmha_fwd_appendkv_traits, + fmha_fwd_appendkv_args, + const ck_tile::stream_config&); + +struct fmha_batch_prefill_traits : public fmha_fwd_traits +{ + ck_tile::BlockAttentionKVCacheMemoryLayoutEnum kv_memory_layout = + ck_tile::BlockAttentionKVCacheMemoryLayoutEnum::VECTORIZED_LAYOUT; + ck_tile::BlockAttentionKVCacheLookupTableEnum kv_lookup_table = + ck_tile::BlockAttentionKVCacheLookupTableEnum::SGLANG_PAGE_TABLE_1D; + int page_size = 1; +}; + +float fmha_batch_prefill(fmha_batch_prefill_traits, + fmha_batch_prefill_args, + const ck_tile::stream_config&); diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_api.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_api.cpp new file mode 100644 index 0000000000..180d4df9ab --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_api.cpp @@ -0,0 +1,418 @@ + +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include + +#include + +#include "fmha_fwd.hpp" + +namespace { +bool get_num_cus(unsigned& num_cus) { + int device; + auto status = hipGetDevice(&device); + if(status != hipSuccess) { + fprintf(stderr, "failed to get device"); + return false; + } + + hipDeviceProp_t props{}; + status = hipGetDeviceProperties(&props, device); + if(status != hipSuccess) { + fprintf(stderr, "failed to get device properties"); + return false; + } + + num_cus = props.multiProcessorCount; + return true; +} + +unsigned get_num_thread_blocks(unsigned batch, unsigned nheads, unsigned max_seqlen_q, unsigned kM0) { + const unsigned num_m_blocks = (max_seqlen_q + kM0 - 1) / kM0; + const unsigned num_n_blocks = 1; // we assume that num_n_blocks is always 1 + + return batch * nheads * num_m_blocks * num_n_blocks; +} +} // namespace + +namespace { +float fmha_fwd_v2([[maybe_unused]] fmha_fwd_traits t, [[maybe_unused]] fmha_fwd_args a, [[maybe_unused]] const ck_tile::stream_config& s) { + float r = -1; + + [[maybe_unused]] const float min_cu_util_rate = 0.8; // minimum CU utilization rate + + unsigned num_cus; + if(!get_num_cus(num_cus)) { + return r; + } + + [[maybe_unused]] auto get_num_blocks = [&](unsigned kM0) { + return get_num_thread_blocks(a.batch, a.nhead_q, a.max_seqlen_q, kM0); + }; + + [[maybe_unused]] const std::string device_name = ck_tile::get_device_name(); + + if(device_name.compare(0, 6, "gfx950") == 0) { + if(t.data_type.compare("fp16") == 0) { + if(t.hdim_q <= 32 && t.hdim_v <= 32) { + if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true) && (true/*fall back to largest tile*/) && ((a.cu_seqlen_k_ptr == nullptr) && (a.seqlen_k != 0 && a.seqlen_k % 64 == 0)) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 128, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, false, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true) && (true/*fall back to largest tile*/) && ((a.cu_seqlen_k_ptr != nullptr) || (a.seqlen_k == 0 || a.seqlen_k % 64 != 0)) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 128, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (a.seqlen_q % 128 == 0) && (true/*fall back to largest tile*/) && ((a.cu_seqlen_k_ptr == nullptr) && (a.seqlen_k != 0 && a.seqlen_k % 64 == 0)) && (a.hdim_q % 32 == 0) && (a.hdim_v % 32 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 128, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, false, false, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 128 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (true /*a.hdim_q % 32 != 0*/) && (true /*a.hdim_v % 32 != 0*/) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 128, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + + } + else if(t.hdim_q <= 64 && t.hdim_v <= 64) { + if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true) && (a.seqlen_q <= 16) && (true) && (a.hdim_q % 64 == 0) && (a.hdim_v % 64 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 16, 32, 64, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, false, false, true, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true) && (a.seqlen_q <= 16) && (true) && (true /*a.hdim_q % 64 != 0*/) && (true /*a.hdim_v % 64 != 0*/) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 16, 32, 64, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, true, true, true, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true) && (a.seqlen_q <= 16) && (true) && (a.hdim_q % 64 == 0) && (a.hdim_v % 64 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 16, 32, 64, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, true, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true) && (a.seqlen_q <= 32) && (true) && (a.hdim_q % 64 == 0) && (a.hdim_v % 64 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 32, 32, 64, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, false, false, true, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true) && (a.seqlen_q <= 32) && (true) && (true /*a.hdim_q % 64 != 0*/) && (true /*a.hdim_v % 64 != 0*/) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 32, 32, 64, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, true, true, true, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true) && (a.seqlen_q <= 32) && (true) && (a.hdim_q % 64 == 0) && (a.hdim_v % 64 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 32, 32, 64, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, true, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true) && (true/*fall back to largest tile*/) && ((a.cu_seqlen_k_ptr == nullptr) && (a.seqlen_k != 0 && a.seqlen_k % 64 == 0)) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 128, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, false, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true) && (true/*fall back to largest tile*/) && ((a.cu_seqlen_k_ptr != nullptr) || (a.seqlen_k == 0 || a.seqlen_k % 64 != 0)) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 128, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (a.seqlen_q % 128 == 0) && (true/*fall back to largest tile*/) && ((a.cu_seqlen_k_ptr == nullptr) && (a.seqlen_k != 0 && a.seqlen_k % 64 == 0)) && (a.hdim_q % 64 == 0) && (a.hdim_v % 64 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 128, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, false, false, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 128 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (true /*a.hdim_q % 64 != 0*/) && (true /*a.hdim_v % 64 != 0*/) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 128, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true) && (true/*fall back to largest tile*/) && (true) && (a.hdim_q % 64 == 0) && (a.hdim_v % 64 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 128, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, false, false, true, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true) && (true/*fall back to largest tile*/) && (true) && (true /*a.hdim_q % 64 != 0*/) && (true /*a.hdim_v % 64 != 0*/) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 128, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, true, true, true, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true) && (true/*fall back to largest tile*/) && (true) && (a.hdim_q % 64 == 0) && (a.hdim_v % 64 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 128, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, true, false, false>; + return fmha_fwd_(s, a); + } + + } + + } + + } + else if(device_name.compare(0, 6, "gfx115") == 0) { + if(t.data_type.compare("fp16") == 0) { + if(t.hdim_q <= 32 && t.hdim_v <= 32) { + if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 32 == 0) && (a.hdim_v % 32 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 64, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 64, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 32 == 0) && (a.hdim_v % 32 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 64, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 64, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + + } + else if(t.hdim_q <= 64 && t.hdim_v <= 64) { + if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 64 == 0) && (a.hdim_v % 64 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 64, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 64, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 64 == 0) && (a.hdim_v % 64 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 64, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 64, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + + } + + } + + } + else if(device_name.compare(0, 5, "gfx11") == 0) { + if(t.data_type.compare("fp16") == 0) { + if(t.hdim_q <= 32 && t.hdim_v <= 32) { + if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 32 == 0) && (a.hdim_v % 32 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 64, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 64, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 32 == 0) && (a.hdim_v % 32 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 64, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 64, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + + } + else if(t.hdim_q <= 64 && t.hdim_v <= 64) { + if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 64 == 0) && (a.hdim_v % 64 == 0) && ((a.max_seqlen_q < 4096) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 64, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((a.max_seqlen_q < 4096) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 64, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 64 == 0) && (a.hdim_v % 64 == 0) && ((a.max_seqlen_q < 4096) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 64, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((a.max_seqlen_q < 4096) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 64, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 128 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 64 == 0) && (a.hdim_v % 64 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 128, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 128 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 128, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 128 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 64 == 0) && (a.hdim_v % 64 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 128, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 128 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 128, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + + } + + } + + } + else if(device_name.compare(0, 5, "gfx12") == 0) { + if(t.data_type.compare("fp16") == 0) { + if(t.hdim_q <= 32 && t.hdim_v <= 32) { + if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 32 == 0) && (a.hdim_v % 32 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 64, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 64, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 32 == 0) && (a.hdim_v % 32 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 64, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 64, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + + } + else if(t.hdim_q <= 64 && t.hdim_v <= 64) { + if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 64 == 0) && (a.hdim_v % 64 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 64, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 64, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 64 == 0) && (a.hdim_v % 64 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 64, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 64 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 64, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + + } + + } + + } + else if(device_name.compare(0, 4, "gfx9") == 0) { + if(t.data_type.compare("fp16") == 0) { + if(t.hdim_q <= 32 && t.hdim_v <= 32) { + if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true) && (true/*fall back to largest tile*/) && ((a.cu_seqlen_k_ptr == nullptr) && (a.seqlen_k != 0 && a.seqlen_k % 64 == 0)) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 128, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, false, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true) && (true/*fall back to largest tile*/) && ((a.cu_seqlen_k_ptr != nullptr) || (a.seqlen_k == 0 || a.seqlen_k % 64 != 0)) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 128, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (a.seqlen_q % 128 == 0) && (true/*fall back to largest tile*/) && ((a.cu_seqlen_k_ptr == nullptr) && (a.seqlen_k != 0 && a.seqlen_k % 64 == 0)) && (a.hdim_q % 32 == 0) && (a.hdim_v % 32 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 128, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, false, false, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 128 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (true /*a.hdim_q % 32 != 0*/) && (true /*a.hdim_v % 32 != 0*/) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<32, FmhaFwdFp16, false, 128, 64, 16, 32, 32, 32, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + + } + else if(t.hdim_q <= 64 && t.hdim_v <= 64) { + if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true) && (true/*fall back to largest tile*/) && ((a.cu_seqlen_k_ptr == nullptr) && (a.seqlen_k != 0 && a.seqlen_k % 64 == 0)) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 128, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, false, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::no_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true) && (true/*fall back to largest tile*/) && ((a.cu_seqlen_k_ptr != nullptr) || (a.seqlen_k == 0 || a.seqlen_k % 64 != 0)) && (a.hdim_q % 8 == 0) && (a.hdim_v % 8 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 128, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (a.seqlen_q % 128 == 0) && (true/*fall back to largest tile*/) && ((a.cu_seqlen_k_ptr == nullptr) && (a.seqlen_k != 0 && a.seqlen_k % 64 == 0)) && (a.hdim_q % 64 == 0) && (a.hdim_v % 64 == 0) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 128, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, false, false, false, false, false>; + return fmha_fwd_(s, a); + } + else if((t.is_group_mode == false) && (t.is_v_rowmajor == true) && (t.has_logits_soft_cap == false) && (t.mask_type == mask_enum::no_mask) && (t.bias_type == bias_enum::elementwise_bias) && (t.has_lse == false) && (t.has_dropout == false) && (t.qscale_type == quant_scale_enum::no_scale) && (t.skip_min_seqlen_q == false) &&(t.has_sink == false) && + (true /*a.seqlen_q % 128 != 0*/) && (true/*fall back to largest tile*/) && (true /*a.seqlen_k % 64 != 0*/) && (true /*a.hdim_q % 64 != 0*/) && (true /*a.hdim_v % 64 != 0*/) && ((true) && (true))) { + using trait_ = fmha_fwd_traits_<64, FmhaFwdFp16, false, 128, 64, 32, 64, 32, 64, true, ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, ck_tile::SimplifiedGenericAttentionMask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + return fmha_fwd_(s, a); + } + + } + + } + + } + + return r; +} +} // namespace + +namespace { +float fmha_fwd_v3([[maybe_unused]] fmha_fwd_traits t, [[maybe_unused]] fmha_fwd_args a, [[maybe_unused]] const ck_tile::stream_config& s) { + float r = -1; + + [[maybe_unused]] const float min_cu_util_rate = 0.8; // minimum CU utilization rate + + unsigned num_cus; + if(!get_num_cus(num_cus)) { + return r; + } + + [[maybe_unused]] auto get_num_blocks = [&](unsigned kM0) { + return get_num_thread_blocks(a.batch, a.nhead_q, a.max_seqlen_q, kM0); + }; + + [[maybe_unused]] const std::string device_name = ck_tile::get_device_name(); + + + return r; +} +} // namespace + +float fmha_fwd(fmha_fwd_traits traits, fmha_fwd_args args, const ck_tile::stream_config& config) { +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wunreachable-code" + if (false) { + float r = fmha_fwd_v3(traits, args, config); + if (r >= 0) return r; + } +#pragma clang diagnostic pop + return fmha_fwd_v2(traits, args, config); +} diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp new file mode 100644 index 0000000000..269e07e526 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx9__) && !defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSAsync< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,128, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, false, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx9__) && !defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp new file mode 100644 index 0000000000..e7620f684a --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSAsync< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,128, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, false, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp new file mode 100644 index 0000000000..188923a1e5 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx9__) && !defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSAsync< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,128, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx9__) && !defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp new file mode 100644 index 0000000000..4dd9e29e33 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSAsync< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,128, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_npad_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_npad_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp new file mode 100644 index 0000000000..024742cd4b --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_npad_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx9__) && !defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + false, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,128, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, false, false, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_npad_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx9__) && !defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_npad_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_npad_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp new file mode 100644 index 0000000000..502d6bd32b --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_npad_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + false, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,128, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, false, false, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_npad_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp new file mode 100644 index 0000000000..a9979075eb --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx9__) && !defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,128, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx9__) && !defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp new file mode 100644 index 0000000000..60eea2125f --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,128, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b128x64x16x32x32x32_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp new file mode 100644 index 0000000000..644fefa301 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#if defined(__HIP_DEVICE_COMPILE__) && (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__) || defined(__gfx11_generic__) || defined(__gfx1200__) || defined(__gfx1201__) || defined(__gfx12_generic__)) +#if !defined(CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK) +#define CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK 1 +#endif +#endif +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSHpad< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,64, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp new file mode 100644 index 0000000000..ba81fe1c7a --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#if defined(__HIP_DEVICE_COMPILE__) && (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__) || defined(__gfx11_generic__) || defined(__gfx1200__) || defined(__gfx1201__) || defined(__gfx12_generic__)) +#if !defined(CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK) +#define CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK 1 +#endif +#endif +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSHpad< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,64, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp new file mode 100644 index 0000000000..744e724cc4 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#if defined(__HIP_DEVICE_COMPILE__) && (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__) || defined(__gfx11_generic__) || defined(__gfx1200__) || defined(__gfx1201__) || defined(__gfx12_generic__)) +#if !defined(CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK) +#define CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK 1 +#endif +#endif +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx12__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSHpad< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,64, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx12__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp new file mode 100644 index 0000000000..dd5507699f --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#if defined(__HIP_DEVICE_COMPILE__) && (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__) || defined(__gfx11_generic__) || defined(__gfx1200__) || defined(__gfx1201__) || defined(__gfx12_generic__)) +#if !defined(CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK) +#define CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK 1 +#endif +#endif +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSHpad< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,64, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp new file mode 100644 index 0000000000..e657bd3940 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#if defined(__HIP_DEVICE_COMPILE__) && (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__) || defined(__gfx11_generic__) || defined(__gfx1200__) || defined(__gfx1201__) || defined(__gfx12_generic__)) +#if !defined(CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK) +#define CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK 1 +#endif +#endif +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSHpad< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,64, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp new file mode 100644 index 0000000000..55e45fd07b --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#if defined(__HIP_DEVICE_COMPILE__) && (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__) || defined(__gfx11_generic__) || defined(__gfx1200__) || defined(__gfx1201__) || defined(__gfx12_generic__)) +#if !defined(CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK) +#define CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK 1 +#endif +#endif +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx12__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSHpad< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,64, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx12__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp new file mode 100644 index 0000000000..c3e3b23905 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,64, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp new file mode 100644 index 0000000000..6dd03008aa --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,64, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp new file mode 100644 index 0000000000..e91d1c66d6 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx12__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,64, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx12__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp new file mode 100644 index 0000000000..ec2daba92c --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,64, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp new file mode 100644 index 0000000000..ef19ce1ccf --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,64, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp new file mode 100644 index 0000000000..fb96d2ae9c --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx12__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 16, 32, 32, 32>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<32, FmhaFwdFp16, false,64, 64, 16, 32, 32, 32, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d32_fp16_batch_b64x64x16x32x32x32_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx12__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_npad_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_npad_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp new file mode 100644 index 0000000000..75d5376d4b --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_npad_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + true, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSAsyncTrload< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + false, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,128, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, false, false, true, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_npad_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_pddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_pddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp new file mode 100644 index 0000000000..e1f69416f0 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_pddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + true, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSAsyncTrload< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + false, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,128, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, true, true, true, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_pddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp new file mode 100644 index 0000000000..6aaaf5f59e --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + true, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSAsyncTrload< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,128, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, true, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp new file mode 100644 index 0000000000..05d662a73e --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx9__) && !defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSAsync< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,128, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, false, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx9__) && !defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp new file mode 100644 index 0000000000..f2a603573d --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSAsync< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,128, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, false, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp new file mode 100644 index 0000000000..344abe80d0 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx9__) && !defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSAsync< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,128, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx9__) && !defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp new file mode 100644 index 0000000000..ad7c3fba0c --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSAsync< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,128, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_async_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_npad_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_npad_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp new file mode 100644 index 0000000000..11c03f42fe --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_npad_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx9__) && !defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + false, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,128, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, false, false, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_npad_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx9__) && !defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_npad_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_npad_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp new file mode 100644 index 0000000000..ac337f5e78 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_npad_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + false, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,128, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, false, false, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_npad_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp new file mode 100644 index 0000000000..b1327c1d74 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx9.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx9__) && !defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,128, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx9__) && !defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp new file mode 100644 index 0000000000..f028a8cc38 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx950.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,128, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r4x1x1_r4x1x1_w32x32x16_w32x32x16_qr_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r8x1x1_r8x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r8x1x1_r8x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp new file mode 100644 index 0000000000..7d4725a52d --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r8x1x1_r8x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#if defined(__HIP_DEVICE_COMPILE__) && (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__) || defined(__gfx11_generic__) || defined(__gfx1200__) || defined(__gfx1201__) || defined(__gfx12_generic__)) +#if !defined(CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK) +#define CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK 1 +#endif +#endif +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<8, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSHpad< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,128, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r8x1x1_r8x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r8x1x1_r8x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r8x1x1_r8x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp new file mode 100644 index 0000000000..5f932d4de0 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r8x1x1_r8x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#if defined(__HIP_DEVICE_COMPILE__) && (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__) || defined(__gfx11_generic__) || defined(__gfx1200__) || defined(__gfx1201__) || defined(__gfx12_generic__)) +#if !defined(CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK) +#define CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK 1 +#endif +#endif +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<8, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSHpad< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,128, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r8x1x1_r8x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r8x1x1_r8x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r8x1x1_r8x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp new file mode 100644 index 0000000000..1f6073a0c4 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r8x1x1_r8x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<8, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,128, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r8x1x1_r8x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r8x1x1_r8x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r8x1x1_r8x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp new file mode 100644 index 0000000000..bc7ea0bb3b --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r8x1x1_r8x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<128, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<8, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,128, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b128x64x32x64x32x64_r8x1x1_r8x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b16x32x64x64x32x64_r1x1x1_r1x1x1_w16x16x32_w16x16x32_qr_async_trload_vr_npad_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b16x32x64x64x32x64_r1x1x1_r1x1x1_w16x16x32_w16x16x32_qr_async_trload_vr_npad_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp new file mode 100644 index 0000000000..7e6954dd21 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b16x32x64x64x32x64_r1x1x1_r1x1x1_w16x16x32_w16x16x32_qr_async_trload_vr_npad_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<16, 32, 64, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 32>, + ck_tile::sequence<1, 1, 1>, + ck_tile::sequence<16, 16, 32>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + true, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSAsyncTrload< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + false, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,16, 32, 64, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, false, false, true, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b16x32x64x64x32x64_r1x1x1_r1x1x1_w16x16x32_w16x16x32_qr_async_trload_vr_npad_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b16x32x64x64x32x64_r1x1x1_r1x1x1_w16x16x32_w16x16x32_qr_async_trload_vr_pddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b16x32x64x64x32x64_r1x1x1_r1x1x1_w16x16x32_w16x16x32_qr_async_trload_vr_pddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp new file mode 100644 index 0000000000..a36118b62d --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b16x32x64x64x32x64_r1x1x1_r1x1x1_w16x16x32_w16x16x32_qr_async_trload_vr_pddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<16, 32, 64, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 32>, + ck_tile::sequence<1, 1, 1>, + ck_tile::sequence<16, 16, 32>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + true, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSAsyncTrload< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + false, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,16, 32, 64, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, true, true, true, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b16x32x64x64x32x64_r1x1x1_r1x1x1_w16x16x32_w16x16x32_qr_async_trload_vr_pddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b16x32x64x64x32x64_r1x1x1_r1x1x1_w16x16x32_w16x16x32_qr_async_trload_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b16x32x64x64x32x64_r1x1x1_r1x1x1_w16x16x32_w16x16x32_qr_async_trload_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp new file mode 100644 index 0000000000..953b19e1b1 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b16x32x64x64x32x64_r1x1x1_r1x1x1_w16x16x32_w16x16x32_qr_async_trload_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<16, 32, 64, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 32>, + ck_tile::sequence<1, 1, 1>, + ck_tile::sequence<16, 16, 32>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + true, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSAsyncTrload< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,16, 32, 64, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, true, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b16x32x64x64x32x64_r1x1x1_r1x1x1_w16x16x32_w16x16x32_qr_async_trload_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b32x32x64x64x32x64_r1x1x1_r1x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_npad_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b32x32x64x64x32x64_r1x1x1_r1x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_npad_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp new file mode 100644 index 0000000000..2da6e8bcd6 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b32x32x64x64x32x64_r1x1x1_r1x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_npad_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<32, 32, 64, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<1, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + true, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSAsyncTrload< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + false, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,32, 32, 64, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, false, false, true, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b32x32x64x64x32x64_r1x1x1_r1x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_npad_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b32x32x64x64x32x64_r1x1x1_r1x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_pddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b32x32x64x64x32x64_r1x1x1_r1x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_pddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp new file mode 100644 index 0000000000..648f8c4388 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b32x32x64x64x32x64_r1x1x1_r1x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_pddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<32, 32, 64, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<1, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + true, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSAsyncTrload< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + false, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,32, 32, 64, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, false, false, true, true, true, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b32x32x64x64x32x64_r1x1x1_r1x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_pddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b32x32x64x64x32x64_r1x1x1_r1x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b32x32x64x64x32x64_r1x1x1_r1x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp new file mode 100644 index 0000000000..a975c1619c --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b32x32x64x64x32x64_r1x1x1_r1x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink_gfx950.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<32, 32, 64, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<32, 32, 16>, + ck_tile::sequence<1, 1, 1>, + ck_tile::sequence<32, 32, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + true, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSAsyncTrload< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,32, 32, 64, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC_TRLOAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, true, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b32x32x64x64x32x64_r1x1x1_r1x1x1_w32x32x16_w32x32x16_qr_async_trload_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_trload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx950__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp new file mode 100644 index 0000000000..62f809583c --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#if defined(__HIP_DEVICE_COMPILE__) && (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__) || defined(__gfx11_generic__) || defined(__gfx1200__) || defined(__gfx1201__) || defined(__gfx12_generic__)) +#if !defined(CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK) +#define CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK 1 +#endif +#endif +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSHpad< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,64, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp new file mode 100644 index 0000000000..d2a009bb92 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#if defined(__HIP_DEVICE_COMPILE__) && (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__) || defined(__gfx11_generic__) || defined(__gfx1200__) || defined(__gfx1201__) || defined(__gfx12_generic__)) +#if !defined(CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK) +#define CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK 1 +#endif +#endif +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSHpad< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,64, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp new file mode 100644 index 0000000000..337a75ad90 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#if defined(__HIP_DEVICE_COMPILE__) && (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__) || defined(__gfx11_generic__) || defined(__gfx1200__) || defined(__gfx1201__) || defined(__gfx12_generic__)) +#if !defined(CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK) +#define CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK 1 +#endif +#endif +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx12__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSHpad< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,64, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx12__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp new file mode 100644 index 0000000000..10d6844a3c --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#if defined(__HIP_DEVICE_COMPILE__) && (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__) || defined(__gfx11_generic__) || defined(__gfx1200__) || defined(__gfx1201__) || defined(__gfx12_generic__)) +#if !defined(CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK) +#define CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK 1 +#endif +#endif +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSHpad< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,64, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp new file mode 100644 index 0000000000..94b45a9586 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#if defined(__HIP_DEVICE_COMPILE__) && (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__) || defined(__gfx11_generic__) || defined(__gfx1200__) || defined(__gfx1201__) || defined(__gfx12_generic__)) +#if !defined(CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK) +#define CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK 1 +#endif +#endif +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSHpad< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,64, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp new file mode 100644 index 0000000000..f4826603a7 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#if defined(__HIP_DEVICE_COMPILE__) && (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__) || defined(__gfx11_generic__) || defined(__gfx1200__) || defined(__gfx1201__) || defined(__gfx12_generic__)) +#if !defined(CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK) +#define CK_TILE_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK 1 +#endif +#endif +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx12__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVSHpad< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, true>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,64, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS_HPAD, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, true, true, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_hpad_vr_psskddv_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx12__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp new file mode 100644 index 0000000000..422de5ae49 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,64, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp new file mode 100644 index 0000000000..e5688428ec --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,64, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp new file mode 100644 index 0000000000..026b1e2881 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx12__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,64, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_bias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx12__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp new file mode 100644 index 0000000000..ec28c1943f --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx11.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,64, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx11__) && !defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp new file mode 100644 index 0000000000..7c25621fe3 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx115.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx115__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,64, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx115__)) diff --git a/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp new file mode 100644 index 0000000000..49eceb7408 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/generated/fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink_gfx12.cpp @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +// auto generated by generate.py +#include "ck_tile/ops/fmha/block/variants.hpp" +#include "fmha_fwd.hpp" + +#include + +#if !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx12__)) + +using fmha_dtype = FmhaFwdFp16; + +using fmha_block_tile = ck_tile::sequence<64, 64, 32, 64, 32, 64>; + +using fmha_shape = ck_tile::TileFmhaShape, + ck_tile::sequence<16, 16, 16>, + ck_tile::sequence<4, 1, 1>, + ck_tile::sequence<16, 16, 16>, + true>; + +using fmha_traits = ck_tile::TileFmhaTraits; + +using fmha_variant = ck_tile::ComposedAttention; + +using fmha_mask = ck_tile::SimplifiedGenericAttentionMask; + +using fmha_pipeline_problem = ck_tile::BlockFmhaPipelineProblem< + typename FmhaFwdTypeConfig::QDataType, + typename FmhaFwdTypeConfig::KDataType, + typename FmhaFwdTypeConfig::VDataType, + typename FmhaFwdTypeConfig::SaccDataType, + typename FmhaFwdTypeConfig::SMPLComputeDataType, + typename FmhaFwdTypeConfig::BiasDataType, + typename FmhaFwdTypeConfig::RandValOutputDataType, + typename FmhaFwdTypeConfig::LSEDataType, + typename FmhaFwdTypeConfig::PDataType, + typename FmhaFwdTypeConfig::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + fmha_shape, + false, + fmha_variant, + fmha_mask, + false, + fmha_traits>; + +using fmha_pipeline = ck_tile::BlockFmhaPipelineQRKSVS< + fmha_pipeline_problem>; + +using fmha_epilogue = + ck_tile::Default2DEpilogue::OaccDataType, + typename FmhaFwdTypeConfig::ODataType, + true, false>>; + +using fmha_kernel = ck_tile::FmhaFwdKernel; + + +using trait = fmha_fwd_traits_<64, FmhaFwdFp16, false,64, 64, 32, 64, 32, 64, true, + ck_tile::BlockFmhaPipelineEnum::QRKSVS, false, fmha_mask, ck_tile::BlockAttentionBiasEnum::NO_BIAS, false, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE, true, true, false, false, false, false, false>; + +template<> +float fmha_fwd_(const ck_tile::stream_config& s, fmha_fwd_args a) +{ + using k_ = fmha_kernel; + if(s.log_level_ > 0) + std::cout << ", fmha_fwd_d64_fp16_batch_b64x64x32x64x32x64_r4x1x1_r4x1x1_w16x16x16_w16x16x16_qr_vr_pssk_nlogits_nbias_nmask_nlse_ndropout_nskip_nqscale_ntrload_nsink" << std::flush; + auto [kargs, grids] = fmha_fwd_create_kargs_and_grids(a); + const dim3 blocks = k_::BlockSize(); + constexpr ck_tile::index_t kBlockPerCu = k_::kBlockPerCu; + return ck_tile::launch_kernel(s, ck_tile::make_kernel(k_{}, grids, blocks, 0, kargs)); +} + +#endif // !defined(__HIP_DEVICE_COMPILE__) || (defined(__gfx12__)) diff --git a/cpp/external/composable_kernel_fmha/mask.hpp b/cpp/external/composable_kernel_fmha/mask.hpp new file mode 100644 index 0000000000..03e1537c5d --- /dev/null +++ b/cpp/external/composable_kernel_fmha/mask.hpp @@ -0,0 +1,203 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include +#include + +#include "ck_tile/core.hpp" +#include "ck_tile/ops/fmha.hpp" + +// keep this in sync with ck_tile::GenericAttentionMaskEnum +enum class mask_enum +{ + no_mask = 0, + mask_top_left, + mask_bottom_right, + window_generic, +}; + +struct mask_info +{ + mask_enum type; + ck_tile::index_t seqlen_q; + ck_tile::index_t seqlen_k; + ck_tile::index_t y, x; + ck_tile::index_t left, right; // FA style SWA left/right + ck_tile::index_t sink; + + void serialize(std::ostream& os) const + { + if(type == mask_enum::no_mask) + os << "n"; + else if(type == mask_enum::mask_top_left) + os << "t(" << left << ":" << right << ")"; + else if(type == mask_enum::mask_bottom_right) + os << "b(" << left << ":" << right << ")"; + else + { + os << "g(" << y << ":" << x << ")"; + } + } + + static mask_info decode(std::string str, ck_tile::index_t seqlen_q, ck_tile::index_t seqlen_k) + { + ck_tile::index_t x_total = seqlen_k; + ck_tile::index_t y_total = seqlen_q; + mask_info tmp; + tmp.seqlen_q = seqlen_q; + tmp.seqlen_k = seqlen_k; + auto found_0 = str.find(':'); + if(found_0 != std::string::npos) + { + std::string t = str.substr(0, found_0); + std::string v = str.substr(found_0 + 1); + if(t == "xt" || t == "xb") + { + // xformer style sliding window attn from top-left + ck_tile::index_t window_size = std::stoi(v); + ck_tile::index_t left_size = -1; + ck_tile::index_t right_size = 0; + ck_tile::index_t sink_size = 0; + if(window_size > 0) + { + left_size = window_size / 2; + right_size = window_size - 1 - left_size; + } + auto r = ck_tile::make_generic_attention_mask_coordinates_from_lr_window( + left_size, right_size, sink_size, y_total, x_total, t == "xt"); + + tmp.type = t == "xt" ? mask_enum::mask_top_left : mask_enum::mask_bottom_right; + tmp.y = r.at(ck_tile::number<0>{}); + tmp.x = r.at(ck_tile::number<1>{}); + tmp.left = left_size; + tmp.right = right_size; + tmp.sink = 0; + } + else if(t == "t" || t == "b" || t == "g") + { + auto found_1 = v.find(","); + if(found_1 == std::string::npos) + { + throw std::invalid_argument("invalid mask value: " + str); + } + tmp.type = mask_enum::window_generic; + ck_tile::index_t v0 = atoi(v.substr(0, found_1).c_str()); + auto found_2 = v.find(',', found_1 + 1); + ck_tile::index_t v1 = 0; + ck_tile::index_t sink = 0; + // ck_tile::index_t v1 = atoi(v.substr(found_1 + 1).c_str()); + // TODO: some validation + if(t == "t") + { + if(found_2 != std::string::npos) + { + v1 = atoi(v.substr(found_1 + 1, found_2 - found_1 - 1).c_str()); + sink = atoi(v.substr(found_2 + 1).c_str()); + } + else + { + v1 = atoi(v.substr(found_1 + 1).c_str()); + sink = 0; + } + tmp.type = mask_enum::mask_top_left; + auto r = ck_tile::make_generic_attention_mask_coordinates_from_lr_window( + v0, v1, sink, y_total, x_total, true); + tmp.y = r.at(ck_tile::number<0>{}); + tmp.x = r.at(ck_tile::number<1>{}); + tmp.left = v0; + tmp.right = v1; + tmp.sink = sink; + } + else if(t == "b") + { + if(found_2 != std::string::npos) + { + v1 = atoi(v.substr(found_1 + 1, found_2 - found_1 - 1).c_str()); + sink = atoi(v.substr(found_2 + 1).c_str()); + } + else + { + v1 = atoi(v.substr(found_1 + 1).c_str()); + sink = 0; + } + tmp.type = mask_enum::mask_bottom_right; + auto r = ck_tile::make_generic_attention_mask_coordinates_from_lr_window( + v0, v1, sink, y_total, x_total, false); + tmp.y = r.at(ck_tile::number<0>{}); + tmp.x = r.at(ck_tile::number<1>{}); + tmp.left = v0; + tmp.right = v1; + tmp.sink = sink; + } + else if(t == "g") + { + tmp.type = mask_enum::window_generic; + tmp.y = v0; + tmp.x = v1; + tmp.left = v0; // TODO: don't use this? + tmp.right = v1; + tmp.sink = 0; + } + } + else + { + throw std::invalid_argument("invalid mask value: " + str); + } + } + else if(str == "0") + { + tmp.type = mask_enum::no_mask; + tmp.left = -1; + tmp.right = -1; + tmp.sink = 0; + } + else if(str == "1" || str == "t") + { + tmp.type = mask_enum::mask_top_left; + tmp.y = seqlen_q; + tmp.x = 1; + tmp.left = -1; + tmp.right = 0; + tmp.sink = 0; + } + else if(str == "2" || str == "b") + { + tmp.type = mask_enum::mask_bottom_right; + tmp.y = seqlen_q; + tmp.x = seqlen_k - seqlen_q + 1; + tmp.left = -1; + tmp.right = 0; + tmp.sink = 0; + } + else + { + throw std::invalid_argument("invalid mask value: " + str); + } + return tmp; + } + + std::size_t get_unmaskarea() const + { + if(type == mask_enum::no_mask) + return static_cast(seqlen_q) * seqlen_k; + std::size_t area = 0; + for(ck_tile::index_t i_y = 0; i_y < seqlen_q; ++i_y) + { + ck_tile::index_t x_start = std::max(-y + i_y + 1, static_cast(0)); + ck_tile::index_t x_end = std::min(i_y + x, seqlen_k); + if(x_end > x_start) + { + area += (x_end - x_start); + } + } + return area; + } + + friend std::ostream& operator<<([[clang::lifetimebound]] std::ostream& os, const mask_info& mi) + { + mi.serialize(os); + return os; + } +}; diff --git a/cpp/external/composable_kernel_fmha/quant.hpp b/cpp/external/composable_kernel_fmha/quant.hpp new file mode 100644 index 0000000000..4b8cd2e9a4 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/quant.hpp @@ -0,0 +1,78 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include +#include +#include "ck_tile/core.hpp" +#include "ck_tile/ops/fmha.hpp" + +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wlifetime-safety-intra-tu-suggestions" + +// keep sync with BlockAttentionQuantScaleEnum +enum class quant_scale_enum +{ + no_scale = 0, + pertensor = 1, + blockscale = 2, + kv_blockscale = 3, // Q per-tensor, K/V per-page block scale + mx = 4, // Microscaling (MX) +}; + +struct quant_scale_info +{ + quant_scale_enum type; + + void serialize(std::ostream& os) const + { + if(type == quant_scale_enum::no_scale) + os << "n"; + else if(type == quant_scale_enum::pertensor) + os << "pt"; + else if(type == quant_scale_enum::blockscale) + os << "bs"; + else if(type == quant_scale_enum::kv_blockscale) + os << "kvbs"; + else if(type == quant_scale_enum::mx) + os << "mx"; + } + + static quant_scale_info decode(std::string str) + { + quant_scale_info info{quant_scale_enum::no_scale}; + if(str == "n" || str == "0") + { + info.type = quant_scale_enum::no_scale; + } + else if(str == "pt" || str == "1") + { + info.type = quant_scale_enum::pertensor; + } + else if(str == "bs" || str == "2") + { + info.type = quant_scale_enum::blockscale; + } + else if(str == "kvbs" || str == "3") + { + info.type = quant_scale_enum::kv_blockscale; + } + else if(str == "mx" || str == "4") + { + info.type = quant_scale_enum::mx; + } + else + { + throw std::invalid_argument("invalid quant scale value: " + str); + } + return info; + } + + friend std::ostream& operator<<(std::ostream& os, const quant_scale_info& qsi) + { + qsi.serialize(os); + return os; + } +}; +#pragma clang diagnostic pop diff --git a/cpp/external/composable_kernel_fmha/rotary.hpp b/cpp/external/composable_kernel_fmha/rotary.hpp new file mode 100644 index 0000000000..a6458c2173 --- /dev/null +++ b/cpp/external/composable_kernel_fmha/rotary.hpp @@ -0,0 +1,89 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include "ck_tile/core.hpp" +#include "ck_tile/host/host_tensor.hpp" + +#include +#include + +#ifndef M_PI // Not there on windows... +#define M_PI 3.141592653589793238462643383279502884 +#endif + +#include +#include +#include +#include +#include + +// keep sync with RotaryEmbeddingEnum +enum class rope_enum +{ + none = 0, + interleaved = 1, + half_rotated = 2, +}; + +template +std::tuple, ck_tile::HostTensor> +generate_rotary_cos_sin(ck_tile::index_t seqlen, + ck_tile::index_t rotary_dim, + std::optional seed = std::nullopt) +{ + // return dummy tensors if we won't apply RoPE at all + if(rotary_dim <= 0) + { + ck_tile::HostTensor dummy({1, 1}); + return std::make_tuple(dummy, dummy); + } + + std::mt19937 random_engine(seed.has_value() ? *seed : std::random_device{}()); + std::uniform_real_distribution generator(0.0f, 1.0f); + + const ck_tile::index_t num_rows = seqlen * 2; + const ck_tile::index_t num_cols = rotary_dim / 2; + + using std::begin, std::end; + + ck_tile::HostTensor angle({num_rows, num_cols}); + std::generate(begin(angle), end(angle), [&] { return generator(random_engine) * 2 * M_PI; }); + + ck_tile::HostTensor cos({num_rows, num_cols}); + std::transform(begin(angle), end(angle), begin(cos), [](float origin_value) { + return ck_tile::type_convert(std::cos(origin_value)); + }); + + ck_tile::HostTensor sin({num_rows, num_cols}); + std::transform(begin(angle), end(angle), begin(sin), [](float origin_value) { + return ck_tile::type_convert(std::sin(origin_value)); + }); + + return std::make_tuple(cos, sin); +} + +template +std::tuple, ck_tile::HostTensor> +slice_rotary_cos_sin(const ck_tile::HostTensor& cos, + const ck_tile::HostTensor& sin, + ck_tile::index_t seqlen_offset, + ck_tile::index_t seqlen) +{ + assert(cos.get_num_of_dimension() == 2 && sin.get_num_of_dimension() == 2); + assert(cos.get_length(0) == sin.get_length(0) && cos.get_length(1) == sin.get_length(1)); + + assert(static_cast(seqlen_offset + seqlen) <= cos.get_length(0)); + + const ck_tile::index_t num_rows = seqlen; + const ck_tile::index_t num_cols = cos.get_length(1); + + ck_tile::HostTensor cos_pt({num_rows, num_cols}); + cos_pt.ForEach([&](auto& self, auto i) { self(i) = cos(i[0] + seqlen_offset, i[1]); }); + + ck_tile::HostTensor sin_pt({num_rows, num_cols}); + sin_pt.ForEach([&](auto& self, auto i) { self(i) = sin(i[0] + seqlen_offset, i[1]); }); + + return std::make_tuple(cos_pt, sin_pt); +} diff --git a/cpp/main.cpp b/cpp/main.cpp index f9c95e09a6..3c2ca3132f 100644 --- a/cpp/main.cpp +++ b/cpp/main.cpp @@ -251,6 +251,13 @@ string Version::getKataGoVersionFullInfo() { out << "Using Metal backend" << endl; #elif defined(USE_OPENCL_BACKEND) out << "Using OpenCL backend" << endl; +#elif defined(USE_ROCM_BACKEND) + out << "Using ROCm backend" << endl; +#if defined(HIP_TARGET_VERSION) +#define STRINGIFY(x) #x +#define STRINGIFY2(x) STRINGIFY(x) + out << "Compiled with HIP runtime version " << STRINGIFY2(HIP_TARGET_VERSION) << endl; +#endif #elif defined(USE_EIGEN_BACKEND) out << "Using Eigen(CPU) backend" << endl; #else @@ -283,6 +290,8 @@ string Version::getGitRevisionWithBackend() { s += "-cuda"; #elif defined(USE_TENSORRT_BACKEND) s += "-trt"; +#elif defined(USE_ROCM_BACKEND) + s += "-rocm"; #elif defined(USE_METAL_BACKEND) s += "-metal"; #elif defined(USE_OPENCL_BACKEND) diff --git a/cpp/neuralnet/rocmbackend.cpp b/cpp/neuralnet/rocmbackend.cpp new file mode 100644 index 0000000000..fc65738040 --- /dev/null +++ b/cpp/neuralnet/rocmbackend.cpp @@ -0,0 +1,3765 @@ +#ifdef USE_ROCM_BACKEND + +#include "../neuralnet/rocmerrorcheck.h" +#include "../neuralnet/rocmincludes.h" + +// Composable Kernel FMHA fused attention support (optional; see cpp/external/composable_kernel_fmha +// and KATAGO_ROCM_HAS_CK_FMHA in CMakeLists.txt). Mirrors the CUDA backend's optional cudnn-frontend +// SDPA path, but CK's fmha_fwd() has no expensive one-time "build plan" step to cache - each call +// directly checks traits/shape compatibility and either executes or returns a negative "unsupported" +// sentinel, so there is no warmup-only tolerance needed the way cudnn_frontend graph building has. +// ck_tile's own arch.hpp has no get_compiler_target() branch for gfx906 or RDNA1 (gfx1010/1011/ +// 1012) - its FMHA kernels were never ported to these architectures (no MFMA/WMMA). Merely +// including its headers while compiling a device pass for one of these archs (as happens in a +// multi-arch fat binary that targets them) hard-fails with "member reference base type 'void' is +// not a structure or union", since get_compiler_target() falls through without a return. Skip CK +// entirely for just these archs' device-compile passes; the plain (non-fused) attention kernel +// still covers them. +#if defined(__gfx906__) || defined(__gfx1010__) || defined(__gfx1011__) || defined(__gfx1012__) + #define KATAGO_ROCM_CK_FMHA_ARCH_OK 0 +#else + #define KATAGO_ROCM_CK_FMHA_ARCH_OK 1 +#endif + +#if KATAGO_ROCM_HAS_CK_FMHA && KATAGO_ROCM_CK_FMHA_ARCH_OK + #include + #include + #include "fmha_fwd.hpp" +#endif + +#include "../neuralnet/rocmhelpers.h" +#include "../neuralnet/rocmutils.h" +#include "../neuralnet/modelversion.h" +#include "../neuralnet/nninterface.h" +#include "../neuralnet/nninputs.h" +#include "../neuralnet/sgfmetadata.h" +#include "../neuralnet/nneval.h" +#include "../neuralnet/desc.h" + +#include "../core/simpleallocator.h" +#include "../core/test.h" + +#include "../external/half-2.2.0/include/half.hpp" + +//------------------------ +#include "../core/using.h" +//------------------------ + +using half_t = half_float::half; + +//Define this to print out some of the intermediate values of the neural net +//#define DEBUG_INTERMEDIATE_VALUES + +void NeuralNet::globalInitialize() { + //Empty for cudnn backend +} + +void NeuralNet::globalCleanup() { + hipDeviceReset(); +} + +struct CudaHandles { + hipblasHandle_t cublas; + miopenHandle_t cudnn; + const int majorComputeCapability; + const int minorComputeCapability; + // Disables the optional CK FMHA fused attention path (see KATAGO_ROCM_HAS_CK_FMHA), falling back + // to the plain attention kernel unconditionally. Off by default. + const bool disableFusedAttention; + + CudaHandles(int major, int minor, bool disableFusedAttention_ = false) + : majorComputeCapability(major), + minorComputeCapability(minor), + disableFusedAttention(disableFusedAttention_) + { + CUBLAS_ERR("CudaHandles",hipblasCreate(&cublas)); + CUDNN_ERR("CudaHandles",miopenCreate(&cudnn)); + } + + ~CudaHandles() { + hipblasDestroy(cublas); + miopenDestroy(cudnn); + } + + static CudaHandles* cudaHandlesTesting() { + const int gpuIdxForThisThread = 0; + hipDeviceProp_t prop; + hipGetDeviceProperties(&prop,gpuIdxForThisThread); + return new CudaHandles(prop.major, prop.minor); + } + + CudaHandles(const CudaHandles&) = delete; + CudaHandles& operator=(const CudaHandles&) = delete; +}; + +//--------------------------------------------------------------------------------- + +template +struct ByBatchSize { + const int maxBatchSize; + T* data; + miopenStatus_t (*destroyFunc)(T); + + ByBatchSize() + : maxBatchSize(0), data(nullptr), destroyFunc(nullptr) + {} + + ByBatchSize( + int maxBatchSize_ + ) : maxBatchSize(maxBatchSize_), data(nullptr), destroyFunc(nullptr) { + data = new T[maxBatchSize]; + } + + ByBatchSize(const ByBatchSize&) = delete; + ByBatchSize& operator=(const ByBatchSize&) = delete; + + ~ByBatchSize() { + if(destroyFunc != nullptr && data != nullptr) { + for(int batchSize = 1; batchSize <= maxBatchSize; batchSize++) { + (*destroyFunc)(data[batchSize-1]); + } + } + if(data != nullptr) { + delete[] data; + data = nullptr; + } + } + T& operator[](int batchSize) { + return data[batchSize-1]; + } + const T& operator[](int batchSize) const { + return data[batchSize-1]; + } +}; + +template +struct ByBatchSizeView { + int maxBatchSize; + T* data; + + ByBatchSizeView() + : maxBatchSize(0), data(nullptr) + {} + + ByBatchSizeView(const ByBatchSize& toView) + : maxBatchSize(toView.maxBatchSize), data(toView.data) + {} + ByBatchSizeView& operator=(const ByBatchSize& toView) { + maxBatchSize = toView.maxBatchSize; + data = toView.data; + } + + ~ByBatchSizeView() { + } + T& operator[](int batchSize) { + return data[batchSize-1]; + } + const T& operator[](int batchSize) const { + return data[batchSize-1]; + } +}; + +//--------------------------------------------------------------------------------- + + +//channels, useFP16, useNHWC +typedef std::tuple CudnnTensorDesc4DKey; + +struct CudnnManager { + const string name; + const int maxBatchSize; + const int nnXLen; + const int nnYLen; + std::map*> tensorDesc4DByBatchSizeByKey; + + CudnnManager(string name_, int maxBatchSize_, int nnXLen_, int nnYLen_) + :name(name_), + maxBatchSize(maxBatchSize_), + nnXLen(nnXLen_), + nnYLen(nnYLen_), + tensorDesc4DByBatchSizeByKey() + { + } + + ~CudnnManager() { + for(auto& iter: tensorDesc4DByBatchSizeByKey) { + delete iter.second; + } + } + + ByBatchSizeView getTensorDesc4DByBatchSize( + int channels, bool useFP16, bool useNHWC + ) { + auto iter = tensorDesc4DByBatchSizeByKey.find({channels, useFP16, useNHWC}); + if(iter != tensorDesc4DByBatchSizeByKey.end()) { + return ByBatchSizeView(*(iter->second)); + } + ByBatchSize* descs = new ByBatchSize(maxBatchSize); + for(int batchSize = 1; batchSize <= maxBatchSize; batchSize++) { + miopenTensorDescriptor_t& desc = (*descs)[batchSize]; + CUDNN_ERR(name.c_str(),miopenCreateTensorDescriptor(&desc)); + CUDNN_ERR(name.c_str(),miopenSet4dTensorDescriptor( + desc, + (useFP16 ? miopenHalf : miopenFloat), + batchSize, + channels, + nnYLen, + nnXLen + )); + } + descs->destroyFunc = miopenDestroyTensorDescriptor; + tensorDesc4DByBatchSizeByKey[{channels, useFP16, useNHWC}] = descs; + return ByBatchSizeView(*descs); + } +}; + +//--------------------------------------------------------------------------------- + +struct ScratchBuffers { + + const size_t batchXYFloatBytes; + const size_t batchFloatBytes; + const size_t batchXYBytes; + const size_t batchBytes; + + SimpleAllocator* allocator; + + // Not scratch, but convenient to have here + void* zeroBuf; + void* oneBuf; + + ScratchBuffers() = delete; + ScratchBuffers(const ScratchBuffers&) = delete; + ScratchBuffers& operator=(const ScratchBuffers&) = delete; + + ScratchBuffers(int maxBatchSize, int nnXLen, int nnYLen, bool useFP16) + : batchXYFloatBytes((size_t)maxBatchSize * nnXLen * nnYLen * sizeof(float)), + batchFloatBytes((size_t)maxBatchSize * sizeof(float)), + batchXYBytes((size_t)maxBatchSize * nnXLen * nnYLen * (useFP16 ? sizeof(half_t) : sizeof(float))), + batchBytes((size_t)maxBatchSize * (useFP16 ? sizeof(half_t) : sizeof(float))) + { + std::function allocateFunc = [](size_t size) { + void* buf; + CUDA_ERR("ScratchBuffers",hipMalloc(&buf, size)); + return buf; + }; + std::function releaseFunc = [](void* buf) { + hipFree(buf); + }; + + allocator = new SimpleAllocator(allocateFunc, releaseFunc); + + CudaUtils::hostMallocZeroOneBufs(zeroBuf, oneBuf, useFP16); + } + ~ScratchBuffers() { + delete allocator; + free(zeroBuf); + free(oneBuf); + } + + size_t getBufSizeXY(int channels) const { + return channels * batchXYBytes; + } + size_t getBufSizeXYFloat(int channels) const { + return channels * batchXYFloatBytes; + } + size_t getBufSizeFloat(int channels) const { + return channels * batchFloatBytes; + } + size_t getBufSize(int channels) const { + return channels * batchBytes; + } + +}; + + +//--------------------------------------------------------------------------------- + +struct ConvLayer { + const string name; + const int inChannels; + const int outChannels; + const int nnXLen; + const int nnYLen; + const int maxBatchSize; + const bool usingFP16; + ByBatchSizeView inputDescriptors; + ByBatchSizeView outputDescriptors; + miopenTensorDescriptor_t filterDescriptor; + miopenConvolutionDescriptor_t convolutionDescriptor; + ByBatchSize* convolutionAlgorithms; //array of one for each batch size + void* filterBuf; + void* accumBuf; // Pre-allocated buffer for residual accumulation (miopenConvolutionForwardImmediate has no beta) + + ConvLayer() = delete; + ConvLayer(const ConvLayer&) = delete; + ConvLayer& operator=(const ConvLayer&) = delete; + + ConvLayer( + CudaHandles* cudaHandles, + CudnnManager* manager, + const ConvLayerDesc* desc, + bool useFP16, + bool useNHWC + ) : ConvLayer(cudaHandles, manager, desc, useFP16, useNHWC, useNHWC) + {} + + ConvLayer( + CudaHandles* cudaHandles, + CudnnManager* manager, + const ConvLayerDesc* desc, + bool useFP16, + bool useNHWCIn, + bool useNHWCOut + ) : + name(desc->name), + inChannels(desc->inChannels), + outChannels(desc->outChannels), + nnXLen(manager->nnXLen), + nnYLen(manager->nnYLen), + maxBatchSize(manager->maxBatchSize), + usingFP16(useFP16) + { + int convYSize = desc->convYSize; + int convXSize = desc->convXSize; + int dilationY = desc->dilationY; + int dilationX = desc->dilationX; + int paddingX = (convXSize / 2) * dilationX; + int paddingY = (convYSize / 2) * dilationY; + + assert(convXSize % 2 == 1); + assert(convYSize % 2 == 1); + + inputDescriptors = manager->getTensorDesc4DByBatchSize(inChannels,useFP16,useNHWCIn); + outputDescriptors = manager->getTensorDesc4DByBatchSize(outChannels,useFP16,useNHWCOut); + int maxBatchSize = manager->maxBatchSize; + + bool filterNHWC = useNHWCOut && dilationY == 1 && dilationX == 1; + + CUDNN_ERR(name.c_str(),miopenCreateTensorDescriptor(&filterDescriptor)); + CUDNN_ERR(name.c_str(),miopenSet4dTensorDescriptor( + filterDescriptor, + (useFP16 ? miopenHalf : miopenFloat), + outChannels, + inChannels, + convYSize, + convXSize + )); + + int yStride = 1; + int xStride = 1; + + + CUDNN_ERR(name.c_str(),miopenCreateConvolutionDescriptor(&convolutionDescriptor)); + CUDNN_ERR(name.c_str(),miopenInitConvolutionDescriptor( + convolutionDescriptor, + miopenConvolution, + paddingY, + paddingX, + yStride, + xStride, + dilationY, + dilationX + )); + if(useFP16) { + int alt = 1; // non‑zero enables alt‑impl on MI2xx+ GPUs + CUDNN_ERR(name.c_str(),miopenSetConvolutionAttribute(convolutionDescriptor,MIOPEN_CONVOLUTION_ATTRIB_FP16_ALT_IMPL,alt)); + } + + convolutionAlgorithms = new ByBatchSize(maxBatchSize); + + for(int batchSize = 1; batchSize <= maxBatchSize; batchSize++) { + const miopenTensorDescriptor_t& inputDescriptor = inputDescriptors[batchSize]; + const miopenTensorDescriptor_t& outputDescriptor = outputDescriptors[batchSize]; + size_t availableAlgoCount = 0; + CUDNN_ERR(name.c_str(),miopenConvolutionForwardGetSolutionCount( + cudaHandles->cudnn, + filterDescriptor, + inputDescriptor, + convolutionDescriptor, + outputDescriptor, + &availableAlgoCount + )); + if(availableAlgoCount <= 0) + throw StringError("miopenConvolutionForwardGetSolutionCount returned 0 algorithms?"); + std::vector solutions(availableAlgoCount); + size_t returnedAlgoCount = 0; + CUDNN_ERR(name.c_str(),miopenConvolutionForwardGetSolution( + cudaHandles->cudnn, + filterDescriptor, + inputDescriptor, + convolutionDescriptor, + outputDescriptor, + availableAlgoCount, + &returnedAlgoCount, + solutions.data() + )); + if(returnedAlgoCount <= 0) + throw StringError("miopenConvolutionForwardGetSolution returned no algorithms?"); + (*convolutionAlgorithms)[batchSize] = solutions[0]; + CUDNN_ERR(name.c_str(),miopenConvolutionForwardCompileSolution( + cudaHandles->cudnn, + filterDescriptor, + inputDescriptor, + convolutionDescriptor, + outputDescriptor, + (*convolutionAlgorithms)[batchSize].solution_id + )); + } + + assert(desc->weights.size() == convYSize * convXSize * inChannels * outChannels); + + if(filterNHWC) { + vector weightsTransposed(desc->weights.size()); + for(int y = 0; y < convYSize; y++) { + for(int x = 0; x < convXSize; x++) { + for(int ic = 0; ic < inChannels; ic++) { + for(int oc = 0; oc < outChannels; oc++) { + weightsTransposed[((oc*convYSize + y)*convXSize + x)*inChannels + ic] = + desc->weights[((oc*inChannels + ic)*convYSize + y)*convXSize + x]; + } + } + } + } + CudaUtils::mallocAndCopyToDevice(name,weightsTransposed,filterBuf,useFP16); + hipDeviceSynchronize(); + } + else + CudaUtils::mallocAndCopyToDevice(name,desc->weights,filterBuf,useFP16); + + // Pre-allocate buffer for accumulate mode (residual skip connections). + // miopenConvolutionForwardImmediate does not support alpha/beta unlike cuDNN, + // so we need to save the output before conv and add it back afterwards. + { + int elemSize = usingFP16 ? sizeof(half_t) : sizeof(float); + size_t accumBytes = (size_t)maxBatchSize * outChannels * nnXLen * nnYLen * elemSize; + CUDA_ERR(name.c_str(), hipMalloc(&accumBuf, accumBytes)); + } + } + + ~ConvLayer() { + hipFree(filterBuf); + hipFree(accumBuf); + miopenDestroyTensorDescriptor(filterDescriptor); + miopenDestroyConvolutionDescriptor(convolutionDescriptor); + delete convolutionAlgorithms; + } + + size_t requiredWorkspaceBytes( + CudaHandles* cudaHandles, + int batchSize + ) const { + size_t workspaceBytes = 0; + CUDNN_ERR(name.c_str(),miopenConvolutionForwardGetSolutionWorkspaceSize( + cudaHandles->cudnn, + filterDescriptor, + inputDescriptors[batchSize], + convolutionDescriptor, + outputDescriptors[batchSize], + (*convolutionAlgorithms)[batchSize].solution_id, + &workspaceBytes + )); + return workspaceBytes; + } + + void apply( + CudaHandles* cudaHandles, + int batchSize, + bool accumulate, + void* inputBuf, + void* outputBuf, + void* workspaceBuf, + size_t workspaceBytes + ) const { + // miopenConvolutionForwardImmediate does NOT support alpha/beta (unlike cuDNN). + // When accumulate=true, we need: outputBuf = conv(inputBuf) + outputBuf (residual skip connection). + // Save outputBuf content to pre-allocated accumBuf, run conv, then add back. + if(accumulate) { + int elemSize = usingFP16 ? sizeof(half) : sizeof(float); + size_t outputBytes = (size_t)batchSize * outChannels * nnXLen * nnYLen * elemSize; + CUDA_ERR(name.c_str(), hipMemcpyAsync(accumBuf, outputBuf, outputBytes, hipMemcpyDeviceToDevice)); + } + + CUDNN_ERR(name.c_str(), miopenConvolutionForwardImmediate( + cudaHandles->cudnn, + filterDescriptor, + filterBuf, + inputDescriptors[batchSize], + inputBuf, + convolutionDescriptor, + outputDescriptors[batchSize], + outputBuf, + workspaceBuf, + workspaceBytes, + (*convolutionAlgorithms)[batchSize].solution_id + )); + + if(accumulate) { + int outputElems = (int)((size_t)batchSize * outChannels * nnXLen * nnYLen); + if(usingFP16) + customCudaAddTensorsInplace((half*)outputBuf, (const half*)accumBuf, outputElems); + else + customCudaAddTensorsInplace((float*)outputBuf, (const float*)accumBuf, outputElems); + } + } + +}; + +//--------------------------------------------------------------------------------- + +struct BatchNormLayer { + const string name; + const int numChannels; + const float epsilon; + const int activation; + const int nnXLen; + const int nnYLen; + + const bool usingFP16; + const bool usingNHWC; + + void* mergedScaleBuf; + void* mergedBiasBuf; + + BatchNormLayer() = delete; + BatchNormLayer(const BatchNormLayer&) = delete; + BatchNormLayer& operator=(const BatchNormLayer&) = delete; + + BatchNormLayer( + CudaHandles* cudaHandles, + const BatchNormLayerDesc* desc, + const ActivationLayerDesc* actDesc, + int nnX, + int nnY, + bool useFP16, + bool useNHWC + ) : + name(desc->name), + numChannels(desc->numChannels), + epsilon(desc->epsilon), + activation(actDesc->activation), + nnXLen(nnX), + nnYLen(nnY), + usingFP16(useFP16), + usingNHWC(useNHWC) + { + (void)cudaHandles; + + assert(desc->mean.size() == numChannels); + assert(desc->variance.size() == numChannels); + assert(desc->scale.size() == numChannels); + assert(desc->bias.size() == numChannels); + assert(desc->mergedScale.size() == numChannels); + assert(desc->mergedBias.size() == numChannels); + CudaUtils::mallocAndCopyToDevice(name,desc->mergedScale,mergedScaleBuf,useFP16); + CudaUtils::mallocAndCopyToDevice(name,desc->mergedBias,mergedBiasBuf,useFP16); + } + ~BatchNormLayer() { + hipFree(mergedScaleBuf); + hipFree(mergedBiasBuf); + } + + void apply( + CudaHandles* cudaHandles, + int batchSize, + void* inputBuf, + const void* maskBuf, //ok to be null + void* outputBuf + ) const { + (void)cudaHandles; + if(!usingFP16) { + if(!usingNHWC) + customCudaApplyCScaleBiasNCHW((const float*)inputBuf,(float*)outputBuf,(const float*)mergedScaleBuf,(const float*)mergedBiasBuf, + (const float*)maskBuf, + batchSize,numChannels,nnXLen*nnYLen,activation); + else + customCudaApplyCScaleBiasNHWC((const float*)inputBuf,(float*)outputBuf,(const float*)mergedScaleBuf,(const float*)mergedBiasBuf, + (const float*)maskBuf, + batchSize,nnXLen*nnYLen,numChannels,activation); + } + else { + if(!usingNHWC) + customCudaApplyCScaleBiasNCHW((const half*)inputBuf,(half*)outputBuf,(const half*)mergedScaleBuf,(const half*)mergedBiasBuf, + (const half*)maskBuf, + batchSize,numChannels,nnXLen*nnYLen,activation); + else + customCudaApplyCScaleBiasNHWC((const half*)inputBuf,(half*)outputBuf,(const half*)mergedScaleBuf,(const half*)mergedBiasBuf, + (const half*)maskBuf, + batchSize,nnXLen*nnYLen,numChannels,activation); + CUDA_ERR(name.c_str(),hipPeekAtLastError()); + } + + } + +}; + + +//--------------------------------------------------------------------------------- + +struct MatMulLayer { + const string name; + const int inChannels; + const int outChannels; + const bool usingFP16; + void* matBuf; + + MatMulLayer() = delete; + MatMulLayer(const MatMulLayer&) = delete; + MatMulLayer& operator=(const MatMulLayer&) = delete; + + MatMulLayer( + CudaHandles* cudaHandles, + const MatMulLayerDesc* desc, + bool useFP16 + ) : + name(desc->name), + inChannels(desc->inChannels), + outChannels(desc->outChannels), + usingFP16(useFP16) + { + (void)cudaHandles; + + if(inChannels > 0 && outChannels > 0) { + assert(desc->weights.size() == inChannels * outChannels); + CudaUtils::mallocAndCopyToDevice(name,desc->weights,matBuf,useFP16); + } + else { + matBuf = NULL; + } + } + + ~MatMulLayer() { + if(inChannels > 0 && outChannels > 0) + hipFree(matBuf); + } + + size_t requiredWorkspaceBytes( + CudaHandles* cudaHandles + ) const { + (void)cudaHandles; + size_t workspaceBytes = 0; + return workspaceBytes; + } + + void apply( + CudaHandles* cudaHandles, + ScratchBuffers* scratch, + int batchSize, + void* inputBuf, + void* outputBuf, + void* workspaceBuf, + size_t workspaceBytes + ) const { + (void)workspaceBuf; + (void)workspaceBytes; + assert(inChannels > 0 && outChannels > 0); + + if(!usingFP16) { + const float alpha = 1.0f; + const float beta = 0.0f; + CUBLAS_ERR(name.c_str(),hipblasSgemm( + cudaHandles->cublas, + HIPBLAS_OP_N, + HIPBLAS_OP_N, + outChannels, + batchSize, + inChannels, + &alpha, + (const float*)matBuf,outChannels, + (const float*)inputBuf,inChannels, + &beta, + (float*)outputBuf,outChannels + )); + } + else { + const hipblasHalf* alpha = (const hipblasHalf*)scratch->oneBuf; + const hipblasHalf* beta = (const hipblasHalf*)scratch->zeroBuf; + CUBLAS_ERR(name.c_str(),hipblasHgemm( + cudaHandles->cublas, + HIPBLAS_OP_N, + HIPBLAS_OP_N, + outChannels, + batchSize, + inChannels, + alpha, + (const hipblasHalf*)matBuf,outChannels, + (const hipblasHalf*)inputBuf,inChannels, + beta, + (hipblasHalf*)outputBuf,outChannels + )); + } + + } + +}; + +//--------------------------------------------------------------------------------- + +struct MatBiasLayer { + const string name; + const int numChannels; + const bool usingFP16; + const int activation; + + void* biasBuf; + + MatBiasLayer() = delete; + MatBiasLayer(const MatBiasLayer&) = delete; + MatBiasLayer& operator=(const MatBiasLayer&) = delete; + + MatBiasLayer( + CudaHandles* cudaHandles, + const MatBiasLayerDesc* desc, + bool useFP16, + int activation_ + ) : + name(desc->name), + numChannels(desc->numChannels), + usingFP16(useFP16), + activation(activation_) + { + (void)cudaHandles; + if(numChannels > 0) { + assert(desc->weights.size() == numChannels); + CudaUtils::mallocAndCopyToDevice(name,desc->weights,biasBuf,useFP16); + } + else + biasBuf = NULL; + } + + ~MatBiasLayer() { + if(numChannels > 0) + hipFree(biasBuf); + } + + void apply( + CudaHandles* cudaHandles, + int batchSize, + void* matBuf + ) const { + (void)cudaHandles; + assert(numChannels > 0); + if(!usingFP16) { + customCudaAddCBiasInplaceNC((float*)matBuf,(const float*)biasBuf,batchSize,numChannels,activation); + CUDA_ERR(name.c_str(),hipPeekAtLastError()); + } + else { + customCudaAddCBiasInplaceNC((half*)matBuf,(const half*)biasBuf,batchSize,numChannels,activation); + CUDA_ERR(name.c_str(),hipPeekAtLastError()); + } + } + +}; + +//--------------------------------------------------------------------------------- + +struct NormActConv { + const BatchNormLayer norm; + const ConvLayer conv; + + const int inChannels; + const int outChannels; + const int nnXLen; + const int nnYLen; + const bool usingFP16; + const bool usingNHWC; + + NormActConv() = delete; + NormActConv(const NormActConv&) = delete; + NormActConv& operator=(const NormActConv&) = delete; + + NormActConv( + CudaHandles* cudaHandles, + CudnnManager* manager, + const BatchNormLayerDesc* normDesc, + const ActivationLayerDesc* actDesc, + const ConvLayerDesc* convDesc, + int nnX, + int nnY, + bool useFP16, + bool useNHWC + ): norm(cudaHandles,normDesc,actDesc,nnX,nnY,useFP16,useNHWC), + conv(cudaHandles,manager,convDesc,useFP16,useNHWC), + inChannels(norm.numChannels), + outChannels(conv.outChannels), + nnXLen(nnX), + nnYLen(nnY), + usingFP16(useFP16), + usingNHWC(useNHWC) + { + assert(norm.numChannels == conv.inChannels); + } + + ~NormActConv() + {} + + size_t requiredWorkspaceBytes( + CudaHandles* cudaHandles, + int batchSize + ) const { + size_t bytes = 0; + size_t b; + b = conv.requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + return bytes; + } + + void apply( + CudaHandles* cudaHandles, + int batchSize, + bool accumulate, + void* inBuf, + void* inScratchBuf, + void* outBuf, + void* maskBuf, + void* workspaceBuf, + size_t workspaceBytes + ) const { + norm.apply(cudaHandles,batchSize,inBuf,maskBuf,inScratchBuf); +#ifdef DEBUG_INTERMEDIATE_VALUES + CudaUtils::debugPrint4D(string("AFTER NORM "), inScratchBuf, batchSize, inChannels, nnXLen, nnYLen, usingNHWC, usingFP16); +#endif + conv.apply(cudaHandles,batchSize,accumulate,inScratchBuf,outBuf,workspaceBuf,workspaceBytes); + } + +}; + + +//--------------------------------------------------------------------------------- + +struct ResidualBlock { + const string name; + const NormActConv normActConv1; + const NormActConv normActConv2; + + ResidualBlock() = delete; + ResidualBlock(const ResidualBlock&) = delete; + ResidualBlock& operator=(const ResidualBlock&) = delete; + + ResidualBlock( + CudaHandles* cudaHandles, + CudnnManager* manager, + const ResidualBlockDesc* desc, + int nnX, + int nnY, + bool useFP16, + bool useNHWC + ): name(desc->name), + normActConv1(cudaHandles,manager,&desc->preBN,&desc->preActivation,&desc->regularConv,nnX,nnY,useFP16,useNHWC), + normActConv2(cudaHandles,manager,&desc->midBN,&desc->midActivation,&desc->finalConv,nnX,nnY,useFP16,useNHWC) + { + } + + ~ResidualBlock() + {} + + size_t requiredWorkspaceBytes( + CudaHandles* cudaHandles, + int batchSize + ) const { + size_t bytes = 0; + size_t b; + b = normActConv1.requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + b = normActConv2.requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + return bytes; + } + + void apply( + CudaHandles* cudaHandles, + ScratchBuffers* scratch, + int batchSize, + void* trunkBuf, + void* trunkScratchBuf, + void* maskBuf, + void* workspaceBuf, + size_t workspaceBytes + ) const { + SizedBuf midIn(scratch->allocator, scratch->getBufSizeXY(normActConv1.outChannels)); + SizedBuf midScratch(scratch->allocator, scratch->getBufSizeXY(normActConv1.outChannels)); + normActConv1.apply(cudaHandles,batchSize,false,trunkBuf,trunkScratchBuf,midIn.buf,maskBuf,workspaceBuf,workspaceBytes); + normActConv2.apply(cudaHandles,batchSize,true,midIn.buf,midScratch.buf,trunkBuf,maskBuf,workspaceBuf,workspaceBytes); + } + +}; + + +//---------------------------------------------------------------------------- + + +struct GlobalPoolingResidualBlock { + const string name; + const BatchNormLayer preBN; + const ConvLayer regularConv; + const ConvLayer gpoolConv; + const BatchNormLayer gpoolBN; + const MatMulLayer gpoolToBiasMul; + const NormActConv normActConv2; + + const int nnXLen; + const int nnYLen; + const int regularChannels; + const int gpoolChannels; + const bool usingFP16; + const bool usingNHWC; + + GlobalPoolingResidualBlock() = delete; + GlobalPoolingResidualBlock(const GlobalPoolingResidualBlock&) = delete; + GlobalPoolingResidualBlock& operator=(const GlobalPoolingResidualBlock&) = delete; + + GlobalPoolingResidualBlock( + CudaHandles* cudaHandles, + CudnnManager* manager, + const GlobalPoolingResidualBlockDesc* desc, + int nnX, + int nnY, + bool useFP16, + bool useNHWC + ): name(desc->name), + preBN(cudaHandles,&desc->preBN,&desc->preActivation,nnX,nnY,useFP16,useNHWC), + regularConv(cudaHandles,manager,&desc->regularConv,useFP16,useNHWC), + gpoolConv(cudaHandles,manager,&desc->gpoolConv,useFP16,useNHWC), + gpoolBN(cudaHandles,&desc->gpoolBN,&desc->gpoolActivation,nnX,nnY,useFP16,useNHWC), + gpoolToBiasMul(cudaHandles,&desc->gpoolToBiasMul,useFP16), + normActConv2(cudaHandles,manager,&desc->midBN,&desc->midActivation,&desc->finalConv,nnX,nnY,useFP16,useNHWC), + nnXLen(nnX), + nnYLen(nnY), + regularChannels(desc->regularConv.outChannels), + gpoolChannels(desc->gpoolConv.outChannels), + usingFP16(useFP16), + usingNHWC(useNHWC) + { + } + + ~GlobalPoolingResidualBlock() { + } + + size_t requiredWorkspaceBytes( + CudaHandles* cudaHandles, + int batchSize + ) const { + size_t bytes = 0; + size_t b; + b = regularConv.requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + b = gpoolConv.requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + b = gpoolToBiasMul.requiredWorkspaceBytes(cudaHandles); + bytes = std::max(bytes,b); + b = normActConv2.requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + b = sizeof(float)*batchSize*gpoolChannels*nnXLen*nnYLen; + bytes = std::max(bytes,b); + return bytes; + } + + void apply( + CudaHandles* cudaHandles, + ScratchBuffers* scratch, + int batchSize, + void* trunkBuf, + void* trunkScratchBuf, + void* maskBuf, + float* maskSumBuf, + void* workspaceBuf, + size_t workspaceBytes + ) const { + SizedBuf regularOut(scratch->allocator, scratch->getBufSizeXY(regularChannels)); + SizedBuf regularScratch(scratch->allocator, scratch->getBufSizeXY(regularChannels)); + SizedBuf gpoolOut(scratch->allocator, scratch->getBufSizeXY(gpoolChannels)); + SizedBuf gpoolOut2(scratch->allocator, scratch->getBufSizeXY(gpoolChannels)); + SizedBuf gpoolConcat(scratch->allocator, scratch->getBufSize(gpoolChannels*3)); + SizedBuf gpoolBias(scratch->allocator, scratch->getBufSize(regularChannels)); + + preBN.apply(cudaHandles,batchSize,trunkBuf,maskBuf,trunkScratchBuf); + regularConv.apply(cudaHandles,batchSize,false,trunkScratchBuf,regularOut.buf,workspaceBuf,workspaceBytes); + gpoolConv.apply(cudaHandles,batchSize,false,trunkScratchBuf,gpoolOut.buf,workspaceBuf,workspaceBytes); + gpoolBN.apply(cudaHandles,batchSize,gpoolOut.buf,maskBuf,gpoolOut2.buf); + + if(!usingFP16) { + if(!usingNHWC) + customCudaPoolRowsGPoolNCHW((const float*)gpoolOut2.buf,(float*)gpoolConcat.buf,batchSize,gpoolChannels,nnXLen*nnYLen,(const float*)maskBuf,maskSumBuf); + else + customCudaPoolRowsGPoolNHWC((const float*)gpoolOut2.buf,(float*)gpoolConcat.buf,batchSize,nnXLen*nnYLen,gpoolChannels,(const float*)maskBuf,maskSumBuf); + } + else { + if(!usingNHWC) + customCudaPoolRowsGPoolNCHW((const half*)gpoolOut2.buf,(half*)gpoolConcat.buf,batchSize,gpoolChannels,nnXLen*nnYLen,(const half*)maskBuf,maskSumBuf); + else + customCudaPoolRowsGPoolNHWC((const half*)gpoolOut2.buf,(half*)gpoolConcat.buf,batchSize,nnXLen*nnYLen,gpoolChannels,(const half*)maskBuf,maskSumBuf); + } + CUDA_ERR(name.c_str(),hipPeekAtLastError()); + + gpoolToBiasMul.apply(cudaHandles,scratch,batchSize,gpoolConcat.buf,gpoolBias.buf,workspaceBuf,workspaceBytes); + + if(!usingFP16) { + if(!usingNHWC) + customCudaAddNCBiasInplaceNCHW((float*)regularOut.buf,(const float*)gpoolBias.buf,batchSize,regularChannels,nnXLen*nnYLen); + else + customCudaAddNCBiasInplaceNHWC((float*)regularOut.buf,(const float*)gpoolBias.buf,batchSize,nnXLen*nnYLen,regularChannels); + } + else { + if(!usingNHWC) + customCudaAddNCBiasInplaceNCHW((half*)regularOut.buf,(const half*)gpoolBias.buf,batchSize,regularChannels,nnXLen*nnYLen); + else + customCudaAddNCBiasInplaceNHWC((half*)regularOut.buf,(const half*)gpoolBias.buf,batchSize,nnXLen*nnYLen,regularChannels); + } + CUDA_ERR(name.c_str(),hipPeekAtLastError()); + + normActConv2.apply(cudaHandles,batchSize,true,regularOut.buf,regularScratch.buf,trunkBuf,maskBuf,workspaceBuf,workspaceBytes); + } + +}; + +//------------------------------------------------------------------------------ + +// Lightweight RMSNorm used inside transformer blocks (weight only, no bias, no spatial modes) +struct TransformerRMSNormLayer { + const string name; + const int numChannels; + const float epsilon; + const bool usingFP16; + void* weightBuf; + void* zeroBetaBuf; + + TransformerRMSNormLayer() = delete; + TransformerRMSNormLayer(const TransformerRMSNormLayer&) = delete; + TransformerRMSNormLayer& operator=(const TransformerRMSNormLayer&) = delete; + + TransformerRMSNormLayer( + CudaHandles* cudaHandles, + const TransformerRMSNormDesc* desc, + bool useFP16 + ) : + name(desc->name), + numChannels(desc->numChannels), + epsilon(desc->epsilon), + usingFP16(useFP16) + { + (void)cudaHandles; + testAssert((int)desc->weight.size() == numChannels); + CudaUtils::mallocAndCopyToDevice(name, desc->weight, weightBuf, useFP16); + vector zeros(numChannels, 0.0f); + CudaUtils::mallocAndCopyToDevice(name + ":zeroBeta", zeros, zeroBetaBuf, useFP16); + } + + ~TransformerRMSNormLayer() { + hipFree(weightBuf); + hipFree(zeroBetaBuf); + } + + // Apply RMSNorm on NHWC data [N, XY, C], applying mask [N, XY] to zero padded positions. + void apply( + CudaHandles* cudaHandles, + int batchSize, + int xySize, + void* inputBuf, + void* outputBuf, + const void* maskBuf + ) const { + (void)cudaHandles; + if(!usingFP16) { + customCudaRMSNormGammaBetaNHWC( + (const float*)inputBuf, (float*)outputBuf, + (const float*)weightBuf, (const float*)zeroBetaBuf, + (const float*)maskBuf, + batchSize, xySize, numChannels, epsilon, ACTIVATION_IDENTITY); + } + else { + customCudaRMSNormGammaBetaNHWC( + (const half*)inputBuf, (half*)outputBuf, + (const half*)weightBuf, (const half*)zeroBetaBuf, + (const half*)maskBuf, + batchSize, xySize, numChannels, epsilon, ACTIVATION_IDENTITY); + } + CUDA_ERR(name.c_str(), hipPeekAtLastError()); + } +}; + +//------------------------------------------------------------------------------ + +struct RMSNormLayer { + const string name; + const int numChannels; + const bool spatial; + const int activation; + const float epsilon; + const int nnXLen; + const int nnYLen; + const bool usingFP16; + const bool usingNHWC; + + void* gammaBuf; + void* betaBuf; + + RMSNormLayer() = delete; + RMSNormLayer(const RMSNormLayer&) = delete; + RMSNormLayer& operator=(const RMSNormLayer&) = delete; + + RMSNormLayer( + CudaHandles* cudaHandles, + const RMSNormLayerDesc* desc, + int act, + int nnX, + int nnY, + bool useFP16, + bool useNHWC + ) : + name(desc->name), + numChannels(desc->numChannels), + spatial(desc->spatial), + activation(act), + epsilon(desc->epsilon), + nnXLen(nnX), + nnYLen(nnY), + usingFP16(useFP16), + usingNHWC(useNHWC) + { + (void)cudaHandles; + testAssert((int)desc->gamma.size() == numChannels); + testAssert((int)desc->beta.size() == numChannels); + // The device kernels apply only RELU/MISH/SILU explicitly and treat anything else as + // identity; guard here so an unsupported kind (e.g. MISH_SCALE8, which applyScale8 can + // produce for non-transformer nets) fails loudly instead of silently skipping activation. + if(activation != ACTIVATION_IDENTITY && activation != ACTIVATION_RELU && + activation != ACTIVATION_MISH && activation != ACTIVATION_SILU) + throw StringError(name + ": RMSNorm layer unsupported activation: " + Global::intToString(activation)); + CudaUtils::mallocAndCopyToDevice(name, desc->gamma, gammaBuf, useFP16); + CudaUtils::mallocAndCopyToDevice(name, desc->beta, betaBuf, useFP16); + } + + ~RMSNormLayer() { + hipFree(gammaBuf); + hipFree(betaBuf); + } + + void apply( + CudaHandles* cudaHandles, + ScratchBuffers* scratch, + int batchSize, + void* inputBuf, + void* outputBuf, + const void* maskBuf, + const float* maskSumBuf + ) const { + (void)cudaHandles; + int xySize = nnXLen * nnYLen; + if(!spatial) { + if(!usingFP16) { + if(!usingNHWC) + customCudaRMSNormGammaBetaNCHW( + (const float*)inputBuf, (float*)outputBuf, (const float*)gammaBuf, (const float*)betaBuf, + (const float*)maskBuf, batchSize, numChannels, xySize, epsilon, activation); + else + customCudaRMSNormGammaBetaNHWC( + (const float*)inputBuf, (float*)outputBuf, (const float*)gammaBuf, (const float*)betaBuf, + (const float*)maskBuf, batchSize, xySize, numChannels, epsilon, activation); + } + else { + if(!usingNHWC) + customCudaRMSNormGammaBetaNCHW( + (const half*)inputBuf, (half*)outputBuf, (const half*)gammaBuf, (const half*)betaBuf, + (const half*)maskBuf, batchSize, numChannels, xySize, epsilon, activation); + else + customCudaRMSNormGammaBetaNHWC( + (const half*)inputBuf, (half*)outputBuf, (const half*)gammaBuf, (const half*)betaBuf, + (const half*)maskBuf, batchSize, xySize, numChannels, epsilon, activation); + } + } + else { + // Scratch buffer for spatial reduction (float regardless of FP16 mode). Holds per-block + // partial sums plus the final reduced value per batch element. + SizedBuf sumSqBuf(scratch->allocator, (size_t)batchSize * CUDA_SPATIAL_RMSNORM_SUMSQ_STRIDE * sizeof(float)); + if(!usingFP16) { + if(!usingNHWC) + customCudaSpatialRMSNormNCHW( + (const float*)inputBuf, (float*)outputBuf, (const float*)gammaBuf, (const float*)betaBuf, + (const float*)maskBuf, maskSumBuf, batchSize, numChannels, xySize, epsilon, activation, (float*)sumSqBuf.buf); + else + customCudaSpatialRMSNormNHWC( + (const float*)inputBuf, (float*)outputBuf, (const float*)gammaBuf, (const float*)betaBuf, + (const float*)maskBuf, maskSumBuf, batchSize, xySize, numChannels, epsilon, activation, (float*)sumSqBuf.buf); + } + else { + if(!usingNHWC) + customCudaSpatialRMSNormNCHW( + (const half*)inputBuf, (half*)outputBuf, (const half*)gammaBuf, (const half*)betaBuf, + (const half*)maskBuf, maskSumBuf, batchSize, numChannels, xySize, epsilon, activation, (float*)sumSqBuf.buf); + else + customCudaSpatialRMSNormNHWC( + (const half*)inputBuf, (half*)outputBuf, (const half*)gammaBuf, (const half*)betaBuf, + (const half*)maskBuf, maskSumBuf, batchSize, xySize, numChannels, epsilon, activation, (float*)sumSqBuf.buf); + } + } + CUDA_ERR(name.c_str(), hipPeekAtLastError()); + } +}; + +//------------------------------------------------------------------------------ + +struct TransformerAttentionBlock { + const string name; + const int numHeads; + const int numKVHeads; + const int qHeadDim; + const int vHeadDim; + const bool useRope; + const bool learnableRope; + const int inChannels; + + const int nnXLen; + const int nnYLen; + const bool usingFP16; + const bool usingNHWC; + + const TransformerRMSNormLayer preLN; + const MatMulLayer qProj; + const MatMulLayer kProj; + const MatMulLayer vProj; + const MatMulLayer outProj; + + // Fixed RoPE: precomputed cos/sin tables on device (NULL for learnable RoPE). + // Learnable RoPE: per-head frequencies on device (ropeFreqsBuf, FP32), cos/sin recomputed in-kernel. + void* ropeCosTable; + void* ropeSinTable; + float* ropeFreqsBuf; + int ropeNumPairs; + int ropeNumKVHeads; + + TransformerAttentionBlock() = delete; + TransformerAttentionBlock(const TransformerAttentionBlock&) = delete; + TransformerAttentionBlock& operator=(const TransformerAttentionBlock&) = delete; + + TransformerAttentionBlock( + CudaHandles* cudaHandles, + const TransformerAttentionDesc* desc, + int nnX, + int nnY, + bool useFP16, + bool useNHWC + ) : + name(desc->name), + numHeads(desc->numHeads), + numKVHeads(desc->numKVHeads), + qHeadDim(desc->qHeadDim), + vHeadDim(desc->vHeadDim), + useRope(desc->useRope), + learnableRope(desc->learnableRope), + inChannels(desc->qProj.inChannels), + nnXLen(nnX), + nnYLen(nnY), + usingFP16(useFP16), + usingNHWC(useNHWC), + preLN(cudaHandles, &desc->preLN, useFP16), + qProj(cudaHandles, &desc->qProj, useFP16), + kProj(cudaHandles, &desc->kProj, useFP16), + vProj(cudaHandles, &desc->vProj, useFP16), + outProj(cudaHandles, &desc->outProj, useFP16), + ropeCosTable(NULL), + ropeSinTable(NULL), + ropeFreqsBuf(NULL), + ropeNumPairs(0), + ropeNumKVHeads(0) + { + if(!useNHWC) { + throw StringError("Transformer blocks with NCHW layout are not yet supported by the ROCm backend"); + } + if(useRope) { + ropeNumPairs = qHeadDim / 2; + ropeNumKVHeads = numKVHeads; + if(learnableRope) { + testAssert(desc->ropeFreqs.size() == (size_t)(numKVHeads * ropeNumPairs * 2)); + void* freqsVoid = NULL; + CudaUtils::mallocAndCopyToDevice(name + ":ropeFreqs", desc->ropeFreqs.data(), (int)desc->ropeFreqs.size(), freqsVoid, false); + ropeFreqsBuf = (float*)freqsVoid; + } + else { + int seqLen = nnXLen * nnYLen; + vector cosTableData; + vector sinTableData; + desc->computeRopeCosSin(nnXLen, nnYLen, seqLen, cosTableData, sinTableData); + CudaUtils::mallocAndCopyToDevice(name + ":ropeCos", cosTableData.data(), (int)cosTableData.size(), ropeCosTable, useFP16); + CudaUtils::mallocAndCopyToDevice(name + ":ropeSin", sinTableData.data(), (int)sinTableData.size(), ropeSinTable, useFP16); + } + } + } + + ~TransformerAttentionBlock() { + if(ropeCosTable != NULL) hipFree(ropeCosTable); + if(ropeSinTable != NULL) hipFree(ropeSinTable); + if(ropeFreqsBuf != NULL) hipFree(ropeFreqsBuf); + } + + size_t requiredWorkspaceBytes( + CudaHandles* cudaHandles, + int batchSize + ) const { + (void)cudaHandles; + (void)batchSize; + return 0; + } + + void apply( + CudaHandles* cudaHandles, + ScratchBuffers* scratch, + int batchSize, + void* trunkBuf, + void* trunkScratchBuf, + void* maskBuf, + float* maskSumBuf, + void* workspaceBuf, + size_t workspaceBytes + ) const { + (void)maskSumBuf; + (void)workspaceBuf; + (void)workspaceBytes; + + int seqLen = nnXLen * nnYLen; + int qTotalDim = numHeads * qHeadDim; + int kTotalDim = numKVHeads * qHeadDim; + int vTotalDim = numKVHeads * vHeadDim; + + // NHWC: trunk is [N, XY, C]. RMSNorm + mask zeroing. + preLN.apply(cudaHandles, batchSize, seqLen, trunkBuf, trunkScratchBuf, maskBuf); + + // Step 2: Q/K/V projections. trunkScratchBuf is [N, XY, C] NHWC = [C, N*seqLen] column-major, + // which matches what MatMulLayer expects as input ([inChannels, batchSize]). + int matBatchSize = batchSize * seqLen; + + SizedBuf qBuf(scratch->allocator, scratch->getBufSizeXY(qTotalDim)); + SizedBuf kBuf(scratch->allocator, scratch->getBufSizeXY(kTotalDim)); + SizedBuf vBuf(scratch->allocator, scratch->getBufSizeXY(vTotalDim)); + + qProj.apply(cudaHandles, scratch, matBatchSize, trunkScratchBuf, qBuf.buf, workspaceBuf, workspaceBytes); + kProj.apply(cudaHandles, scratch, matBatchSize, trunkScratchBuf, kBuf.buf, workspaceBuf, workspaceBytes); + vProj.apply(cudaHandles, scratch, matBatchSize, trunkScratchBuf, vBuf.buf, workspaceBuf, workspaceBytes); + + // Step 3: Apply RoPE to Q and K. + if(useRope) { + if(learnableRope) { + if(!usingFP16) { + customCudaApplyRoPELearnableRecompute((float*)qBuf.buf, ropeFreqsBuf, + batchSize, seqLen, numHeads, numKVHeads, qHeadDim, ropeNumPairs, nnXLen); + customCudaApplyRoPELearnableRecompute((float*)kBuf.buf, ropeFreqsBuf, + batchSize, seqLen, numKVHeads, numKVHeads, qHeadDim, ropeNumPairs, nnXLen); + } + else { + customCudaApplyRoPELearnableRecompute((half*)qBuf.buf, ropeFreqsBuf, + batchSize, seqLen, numHeads, numKVHeads, qHeadDim, ropeNumPairs, nnXLen); + customCudaApplyRoPELearnableRecompute((half*)kBuf.buf, ropeFreqsBuf, + batchSize, seqLen, numKVHeads, numKVHeads, qHeadDim, ropeNumPairs, nnXLen); + } + } + else { + if(!usingFP16) { + customCudaApplyRoPE((float*)qBuf.buf, (const float*)ropeCosTable, (const float*)ropeSinTable, + batchSize, seqLen, numHeads, numKVHeads, qHeadDim, ropeNumPairs, learnableRope); + customCudaApplyRoPE((float*)kBuf.buf, (const float*)ropeCosTable, (const float*)ropeSinTable, + batchSize, seqLen, numKVHeads, numKVHeads, qHeadDim, ropeNumPairs, learnableRope); + } + else { + customCudaApplyRoPE((half*)qBuf.buf, (const half*)ropeCosTable, (const half*)ropeSinTable, + batchSize, seqLen, numHeads, numKVHeads, qHeadDim, ropeNumPairs, learnableRope); + customCudaApplyRoPE((half*)kBuf.buf, (const half*)ropeCosTable, (const half*)ropeSinTable, + batchSize, seqLen, numKVHeads, numKVHeads, qHeadDim, ropeNumPairs, learnableRope); + } + } + CUDA_ERR(name.c_str(), hipPeekAtLastError()); + } + + // Step 4: Scaled dot-product attention. When enabled and available, try the CK FMHA fused path + // first (only supports FP16, matching the CUDA backend's cudnn-frontend SDPA path, which is + // also FP16-only); otherwise, or if CK reports the shape/traits as unsupported at runtime, fall + // back to the plain online-softmax kernel, which is fully adequate at KataGo's sequence lengths. + SizedBuf attnOutBuf(scratch->allocator, scratch->getBufSizeXY(numHeads * vHeadDim)); + + bool usedFusedAttention = false; +#if KATAGO_ROCM_HAS_CK_FMHA && KATAGO_ROCM_CK_FMHA_ARCH_OK + if(usingFP16 && !cudaHandles->disableFusedAttention) { + bool hasMask = (maskBuf != NULL); + + // CK's fused path takes a pre-materialized additive bias rather than a raw per-position + // mask; build a [B, S, S] bias broadcast over heads (matches the CUDA backend's approach for + // its cudnn-frontend graph SDPA path). + SizedBuf biasBuf(scratch->allocator, hasMask ? (size_t)batchSize * seqLen * seqLen * sizeof(half) : 1); + if(hasMask) { + customCudaMaskToAttnBiasFull((const half*)maskBuf, (half*)biasBuf.buf, batchSize, seqLen); + } + + fmha_fwd_traits traits; + traits.hdim_q = qHeadDim; + traits.hdim_v = vHeadDim; + traits.data_type = "fp16"; + traits.is_group_mode = false; + traits.is_v_rowmajor = true; + traits.has_logits_soft_cap = false; + traits.mask_type = mask_enum::no_mask; + traits.bias_type = hasMask ? bias_enum::elementwise_bias : bias_enum::no_bias; + traits.has_lse = false; + traits.has_dropout = false; + traits.qscale_type = quant_scale_enum::no_scale; + traits.skip_min_seqlen_q = false; + traits.has_sink = false; + + // Physical layout for Q/K/V/O is [N, S, H, D] (i_perm=false, o_perm=false), V is row-major + // ([N, S, H, Dv], is_v_rowmajor=true) - all matching the BSHD buffers MatMulLayer produces. + fmha_fwd_args args; + memset(&args, 0, sizeof(args)); + args.q_ptr = qBuf.buf; + args.k_ptr = kBuf.buf; + args.v_ptr = vBuf.buf; + args.bias_ptr = hasMask ? biasBuf.buf : nullptr; + args.o_ptr = attnOutBuf.buf; + args.seqlen_q = seqLen; + args.seqlen_k = seqLen; + args.batch = batchSize; + args.max_seqlen_q = seqLen; + args.hdim_q = qHeadDim; + args.hdim_v = vHeadDim; + args.nhead_q = numHeads; + args.nhead_k = numKVHeads; + args.scale_s = 1.0f / sqrtf((float)qHeadDim); + args.logits_soft_cap = 0.0f; + args.stride_q = (ck_tile::index_t)numHeads * qHeadDim; + args.stride_k = (ck_tile::index_t)numKVHeads * qHeadDim; + args.stride_v = (ck_tile::index_t)numKVHeads * vHeadDim; + args.stride_bias = hasMask ? seqLen : 0; + args.stride_o = (ck_tile::index_t)numHeads * vHeadDim; + args.nhead_stride_q = qHeadDim; + args.nhead_stride_k = qHeadDim; + args.nhead_stride_v = vHeadDim; + args.nhead_stride_bias = 0; // broadcast the bias over heads + args.nhead_stride_o = vHeadDim; + args.batch_stride_q = (ck_tile::index_t)numHeads * seqLen * qHeadDim; + args.batch_stride_k = (ck_tile::index_t)numKVHeads * seqLen * qHeadDim; + args.batch_stride_v = (ck_tile::index_t)numKVHeads * seqLen * vHeadDim; + args.batch_stride_bias = hasMask ? (ck_tile::index_t)seqLen * seqLen : 0; + args.batch_stride_o = (ck_tile::index_t)numHeads * seqLen * vHeadDim; + args.window_size_left = -1; + args.window_size_right = -1; + args.mask_type = static_cast(mask_enum::no_mask); + args.min_seqlen_q = 0; + args.p_drop = 0.0f; + args.s_randval = false; + args.drop_seed_offset = std::pair{0, 0}; + + float ckResult = fmha_fwd(traits, args, ck_tile::stream_config{}); + if(ckResult >= 0.0f) { + usedFusedAttention = true; + } + } +#endif + + if(!usedFusedAttention) { + if(!usingFP16) { + customCudaFlashAttention( + (const float*)qBuf.buf, (const float*)kBuf.buf, (const float*)vBuf.buf, + (const float*)maskBuf, (float*)attnOutBuf.buf, + batchSize, seqLen, numHeads, numKVHeads, qHeadDim, vHeadDim); + } + else { + customCudaFlashAttention( + (const half*)qBuf.buf, (const half*)kBuf.buf, (const half*)vBuf.buf, + (const half*)maskBuf, (half*)attnOutBuf.buf, + batchSize, seqLen, numHeads, numKVHeads, qHeadDim, vHeadDim); + } + CUDA_ERR(name.c_str(), hipPeekAtLastError()); + } + + // Step 5: Output projection. + outProj.apply(cudaHandles, scratch, matBatchSize, attnOutBuf.buf, trunkScratchBuf, workspaceBuf, workspaceBytes); + + // Step 6: Residual addition: trunk += trunkScratch * mask + if(!usingFP16) { + customCudaMaskedResidualAddNHWC((float*)trunkBuf, (const float*)trunkScratchBuf, (const float*)maskBuf, batchSize, seqLen, inChannels); + } + else { + customCudaMaskedResidualAddNHWC((half*)trunkBuf, (const half*)trunkScratchBuf, (const half*)maskBuf, batchSize, seqLen, inChannels); + } + CUDA_ERR(name.c_str(), hipPeekAtLastError()); + } +}; + +//------------------------------------------------------------------------------ + +struct TransformerFFNBlock { + const string name; + const int numChannels; + const int ffnChannels; + const bool useSwiGLU; + + const int nnXLen; + const int nnYLen; + const bool usingFP16; + const bool usingNHWC; + + const TransformerRMSNormLayer preLN; + const MatMulLayer linear1; + std::unique_ptr linearGate; + const MatMulLayer linear2; + + TransformerFFNBlock() = delete; + TransformerFFNBlock(const TransformerFFNBlock&) = delete; + TransformerFFNBlock& operator=(const TransformerFFNBlock&) = delete; + + TransformerFFNBlock( + CudaHandles* cudaHandles, + const TransformerFFNDesc* desc, + int nnX, + int nnY, + bool useFP16, + bool useNHWC + ) : + name(desc->name), + numChannels(desc->numChannels), + ffnChannels(desc->ffnChannels), + useSwiGLU(desc->useSwiGLU), + nnXLen(nnX), + nnYLen(nnY), + usingFP16(useFP16), + usingNHWC(useNHWC), + preLN(cudaHandles, &desc->preLN, useFP16), + linear1(cudaHandles, &desc->linear1, useFP16), + linear2(cudaHandles, &desc->linear2, useFP16) + { + if(!useSwiGLU) { + throw StringError("Non-SwiGLU transformer FFN is not yet supported in ROCm backend"); + } + linearGate = std::make_unique(cudaHandles, &desc->linearGate, useFP16); + if(!useNHWC) { + throw StringError("Transformer blocks with NCHW layout are not yet supported by the ROCm backend"); + } + } + + ~TransformerFFNBlock() + {} + + size_t requiredWorkspaceBytes( + CudaHandles* cudaHandles, + int batchSize + ) const { + (void)cudaHandles; + (void)batchSize; + return 0; + } + + void apply( + CudaHandles* cudaHandles, + ScratchBuffers* scratch, + int batchSize, + void* trunkBuf, + void* trunkScratchBuf, + void* maskBuf, + float* maskSumBuf, + void* workspaceBuf, + size_t workspaceBytes + ) const { + (void)maskSumBuf; + + int seqLen = nnXLen * nnYLen; + int matBatchSize = batchSize * seqLen; + + preLN.apply(cudaHandles, batchSize, seqLen, trunkBuf, trunkScratchBuf, maskBuf); + + SizedBuf hiddenBuf(scratch->allocator, scratch->getBufSizeXY(ffnChannels)); + SizedBuf gateBuf(scratch->allocator, scratch->getBufSizeXY(ffnChannels)); + + linear1.apply(cudaHandles, scratch, matBatchSize, trunkScratchBuf, hiddenBuf.buf, workspaceBuf, workspaceBytes); + linearGate->apply(cudaHandles, scratch, matBatchSize, trunkScratchBuf, gateBuf.buf, workspaceBuf, workspaceBytes); + + int totalElts = matBatchSize * ffnChannels; + if(!usingFP16) { + customCudaSwiGLU((const float*)hiddenBuf.buf, (const float*)gateBuf.buf, (float*)hiddenBuf.buf, totalElts); + } + else { + customCudaSwiGLU((const half*)hiddenBuf.buf, (const half*)gateBuf.buf, (half*)hiddenBuf.buf, totalElts); + } + CUDA_ERR(name.c_str(), hipPeekAtLastError()); + + linear2.apply(cudaHandles, scratch, matBatchSize, hiddenBuf.buf, trunkScratchBuf, workspaceBuf, workspaceBytes); + + if(!usingFP16) { + customCudaMaskedResidualAddNHWC((float*)trunkBuf, (const float*)trunkScratchBuf, (const float*)maskBuf, batchSize, seqLen, numChannels); + } + else { + customCudaMaskedResidualAddNHWC((half*)trunkBuf, (const half*)trunkScratchBuf, (const half*)maskBuf, batchSize, seqLen, numChannels); + } + CUDA_ERR(name.c_str(), hipPeekAtLastError()); + } +}; + +//------------------------------------------------------------------------------ + +struct BlockStack { + const int numBlocks; + const int trunkNumChannels; + const int nnXLen; + const int nnYLen; + const bool usingFP16; + const bool usingNHWC; + vector> blocks; + + BlockStack() = delete; + BlockStack(const BlockStack&) = delete; + BlockStack& operator=(const BlockStack&) = delete; + + BlockStack( + CudaHandles* cudaHandles, + CudnnManager* manager, + int nBlocks, + int trunkChannels, + const std::vector>& descBlocks, + int nnX, + int nnY, + bool useFP16, + bool useNHWC + ); + ~BlockStack(); + + size_t requiredWorkspaceBytes( + CudaHandles* cudaHandles, + int batchSize + ) const; + + void apply( + CudaHandles* cudaHandles, + ScratchBuffers* scratch, + int batchSize, + void* maskBuf, + float* maskSumBuf, + void* trunkBuf, + void* trunkScratchBuf, + void* workspaceBuf, + size_t workspaceBytes + ) const; + +}; + +//------------------------------------------------------------------------------ + +struct NestedBottleneckResidualBlock { + const string name; + const NormActConv normActConv1; + const BlockStack blocks; + const NormActConv normActConv2; + + NestedBottleneckResidualBlock() = delete; + NestedBottleneckResidualBlock(const NestedBottleneckResidualBlock&) = delete; + NestedBottleneckResidualBlock& operator=(const NestedBottleneckResidualBlock&) = delete; + + NestedBottleneckResidualBlock( + CudaHandles* cudaHandles, + CudnnManager* manager, + const NestedBottleneckResidualBlockDesc* desc, + int nnX, + int nnY, + bool useFP16, + bool useNHWC + ): name(desc->name), + normActConv1(cudaHandles,manager,&desc->preBN,&desc->preActivation,&desc->preConv,nnX,nnY,useFP16,useNHWC), + blocks(cudaHandles,manager,desc->numBlocks,desc->preConv.outChannels,desc->blocks,nnX,nnY,useFP16,useNHWC), + normActConv2(cudaHandles,manager,&desc->postBN,&desc->postActivation,&desc->postConv,nnX,nnY,useFP16,useNHWC) + { + } + + ~NestedBottleneckResidualBlock() + {} + + size_t requiredWorkspaceBytes( + CudaHandles* cudaHandles, + int batchSize + ) const { + size_t bytes = 0; + size_t b; + b = normActConv1.requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + b = blocks.requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + b = normActConv2.requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + return bytes; + } + + void apply( + CudaHandles* cudaHandles, + ScratchBuffers* scratch, + int batchSize, + void* trunkBuf, + void* trunkScratchBuf, + void* maskBuf, + float* maskSumBuf, + void* workspaceBuf, + size_t workspaceBytes + ) const { + SizedBuf mid(scratch->allocator, scratch->getBufSizeXY(normActConv1.outChannels)); + SizedBuf midScratch(scratch->allocator, scratch->getBufSizeXY(normActConv1.outChannels)); + assert(normActConv1.outChannels == normActConv2.inChannels); + normActConv1.apply(cudaHandles,batchSize,false,trunkBuf,trunkScratchBuf,mid.buf,maskBuf,workspaceBuf,workspaceBytes); + blocks.apply( + cudaHandles, + scratch, + batchSize, + maskBuf, + maskSumBuf, + mid.buf, + midScratch.buf, + workspaceBuf, + workspaceBytes + ); + normActConv2.apply(cudaHandles,batchSize,true,mid.buf,midScratch.buf,trunkBuf,maskBuf,workspaceBuf,workspaceBytes); + } + +}; + +//------------------------------------------------------------------------------ + +BlockStack::BlockStack( + CudaHandles* cudaHandles, + CudnnManager* manager, + int nBlocks, + int trunkChannels, + const std::vector>& descBlocks, + int nnX, + int nnY, + bool useFP16, + bool useNHWC +) : + numBlocks(nBlocks), + trunkNumChannels(trunkChannels), + nnXLen(nnX), + nnYLen(nnY), + usingFP16(useFP16), + usingNHWC(useNHWC) +{ + assert(numBlocks == descBlocks.size()); + for(int i = 0; irequiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + } + else if(blocks[i].first == GLOBAL_POOLING_BLOCK_KIND) { + GlobalPoolingResidualBlock* block = (GlobalPoolingResidualBlock*)blocks[i].second.get(); + b = block->requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + } + else if(blocks[i].first == NESTED_BOTTLENECK_BLOCK_KIND) { + NestedBottleneckResidualBlock* block = (NestedBottleneckResidualBlock*)blocks[i].second.get(); + b = block->requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + } + else if(blocks[i].first == TRANSFORMER_ATTENTION_BLOCK_KIND) { + TransformerAttentionBlock* block = (TransformerAttentionBlock*)blocks[i].second.get(); + b = block->requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + } + else if(blocks[i].first == TRANSFORMER_FFN_BLOCK_KIND) { + TransformerFFNBlock* block = (TransformerFFNBlock*)blocks[i].second.get(); + b = block->requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + } + else { + ASSERT_UNREACHABLE; + } + } + return bytes; +} + +void BlockStack::apply( + CudaHandles* cudaHandles, + ScratchBuffers* scratch, + int batchSize, + void* maskBuf, + float* maskSumBuf, + void* trunkBuf, + void* trunkScratchBuf, + void* workspaceBuf, + size_t workspaceBytes +) const { + + for(int i = 0; iapply( + cudaHandles, + scratch, + batchSize, + trunkBuf, + trunkScratchBuf, + maskBuf, + workspaceBuf, + workspaceBytes + ); + } + else if(blocks[i].first == GLOBAL_POOLING_BLOCK_KIND) { + GlobalPoolingResidualBlock* block = (GlobalPoolingResidualBlock*)blocks[i].second.get(); + block->apply( + cudaHandles, + scratch, + batchSize, + trunkBuf, + trunkScratchBuf, + maskBuf, + maskSumBuf, + workspaceBuf, + workspaceBytes + ); + } + else if(blocks[i].first == NESTED_BOTTLENECK_BLOCK_KIND) { + NestedBottleneckResidualBlock* block = (NestedBottleneckResidualBlock*)blocks[i].second.get(); + block->apply( + cudaHandles, + scratch, + batchSize, + trunkBuf, + trunkScratchBuf, + maskBuf, + maskSumBuf, + workspaceBuf, + workspaceBytes + ); + } + else if(blocks[i].first == TRANSFORMER_ATTENTION_BLOCK_KIND) { + TransformerAttentionBlock* block = (TransformerAttentionBlock*)blocks[i].second.get(); + block->apply( + cudaHandles, + scratch, + batchSize, + trunkBuf, + trunkScratchBuf, + maskBuf, + maskSumBuf, + workspaceBuf, + workspaceBytes + ); + } + else if(blocks[i].first == TRANSFORMER_FFN_BLOCK_KIND) { + TransformerFFNBlock* block = (TransformerFFNBlock*)blocks[i].second.get(); + block->apply( + cudaHandles, + scratch, + batchSize, + trunkBuf, + trunkScratchBuf, + maskBuf, + maskSumBuf, + workspaceBuf, + workspaceBytes + ); + } + else { + ASSERT_UNREACHABLE; + } + } +} +//------------------------------------------------------------------------------ + +struct SGFMetadataEncoder { + const string name; + + const bool usingFP16; + + const MatMulLayer mul1; + const MatBiasLayer bias1; + const MatMulLayer mul2; + const MatBiasLayer bias2; + const MatMulLayer mul3; + + SGFMetadataEncoder() = delete; + SGFMetadataEncoder(const SGFMetadataEncoder&) = delete; + SGFMetadataEncoder& operator=(const SGFMetadataEncoder&) = delete; + + SGFMetadataEncoder( + CudaHandles* cudaHandles, + const SGFMetadataEncoderDesc* desc, + bool useFP16 + ) : + name(desc->name), + usingFP16(useFP16), + mul1(cudaHandles,&desc->mul1,useFP16), + bias1(cudaHandles,&desc->bias1,useFP16,desc->act1.activation), + mul2(cudaHandles,&desc->mul2,useFP16), + bias2(cudaHandles,&desc->bias2,useFP16,desc->act2.activation), + mul3(cudaHandles,&desc->mul3,useFP16) + { + } + + ~SGFMetadataEncoder() + { + } + + size_t requiredWorkspaceBytes( + CudaHandles* cudaHandles, + int batchSize + ) const { + (void)batchSize; + size_t bytes = 0; + size_t b; + + b = mul1.requiredWorkspaceBytes(cudaHandles); + bytes = std::max(bytes,b); + b = mul2.requiredWorkspaceBytes(cudaHandles); + bytes = std::max(bytes,b); + b = mul3.requiredWorkspaceBytes(cudaHandles); + bytes = std::max(bytes,b); + + return bytes; + } + + void apply( + CudaHandles* cudaHandles, + ScratchBuffers* scratch, + int batchSize, + void* inputBuf, + void* outputBuf, + void* workspaceBuf, + size_t workspaceBytes + ) const { + SizedBuf internalBuf1(scratch->allocator, scratch->getBufSizeFloat(std::max(mul1.outChannels,mul2.outChannels))); + SizedBuf internalBuf2(scratch->allocator, scratch->getBufSizeFloat(std::max(mul1.outChannels,mul2.outChannels))); + + mul1.apply(cudaHandles,scratch,batchSize,inputBuf,internalBuf1.buf,workspaceBuf,workspaceBytes); + bias1.apply(cudaHandles,batchSize,internalBuf1.buf); + mul2.apply(cudaHandles,scratch,batchSize,internalBuf1.buf,internalBuf2.buf,workspaceBuf,workspaceBytes); + bias2.apply(cudaHandles,batchSize,internalBuf2.buf); + mul3.apply(cudaHandles,scratch,batchSize,internalBuf2.buf,outputBuf,workspaceBuf,workspaceBytes); + } + +}; + + +//---------------------------------------------------------------------------- + +struct Trunk { + const string name; + const int modelVersion; + const int numBlocks; + const int trunkNumChannels; + const int trunkNormKind; + + const int nnXLen; + const int nnYLen; + const bool usingFP16; + const bool usingNHWC; + + std::unique_ptr initialConv; + std::unique_ptr initialMatMul; + std::unique_ptr sgfMetadataEncoder; + const BlockStack blocks; + std::unique_ptr trunkTipBN; + std::unique_ptr trunkTipRMSNorm; + + Trunk() = delete; + Trunk(const Trunk&) = delete; + Trunk& operator=(const Trunk&) = delete; + + Trunk( + CudaHandles* cudaHandles, + CudnnManager* manager, + const TrunkDesc* desc, + int nnX, + int nnY, + bool inputsUseNHWC, + bool useFP16, + bool useNHWC + ) : + name(desc->name), + modelVersion(desc->modelVersion), + numBlocks(desc->numBlocks), + trunkNumChannels(desc->trunkNumChannels), + trunkNormKind(desc->trunkNormKind), + nnXLen(nnX), + nnYLen(nnY), + usingFP16(useFP16), + usingNHWC(useNHWC), + blocks(cudaHandles,manager,desc->numBlocks,desc->trunkNumChannels,desc->blocks,nnX,nnY,useFP16,useNHWC) + { + int midNumChannels = desc->midNumChannels; + int regularNumChannels = desc->regularNumChannels; + int gpoolNumChannels = desc->gpoolNumChannels; + + int maxBatchSize = manager->maxBatchSize; + CudaUtils::checkBufferSize(maxBatchSize,nnXLen,nnYLen,trunkNumChannels); + CudaUtils::checkBufferSize(maxBatchSize,nnXLen,nnYLen,midNumChannels); + CudaUtils::checkBufferSize(maxBatchSize,nnXLen,nnYLen,regularNumChannels); + CudaUtils::checkBufferSize(maxBatchSize,nnXLen,nnYLen,gpoolNumChannels); + + initialConv = std::make_unique(cudaHandles,manager,&desc->initialConv,useFP16,inputsUseNHWC,useNHWC); + initialMatMul = std::make_unique(cudaHandles,&desc->initialMatMul,useFP16); + if(desc->metaEncoderVersion > 0) { + sgfMetadataEncoder = std::make_unique(cudaHandles,&desc->sgfMetadataEncoder,useFP16); + testAssert(sgfMetadataEncoder->mul3.outChannels == initialMatMul->outChannels); + } + + if(desc->trunkNormKind == TRUNK_NORM_KIND_STANDARD) { + trunkTipBN = std::make_unique(cudaHandles,&desc->trunkTipBN,&desc->trunkTipActivation,nnXLen,nnYLen,useFP16,useNHWC); + } + else if(desc->trunkNormKind == TRUNK_NORM_KIND_RMSNORM) { + trunkTipRMSNorm = std::make_unique(cudaHandles,&desc->trunkTipRMSNorm,desc->trunkTipActivation.activation,nnXLen,nnYLen,useFP16,useNHWC); + } + else { + throw StringError("Unsupported trunk norm kind: " + Global::intToString(desc->trunkNormKind)); + } + assert(desc->blocks.size() == numBlocks); + } + + ~Trunk() + { + } + + size_t requiredWorkspaceBytes( + CudaHandles* cudaHandles, + int batchSize + ) const { + size_t bytes = 0; + size_t b; + + b = initialConv->requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + + b = initialMatMul->requiredWorkspaceBytes(cudaHandles); + bytes = std::max(bytes,b); + + if(sgfMetadataEncoder != nullptr) { + b = sgfMetadataEncoder->requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + } + + b = blocks.requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + return bytes; + } + + void apply( + CudaHandles* cudaHandles, + ScratchBuffers* scratch, + int batchSize, + void* inputBuf, + void* inputGlobalBuf, + void* inputMetaBuf, + void* maskBuf, + float* maskSumBuf, + void* trunkBuf, + void* workspaceBuf, + size_t workspaceBytes + ) const { + + SizedBuf trunkScratch(scratch->allocator, scratch->getBufSizeXY(trunkNumChannels)); + + //Feed the conv into trunkScratch.buf, not trunkBuf + initialConv->apply(cudaHandles,batchSize,false,inputBuf,trunkScratch.buf,workspaceBuf,workspaceBytes); + + #ifdef DEBUG_INTERMEDIATE_VALUES + CudaUtils::debugPrint4D(string("After initial conv"), trunkScratch.buf, batchSize, trunkNumChannels, nnXLen, nnYLen, usingNHWC, usingFP16); + #endif + + //Feed the matmul into trunkBuf + initialMatMul->apply(cudaHandles,scratch,batchSize,inputGlobalBuf,trunkBuf,workspaceBuf,workspaceBytes); + //Then accumulate it into trunkScratch.buf, broadcasting during the process + if(!usingFP16) { + if(!usingNHWC) + customCudaAddNCBiasInplaceNCHW((float*)trunkScratch.buf,(const float*)trunkBuf,batchSize,trunkNumChannels,nnXLen*nnYLen); + else + customCudaAddNCBiasInplaceNHWC((float*)trunkScratch.buf,(const float*)trunkBuf,batchSize,nnXLen*nnYLen,trunkNumChannels); + } + else { + if(!usingNHWC) + customCudaAddNCBiasInplaceNCHW((half*)trunkScratch.buf,(const half*)trunkBuf,batchSize,trunkNumChannels,nnXLen*nnYLen); + else + customCudaAddNCBiasInplaceNHWC((half*)trunkScratch.buf,(const half*)trunkBuf,batchSize,nnXLen*nnYLen,trunkNumChannels); + } + CUDA_ERR(name.c_str(),hipPeekAtLastError()); + + if(sgfMetadataEncoder != nullptr) { + testAssert(inputMetaBuf != NULL); + //Feed the result into trunkBuf + sgfMetadataEncoder->apply(cudaHandles,scratch,batchSize,inputMetaBuf,trunkBuf,workspaceBuf,workspaceBytes); + //Then accumulate it into trunkScratch.buf, broadcasting during the process + if(!usingFP16) { + if(!usingNHWC) + customCudaAddNCBiasInplaceNCHW((float*)trunkScratch.buf,(const float*)trunkBuf,batchSize,trunkNumChannels,nnXLen*nnYLen); + else + customCudaAddNCBiasInplaceNHWC((float*)trunkScratch.buf,(const float*)trunkBuf,batchSize,nnXLen*nnYLen,trunkNumChannels); + } + else { + if(!usingNHWC) + customCudaAddNCBiasInplaceNCHW((half*)trunkScratch.buf,(const half*)trunkBuf,batchSize,trunkNumChannels,nnXLen*nnYLen); + else + customCudaAddNCBiasInplaceNHWC((half*)trunkScratch.buf,(const half*)trunkBuf,batchSize,nnXLen*nnYLen,trunkNumChannels); + } + CUDA_ERR(name.c_str(),hipPeekAtLastError()); + } + else { + testAssert(inputMetaBuf == NULL); + } + + //Flip trunkBuf and trunkScratch.buf so that the result gets accumulated in trunkScratch.buf + blocks.apply( + cudaHandles, + scratch, + batchSize, + maskBuf, + maskSumBuf, + trunkScratch.buf, + trunkBuf, + workspaceBuf, + workspaceBytes + ); + + //And now with the final norm port it from trunkScratch.buf to trunkBuf. + if(trunkNormKind == TRUNK_NORM_KIND_STANDARD) { + trunkTipBN->apply(cudaHandles,batchSize,trunkScratch.buf,maskBuf,trunkBuf); + } + else { + trunkTipRMSNorm->apply(cudaHandles,scratch,batchSize,trunkScratch.buf,trunkBuf,maskBuf,maskSumBuf); + } + + #ifdef DEBUG_INTERMEDIATE_VALUES + CudaUtils::debugPrint4D(string("Trunk tip"), trunkBuf, batchSize, trunkNumChannels, nnXLen, nnYLen, usingNHWC, usingFP16); + #endif + } + +}; + +//------------------------------------------------------------------------------ + +static void fillMaskFloatBufAndMaskSumBuf(void* maskBuf, float*& maskFloatBuf, float*& maskSumBuf, bool usingFP16, int batchSize, int nnXLen, int nnYLen) { + if(!usingFP16) { + maskFloatBuf = (float*)maskBuf; + customCudaPoolRowsSumNCHW((const float*)maskFloatBuf,maskSumBuf,batchSize,1,nnXLen*nnYLen,1.0); + CUDA_ERR("sumMask",hipPeekAtLastError()); + } + else { + customCudaCopyFromHalf((const half*)maskBuf,maskFloatBuf,batchSize*nnXLen*nnYLen); + CUDA_ERR("copyMaskFromHalf",hipPeekAtLastError()); + customCudaPoolRowsSumNCHW((const float*)maskFloatBuf,maskSumBuf,batchSize,1,nnXLen*nnYLen,1.0); + CUDA_ERR("sumMask",hipPeekAtLastError()); + } +} + + +//------------------------------------------------------------------------------ + +struct PolicyHead { + const string name; + const int modelVersion; + const int nnXLen; + const int nnYLen; + const int p1Channels; + const int g1Channels; + const int p2Channels; + const bool usingFP16; + const bool usingNHWC; + + const ConvLayer p1Conv; + const ConvLayer g1Conv; + const BatchNormLayer g1BN; + const MatMulLayer gpoolToBiasMul; + const BatchNormLayer p1BN; + const ConvLayer p2Conv; + const MatMulLayer gpoolToPassMul; + const MatBiasLayer gpoolToPassBias; + const MatMulLayer gpoolToPassMul2; + + PolicyHead() = delete; + PolicyHead(const PolicyHead&) = delete; + PolicyHead& operator=(const PolicyHead&) = delete; + + PolicyHead( + CudaHandles* cudaHandles, + CudnnManager* manager, + const PolicyHeadDesc* desc, + int nnX, + int nnY, + bool useFP16, + bool useNHWC + ) : + name(desc->name), + modelVersion(desc->modelVersion), + nnXLen(nnX), + nnYLen(nnY), + p1Channels(desc->p1Conv.outChannels), + g1Channels(desc->g1Conv.outChannels), + p2Channels(desc->p2Conv.outChannels), + usingFP16(useFP16), + usingNHWC(useNHWC), + p1Conv(cudaHandles,manager,&desc->p1Conv,useFP16,useNHWC), + g1Conv(cudaHandles,manager,&desc->g1Conv,useFP16,useNHWC), + g1BN(cudaHandles,&desc->g1BN,&desc->g1Activation,nnX,nnY,useFP16,useNHWC), + gpoolToBiasMul(cudaHandles,&desc->gpoolToBiasMul,false), + p1BN(cudaHandles,&desc->p1BN,&desc->p1Activation,nnX,nnY,false,useNHWC), + p2Conv(cudaHandles,manager,&desc->p2Conv,false,useNHWC), + gpoolToPassMul(cudaHandles,&desc->gpoolToPassMul,false), + gpoolToPassBias(cudaHandles,&desc->gpoolToPassBias,false,desc->passActivation.activation), + gpoolToPassMul2(cudaHandles,&desc->gpoolToPassMul2,false) + { + } + + ~PolicyHead() + { + } + + size_t requiredWorkspaceBytes( + CudaHandles* cudaHandles, + int batchSize + ) const { + size_t bytes = 0; + size_t b; + + b = p1Conv.requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + b = g1Conv.requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + b = gpoolToBiasMul.requiredWorkspaceBytes(cudaHandles); + bytes = std::max(bytes,b); + b = p2Conv.requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + b = gpoolToPassMul.requiredWorkspaceBytes(cudaHandles); + bytes = std::max(bytes,b); + b = gpoolToPassMul2.requiredWorkspaceBytes(cudaHandles); + bytes = std::max(bytes,b); + b = sizeof(float)*batchSize*g1Channels*nnXLen*nnYLen; + bytes = std::max(bytes,b); + + return bytes; + } + + void apply( + CudaHandles* cudaHandles, + ScratchBuffers* scratch, + int batchSize, + void* maskBuf, + float* maskFloatBuf, + float* maskSumBuf, + void* trunkBuf, + float* policyPassBuf, + float* policyBuf, + void* workspaceBuf, + size_t workspaceBytes + ) const { + + SizedBuf p1Out(scratch->allocator, scratch->getBufSizeXYFloat(p1Channels)); //Need to hold floats, not just halfs + SizedBuf p1Out2(scratch->allocator, scratch->getBufSizeXYFloat(p1Channels)); //Need to hold floats, not just halfs + SizedBuf g1Out(scratch->allocator, scratch->getBufSizeXY(g1Channels)); + SizedBuf g1Out2(scratch->allocator, scratch->getBufSizeXY(g1Channels)); + SizedBuf g1Concat(scratch->allocator, scratch->getBufSizeFloat(g1Channels*3)); + SizedBuf g1Bias(scratch->allocator, scratch->getBufSizeFloat(p1Channels)); + SizedBuf p1Pass(scratch->allocator, scratch->getBufSizeFloat(p1Channels)); + + p1Conv.apply(cudaHandles,batchSize,false,trunkBuf,p1Out.buf,workspaceBuf,workspaceBytes); + g1Conv.apply(cudaHandles,batchSize,false,trunkBuf,g1Out.buf,workspaceBuf,workspaceBytes); + g1BN.apply(cudaHandles,batchSize,g1Out.buf,maskBuf,g1Out2.buf); + + if(!usingFP16) { + if(!usingNHWC) + customCudaPoolRowsGPoolNCHW((const float*)g1Out2.buf,(float*)g1Concat.buf,batchSize,g1Channels,nnXLen*nnYLen,maskFloatBuf,maskSumBuf); + else + customCudaPoolRowsGPoolNHWC((const float*)g1Out2.buf,(float*)g1Concat.buf,batchSize,nnXLen*nnYLen,g1Channels,maskFloatBuf,maskSumBuf); + CUDA_ERR(name.c_str(),hipPeekAtLastError()); + } + else { + customCudaCopyFromHalf((const half*)g1Out2.buf,(float*)workspaceBuf,batchSize*g1Channels*nnXLen*nnYLen); + CUDA_ERR(name.c_str(),hipPeekAtLastError()); + if(!usingNHWC) + customCudaPoolRowsGPoolNCHW((const float*)workspaceBuf,(float*)g1Concat.buf,batchSize,g1Channels,nnXLen*nnYLen,maskFloatBuf,maskSumBuf); + else + customCudaPoolRowsGPoolNHWC((const float*)workspaceBuf,(float*)g1Concat.buf,batchSize,nnXLen*nnYLen,g1Channels,maskFloatBuf,maskSumBuf); + CUDA_ERR(name.c_str(),hipPeekAtLastError()); + } + + gpoolToBiasMul.apply(cudaHandles,scratch,batchSize,g1Concat.buf,g1Bias.buf,workspaceBuf,workspaceBytes); + + #ifdef DEBUG_INTERMEDIATE_VALUES + CudaUtils::debugPrint4D(string("p1 pre-gpool-sum"), p1Out.buf, batchSize, p1Channels, nnXLen, nnYLen, usingNHWC, usingFP16); + CudaUtils::debugPrint4D(string("g1 pre-gpool"), g1Out.buf, batchSize, g1Channels, nnXLen, nnYLen, usingNHWC, usingFP16); + CudaUtils::debugPrint2D(string("g1 pooled"), g1Concat.buf, batchSize, g1Channels*3, false); + CudaUtils::debugPrint2D(string("g1 biases"), g1Bias.buf, batchSize, p1Channels, false); + #endif + + float* p1OutBufA; + float* p1OutBufB; + if(!usingFP16) { + p1OutBufA = (float*)p1Out.buf; + p1OutBufB = (float*)p1Out2.buf; + } + else { + customCudaCopyFromHalf((const half*)p1Out.buf,(float*)p1Out2.buf,batchSize*p1Channels*nnXLen*nnYLen); + CUDA_ERR(name.c_str(),hipPeekAtLastError()); + p1OutBufA = (float*)p1Out2.buf; + p1OutBufB = (float*)p1Out.buf; + } + + if(!usingNHWC) + customCudaAddNCBiasInplaceNCHW(p1OutBufA,(float*)g1Bias.buf,batchSize,p1Channels,nnXLen*nnYLen); + else + customCudaAddNCBiasInplaceNHWC(p1OutBufA,(float*)g1Bias.buf,batchSize,nnXLen*nnYLen,p1Channels); + CUDA_ERR(name.c_str(),hipPeekAtLastError()); + + p1BN.apply(cudaHandles,batchSize,p1OutBufA,maskFloatBuf,p1OutBufB); + p2Conv.apply(cudaHandles,batchSize,false,p1OutBufB,(float*)policyBuf,workspaceBuf,workspaceBytes); + + if(modelVersion >= 15) { + gpoolToPassMul.apply(cudaHandles,scratch,batchSize,g1Concat.buf,p1Pass.buf,workspaceBuf,workspaceBytes); + gpoolToPassBias.apply(cudaHandles,batchSize,p1Pass.buf); + gpoolToPassMul2.apply(cudaHandles,scratch,batchSize,p1Pass.buf,policyPassBuf,workspaceBuf,workspaceBytes); + } + else { + gpoolToPassMul.apply(cudaHandles,scratch,batchSize,g1Concat.buf,policyPassBuf,workspaceBuf,workspaceBytes); + } + + #ifdef DEBUG_INTERMEDIATE_VALUES + CudaUtils::debugPrint4D(string("p1 after-gpool-sum"), p1OutBufA, batchSize, p1Channels, nnXLen, nnYLen, usingNHWC, false); + CudaUtils::debugPrint2D(string("policypass"), policyPassBuf, batchSize, 1, false); + CudaUtils::debugPrint4D(string("policy"), policyBuf, batchSize, p2Channels, nnXLen, nnYLen, usingNHWC, false); + #endif + + } + +}; + +//------------------------------------------------------------------------------ + +struct ValueHead { + const string name; + const int modelVersion; + const int nnXLen; + const int nnYLen; + const int v1Channels; + const int v2Channels; + const int valueChannels; + const int scoreValueChannels; + const int ownershipChannels; + const bool usingFP16; + const bool usingNHWC; + + const ConvLayer v1Conv; + const BatchNormLayer v1BN; + const MatMulLayer v2Mul; + const MatBiasLayer v2Bias; + const MatMulLayer v3Mul; + const MatBiasLayer v3Bias; + const MatMulLayer sv3Mul; + const MatBiasLayer sv3Bias; + const ConvLayer vOwnershipConv; + + ValueHead() = delete; + ValueHead(const ValueHead&) = delete; + ValueHead& operator=(const ValueHead&) = delete; + + ValueHead( + CudaHandles* cudaHandles, + CudnnManager* manager, + const ValueHeadDesc* desc, + int nnX, + int nnY, + bool useFP16, + bool useNHWC + ) : + name(desc->name), + modelVersion(desc->modelVersion), + nnXLen(nnX), + nnYLen(nnY), + v1Channels(desc->v1Conv.outChannels), + v2Channels(desc->v2Mul.outChannels), + valueChannels(desc->v3Mul.outChannels), + scoreValueChannels(desc->sv3Mul.outChannels), + ownershipChannels(desc->vOwnershipConv.outChannels), + usingFP16(useFP16), + usingNHWC(useNHWC), + v1Conv(cudaHandles,manager,&desc->v1Conv,useFP16,useNHWC), + v1BN(cudaHandles,&desc->v1BN,&desc->v1Activation,nnX,nnY,useFP16,useNHWC), + v2Mul(cudaHandles,&desc->v2Mul,false), + v2Bias(cudaHandles,&desc->v2Bias,false,desc->v2Activation.activation), + v3Mul(cudaHandles,&desc->v3Mul,false), + v3Bias(cudaHandles,&desc->v3Bias,false,ACTIVATION_IDENTITY), + sv3Mul(cudaHandles,&desc->sv3Mul,false), + sv3Bias(cudaHandles,&desc->sv3Bias,false,ACTIVATION_IDENTITY), + vOwnershipConv(cudaHandles,manager,&desc->vOwnershipConv,useFP16,useNHWC) + { + } + + ~ValueHead() + { + } + + size_t requiredWorkspaceBytes( + CudaHandles* cudaHandles, + int batchSize + ) const { + size_t bytes = 0; + size_t b; + + b = v1Conv.requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + b = v2Mul.requiredWorkspaceBytes(cudaHandles); + bytes = std::max(bytes,b); + b = v3Mul.requiredWorkspaceBytes(cudaHandles); + bytes = std::max(bytes,b); + b = sizeof(float)*batchSize*v1Channels*nnXLen*nnYLen; + bytes = std::max(bytes,b); + + b = sv3Mul.requiredWorkspaceBytes(cudaHandles); + bytes = std::max(bytes,b); + b = vOwnershipConv.requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + b = sizeof(float)*batchSize*ownershipChannels*nnXLen*nnYLen; + bytes = std::max(bytes,b); + + return bytes; + } + + + void apply( + CudaHandles* cudaHandles, + ScratchBuffers* scratch, + int batchSize, + void* maskBuf, + float* maskSumBuf, + void* trunkBuf, + float* valueBuf, + float* scoreValueBuf, + void* ownershipBuf, + void* workspaceBuf, + size_t workspaceBytes + ) const { + SizedBuf v1Out(scratch->allocator, scratch->getBufSizeXY(v1Channels)); + SizedBuf v1Out2(scratch->allocator, scratch->getBufSizeXY(v1Channels)); + SizedBuf v1Mean(scratch->allocator, scratch->getBufSizeFloat(v1Channels*3)); + SizedBuf v2Out(scratch->allocator, scratch->getBufSizeFloat(v2Channels)); + SizedBuf ownershipScratch(scratch->allocator, scratch->getBufSizeXYFloat(ownershipChannels)); + + v1Conv.apply(cudaHandles,batchSize,false,trunkBuf,v1Out.buf,workspaceBuf,workspaceBytes); + v1BN.apply(cudaHandles,batchSize,v1Out.buf,maskBuf,v1Out2.buf); + + void* bufToBePooled = v1Out2.buf; + if(usingFP16) { + customCudaCopyFromHalf((const half*)v1Out2.buf,(float*)workspaceBuf,batchSize*v1Channels*nnXLen*nnYLen); + CUDA_ERR(name.c_str(),hipPeekAtLastError()); + bufToBePooled = workspaceBuf; + } + + if(!usingNHWC) + customCudaValueHeadPoolNCHW((float*)bufToBePooled,(float*)v1Mean.buf,batchSize,v1Channels,nnXLen*nnYLen,maskSumBuf); + else + customCudaValueHeadPoolNHWC((const float*)bufToBePooled,(float*)v1Mean.buf,batchSize,nnXLen*nnYLen,v1Channels,maskSumBuf); + CUDA_ERR(name.c_str(),hipPeekAtLastError()); + + v2Mul.apply(cudaHandles,scratch,batchSize,v1Mean.buf,v2Out.buf,workspaceBuf,workspaceBytes); + v2Bias.apply(cudaHandles,batchSize,v2Out.buf); + v3Mul.apply(cudaHandles,scratch,batchSize,v2Out.buf,valueBuf,workspaceBuf,workspaceBytes); + v3Bias.apply(cudaHandles,batchSize,valueBuf); + + sv3Mul.apply(cudaHandles,scratch,batchSize,v2Out.buf,scoreValueBuf,workspaceBuf,workspaceBytes); + sv3Bias.apply(cudaHandles,batchSize,scoreValueBuf); + + #ifdef DEBUG_INTERMEDIATE_VALUES + CudaUtils::debugPrint4D(string("v1"), v1Out.buf, batchSize, v1Channels, nnXLen, nnYLen, usingNHWC, usingFP16); + CudaUtils::debugPrint2D(string("v1 pooled"), v1Mean.buf, batchSize, v1Channels, false); + CudaUtils::debugPrint2D(string("v2"), v2Out.buf, batchSize, v1Channels, false); + #endif + + if(!usingFP16) { + vOwnershipConv.apply(cudaHandles,batchSize,false,v1Out2.buf,ownershipBuf,workspaceBuf,workspaceBytes); + } + else { + vOwnershipConv.apply(cudaHandles,batchSize,false,v1Out2.buf,ownershipScratch.buf,workspaceBuf,workspaceBytes); + customCudaCopyFromHalf((const half*)ownershipScratch.buf,(float*)ownershipBuf,batchSize*ownershipChannels*nnXLen*nnYLen); + CUDA_ERR("vOwnership copy",hipPeekAtLastError()); + } + + } + +}; + +//------------------------------------------------------------------------------ + +struct Model { + const string name; + const int modelVersion; + const int maxBatchSize; + const int nnXLen; + const int nnYLen; + const int numInputChannels; + const int numInputGlobalChannels; + const int numInputMetaChannels; + const int numPolicyChannels; + const int numValueChannels; + const int numScoreValueChannels; + const int numOwnershipChannels; + const bool usingFP16; + const bool usingNHWC; + const bool inputsUsingNHWC; + + std::unique_ptr trunk; + std::unique_ptr policyHead; + std::unique_ptr valueHead; + std::unique_ptr manager; + + Model() = delete; + Model(const Model&) = delete; + Model& operator=(const Model&) = delete; + + Model( + CudaHandles* cudaHandles, + const ModelDesc* desc, + int maxBatchSz, + int nnX, + int nnY, + bool inputsUseNHWC, + bool useFP16, + bool useNHWC + ) : + name(desc->name), + modelVersion(desc->modelVersion), + maxBatchSize(maxBatchSz), + nnXLen(nnX), + nnYLen(nnY), + numInputChannels(desc->numInputChannels), + numInputGlobalChannels(desc->numInputGlobalChannels), + numInputMetaChannels(desc->numInputMetaChannels), + numPolicyChannels(desc->numPolicyChannels), + numValueChannels(desc->numValueChannels), + numScoreValueChannels(desc->numScoreValueChannels), + numOwnershipChannels(desc->numOwnershipChannels), + usingFP16(useFP16), + usingNHWC(useNHWC), + inputsUsingNHWC(inputsUseNHWC) + { + if(nnXLen > NNPos::MAX_BOARD_LEN) + throw StringError(Global::strprintf("nnXLen (%d) is greater than NNPos::MAX_BOARD_LEN (%d)", + nnXLen, NNPos::MAX_BOARD_LEN + )); + if(nnYLen > NNPos::MAX_BOARD_LEN) + throw StringError(Global::strprintf("nnYLen (%d) is greater than NNPos::MAX_BOARD_LEN (%d)", + nnYLen, NNPos::MAX_BOARD_LEN + )); + + int numFeatures = NNModelVersion::getNumSpatialFeatures(modelVersion); + if(numInputChannels != numFeatures) + throw StringError(Global::strprintf("Neural net numInputChannels (%d) was not the expected number based on version (%d)", + numInputChannels, numFeatures + )); + int numGlobalFeatures = NNModelVersion::getNumGlobalFeatures(modelVersion); + if(numInputGlobalChannels != numGlobalFeatures) + throw StringError(Global::strprintf("Neural net numInputGlobalChannels (%d) was not the expected number based on version (%d)", + numInputGlobalChannels, numGlobalFeatures + )); + if(numInputMetaChannels > 0) { + if(numInputMetaChannels != SGFMetadata::METADATA_INPUT_NUM_CHANNELS) + throw StringError(Global::strprintf("Neural net numInputMetaChannels (%d) was not the expected number (%d)", + numInputMetaChannels, SGFMetadata::METADATA_INPUT_NUM_CHANNELS + )); + } + + CudaUtils::checkBufferSize(maxBatchSize,nnXLen,nnYLen,numInputChannels); + CudaUtils::checkBufferSize(maxBatchSize,nnXLen,nnYLen,numInputGlobalChannels); + CudaUtils::checkBufferSize(maxBatchSize,nnXLen,nnYLen,numInputMetaChannels); + CudaUtils::checkBufferSize(maxBatchSize,nnXLen,nnYLen,numPolicyChannels); + CudaUtils::checkBufferSize(maxBatchSize,nnXLen,nnYLen,numValueChannels); + CudaUtils::checkBufferSize(maxBatchSize,nnXLen,nnYLen,numScoreValueChannels); + CudaUtils::checkBufferSize(maxBatchSize,nnXLen,nnYLen,numOwnershipChannels); + + manager = std::make_unique(name, maxBatchSize, nnXLen, nnYLen); + trunk = std::make_unique(cudaHandles,manager.get(),&desc->trunk,nnXLen,nnYLen,inputsUseNHWC,useFP16,useNHWC); + policyHead = std::make_unique(cudaHandles,manager.get(),&desc->policyHead,nnXLen,nnYLen,useFP16,useNHWC); + valueHead = std::make_unique(cudaHandles,manager.get(),&desc->valueHead,nnXLen,nnYLen,useFP16,useNHWC); + } + + ~Model() + { + } + + size_t requiredWorkspaceBytes( + CudaHandles* cudaHandles, + int batchSize + ) const { + size_t bytes = 0; + size_t b; + + b = trunk->requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + b = policyHead->requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + b = valueHead->requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + + return bytes; + } + + void apply( + CudaHandles* cudaHandles, + ScratchBuffers* scratch, + int batchSize, + bool requireExactNNLen, + + void* inputBuf, + void* inputGlobalBuf, + void* inputMetaBuf, + + float* policyPassBuf, + float* policyBuf, + + float* valueBuf, + float* scoreValueBuf, + void* ownershipBuf, + + void* workspaceBuf, + size_t workspaceBytes + ) const { + SizedBuf mask(scratch->allocator, scratch->getBufSizeXY(1)); + SizedBuf maskFloat(scratch->allocator, scratch->getBufSizeXYFloat(1)); + SizedBuf maskSum(scratch->allocator, scratch->getBufSizeFloat(1)); + + void* maskBuf = mask.buf; + float* maskFloatBuf = (float*)maskFloat.buf; + float* maskSumBuf = (float*)maskSum.buf; + + if(!usingFP16) { + if(inputsUsingNHWC) + customCudaChannel0ExtractNHWC((const float*)inputBuf, (float*)maskBuf, batchSize, nnXLen*nnYLen, numInputChannels); + else + customCudaChannel0ExtractNCHW((const float*)inputBuf, (float*)maskBuf, batchSize, numInputChannels, nnXLen*nnYLen); + CUDA_ERR("modelExtractMask",hipPeekAtLastError()); + } + else { + if(inputsUsingNHWC) + customCudaChannel0ExtractNHWC((const half*)inputBuf, (half*)maskBuf, batchSize, nnXLen*nnYLen, numInputChannels); + else + customCudaChannel0ExtractNCHW((const half*)inputBuf, (half*)maskBuf, batchSize, numInputChannels, nnXLen*nnYLen); + CUDA_ERR("modelExtractMask",hipPeekAtLastError()); + } + + fillMaskFloatBufAndMaskSumBuf(maskBuf,maskFloatBuf,maskSumBuf,usingFP16,batchSize,nnXLen,nnYLen); + + //Don't do any masking if we know the board is exactly the desired size + if(requireExactNNLen) { + //Set to NULL to signal downstream that this buf doesn't need to be used + maskBuf = NULL; + maskFloatBuf = NULL; + //The global pooling structures need this no matter what, for normalizing based on this and its sqrt. + //maskSumBuf = NULL; + } + + #ifdef DEBUG_INTERMEDIATE_VALUES + CudaUtils::debugPrint4D(string("Initial bin features"), inputBuf, batchSize, trunk->initialConv->inChannels, nnXLen, nnYLen, inputsUsingNHWC, usingFP16); + CudaUtils::debugPrint2D(string("Initial global features"), inputGlobalBuf, batchSize, trunk->initialMatMul->inChannels, usingFP16); + if(trunk->sgfMetadataEncoder != nullptr) { + assert(inputMetaBuf != NULL); + CudaUtils::debugPrint2D(string("Initial meta features"), inputMetaBuf, batchSize, trunk->sgfMetadataEncoder->mul1.inChannels, usingFP16); + } + #endif + + SizedBuf trunkBuf(scratch->allocator, scratch->getBufSizeXY(trunk->trunkNumChannels)); + + trunk->apply( + cudaHandles, + scratch, + batchSize, + inputBuf, + inputGlobalBuf, + inputMetaBuf, + maskBuf, + maskSumBuf, + trunkBuf.buf, + workspaceBuf, + workspaceBytes + ); + policyHead->apply( + cudaHandles, + scratch, + batchSize, + maskBuf, + maskFloatBuf, + maskSumBuf, + trunkBuf.buf, + policyPassBuf, + policyBuf, + workspaceBuf, + workspaceBytes + ); + valueHead->apply( + cudaHandles, + scratch, + batchSize, + maskBuf, + maskSumBuf, + trunkBuf.buf, + valueBuf, + scoreValueBuf, + ownershipBuf, + workspaceBuf, + workspaceBytes + ); + } + +}; + + +//------------------------------------------------------------------------------ + +struct LoadedModel { + ModelDesc modelDesc; + + LoadedModel(const string& fileName, const string& expectedSha256) { + ModelDesc::loadFromFileMaybeGZipped(fileName,modelDesc,expectedSha256); + modelDesc.applyScale8ToReduceActivations(); + } + + LoadedModel() = delete; + LoadedModel(const LoadedModel&) = delete; + LoadedModel& operator=(const LoadedModel&) = delete; +}; + +LoadedModel* NeuralNet::loadModelFile(const string& file, const string& expectedSha256) { + LoadedModel* loadedModel = new LoadedModel(file,expectedSha256); + return loadedModel; +} + +void NeuralNet::freeLoadedModel(LoadedModel* loadedModel) { + delete loadedModel; +} + +const ModelDesc& NeuralNet::getModelDesc(const LoadedModel* loadedModel) { + return loadedModel->modelDesc; +} + +//------------------------------------------------------------------------------ + +struct Buffers { + //All of these are device pointers + + float* inputBufFloat; + void* inputBuf; + float* inputGlobalBufFloat; + void* inputGlobalBuf; + float* inputMetaBufFloat; + void* inputMetaBuf; + size_t inputBufBytesFloat; + size_t inputBufBytes; + size_t inputGlobalBufBytesFloat; + size_t inputGlobalBufBytes; + size_t inputMetaBufBytesFloat; + size_t inputMetaBufBytes; + + float* policyPassBuf; + size_t policyPassBufBytes; + float* policyBuf; + size_t policyBufBytes; + + float* valueBuf; + size_t valueBufBytes; + float* scoreValueBuf; + size_t scoreValueBufBytes; + void* ownershipBuf; + size_t ownershipBufBytes; + + void* workspaceBuf; + size_t workspaceBytes; + + Buffers() = delete; + Buffers(const Buffers&) = delete; + Buffers& operator=(const Buffers&) = delete; + + Buffers(CudaHandles* cudaHandles, const Model& m, const ScratchBuffers& scratch) { + size_t batchXYFloatBytes = (size_t)scratch.batchXYFloatBytes; + size_t batchFloatBytes = (size_t)scratch.batchFloatBytes; + size_t batchXYBytes = (size_t)scratch.batchXYBytes; + size_t batchBytes = (size_t)scratch.batchBytes; + + inputBufBytesFloat = m.numInputChannels * batchXYFloatBytes; + inputBufBytes = m.numInputChannels * batchXYBytes; + inputGlobalBufBytesFloat = m.numInputGlobalChannels * batchFloatBytes; + inputGlobalBufBytes = m.numInputGlobalChannels * batchBytes; + inputMetaBufBytesFloat = m.numInputMetaChannels * batchFloatBytes; + inputMetaBufBytes = m.numInputMetaChannels * batchBytes; + + CUDA_ERR("Buffers",hipMalloc(reinterpret_cast(&inputBufFloat), inputBufBytesFloat)); + CUDA_ERR("Buffers",hipMalloc(&inputBuf, inputBufBytes)); + CUDA_ERR("Buffers",hipMalloc(reinterpret_cast(&inputGlobalBufFloat), inputGlobalBufBytesFloat)); + CUDA_ERR("Buffers",hipMalloc(&inputGlobalBuf, inputGlobalBufBytes)); + if(m.numInputMetaChannels > 0) { + CUDA_ERR("Buffers",hipMalloc(reinterpret_cast(&inputMetaBufFloat), inputMetaBufBytesFloat)); + CUDA_ERR("Buffers",hipMalloc(&inputMetaBuf, inputMetaBufBytes)); + } + else { + inputMetaBufFloat = NULL; + inputMetaBuf = NULL; + } + + if(m.modelVersion >= 17) + testAssert(m.policyHead->p2Channels == 2 || m.policyHead->p2Channels == 4); + else if(m.modelVersion >= 16) + testAssert(m.policyHead->p2Channels == 4); + else if(m.modelVersion >= 12) + testAssert(m.policyHead->p2Channels == 2); + else + testAssert(m.policyHead->p2Channels == 1); + + policyPassBufBytes = m.policyHead->p2Channels * batchFloatBytes; + CUDA_ERR("Buffers",hipMalloc(reinterpret_cast(&policyPassBuf), policyPassBufBytes)); + policyBufBytes = m.policyHead->p2Channels * batchXYFloatBytes; + CUDA_ERR("Buffers",hipMalloc(reinterpret_cast(&policyBuf), policyBufBytes)); + + valueBufBytes = m.valueHead->valueChannels * batchFloatBytes; + CUDA_ERR("Buffers",hipMalloc(reinterpret_cast(&valueBuf), valueBufBytes)); + + scoreValueBufBytes = m.valueHead->scoreValueChannels * batchFloatBytes; + CUDA_ERR("Buffers",hipMalloc(reinterpret_cast(&scoreValueBuf), scoreValueBufBytes)); + + //This buf is used for both an intermdiate fp16 result in fp16 mode, and ALSO the final fp32 output, so always must be fp32-sized + ownershipBufBytes = m.valueHead->ownershipChannels * batchXYFloatBytes; + CUDA_ERR("Buffers",hipMalloc(&ownershipBuf, ownershipBufBytes)); + + //In theory the requiredWorkspaceBytes calls could give us values non-monotone in batch size + //such as if the convolution algorithm changes between batch size 1 and larger. + //So we call it for all the batch sizes. + size_t bytes = 0; + size_t b; + for(int batchSize = 1; batchSize <= m.maxBatchSize; batchSize++) { + b = m.requiredWorkspaceBytes(cudaHandles,batchSize); + bytes = std::max(bytes,b); + } + + CUDA_ERR("Buffers",hipMalloc(&workspaceBuf, bytes)); + workspaceBytes = bytes; + } + + ~Buffers() { + hipFree(inputBufFloat); + hipFree(inputBuf); + hipFree(inputGlobalBufFloat); + hipFree(inputGlobalBuf); + if(inputMetaBufFloat != NULL) + hipFree(inputMetaBufFloat); + if(inputMetaBuf != NULL) + hipFree(inputMetaBuf); + + hipFree(policyPassBuf); + hipFree(policyBuf); + + hipFree(valueBuf); + hipFree(scoreValueBuf); + hipFree(ownershipBuf); + + hipFree(workspaceBuf); + } + +}; + +//------------------------------------------------------------------------------ + +struct ComputeContext { + int nnXLen; + int nnYLen; + enabled_t useFP16Mode; + enabled_t useNHWCMode; + bool disableFusedAttention; +}; + +ComputeContext* NeuralNet::createComputeContext( + const std::vector& gpuIdxs, + Logger* logger, + int nnXLen, + int nnYLen, + const string& homeDataDirOverride, + enabled_t useFP16Mode, + const LoadedModel* loadedModel, + ConfigParser& cfg +) { + (void)gpuIdxs; + (void)logger; + (void)homeDataDirOverride; + (void)loadedModel; + + // ROCm-specific NHWC override, read directly off cfg (mirrors cudaUseNHWC in the CUDA backend). + enabled_t useNHWCMode = + cfg.contains("rocmUseNHWC") ? cfg.getEnabled("rocmUseNHWC") : enabled_t::Auto; + // Disables the optional CK FMHA fused attention path (see KATAGO_ROCM_HAS_CK_FMHA); mirrors the + // CUDA backend's cudaDisableGraphSDPA. Only meaningful for transformer models. + bool disableFusedAttention = + cfg.contains("rocmDisableFusedAttention") ? cfg.getBool("rocmDisableFusedAttention") : false; + + ComputeContext* context = new ComputeContext(); + context->nnXLen = nnXLen; + context->nnYLen = nnYLen; + context->useFP16Mode = useFP16Mode; + context->useNHWCMode = useNHWCMode; + context->disableFusedAttention = disableFusedAttention; + return context; +} + +void NeuralNet::freeComputeContext(ComputeContext* computeContext) { + delete computeContext; +} + +//------------------------------------------------------------------------------ + +struct ComputeHandle { + std::unique_ptr cudaHandles; + std::unique_ptr model; + std::unique_ptr scratch; + std::unique_ptr buffers; + const bool usingFP16; + const int nnXLen; + const int nnYLen; + const bool requireExactNNLen; + const bool inputsUseNHWC; + const bool usingNHWC; + + ComputeHandle( + const ComputeContext* context, + const LoadedModel* loadedModel, + int majorComputeCapability, + int minorComputeCapability, + int maxBatchSize, + bool requireExactNNLen_, + bool inputsUseNHWC_, + bool useFP16, + bool useNHWC + ) : + usingFP16(useFP16), + nnXLen(context->nnXLen), + nnYLen(context->nnYLen), + requireExactNNLen(requireExactNNLen_), + inputsUseNHWC(inputsUseNHWC_), + usingNHWC(useNHWC) + { + cudaHandles = std::make_unique(majorComputeCapability,minorComputeCapability,context->disableFusedAttention); + model = std::make_unique( + cudaHandles.get(), &(loadedModel->modelDesc), maxBatchSize, + nnXLen, nnYLen, inputsUseNHWC, useFP16, useNHWC + ); + scratch = std::make_unique(maxBatchSize, nnXLen, nnYLen, useFP16); + buffers = std::make_unique(cudaHandles.get(), *model, *scratch); + + //Synchronize after creating buffers and copying all the weights, just in case + CUDA_ERR("ComputeHandle", hipDeviceSynchronize()); + } + ~ComputeHandle() { + } + + ComputeHandle() = delete; + ComputeHandle(const ComputeHandle&) = delete; + ComputeHandle& operator=(const ComputeHandle&) = delete; +}; + +ComputeHandle* NeuralNet::createComputeHandle( + ComputeContext* context, + const LoadedModel* loadedModel, + Logger* logger, + int maxBatchSize, + bool requireExactNNLen, + bool inputsUseNHWC, + int gpuIdxForThisThread, + int serverThreadIdx +) { + //Use whatever CUDA believes GPU 0 to be. + if(gpuIdxForThisThread == -1) + gpuIdxForThisThread = 0; + + CUDA_ERR("createComputeHandle",hipSetDevice(gpuIdxForThisThread)); + + hipDeviceProp_t prop; + hipGetDeviceProperties(&prop,gpuIdxForThisThread); + + bool useFP16 = false; + bool useNHWC = false; + if(context->useFP16Mode == enabled_t::True || context->useFP16Mode == enabled_t::Auto) + useFP16 = true; + if(context->useNHWCMode == enabled_t::True) + useNHWC = true; + + // The transformer block implementation (attention/RoPE/FFN) only supports NHWC, since its channel + // projections assume the channel dim is contiguous per spatial position. Force NHWC for transformer + // models regardless of the useNHWCMode decision above, matching the CUDA backend's behavior. + if(!useNHWC && loadedModel->modelDesc.trunk.hasAnyTransformerBlocks()) { + if(context->useNHWCMode == enabled_t::False) + throw StringError("ROCm backend: transformer models require NHWC, but rocmUseNHWC=false was set"); + useNHWC = true; + } + + if(logger != NULL) { + logger->write( + "ROCm backend thread " + Global::intToString(serverThreadIdx) + ": Found GPU " + string(prop.name) + + " memory " + Global::uint64ToString(prop.totalGlobalMem) + + " compute capability major " + Global::intToString(prop.major) + + " minor " + Global::intToString(prop.minor) + ); + logger->write( + "ROCm backend thread " + Global::intToString(serverThreadIdx) + ": Model version " + Global::intToString(loadedModel->modelDesc.modelVersion) + + " useFP16 = " + Global::boolToString(useFP16) + + " useNHWC = " + Global::boolToString(useNHWC) + ); + logger->write( + "ROCm backend thread " + Global::intToString(serverThreadIdx) + ": Model name: " + loadedModel->modelDesc.name + ); + logger->write( + "MIOpen finding convolution algorithms for GPU " + string(prop.name) + ". This may take a while, please wait......" + ); + } + + ComputeHandle* gpuHandle = new ComputeHandle( + context,loadedModel,prop.major,prop.minor,maxBatchSize,requireExactNNLen,inputsUseNHWC,useFP16,useNHWC + ); + return gpuHandle; +} + +void NeuralNet::freeComputeHandle(ComputeHandle* gpuHandle) { + delete gpuHandle; +} + +bool NeuralNet::isUsingFP16(const ComputeHandle* handle) { + return handle->usingFP16; +} + +bool NeuralNet::setIsWarmup(const ComputeHandle* handle, bool isWarmup) { + (void)handle; + (void)isWarmup; + return false; +} + +//------------------------------------------------------------------------------ + +void NeuralNet::printDevices() { + int numDevices = 0; + hipGetDeviceCount(&numDevices); + for(int i = 0; imodelDesc; + + maxBatchSize = maxBatchSz; + singleInputElts = (size_t)m.numInputChannels * nnXLen * nnYLen; + singleInputBytes = (size_t)m.numInputChannels * nnXLen * nnYLen * sizeof(float); + singleInputGlobalElts = (size_t)m.numInputGlobalChannels; + singleInputGlobalBytes = (size_t)m.numInputGlobalChannels * sizeof(float); + singleInputMetaElts = (size_t)m.numInputMetaChannels; + singleInputMetaBytes = (size_t)m.numInputMetaChannels * sizeof(float); + singlePolicyPassResultElts = (size_t)(m.numPolicyChannels); + singlePolicyPassResultBytes = (size_t)(m.numPolicyChannels) * sizeof(float); + singlePolicyResultElts = (size_t)(m.numPolicyChannels * nnXLen * nnYLen); + singlePolicyResultBytes = (size_t)(m.numPolicyChannels * nnXLen * nnYLen) * sizeof(float); + singleValueResultElts = (size_t)m.numValueChannels; + singleValueResultBytes = (size_t)m.numValueChannels * sizeof(float); + singleScoreValueResultElts = (size_t)m.numScoreValueChannels; + singleScoreValueResultBytes = (size_t)m.numScoreValueChannels * sizeof(float); + singleOwnershipResultElts = (size_t)m.numOwnershipChannels * nnXLen * nnYLen; + singleOwnershipResultBytes = (size_t)m.numOwnershipChannels * nnXLen * nnYLen * sizeof(float); + + assert(NNModelVersion::getNumSpatialFeatures(m.modelVersion) == m.numInputChannels); + assert(NNModelVersion::getNumGlobalFeatures(m.modelVersion) == m.numInputGlobalChannels); + if(m.numInputMetaChannels > 0) { + assert(SGFMetadata::METADATA_INPUT_NUM_CHANNELS == m.numInputMetaChannels); + } + + userInputBufferBytes = (size_t)m.numInputChannels * maxBatchSize * nnXLen * nnYLen * sizeof(float); + userInputGlobalBufferBytes = (size_t)m.numInputGlobalChannels * maxBatchSize * sizeof(float); + userInputMetaBufferBytes = (size_t)m.numInputMetaChannels * maxBatchSize * sizeof(float); + policyPassResultBufferBytes = (size_t)maxBatchSize * m.numPolicyChannels * sizeof(float); + policyResultBufferBytes = (size_t)maxBatchSize * m.numPolicyChannels * nnXLen * nnYLen * sizeof(float); + valueResultBufferBytes = (size_t)maxBatchSize * m.numValueChannels * sizeof(float); + scoreValueResultBufferBytes = (size_t)maxBatchSize * m.numScoreValueChannels * sizeof(float); + ownershipResultBufferBytes = (size_t)maxBatchSize * nnXLen * nnYLen * m.numOwnershipChannels * sizeof(float); + + userInputBuffer = new float[(size_t)m.numInputChannels * maxBatchSize * nnXLen * nnYLen]; + userInputGlobalBuffer = new float[(size_t)m.numInputGlobalChannels * maxBatchSize]; + if(m.numInputMetaChannels > 0) + userInputMetaBuffer = new float[(size_t)m.numInputMetaChannels * maxBatchSize]; + else + userInputMetaBuffer = NULL; + + policyPassResults = new float[(size_t)maxBatchSize * m.numPolicyChannels]; + policyResults = new float[(size_t)maxBatchSize * m.numPolicyChannels * nnXLen * nnYLen]; + valueResults = new float[(size_t)maxBatchSize * m.numValueChannels]; + + scoreValueResults = new float[(size_t)maxBatchSize * m.numScoreValueChannels]; + ownershipResults = new float[(size_t)maxBatchSize * nnXLen * nnYLen * m.numOwnershipChannels]; + } + + ~InputBuffers() { + delete[] userInputBuffer; + delete[] userInputGlobalBuffer; + if(userInputMetaBuffer != NULL) + delete[] userInputMetaBuffer; + delete[] policyPassResults; + delete[] policyResults; + delete[] valueResults; + delete[] scoreValueResults; + delete[] ownershipResults; + } + + InputBuffers() = delete; + InputBuffers(const InputBuffers&) = delete; + InputBuffers& operator=(const InputBuffers&) = delete; + +}; + +InputBuffers* NeuralNet::createInputBuffers(const LoadedModel* loadedModel, int maxBatchSize, int nnXLen, int nnYLen) { + return new InputBuffers(loadedModel,maxBatchSize,nnXLen,nnYLen); +} +void NeuralNet::freeInputBuffers(InputBuffers* inputBuffers) { + delete inputBuffers; +} + +//--------------------------------------------------------------------------------------- + + +void NeuralNet::getOutput( + ComputeHandle* gpuHandle, + InputBuffers* inputBuffers, + int numBatchEltsFilled, + NNResultBuf** inputBufs, + vector& outputs +) { + assert(numBatchEltsFilled <= inputBuffers->maxBatchSize); + assert(numBatchEltsFilled > 0); + const int batchSize = numBatchEltsFilled; + const int nnXLen = gpuHandle->nnXLen; + const int nnYLen = gpuHandle->nnYLen; + const int modelVersion = gpuHandle->model->modelVersion; + + const int numSpatialFeatures = NNModelVersion::getNumSpatialFeatures(modelVersion); + const int numGlobalFeatures = NNModelVersion::getNumGlobalFeatures(modelVersion); + const int numMetaFeatures = inputBuffers->singleInputMetaElts; + assert(numSpatialFeatures == gpuHandle->model->numInputChannels); + assert(numSpatialFeatures * nnXLen * nnYLen == inputBuffers->singleInputElts); + assert(numGlobalFeatures == inputBuffers->singleInputGlobalElts); + const int numPolicyChannels = gpuHandle->model->numPolicyChannels; + + for(int nIdx = 0; nIdxuserInputBuffer + (inputBuffers->singleInputElts * nIdx); + float* rowGlobalInput = inputBuffers->userInputGlobalBuffer + (inputBuffers->singleInputGlobalElts * nIdx); + float* rowMetaInput = inputBuffers->userInputMetaBuffer + (inputBuffers->singleInputMetaElts * nIdx); + + const float* rowGlobal = inputBufs[nIdx]->rowGlobalBuf.data(); + const float* rowSpatial = inputBufs[nIdx]->rowSpatialBuf.data(); + const float* rowMeta = inputBufs[nIdx]->rowMetaBuf.data(); + bool hasRowMeta = inputBufs[nIdx]->hasRowMeta; + std::copy(rowGlobal,rowGlobal+numGlobalFeatures,rowGlobalInput); + if(numMetaFeatures > 0) { + testAssert(rowMeta != NULL); + testAssert(hasRowMeta); + std::copy(rowMeta,rowMeta+numMetaFeatures,rowMetaInput); + } + else { + testAssert(!hasRowMeta); + } + SymmetryHelpers::copyInputsWithSymmetry(rowSpatial, rowSpatialInput, 1, nnYLen, nnXLen, numSpatialFeatures, gpuHandle->inputsUseNHWC, inputBufs[nIdx]->symmetry); + } + + Buffers* buffers = gpuHandle->buffers.get(); + ScratchBuffers* scratch = gpuHandle->scratch.get(); + + if(!gpuHandle->usingFP16) { + assert(inputBuffers->userInputBufferBytes == buffers->inputBufBytes); + assert(inputBuffers->userInputGlobalBufferBytes == buffers->inputGlobalBufBytes); + assert(inputBuffers->userInputMetaBufferBytes == buffers->inputMetaBufBytes); + assert(inputBuffers->policyPassResultBufferBytes == buffers->policyPassBufBytes); + assert(inputBuffers->policyResultBufferBytes == buffers->policyBufBytes); + assert(inputBuffers->valueResultBufferBytes == buffers->valueBufBytes); + assert(inputBuffers->singleInputBytes == inputBuffers->singleInputElts*4); + assert(inputBuffers->singleInputGlobalBytes == inputBuffers->singleInputGlobalElts*4); + assert(inputBuffers->singleInputMetaBytes == inputBuffers->singleInputMetaElts*4); + assert(inputBuffers->singlePolicyPassResultElts == numPolicyChannels); + assert(inputBuffers->singlePolicyPassResultBytes == numPolicyChannels * sizeof(float)); + assert(inputBuffers->singlePolicyResultElts == numPolicyChannels*nnXLen*nnYLen); + assert(inputBuffers->singlePolicyResultBytes == numPolicyChannels*nnXLen*nnYLen * sizeof(float)); + assert(inputBuffers->scoreValueResultBufferBytes == buffers->scoreValueBufBytes); + assert(inputBuffers->ownershipResultBufferBytes == buffers->ownershipBufBytes); + assert(inputBuffers->singleOwnershipResultElts == nnXLen*nnYLen); + assert(inputBuffers->singleOwnershipResultBytes == nnXLen*nnYLen * sizeof(float)); + + CUDA_ERR("getOutput",hipMemcpy(buffers->inputBuf, inputBuffers->userInputBuffer, inputBuffers->singleInputBytes*batchSize, hipMemcpyHostToDevice)); + CUDA_ERR("getOutput",hipMemcpy(buffers->inputGlobalBuf, inputBuffers->userInputGlobalBuffer, inputBuffers->singleInputGlobalBytes*batchSize, hipMemcpyHostToDevice)); + if(numMetaFeatures > 0) { + CUDA_ERR("getOutput",hipMemcpy(buffers->inputMetaBuf, inputBuffers->userInputMetaBuffer, inputBuffers->singleInputMetaBytes*batchSize, hipMemcpyHostToDevice)); + } + } + else { + assert(inputBuffers->userInputBufferBytes == buffers->inputBufBytesFloat); + assert(inputBuffers->userInputGlobalBufferBytes == buffers->inputGlobalBufBytesFloat); + assert(inputBuffers->userInputMetaBufferBytes == buffers->inputMetaBufBytesFloat); + assert(inputBuffers->policyResultBufferBytes == buffers->policyBufBytes); + assert(inputBuffers->valueResultBufferBytes == buffers->valueBufBytes); + assert(inputBuffers->userInputBufferBytes == buffers->inputBufBytes*2); + assert(inputBuffers->userInputGlobalBufferBytes == buffers->inputGlobalBufBytes*2); + assert(inputBuffers->userInputMetaBufferBytes == buffers->inputMetaBufBytes*2); + assert(inputBuffers->singleInputBytes == inputBuffers->singleInputElts*4); + assert(inputBuffers->singleInputGlobalBytes == inputBuffers->singleInputGlobalElts*4); + assert(inputBuffers->singleInputMetaBytes == inputBuffers->singleInputMetaElts*4); + assert(inputBuffers->singlePolicyPassResultElts == numPolicyChannels); + assert(inputBuffers->singlePolicyPassResultBytes == numPolicyChannels * sizeof(float)); + assert(inputBuffers->singlePolicyResultElts == numPolicyChannels*nnXLen*nnYLen); + assert(inputBuffers->singlePolicyResultBytes == numPolicyChannels*nnXLen*nnYLen * sizeof(float)); + assert(inputBuffers->scoreValueResultBufferBytes == buffers->scoreValueBufBytes); + assert(inputBuffers->ownershipResultBufferBytes == buffers->ownershipBufBytes); + assert(inputBuffers->singleOwnershipResultElts == nnXLen*nnYLen); + assert(inputBuffers->singleOwnershipResultBytes == nnXLen*nnYLen * sizeof(float)); + + CUDA_ERR("getOutput",hipMemcpy(buffers->inputBufFloat, inputBuffers->userInputBuffer, inputBuffers->singleInputBytes*batchSize, hipMemcpyHostToDevice)); + CUDA_ERR("getOutput",hipMemcpy(buffers->inputGlobalBufFloat, inputBuffers->userInputGlobalBuffer, inputBuffers->singleInputGlobalBytes*batchSize, hipMemcpyHostToDevice)); + if(numMetaFeatures > 0) { + CUDA_ERR("getOutput",hipMemcpy(buffers->inputMetaBufFloat, inputBuffers->userInputMetaBuffer, inputBuffers->singleInputMetaBytes*batchSize, hipMemcpyHostToDevice)); + } + + customCudaCopyToHalf((const float*)buffers->inputBufFloat,(half*)buffers->inputBuf,inputBuffers->singleInputElts*batchSize); + CUDA_ERR("getOutput",hipPeekAtLastError()); + customCudaCopyToHalf((const float*)buffers->inputGlobalBufFloat,(half*)buffers->inputGlobalBuf,inputBuffers->singleInputGlobalElts*batchSize); + CUDA_ERR("getOutput",hipPeekAtLastError()); + if(numMetaFeatures > 0) { + customCudaCopyToHalf((const float*)buffers->inputMetaBufFloat,(half*)buffers->inputMetaBuf,inputBuffers->singleInputMetaElts*batchSize); + CUDA_ERR("getOutput",hipPeekAtLastError()); + } + } + + gpuHandle->model->apply( + gpuHandle->cudaHandles.get(), + scratch, + batchSize, + gpuHandle->requireExactNNLen, + + buffers->inputBuf, + buffers->inputGlobalBuf, + buffers->inputMetaBuf, + + buffers->policyPassBuf, + buffers->policyBuf, + + buffers->valueBuf, + buffers->scoreValueBuf, + buffers->ownershipBuf, + + buffers->workspaceBuf, + buffers->workspaceBytes + ); + + CUDA_ERR("getOutput",hipMemcpy(inputBuffers->policyPassResults, buffers->policyPassBuf, inputBuffers->singlePolicyPassResultBytes*batchSize, hipMemcpyDeviceToHost)); + CUDA_ERR("getOutput",hipMemcpy(inputBuffers->policyResults, buffers->policyBuf, inputBuffers->singlePolicyResultBytes*batchSize, hipMemcpyDeviceToHost)); + CUDA_ERR("getOutput",hipMemcpy(inputBuffers->valueResults, buffers->valueBuf, inputBuffers->singleValueResultBytes*batchSize, hipMemcpyDeviceToHost)); + CUDA_ERR("getOutput",hipMemcpy(inputBuffers->scoreValueResults, buffers->scoreValueBuf, inputBuffers->singleScoreValueResultBytes*batchSize, hipMemcpyDeviceToHost)); + CUDA_ERR("getOutput",hipMemcpy(inputBuffers->ownershipResults, buffers->ownershipBuf, inputBuffers->singleOwnershipResultBytes*batchSize, hipMemcpyDeviceToHost)); + + assert(outputs.size() == batchSize); + + float policyProbsTmp[NNPos::MAX_NN_POLICY_SIZE]; + + for(int row = 0; row < batchSize; row++) { + NNOutput* output = outputs[row]; + assert(output->nnXLen == nnXLen); + assert(output->nnYLen == nnYLen); + float policyOptimism = (float)inputBufs[row]->policyOptimism; + + const float* policyPassSrcBuf = inputBuffers->policyPassResults + row * numPolicyChannels; + const float* policySrcBuf = inputBuffers->policyResults + row * numPolicyChannels * nnXLen * nnYLen; + float* policyProbs = output->policyProbs; + + // These are in logits, the client does the postprocessing to turn them into + // policy probabilities and white game outcome probabilities + // Also we don't fill in the nnHash here either + // Handle version >= 12 policy optimism + if(numPolicyChannels == 2 || (numPolicyChannels == 4 && modelVersion >= 16)) { + if(gpuHandle->usingNHWC) { + for(int i = 0; isymmetry); + policyProbs[nnXLen*nnYLen] = policyPassSrcBuf[0] + (policyPassSrcBuf[1] - policyPassSrcBuf[0]) * policyOptimism; + } + else { + for(int i = 0; isymmetry); + policyProbs[nnXLen*nnYLen] = policyPassSrcBuf[0] + (policyPassSrcBuf[1] - policyPassSrcBuf[0]) * policyOptimism; + } + } + else { + assert(numPolicyChannels == 1); + SymmetryHelpers::copyOutputsWithSymmetry(policySrcBuf, policyProbs, 1, nnYLen, nnXLen, inputBufs[row]->symmetry); + policyProbs[nnXLen*nnYLen] = policyPassSrcBuf[0]; + } + + int numValueChannels = gpuHandle->model->numValueChannels; + assert(numValueChannels == 3); + output->whiteWinProb = inputBuffers->valueResults[row * numValueChannels]; + output->whiteLossProb = inputBuffers->valueResults[row * numValueChannels + 1]; + output->whiteNoResultProb = inputBuffers->valueResults[row * numValueChannels + 2]; + + //As above, these are NOT actually from white's perspective, but rather the player to move. + //As usual the client does the postprocessing. + if(output->whiteOwnerMap != NULL) { + const float* ownershipSrcBuf = inputBuffers->ownershipResults + row * nnXLen * nnYLen; + assert(gpuHandle->model->numOwnershipChannels == 1); + SymmetryHelpers::copyOutputsWithSymmetry(ownershipSrcBuf, output->whiteOwnerMap, 1, nnYLen, nnXLen, inputBufs[row]->symmetry); + } + + if(modelVersion >= 9) { + int numScoreValueChannels = gpuHandle->model->numScoreValueChannels; + assert(numScoreValueChannels == 6); + output->whiteScoreMean = inputBuffers->scoreValueResults[row * numScoreValueChannels]; + output->whiteScoreMeanSq = inputBuffers->scoreValueResults[row * numScoreValueChannels + 1]; + output->whiteLead = inputBuffers->scoreValueResults[row * numScoreValueChannels + 2]; + output->varTimeLeft = inputBuffers->scoreValueResults[row * numScoreValueChannels + 3]; + output->shorttermWinlossError = inputBuffers->scoreValueResults[row * numScoreValueChannels + 4]; + output->shorttermScoreError = inputBuffers->scoreValueResults[row * numScoreValueChannels + 5]; + } + else if(modelVersion >= 8) { + int numScoreValueChannels = gpuHandle->model->numScoreValueChannels; + assert(numScoreValueChannels == 4); + output->whiteScoreMean = inputBuffers->scoreValueResults[row * numScoreValueChannels]; + output->whiteScoreMeanSq = inputBuffers->scoreValueResults[row * numScoreValueChannels + 1]; + output->whiteLead = inputBuffers->scoreValueResults[row * numScoreValueChannels + 2]; + output->varTimeLeft = inputBuffers->scoreValueResults[row * numScoreValueChannels + 3]; + output->shorttermWinlossError = 0; + output->shorttermScoreError = 0; + } + else if(modelVersion >= 4) { + int numScoreValueChannels = gpuHandle->model->numScoreValueChannels; + assert(numScoreValueChannels == 2); + output->whiteScoreMean = inputBuffers->scoreValueResults[row * numScoreValueChannels]; + output->whiteScoreMeanSq = inputBuffers->scoreValueResults[row * numScoreValueChannels + 1]; + output->whiteLead = output->whiteScoreMean; + output->varTimeLeft = 0; + output->shorttermWinlossError = 0; + output->shorttermScoreError = 0; + } + else if(modelVersion >= 3) { + int numScoreValueChannels = gpuHandle->model->numScoreValueChannels; + assert(numScoreValueChannels == 1); + output->whiteScoreMean = inputBuffers->scoreValueResults[row * numScoreValueChannels]; + //Version 3 neural nets don't have any second moment output, implicitly already folding it in, so we just use the mean squared + output->whiteScoreMeanSq = output->whiteScoreMean * output->whiteScoreMean; + output->whiteLead = output->whiteScoreMean; + output->varTimeLeft = 0; + output->shorttermWinlossError = 0; + output->shorttermScoreError = 0; + } + else { + ASSERT_UNREACHABLE; + } + } + +} + +//TESTING ---------------------------------------------------------------------------------- + + +bool NeuralNet::testEvaluateConv( + const ConvLayerDesc* desc, + int desiredBatchSize, + int nnXLen, + int nnYLen, + bool useFP16, + bool useNHWC, + const vector& inputBuffer, + vector& outputBuffer +) { + hipDeviceSynchronize(); + CudaHandles* cudaHandles = CudaHandles::cudaHandlesTesting(); + + size_t numInputFloats = (size_t)desiredBatchSize * nnXLen * nnYLen * desc->inChannels; + size_t numOutputFloats = (size_t)desiredBatchSize * nnXLen * nnYLen * desc->outChannels; + if(numInputFloats != inputBuffer.size()) + throw StringError("testEvaluateConv: unexpected input buffer size"); + + void* deviceInput; + void* deviceOutput; + CudaUtils::mallocAndCopyToDevice("deviceInput", inputBuffer.data(), numInputFloats, deviceInput, useFP16); + CudaUtils::mallocOnDevice("deviceOutput", numOutputFloats, deviceOutput, useFP16); + + int maxBatchSize = desiredBatchSize; + + CudnnManager* manager = new CudnnManager("manager",maxBatchSize,nnXLen,nnYLen); + ConvLayer* convLayer = new ConvLayer(cudaHandles,manager,desc,useFP16,useNHWC); + + size_t workspaceBytes = + convLayer->requiredWorkspaceBytes(cudaHandles,desiredBatchSize); + void* deviceWorkspace; + CUDA_ERR("deviceWorkspace",hipMalloc(&deviceWorkspace, workspaceBytes)); + + + bool accumulate = false; + convLayer->apply( + cudaHandles, + desiredBatchSize, + accumulate, + deviceInput, + deviceOutput, + deviceWorkspace, + workspaceBytes + ); + + outputBuffer.resize(numOutputFloats); + CudaUtils::expensiveCopyFromDevice("copyResultsToHost", outputBuffer.data(), numOutputFloats, deviceOutput, useFP16); + + hipFree(deviceWorkspace); + + delete convLayer; + delete manager; + hipFree(deviceInput); + hipFree(deviceOutput); + delete cudaHandles; + + return true; +} + + +bool NeuralNet::testEvaluateBatchNorm( + const BatchNormLayerDesc* desc, + int desiredBatchSize, + int nnXLen, + int nnYLen, + bool useFP16, + bool useNHWC, + const vector& inputBuffer, + const vector& maskBuffer, + vector& outputBuffer +) { + hipDeviceSynchronize(); + CudaHandles* cudaHandles = CudaHandles::cudaHandlesTesting(); + + size_t numInputFloats = (size_t)desiredBatchSize * nnXLen * nnYLen * desc->numChannels; + size_t numMaskFloats = (size_t)desiredBatchSize * nnXLen * nnYLen; + size_t numOutputFloats = (size_t)desiredBatchSize * nnXLen * nnYLen * desc->numChannels; + if(numInputFloats != inputBuffer.size()) + throw StringError("testEvaluateBatchNorm: unexpected input buffer size"); + if(numMaskFloats != maskBuffer.size()) + throw StringError("testEvaluateBatchNorm: unexpected mask buffer size"); + + ActivationLayerDesc actDesc; + actDesc.activation = ACTIVATION_IDENTITY; + + void* deviceInput; + void* deviceMask; + void* deviceOutput; + CudaUtils::mallocAndCopyToDevice("deviceInput", inputBuffer.data(), numInputFloats, deviceInput, useFP16); + CudaUtils::mallocAndCopyToDevice("deviceMask", maskBuffer.data(), numMaskFloats, deviceMask, useFP16); + CudaUtils::mallocOnDevice("deviceOutput", numOutputFloats, deviceOutput, useFP16); + + BatchNormLayer* batchNormLayer = new BatchNormLayer(cudaHandles,desc,&actDesc,nnXLen,nnYLen,useFP16,useNHWC); + + batchNormLayer->apply( + cudaHandles, + desiredBatchSize, + deviceInput, + deviceMask, + deviceOutput + ); + + outputBuffer.resize(numOutputFloats); + CudaUtils::expensiveCopyFromDevice("copyResultsToHost", outputBuffer.data(), numOutputFloats, deviceOutput, useFP16); + + delete batchNormLayer; + + hipFree(deviceInput); + hipFree(deviceMask); + hipFree(deviceOutput); + delete cudaHandles; + + return true; +} + + +bool NeuralNet::testEvaluateResidualBlock( + const ResidualBlockDesc* desc, + int desiredBatchSize, + int nnXLen, + int nnYLen, + bool useFP16, + bool useNHWC, + const vector& inputBuffer, + const vector& maskBuffer, + vector& outputBuffer +) { + hipDeviceSynchronize(); + CudaHandles* cudaHandles = CudaHandles::cudaHandlesTesting(); + + size_t numInputFloats = (size_t)desiredBatchSize * nnXLen * nnYLen * desc->preBN.numChannels; + size_t numMaskFloats = (size_t)desiredBatchSize * nnXLen * nnYLen; + size_t numOutputFloats = (size_t)desiredBatchSize * nnXLen * nnYLen * desc->finalConv.outChannels; + if(numInputFloats != inputBuffer.size()) + throw StringError("testEvaluateResidualBlock: unexpected input buffer size"); + if(numMaskFloats != maskBuffer.size()) + throw StringError("testEvaluateResidualBlock: unexpected mask buffer size"); + + ScratchBuffers* scratch = new ScratchBuffers(desiredBatchSize, nnXLen, nnYLen, useFP16); + + void* deviceInput; + void* deviceMask; + void* deviceScratch; + CudaUtils::mallocAndCopyToDevice("deviceInput", inputBuffer.data(), numInputFloats, deviceInput, useFP16); + CudaUtils::mallocAndCopyToDevice("deviceMask", maskBuffer.data(), numMaskFloats, deviceMask, useFP16); + CudaUtils::mallocOnDevice("deviceScratch", numInputFloats, deviceScratch, useFP16); + + int maxBatchSize = desiredBatchSize; + + CudnnManager* manager = new CudnnManager("manager",maxBatchSize,nnXLen,nnYLen); + ResidualBlock* residualBlock = new ResidualBlock(cudaHandles,manager,desc,nnXLen,nnYLen,useFP16,useNHWC); + + size_t workspaceBytes = + residualBlock->requiredWorkspaceBytes(cudaHandles,desiredBatchSize); + void* deviceWorkspace; + CUDA_ERR("deviceWorkspace",hipMalloc(&deviceWorkspace, workspaceBytes)); + + residualBlock->apply( + cudaHandles, + scratch, + desiredBatchSize, + deviceInput, + deviceScratch, + deviceMask, + deviceWorkspace, + workspaceBytes + ); + + outputBuffer.resize(numOutputFloats); + CudaUtils::expensiveCopyFromDevice("copyResultsToHost", outputBuffer.data(), numOutputFloats, deviceInput, useFP16); + + hipFree(deviceWorkspace); + + delete residualBlock; + delete manager; + hipFree(deviceInput); + hipFree(deviceMask); + hipFree(deviceScratch); + delete scratch; + delete cudaHandles; + + return true; +} + +bool NeuralNet::testEvaluateGlobalPoolingResidualBlock( + const GlobalPoolingResidualBlockDesc* desc, + int desiredBatchSize, + int nnXLen, + int nnYLen, + bool useFP16, + bool useNHWC, + const vector& inputBuffer, + const vector& maskBuffer, + vector& outputBuffer +) { + hipDeviceSynchronize(); + CudaHandles* cudaHandles = CudaHandles::cudaHandlesTesting(); + + size_t numInputFloats = (size_t)desiredBatchSize * nnXLen * nnYLen * desc->preBN.numChannels; + size_t numMaskFloats = (size_t)desiredBatchSize * nnXLen * nnYLen; + size_t numMaskSumFloats = (size_t)desiredBatchSize; + size_t numOutputFloats = (size_t)desiredBatchSize * nnXLen * nnYLen * desc->finalConv.outChannels; + + if(numInputFloats != inputBuffer.size()) + throw StringError("testEvaluateGlobalPoolingResidualBlock: unexpected input buffer size"); + if(numMaskFloats != maskBuffer.size()) + throw StringError("testEvaluateGlobalPoolingResidualBlock: unexpected mask buffer size"); + + ScratchBuffers* scratch = new ScratchBuffers(desiredBatchSize, nnXLen, nnYLen, useFP16); + + void* deviceInput; + void* deviceMask; + float* deviceMaskFloatOrig; + float* deviceMaskFloat; + float* deviceMaskSum; + void* deviceScratch; + + CudaUtils::mallocAndCopyToDevice("deviceInput", inputBuffer.data(), numInputFloats, deviceInput, useFP16); + CudaUtils::mallocAndCopyToDevice("deviceMask", maskBuffer.data(), numMaskFloats, deviceMask, useFP16); + CUDA_ERR("deviceMaskFloat",hipMalloc(reinterpret_cast(&deviceMaskFloat), numMaskFloats * sizeof(float))); + CUDA_ERR("deviceMaskSum",hipMalloc(reinterpret_cast(&deviceMaskSum), numMaskSumFloats * sizeof(float))); + deviceMaskFloatOrig = deviceMaskFloat; + CudaUtils::mallocOnDevice("deviceScratch", numInputFloats, deviceScratch, useFP16); + + fillMaskFloatBufAndMaskSumBuf(deviceMask, deviceMaskFloat, deviceMaskSum, useFP16, desiredBatchSize, nnXLen, nnYLen); + + int maxBatchSize = desiredBatchSize; + + CudnnManager* manager = new CudnnManager("manager",maxBatchSize,nnXLen,nnYLen); + GlobalPoolingResidualBlock* residualBlock = new GlobalPoolingResidualBlock( + cudaHandles,manager,desc,nnXLen,nnYLen,useFP16,useNHWC + ); + + size_t workspaceBytes = + residualBlock->requiredWorkspaceBytes( + cudaHandles,desiredBatchSize + ); + + void* deviceWorkspace; + CUDA_ERR("deviceWorkspace",hipMalloc(&deviceWorkspace, workspaceBytes)); + + residualBlock->apply( + cudaHandles, + scratch, + desiredBatchSize, + deviceInput, + deviceScratch, + deviceMask, + deviceMaskSum, + deviceWorkspace, + workspaceBytes + ); + + outputBuffer.resize(numOutputFloats); + CudaUtils::expensiveCopyFromDevice("copyResultsToHost", outputBuffer.data(), numOutputFloats, deviceInput, useFP16); + + hipFree(deviceWorkspace); + + delete residualBlock; + delete manager; + + hipFree(deviceInput); + hipFree(deviceMask); + hipFree(deviceMaskFloatOrig); + hipFree(deviceMaskSum); + hipFree(deviceScratch); + delete scratch; + delete cudaHandles; + + return true; +} + + +#endif // USE_ROCM_BACKEND diff --git a/cpp/neuralnet/rocmerrorcheck.h b/cpp/neuralnet/rocmerrorcheck.h new file mode 100644 index 0000000000..8cb214ad7e --- /dev/null +++ b/cpp/neuralnet/rocmerrorcheck.h @@ -0,0 +1,59 @@ +#ifndef NEURALNET_ROCMERRORCHECK_H_ +#define NEURALNET_ROCMERRORCHECK_H_ + +#include "../neuralnet/rocmincludes.h" +#include "../core/global.h" + +// ---------- HIP runtime ---------- +static inline void checkCudaError(hipError_t status, + const char* opName, + const char* file, + const char* func, + int line) { + if(status != hipSuccess) + throw StringError(std::string("HIP Error @") + opName + " " + + file + ":" + func + ":" + Global::intToString(line) + + " : " + hipGetErrorString(status)); +} +#define CUDA_ERR(opName,x) checkCudaError((x),opName,__FILE__,#x,__LINE__) + +// ---------- hipBLAS ---------- +static inline const char* cublasGetErrorString(hipblasStatus_t s) { + switch(s) { + case HIPBLAS_STATUS_SUCCESS: return "HIPBLAS_STATUS_SUCCESS"; + case HIPBLAS_STATUS_ALLOC_FAILED: return "HIPBLAS_STATUS_ALLOC_FAILED"; + case HIPBLAS_STATUS_MAPPING_ERROR: return "HIPBLAS_STATUS_MAPPING_ERROR"; + case HIPBLAS_STATUS_EXECUTION_FAILED: return "HIPBLAS_STATUS_EXECUTION_FAILED"; + case HIPBLAS_STATUS_INTERNAL_ERROR: return "HIPBLAS_STATUS_INTERNAL_ERROR"; + case HIPBLAS_STATUS_INVALID_VALUE: return "HIPBLAS_STATUS_INVALID_VALUE"; + case HIPBLAS_STATUS_NOT_INITIALIZED: return "HIPBLAS_STATUS_NOT_INITIALIZED"; + case HIPBLAS_STATUS_NOT_SUPPORTED: return "HIPBLAS_STATUS_NOT_SUPPORTED"; + default: return "HIPBLAS_STATUS_UNKNOWN"; + } +} +static inline void checkCublasError(hipblasStatus_t status, + const char* opName, + const char* file, + const char* func, + int line) { + if(status != HIPBLAS_STATUS_SUCCESS) + throw StringError(std::string("hipBLAS Error @") + opName + " " + + file + ":" + func + ":" + Global::intToString(line) + + " : " + cublasGetErrorString(status)); +} +#define CUBLAS_ERR(opName,x) checkCublasError((x),opName,__FILE__,#x,__LINE__) + +// ---------- MIOpen ---------- +static inline void checkCudnnError(miopenStatus_t status, + const char* opName, + const char* file, + const char* func, + int line) { + if(status != miopenStatusSuccess) + throw StringError(std::string("MIOpen Error @") + opName + " " + + file + ":" + func + ":" + Global::intToString(line) + + " : " + miopenGetErrorString(status)); +} +#define CUDNN_ERR(opName,x) checkCudnnError((x),opName,__FILE__,#x,__LINE__) + +#endif // NEURALNET_ROCMERRORCHECK_H_ diff --git a/cpp/neuralnet/rocmhelpers.h b/cpp/neuralnet/rocmhelpers.h new file mode 100644 index 0000000000..086be892c6 --- /dev/null +++ b/cpp/neuralnet/rocmhelpers.h @@ -0,0 +1,156 @@ +#ifndef NEURALNET_ROCMHELPERS_H_ +#define NEURALNET_ROCMHELPERS_H_ + +#include "../neuralnet/rocmincludes.h" +#include "../neuralnet/activations.h" + +//Given two tensors with shapes inA: [n,cA,h,w] and inB: [n,cB,h,w], that are on the GPU +//Copy them into a single tensor out: [n,cA+cB,h,w] that is also allocated on the gpu +void customCudaChannelConcat(const float* inA, const float* inB, float* out, int chwA, int chwB, int n); +void customCudaChannelConcat(const half* inA, const half* inB, half* out, int chwA, int chwB, int n); + +//Given a tensor [n,c,hw], extract out channel 0 to [n,hw] +void customCudaChannel0ExtractNCHW(const float* in, float* out, int n, int c, int hw); +void customCudaChannel0ExtractNCHW(const half* in, half* out, int n, int c, int hw); +//Given a tensor [n,hw,c], extract out channel 0 to [n,hw] +void customCudaChannel0ExtractNHWC(const float* in, float* out, int n, int hw, int c); +void customCudaChannel0ExtractNHWC(const half* in, half* out, int n, int hw, int c); + +//Given an input tensor and an output buffer of shape [n,c], fill output buffer with sum or max over c. +void customCudaPoolRowsSumNCHW(const float* in, float* out, int nSize, int cSize, int xySize, float scaleSum); +void customCudaPoolRowsSumNHWC(const float* in, float* out, int nSize, int xySize, int cSize, float scaleSum); + +//Specialized operations for value head and general global pooling. Same as the other pooling, but fusedly fills +//an output buffer of shape [n,c*3]. +void customCudaValueHeadPoolNCHW(const float* in, float* out, int nSize, int cSize, int xySize, const float* maskSum); +void customCudaValueHeadPoolNHWC(const float* in, float* out, int nSize, int xySize, int cSize, const float* maskSum); +void customCudaPoolRowsGPoolNCHW(const float* in, float* out, int nSize, int cSize, int xySize, const float* mask, const float* maskSum); +void customCudaPoolRowsGPoolNHWC(const float* in, float* out, int nSize, int xySize, int cSize, const float* mask, const float* maskSum); +void customCudaPoolRowsGPoolNCHW(const half* in, half* out, int nSize, int cSize, int xySize, const half* mask, const float* maskSum); +void customCudaPoolRowsGPoolNHWC(const half* in, half* out, int nSize, int xySize, int cSize, const half* mask, const float* maskSum); + +void customCudaCopyToHalf(const float* in, half* out, int n); +void customCudaCopyFromHalf(const half* in, float* out, int n); + +//Given a tensor, add another tensor element-wise to it (same shape). +void customCudaAddTensorsInplace(float* buf, const float* toAdd, int n); +void customCudaAddTensorsInplace(half* buf, const half* toAdd, int n); +//Given a tensor, add another tensor to it. +void customCudaAddTensorInplace(half* buf, const half* biases, int n); +//Given an input with shape [n,c] and biases of shape [c], add the biases in-place. +void customCudaAddCBiasInplaceNC(float* buf, const float* biases, int n, int c, int activation); +void customCudaAddCBiasInplaceNC(half* buf, const half* biases, int n, int c, int activation); +//Given an input with shape [n,c,xy] and biases of shape [n,c], add the biases in-place. +void customCudaAddNCBiasInplaceNCHW(float *buf, const float* biases, int nSize, int cSize, int xySize); +void customCudaAddNCBiasInplaceNCHW(half *buf, const half* biases, int nSize, int cSize, int xySize); +//Given an input with shape [n,xy,c] and biases of shape [n,c], add the biases in-place. +void customCudaAddNCBiasInplaceNHWC(float *buf, const float* biases, int nSize, int xySize, int cSize); +void customCudaAddNCBiasInplaceNHWC(half *buf, const half* biases, int nSize, int xySize, int cSize); + +//Given an input with shape [n,c,xy] and scale and biases of shape [c], multiply by scale and add the biases +//Optionally also apply an activation. +//Optionally also multiply by mask (can be null), with shape [n,xy] +void customCudaApplyCScaleBiasNCHW(const float* in, float* out, const float* scale, const float* biases, const float* mask, int n, int c, int xy, int activation); +void customCudaApplyCScaleBiasNCHW(const half* in, half* out, const half* scale, const half* biases, const half* mask, int n, int c, int xy, int activation); +//Given an input with shape [n,xy,c] and scale and biases of shape [c], multiply by scale and add the biases +//Optionally also apply relu. +//Optionally also multiply by mask (can be null), with shape [n,xy] +void customCudaApplyCScaleBiasNHWC(const float* in, float* out, const float* scale, const float* biases, const float* mask, int n, int xy, int c, int activation); +void customCudaApplyCScaleBiasNHWC(const half* in, half* out, const half* scale, const half* biases, const half* mask, int n, int xy, int c, int activation); + +//============================================================================================== +// Transformer support kernels +//============================================================================================== + +//Apply rotary position embeddings in-place to a BSHD-laid-out Q or K buffer. +//buf: [totalDim, seqLen*batchSize] column-major (totalDim = numBufHeads*qHeadDim). +//Fixed RoPE: cosTable/sinTable are (numPairs, seqLen) if !learnableRope, else (numKVHeads, numPairs, seqLen). +void customCudaApplyRoPE( + float* buf, const float* cosTable, const float* sinTable, + int batchSize, int seqLen, int numBufHeads, int numKVHeads, int qHeadDim, int numPairs, bool learnableRope +); +void customCudaApplyRoPE( + half* buf, const half* cosTable, const half* sinTable, + int batchSize, int seqLen, int numBufHeads, int numKVHeads, int qHeadDim, int numPairs, bool learnableRope +); +//Table-free learnable RoPE: recomputes cos/sin in-kernel from per-head frequencies (numKVHeads, numPairs, 2) flattened. +void customCudaApplyRoPELearnableRecompute( + float* buf, const float* freqs, + int batchSize, int seqLen, int numBufHeads, int numKVHeads, int qHeadDim, int numPairs, int nnXLen +); +void customCudaApplyRoPELearnableRecompute( + half* buf, const float* freqs, + int batchSize, int seqLen, int numBufHeads, int numKVHeads, int qHeadDim, int numPairs, int nnXLen +); + +//Scaled dot product attention (online-softmax, tiled). Q/K/V/output are BSHD row-major. +//mask (can be null) is [batchSize, seqLen]. +void customCudaFlashAttention( + const float* Q, const float* K, const float* V, const float* mask, float* output, + int batchSize, int seqLen, int numHeads, int numKVHeads, int qHeadDim, int vHeadDim +); +void customCudaFlashAttention( + const half* Q, const half* K, const half* V, const half* mask, half* output, + int batchSize, int seqLen, int numHeads, int numKVHeads, int qHeadDim, int vHeadDim +); + +//SwiGLU: out[i] = SiLU(a[i]) * b[i] +void customCudaSwiGLU(const float* a, const float* b, float* out, int size); +void customCudaSwiGLU(const half* a, const half* b, half* out, int size); + +//Convert mask [batchSize, seqLen] (0/1) into a fully-materialized additive attention bias of shape +//[batchSize, seqLen, seqLen]: bias[b,q,k] = (mask[b,k] != 0 ? 0 : -3e4). Used to feed KataGo's mask +//into the CK FMHA fused attention path, which (unlike the plain kernel) needs the bias +//pre-materialized rather than taking the raw per-position mask directly. +//See the comment in rocmhelpers.hip for why this bias value for the mask. +void customCudaMaskToAttnBiasFull(const float* mask, float* outBias, int batchSize, int seqLen); +void customCudaMaskToAttnBiasFull(const half* mask, half* outBias, int batchSize, int seqLen); + +//Masked residual add: trunk[i] += residual[i] * mask[spatial_idx]. mask can be null (treated as all ones). +//NCHW: trunk/residual [n,c,xy], mask [n,xy]. NHWC: trunk/residual [n,xy,c], mask [n,xy]. +void customCudaMaskedResidualAddNCHW(float* trunk, const float* residual, const float* mask, int nSize, int cSize, int xySize); +void customCudaMaskedResidualAddNCHW(half* trunk, const half* residual, const half* mask, int nSize, int cSize, int xySize); +void customCudaMaskedResidualAddNHWC(float* trunk, const float* residual, const float* mask, int nSize, int xySize, int cSize); +void customCudaMaskedResidualAddNHWC(half* trunk, const half* residual, const half* mask, int nSize, int xySize, int cSize); + +//RMSNorm with gamma/beta/activation, non-spatial mode (for transformer pre-norm and trunk tip). +//NHWC: input/output [n,xy,c], gamma/beta [c], mask [n,xy] (can be null). +void customCudaRMSNormGammaBetaNHWC( + const float* in, float* out, const float* gamma, const float* beta, const float* mask, + int nSize, int xySize, int cSize, float epsilon, int activation +); +void customCudaRMSNormGammaBetaNHWC( + const half* in, half* out, const half* gamma, const half* beta, const half* mask, + int nSize, int xySize, int cSize, float epsilon, int activation +); +//NCHW: input/output [n,c,xy], gamma/beta [c], mask [n,xy] (can be null). +void customCudaRMSNormGammaBetaNCHW( + const float* in, float* out, const float* gamma, const float* beta, const float* mask, + int nSize, int cSize, int xySize, float epsilon, int activation +); +void customCudaRMSNormGammaBetaNCHW( + const half* in, half* out, const half* gamma, const half* beta, const half* mask, + int nSize, int cSize, int xySize, float epsilon, int activation +); + +//sumSqBuf must be a scratch buffer of size nSize * CUDA_SPATIAL_RMSNORM_SUMSQ_STRIDE floats. +//Spatial RMSNorm: normalizes over all C*H*W per batch element (rather than per-position over C). +#define CUDA_SPATIAL_RMSNORM_SUMSQ_STRIDE 9 // SPATIAL_RMSNORM_BLOCKS_PER_BATCH (8) partials + 1 final +void customCudaSpatialRMSNormNHWC( + const float* in, float* out, const float* gamma, const float* beta, const float* mask, const float* maskSum, + int nSize, int xySize, int cSize, float epsilon, int activation, float* sumSqBuf +); +void customCudaSpatialRMSNormNHWC( + const half* in, half* out, const half* gamma, const half* beta, const half* mask, const float* maskSum, + int nSize, int xySize, int cSize, float epsilon, int activation, float* sumSqBuf +); +void customCudaSpatialRMSNormNCHW( + const float* in, float* out, const float* gamma, const float* beta, const float* mask, const float* maskSum, + int nSize, int cSize, int xySize, float epsilon, int activation, float* sumSqBuf +); +void customCudaSpatialRMSNormNCHW( + const half* in, half* out, const half* gamma, const half* beta, const half* mask, const float* maskSum, + int nSize, int cSize, int xySize, float epsilon, int activation, float* sumSqBuf +); + +#endif // NEURALNET_ROCMHELPERS_H_ diff --git a/cpp/neuralnet/rocmhelpers.hip b/cpp/neuralnet/rocmhelpers.hip new file mode 100644 index 0000000000..320963b7cb --- /dev/null +++ b/cpp/neuralnet/rocmhelpers.hip @@ -0,0 +1,3324 @@ +#include "hip/hip_runtime.h" + +#include "../neuralnet/rocmhelpers.h" + +#include +#include + +#if defined(__HIP_ARCH_HAS_FP16__) || (defined(__HIP_DEVICE_COMPILE__) && (__HIP_ARCH_GFX803__ || __HIP_ARCH_GFX900__)) +#define HIP_SUPPORTS_FP16 +#endif + +//TODO maybe tune this number, it varies by GPU +static const int targetNumThreads = 512; + +// The custom kernels below compute flattened element indices in 32-bit int. Guard the launchers: +// throw loudly if the total element count does not fit in a positive int, rather than silently +// overflowing into an out-of-bounds access. +static void checkBufferIndexFitsInt(int64_t a, int64_t b, int64_t c, const char* whatKernel) { + int64_t total = a * b * c; + if(total >= (int64_t)2147483647) + throw std::runtime_error( + std::string(whatKernel) + ": total element count " + std::to_string(total) + + " exceeds the 32-bit index limit used by this kernel"); +} + +void splitThreadsAcrossDim01(int dim0Size, int dim1Size, int& threads0, int& blocks0, int& threads1, int& blocks1) { + if(dim0Size > targetNumThreads) { + threads0 = targetNumThreads/2; + blocks0 = (dim0Size + threads0 - 1) / threads0; + threads1 = 1; + blocks1 = dim1Size; + } + else if(dim0Size > targetNumThreads/2) { + threads0 = dim0Size; + blocks0 = 1; + threads1 = 1; + blocks1 = dim1Size; + } + else { + threads0 = dim0Size; + blocks0 = 1; + threads1 = targetNumThreads / dim0Size; + blocks1 = (dim1Size + threads1 - 1) / threads1; + } +} + +__forceinline__ __device__ float mishf(float a) { + return a * tanhf(a < 20.0f ? log1pf(expf(a)) : a); +} +__forceinline__ __device__ float mishf_scale8(float a) { + return a < 2.5f ? a * tanhf(log1pf(expf(a*8.0f))) : a; +} +__forceinline__ __device__ float siluf(float a) { + return a / (1.0f + expf(-a)); +} + +#ifdef HIP_SUPPORTS_FP16 +__forceinline__ __device__ half mishh(half h) { + float a = __half2float(h); + return __float2half(a * tanhf(a < 20.0f ? log1pf(expf(a)) : a)); +} +__forceinline__ __device__ half mishh_scale8(half h) { + float a = __half2float(h); + return __float2half(a < 2.5f ? a * tanhf(log1pf(expf(a*8.0f))) : a); +} +__forceinline__ __device__ half siluh(half h) { + float a = __half2float(h); + return __float2half(a / (1.0f + expf(-a))); +} +#endif + +//-------------------------------------------------------------------------------------------------------------- + +template +__global__ +void channelConcatKernel( + const T* inA, + const T* inB, + T* out, + int chwA, + int chwB, + int numBlocksA, + int numBlocksB, + int n +) { + if(blockIdx.x < numBlocksA) { + int index = blockIdx.x * blockDim.x + threadIdx.x; + if(index < chwA) { + int nchwA = n*chwA; + int chwOut = (chwA+chwB); + + int aIdx = index; + int outIdx = index; + while(aIdx < nchwA) { + out[outIdx] = inA[aIdx]; + aIdx += chwA; + outIdx += chwOut; + } + } + } + else { + int index = (blockIdx.x - numBlocksA) * blockDim.x + threadIdx.x; + if(index < chwB) { + int nchwB = n*chwB; + int chwOut = (chwA+chwB); + + int bIdx = index; + int outIdx = chwA+index; + while(bIdx < nchwB) { + out[outIdx] = inB[bIdx]; + bIdx += chwB; + outIdx += chwOut; + } + } + } +} + +template +void customCudaChannelConcatTemplate(const T* inA, const T* inB, T* out, int chwA, int chwB, int n) { + int blockSize = targetNumThreads; + int numBlocksA = (chwA + blockSize-1) / blockSize; + int numBlocksB = (chwB + blockSize-1) / blockSize; + int numBlocks = numBlocksA + numBlocksB; + channelConcatKernel<<>>(inA,inB,out,chwA,chwB,numBlocksA,numBlocksB,n); +} +template void customCudaChannelConcatTemplate(const float* inA, const float* inB, float* out, int chwA, int chwB, int n); +template void customCudaChannelConcatTemplate(const half* inA, const half* inB, half* out, int chwA, int chwB, int n); + +void customCudaChannelConcat(const float* inA, const float* inB, float* out, int chwA, int chwB, int n) { + customCudaChannelConcatTemplate(inA,inB,out,chwA,chwB,n); +} +void customCudaChannelConcat(const half* inA, const half* inB, half* out, int chwA, int chwB, int n) { + customCudaChannelConcatTemplate(inA,inB,out,chwA,chwB,n); +} + +//-------------------------------------------------------------------------------------------------------------- + +template +__global__ +void extractChannel0KernelNHWC(const T *in, T* out, int nhwSize, int cSize) +{ + int nhwIdx = blockIdx.x * blockDim.x + threadIdx.x; + if(nhwIdx < nhwSize) { + out[nhwIdx] = in[nhwIdx*cSize]; + } +} +template +void customCudaChannel0ExtractNHWCTemplate(const T *in, T* out, int n, int hw, int c) { + int nhw = n*hw; + int blockSize = targetNumThreads; + int numBlocks = (nhw+blockSize-1)/blockSize; + extractChannel0KernelNHWC<<>>(in,out,nhw,c); +} + +template +__global__ +void extractChannel0KernelNCHW(const T *in, T* out, int nSize, int cSize, int hwSize) +{ + int hwIdx = blockIdx.x * blockDim.x + threadIdx.x; + int nIdx = blockIdx.y * blockDim.y + threadIdx.y; + if(hwIdx < hwSize && nIdx < nSize) { + out[nIdx * hwSize + hwIdx] = in[nIdx * cSize * hwSize + hwIdx]; + } +} +template +void customCudaChannel0ExtractNCHWTemplate(const T *in, T* out, int nSize, int cSize, int hwSize) { + int hwThreads; + int hwBlocks; + int nThreads; + int nBlocks; + splitThreadsAcrossDim01(hwSize, nSize, hwThreads, hwBlocks, nThreads, nBlocks); + + if(nBlocks > 65536) + throw std::runtime_error("customCudaChannel0ExtractNCHW: nSize too large given hwSize"); + + dim3 grid(hwBlocks,nBlocks,1); + dim3 threads(hwThreads,nThreads,1); + extractChannel0KernelNCHW<<>>(in,out,nSize,cSize,hwSize); +} + +void customCudaChannel0ExtractNCHW(const float* in, float* out, int n, int c, int hw) { + customCudaChannel0ExtractNCHWTemplate(in,out,n,c,hw); +} +void customCudaChannel0ExtractNCHW(const half* in, half* out, int n, int c, int hw) { + customCudaChannel0ExtractNCHWTemplate(in,out,n,c,hw); +} +void customCudaChannel0ExtractNHWC(const float* in, float* out, int n, int hw, int c) { + customCudaChannel0ExtractNHWCTemplate(in,out,n,hw,c); +} +void customCudaChannel0ExtractNHWC(const half* in, half* out, int n, int hw, int c) { + customCudaChannel0ExtractNHWCTemplate(in,out,n,hw,c); +} + +//-------------------------------------------------------------------------------------------------------------- + +__global__ +void sumChannelsNCHWKernel(const float* in, float* out, int cSize, int xySize, float scaleSum) +{ + extern __shared__ float sumPoolNCHWShared[]; + int xyId = threadIdx.x; + int xyBlockDim = blockDim.x; + int cId = threadIdx.y; + int cBlockDim = blockDim.y; + int cIdx = blockIdx.y * cBlockDim + cId; + int nIdx = blockIdx.z; + + int xycSize = xySize*cSize; + int sharedIdx = xyId + cId * xyBlockDim; + + float acc = 0.0f; + if(cIdx < cSize) { + int xyIdx = xyId; + while(xyIdx < xySize) { + acc += in[xyIdx + cIdx * xySize + nIdx * xycSize]; + xyIdx += xyBlockDim; + } + sumPoolNCHWShared[sharedIdx] = acc; + } + __syncthreads(); + + for(int s = xyBlockDim>>1; s > 0; s >>= 1) { + if(xyId < s) { + sumPoolNCHWShared[sharedIdx] += sumPoolNCHWShared[sharedIdx + s]; + } + __syncthreads(); + } + if(xyId == 0 && cIdx < cSize) + out[cIdx + nIdx * cSize] = sumPoolNCHWShared[sharedIdx] * scaleSum; +} +__global__ +void valueHeadPoolChannelsNCHWKernel(const float* in, float* out, int nSize, int cSize, int xySize, const float* maskSum) +{ + extern __shared__ float sumPoolNCHWShared[]; + int xyId = threadIdx.x; + int xyBlockDim = blockDim.x; + int cId = threadIdx.y; + int cBlockDim = blockDim.y; + int cIdx = blockIdx.y * cBlockDim + cId; + int nIdx = blockIdx.z; + + int xycSize = xySize*cSize; + int sharedIdx = xyId + cId * xyBlockDim; + + float acc = 0.0f; + if(cIdx < cSize) { + int xyIdx = xyId; + while(xyIdx < xySize) { + acc += in[xyIdx + cIdx * xySize + nIdx * xycSize]; + xyIdx += xyBlockDim; + } + sumPoolNCHWShared[sharedIdx] = acc; + } + __syncthreads(); + + for(int s = xyBlockDim>>1; s > 0; s >>= 1) { + if(xyId < s) { + sumPoolNCHWShared[sharedIdx] += sumPoolNCHWShared[sharedIdx + s]; + } + __syncthreads(); + } + if(xyId == 0 && cIdx < cSize) { + float sum = sumPoolNCHWShared[sharedIdx]; + float div = maskSum[nIdx]; + float sqrtdiv = sqrt(div); + float mean = sum/div; + out[cIdx + nIdx * cSize*3] = mean; + out[cIdx + nIdx * cSize*3 + cSize] = mean * (sqrtdiv - 14.0f) * 0.1f; + out[cIdx + nIdx * cSize*3 + cSize*2] = mean * ((sqrtdiv - 14.0f) * (sqrtdiv - 14.0f) * 0.01f - 0.1f); + } +} +__global__ +void gPoolChannelsNCHWKernel(const float* in, float* out, int cSize, int xySize, const float* maskSum, int sharedMemElts) +{ + extern __shared__ float poolNCHWShared[]; + float* sumShared = (float*)poolNCHWShared; + float* maxShared = (float*)poolNCHWShared + sharedMemElts; + + int xyId = threadIdx.x; + int xyBlockDim = blockDim.x; + int cId = threadIdx.y; + int cBlockDim = blockDim.y; + int cIdx = blockIdx.y * cBlockDim + cId; + int nIdx = blockIdx.z; + + int xycSize = xySize*cSize; + int sharedIdx = xyId + cId * xyBlockDim; + + if(cIdx < cSize) { + float accSum = 0.0f; + float accMax = -1.0f; + int xyIdx = xyId; + while(xyIdx < xySize) { + float a = in[xyIdx + cIdx * xySize + nIdx * xycSize]; + accSum += a; + accMax = fmaxf(accMax, a); + xyIdx += xyBlockDim; + } + sumShared[sharedIdx] = accSum; + maxShared[sharedIdx] = accMax; + } + __syncthreads(); + + for(int s = xyBlockDim>>1; s > 0; s >>= 1) { + if(xyId < s) { + sumShared[sharedIdx] += sumShared[sharedIdx + s]; + maxShared[sharedIdx] = fmaxf(maxShared[sharedIdx], maxShared[sharedIdx + s]); + } + __syncthreads(); + } + if(xyId == 0 && cIdx < cSize) { + float sum = sumShared[sharedIdx]; + float div = maskSum[nIdx]; + float sqrtdiv = sqrt(div); + float mean = sum/div; + + out[cIdx + nIdx * (cSize*3)] = mean; + out[cIdx + nIdx * (cSize*3) + cSize] = mean * (sqrtdiv - 14.0f) * 0.1f; + out[cIdx + nIdx * (cSize*3) + cSize*2] = maxShared[sharedIdx]; + } +} +__global__ +void gPoolChannelsNCHWMaskKernel(const float* in, float* out, int cSize, int xySize, const float* mask, const float* maskSum, int sharedMemElts) +{ + extern __shared__ float poolNCHWShared[]; + float* sumShared = (float*)poolNCHWShared; + float* maxShared = (float*)poolNCHWShared + sharedMemElts; + + int xyId = threadIdx.x; + int xyBlockDim = blockDim.x; + int cId = threadIdx.y; + int cBlockDim = blockDim.y; + int cIdx = blockIdx.y * cBlockDim + cId; + int nIdx = blockIdx.z; + + int xycSize = xySize*cSize; + int sharedIdx = xyId + cId * xyBlockDim; + + if(cIdx < cSize) { + float accSum = 0.0f; + float accMax = -1.0f; + int xyIdx = xyId; + while(xyIdx < xySize) { + float a = in[xyIdx + cIdx * xySize + nIdx * xycSize]; + accSum += a; + // Init to -1.0 above and + mask - 1.0 is because it will effectively make all padded space into -1.0 + // which is lower than the lowest value that any current activation function will produce. + // so the max over all valid spaces will the same as the mask over all spaces including padding + // We're relying on all padded space being equal to 0 because this gpool only ever follows a BN+Activate with a mask. + accMax = fmaxf(accMax, a + (mask[xyIdx + nIdx * xySize] - 1.0f)); + xyIdx += xyBlockDim; + } + sumShared[sharedIdx] = accSum; + maxShared[sharedIdx] = accMax; + } + __syncthreads(); + + for(int s = xyBlockDim>>1; s > 0; s >>= 1) { + if(xyId < s) { + sumShared[sharedIdx] += sumShared[sharedIdx + s]; + maxShared[sharedIdx] = fmaxf(maxShared[sharedIdx], maxShared[sharedIdx + s]); + } + __syncthreads(); + } + if(xyId == 0 && cIdx < cSize) { + float sum = sumShared[sharedIdx]; + float div = maskSum[nIdx]; + float sqrtdiv = sqrt(div); + float mean = sum/div; + + out[cIdx + nIdx * (cSize*3)] = mean; + out[cIdx + nIdx * (cSize*3) + cSize] = mean * (sqrtdiv - 14.0f) * 0.1f; + out[cIdx + nIdx * (cSize*3) + cSize*2] = maxShared[sharedIdx]; + } +} + +void customCudaPoolRowsSumNCHW(const float* in, float* out, int nSize, int cSize, int xySize, float scaleSum) { + if(nSize > 65536) + throw std::runtime_error("customCudaPoolRowsSumNCHW: nSize too large"); + if(cSize > 65536) + throw std::runtime_error("customCudaPoolRowsSumNCHW: cSize too large"); + + //Use up as many threads as possible along the xy dimension. + int xyThreads = 1; + while(xyThreads < targetNumThreads && xyThreads < xySize/2) + xyThreads *= 2; + + //Distribute the extra threads along the c dimension. + int cThreads = (targetNumThreads < xyThreads) ? 1 : (targetNumThreads / xyThreads); + int cBlocks = (cSize + cThreads - 1) / cThreads; + + //We need one shared memory spot per thread + int sharedMemSize = sizeof(float) * cThreads * xyThreads; + + dim3 grid(1,cBlocks,nSize); + dim3 threads(xyThreads,cThreads,1); + sumChannelsNCHWKernel<<>>(in,out,cSize,xySize,scaleSum); +} +void customCudaValueHeadPoolNCHW(const float* in, float* out, int nSize, int cSize, int xySize, const float* maskSum) { + if(nSize > 65536) + throw std::runtime_error("customCudaValueHeadPoolNCHW: nSize too large"); + if(cSize > 65536) + throw std::runtime_error("customCudaValueHeadPoolNCHW: cSize too large"); + + //Use up as many threads as possible along the xy dimension. + int xyThreads = 1; + while(xyThreads < targetNumThreads && xyThreads < xySize/2) + xyThreads *= 2; + + //Distribute the extra threads along the c dimension. + int cThreads = (targetNumThreads < xyThreads) ? 1 : (targetNumThreads / xyThreads); + int cBlocks = (cSize + cThreads - 1) / cThreads; + + //We need one shared memory spot per thread + int sharedMemSize = sizeof(float) * cThreads * xyThreads; + + dim3 grid(1,cBlocks,nSize); + dim3 threads(xyThreads,cThreads,1); + valueHeadPoolChannelsNCHWKernel<<>>(in,out,nSize,cSize,xySize,maskSum); +} +void customCudaPoolRowsGPoolNCHW(const float* in, float* out, int nSize, int cSize, int xySize, const float* mask, const float* maskSum) { + if(nSize > 65536) + throw std::runtime_error("customCudaPoolRowsGPoolNCHW: nSize too large"); + if(cSize > 65536) + throw std::runtime_error("customCudaPoolRowsGPoolNCHW: cSize too large"); + + //Use up as many threads as possible along the xy dimension. + int xyThreads = 1; + while(xyThreads < targetNumThreads && xyThreads < xySize/2) + xyThreads *= 2; + + //Distribute the extra threads along the c dimension. + int cThreads = (targetNumThreads < xyThreads) ? 1 : (targetNumThreads / xyThreads); + int cBlocks = (cSize + cThreads - 1) / cThreads; + + //We need one shared memory spot per thread, and then we double it because we need both sum and max. + //We also make sure it's a power of two to address any alignment concerns. + int sharedMemElts = 128; + while(sharedMemElts < cThreads * xyThreads) + sharedMemElts *= 2; + int sharedMemSize = sizeof(float) * sharedMemElts * 2; + + dim3 grid(1,cBlocks,nSize); + dim3 threads(xyThreads,cThreads,1); + if(mask != NULL) + gPoolChannelsNCHWMaskKernel<<>>(in,out,cSize,xySize,mask,maskSum,sharedMemElts); + else + gPoolChannelsNCHWKernel<<>>(in,out,cSize,xySize,maskSum,sharedMemElts); +} + +//-------------------------------------------------------------------------------------------------------------- + +__global__ +void gPoolChannelsNCHWHalfKernel(const half* in, half* out, int cSize, int xySize, const float* maskSum, int sharedMemElts) +{ +#ifdef HIP_SUPPORTS_FP16 + extern __shared__ float poolNCHWShared[]; + float* sumShared = (float*)poolNCHWShared; + float* maxShared = (float*)poolNCHWShared + sharedMemElts; + + int xyId = threadIdx.x; + int xyBlockDim = blockDim.x; + int cId = threadIdx.y; + int cBlockDim = blockDim.y; + int cIdx = blockIdx.y * cBlockDim + cId; + int nIdx = blockIdx.z; + + int xycSize = xySize*cSize; + int sharedIdx = xyId + cId * xyBlockDim; + + if(cIdx < cSize) { + float accSum = 0.0f; + float accMax = -1.0f; + int xyIdx = xyId; + while(xyIdx < xySize) { + float a = __half2float(in[xyIdx + cIdx * xySize + nIdx * xycSize]); + accSum += a; + accMax = fmaxf(accMax, a); + xyIdx += xyBlockDim; + } + sumShared[sharedIdx] = accSum; + maxShared[sharedIdx] = accMax; + } + __syncthreads(); + + for(int s = xyBlockDim>>1; s > 0; s >>= 1) { + if(xyId < s) { + sumShared[sharedIdx] += sumShared[sharedIdx + s]; + maxShared[sharedIdx] = fmaxf(maxShared[sharedIdx], maxShared[sharedIdx + s]); + } + __syncthreads(); + } + if(xyId == 0 && cIdx < cSize) { + float sum = sumShared[sharedIdx]; + float div = maskSum[nIdx]; + float sqrtdiv = sqrt(div); + float mean = sum/div; + + out[cIdx + nIdx * (cSize*3)] = __float2half(mean); + out[cIdx + nIdx * (cSize*3) + cSize] = __float2half(mean * (sqrtdiv - 14.0f) * 0.1f); + out[cIdx + nIdx * (cSize*3) + cSize*2] = __float2half(maxShared[sharedIdx]); + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void gPoolChannelsNCHWHalfMaskKernel(const half* in, half* out, int cSize, int xySize, const half* mask, const float* maskSum, int sharedMemElts) +{ +#ifdef HIP_SUPPORTS_FP16 + extern __shared__ float poolNCHWShared[]; + float* sumShared = (float*)poolNCHWShared; + float* maxShared = (float*)poolNCHWShared + sharedMemElts; + + int xyId = threadIdx.x; + int xyBlockDim = blockDim.x; + int cId = threadIdx.y; + int cBlockDim = blockDim.y; + int cIdx = blockIdx.y * cBlockDim + cId; + int nIdx = blockIdx.z; + + int xycSize = xySize*cSize; + int sharedIdx = xyId + cId * xyBlockDim; + + if(cIdx < cSize) { + float accSum = 0.0f; + float accMax = -1.0f; + int xyIdx = xyId; + while(xyIdx < xySize) { + float a = __half2float(in[xyIdx + cIdx * xySize + nIdx * xycSize]); + accSum += a; + // Init to -1.0 above and + mask - 1.0 is because it will effectively make all padded space into -1.0 + // which is lower than the lowest value that any current activation function will produce. + // so the max over all valid spaces will the same as the mask over all spaces including padding + accMax = fmaxf(accMax, a + (__half2float(mask[xyIdx + nIdx * xySize]) - 1.0f)); + xyIdx += xyBlockDim; + } + sumShared[sharedIdx] = accSum; + maxShared[sharedIdx] = accMax; + } + __syncthreads(); + + for(int s = xyBlockDim>>1; s > 0; s >>= 1) { + if(xyId < s) { + sumShared[sharedIdx] += sumShared[sharedIdx + s]; + maxShared[sharedIdx] = fmaxf(maxShared[sharedIdx], maxShared[sharedIdx + s]); + } + __syncthreads(); + } + if(xyId == 0 && cIdx < cSize) { + float sum = sumShared[sharedIdx]; + float div = maskSum[nIdx]; + float sqrtdiv = sqrt(div); + float mean = sum/div; + + out[cIdx + nIdx * (cSize*3)] = __float2half(mean); + out[cIdx + nIdx * (cSize*3) + cSize] = __float2half(mean * (sqrtdiv - 14.0f) * 0.1f); + out[cIdx + nIdx * (cSize*3) + cSize*2] = __float2half(maxShared[sharedIdx]); + } +#else + //Do nothing, FP16 not supported +#endif +} + +void customCudaPoolRowsGPoolNCHW(const half* in, half* out, int nSize, int cSize, int xySize, const half* mask, const float* maskSum) { + if(nSize > 65536) + throw std::runtime_error("customCudaPoolRowsGPoolNCHW: nSize too large"); + if(cSize > 65536) + throw std::runtime_error("customCudaPoolRowsGPoolNCHW: cSize too large"); + + //Use up as many threads as possible along the xy dimension. + int xyThreads = 1; + while(xyThreads < targetNumThreads && xyThreads < xySize/2) + xyThreads *= 2; + + //Distribute the extra threads along the c dimension. + int cThreads = (targetNumThreads < xyThreads) ? 1 : (targetNumThreads / xyThreads); + int cBlocks = (cSize + cThreads - 1) / cThreads; + + //We need one shared memory spot per thread, and then we double it because we need both sum and max. + //We also make sure it's a power of two to address any alignment concerns. + int sharedMemElts = 128; + while(sharedMemElts < cThreads * xyThreads) + sharedMemElts *= 2; + int sharedMemSize = sizeof(float) * sharedMemElts * 2; + + dim3 grid(1,cBlocks,nSize); + dim3 threads(xyThreads,cThreads,1); + if(mask != NULL) + gPoolChannelsNCHWHalfMaskKernel<<>>(in,out,cSize,xySize,mask,maskSum,sharedMemElts); + else + gPoolChannelsNCHWHalfKernel<<>>(in,out,cSize,xySize,maskSum,sharedMemElts); +} + + + +//-------------------------------------------------------------------------------------------------------------- + +__global__ +void sumChannelsNHWCKernel(const float* in, float* out, int xySize, int cSize, float scaleSum) +{ + extern __shared__ float sumPoolNHWCShared[]; + int cId = threadIdx.x; + int cBlockDim = blockDim.x; + int xyId = threadIdx.y; + int xyBlockDim = blockDim.y; + + int cIdx = blockIdx.x * cBlockDim + cId; + int nIdx = blockIdx.z; + int sharedIdx = cId + cBlockDim * xyId; + int xycSize = xySize*cSize; + + sumPoolNHWCShared[sharedIdx] = 0; + + if(cIdx < cSize) { + int xyIdx = xyId; + while(xyIdx < xySize) { + sumPoolNHWCShared[sharedIdx] += in[cIdx + xyIdx * cSize + nIdx * xycSize]; + xyIdx += xyBlockDim; + } + } + __syncthreads(); + + for(int s = xyBlockDim>>1; s > 0; s >>= 1) { + if(xyId < s) { + sumPoolNHWCShared[sharedIdx] += sumPoolNHWCShared[sharedIdx + cBlockDim * s]; + } + __syncthreads(); + } + if(xyId == 0 && cIdx < cSize) + out[cIdx + nIdx * cSize] = sumPoolNHWCShared[sharedIdx] * scaleSum; +} +__global__ +void valueHeadPoolChannelsNHWCKernel(const float* in, float* out, int nSize, int xySize, int cSize, const float* maskSum) +{ + extern __shared__ float sumPoolNHWCShared[]; + int cId = threadIdx.x; + int cBlockDim = blockDim.x; + int xyId = threadIdx.y; + int xyBlockDim = blockDim.y; + + int cIdx = blockIdx.x * cBlockDim + cId; + int nIdx = blockIdx.z; + int sharedIdx = cId + cBlockDim * xyId; + int xycSize = xySize*cSize; + + sumPoolNHWCShared[sharedIdx] = 0; + + if(cIdx < cSize) { + int xyIdx = xyId; + while(xyIdx < xySize) { + sumPoolNHWCShared[sharedIdx] += in[cIdx + xyIdx * cSize + nIdx * xycSize]; + xyIdx += xyBlockDim; + } + } + __syncthreads(); + + for(int s = xyBlockDim>>1; s > 0; s >>= 1) { + if(xyId < s) { + sumPoolNHWCShared[sharedIdx] += sumPoolNHWCShared[sharedIdx + cBlockDim * s]; + } + __syncthreads(); + } + if(xyId == 0 && cIdx < cSize) { + float sum = sumPoolNHWCShared[sharedIdx]; + float div = maskSum[nIdx]; + float sqrtdiv = sqrt(div); + float mean = sum/div; + out[cIdx + nIdx * cSize*3] = mean; + out[cIdx + nIdx * cSize*3 + cSize] = mean * (sqrtdiv - 14.0f) * 0.1f; + out[cIdx + nIdx * cSize*3 + cSize*2] = mean * ((sqrtdiv - 14.0f) * (sqrtdiv - 14.0f) * 0.01f - 0.1f); + } +} +__global__ +void gPoolChannelsNHWCKernel(const float* in, float* out, int xySize, int cSize, const float* maskSum, int sharedMemElts) +{ + extern __shared__ float poolNHWCShared[]; + float* sumShared = (float*)poolNHWCShared; + float* maxShared = (float*)poolNHWCShared + sharedMemElts; + + int cId = threadIdx.x; + int cBlockDim = blockDim.x; + int xyId = threadIdx.y; + int xyBlockDim = blockDim.y; + + int cIdx = blockIdx.x * cBlockDim + cId; + int nIdx = blockIdx.z; + int sharedIdx = cId + cBlockDim * xyId; + int xycSize = xySize*cSize; + + sumShared[sharedIdx] = 0; + maxShared[sharedIdx] = -1.0f; + + if(cIdx < cSize) { + int xyIdx = xyId; + while(xyIdx < xySize) { + float a = in[cIdx + xyIdx * cSize + nIdx * xycSize]; + sumShared[sharedIdx] += a; + maxShared[sharedIdx] = fmaxf(maxShared[sharedIdx], a); + xyIdx += xyBlockDim; + } + } + __syncthreads(); + + for(int s = xyBlockDim>>1; s > 0; s >>= 1) { + if(xyId < s) { + sumShared[sharedIdx] += sumShared[sharedIdx + cBlockDim * s]; + maxShared[sharedIdx] = fmaxf(maxShared[sharedIdx],maxShared[sharedIdx + cBlockDim * s]); + } + __syncthreads(); + } + if(xyId == 0 && cIdx < cSize) { + float sum = sumShared[sharedIdx]; + float div = maskSum[nIdx]; + float sqrtdiv = sqrt(div); + float mean = sum/div; + + out[cIdx + nIdx * (cSize*3)] = mean; + out[cIdx + nIdx * (cSize*3) + cSize] = mean * (sqrtdiv - 14.0f) * 0.1f; + out[cIdx + nIdx * (cSize*3) + cSize*2] = maxShared[sharedIdx]; + } +} +__global__ +void gPoolChannelsNHWCMaskKernel(const float* in, float* out, int xySize, int cSize, const float* mask, const float* maskSum, int sharedMemElts) +{ + extern __shared__ float poolNHWCShared[]; + float* sumShared = (float*)poolNHWCShared; + float* maxShared = (float*)poolNHWCShared + sharedMemElts; + + int cId = threadIdx.x; + int cBlockDim = blockDim.x; + int xyId = threadIdx.y; + int xyBlockDim = blockDim.y; + + int cIdx = blockIdx.x * cBlockDim + cId; + int nIdx = blockIdx.z; + int sharedIdx = cId + cBlockDim * xyId; + int xycSize = xySize*cSize; + + sumShared[sharedIdx] = 0; + maxShared[sharedIdx] = -1.0f; + + if(cIdx < cSize) { + int xyIdx = xyId; + while(xyIdx < xySize) { + float a = in[cIdx + xyIdx * cSize + nIdx * xycSize]; + sumShared[sharedIdx] += a; + // Init to -1.0 above and + mask - 1.0 is because it will effectively make all padded space into -1.0 + // which is lower than the lowest value that any current activation function will produce. + // so the max over all valid spaces will the same as the mask over all spaces including padding + maxShared[sharedIdx] = fmaxf(maxShared[sharedIdx], a + (mask[xyIdx + nIdx * xySize] - 1.0f)); + xyIdx += xyBlockDim; + } + } + __syncthreads(); + + for(int s = xyBlockDim>>1; s > 0; s >>= 1) { + if(xyId < s) { + sumShared[sharedIdx] += sumShared[sharedIdx + cBlockDim * s]; + maxShared[sharedIdx] = fmaxf(maxShared[sharedIdx],maxShared[sharedIdx + cBlockDim * s]); + } + __syncthreads(); + } + if(xyId == 0 && cIdx < cSize) { + float sum = sumShared[sharedIdx]; + float div = maskSum[nIdx]; + float sqrtdiv = sqrt(div); + float mean = sum/div; + + out[cIdx + nIdx * (cSize*3)] = mean; + out[cIdx + nIdx * (cSize*3) + cSize] = mean * (sqrtdiv - 14.0f) * 0.1f; + out[cIdx + nIdx * (cSize*3) + cSize*2] = maxShared[sharedIdx]; + } +} + + +void customCudaPoolRowsSumNHWC(const float* in, float* out, int nSize, int xySize, int cSize, float scaleSum) { + if(nSize > 65536) + throw std::runtime_error("customCudaPoolRowsSumNHWC: nSize too large"); + if(cSize > 65536) + throw std::runtime_error("customCudaPoolRowsSumNHWC: cSize too large"); + + //Use up to two warps worth of threads along the channel dimension, which is the + //most compact + int cThreads = 1; + while(cThreads < 64 && cThreads < cSize/2) + cThreads *= 2; + int cBlocks = (cSize + cThreads - 1) / cThreads; + + //Distribute the extra threads to perform parallel reduction along the xy dimension. + int xyThreads = (targetNumThreads < cThreads) ? 1 : (targetNumThreads / cThreads); + + //We need one shared memory spot per thread + int sharedMemSize = sizeof(float) * cThreads * xyThreads; + + dim3 grid(cBlocks,1,nSize); + dim3 threads(cThreads,xyThreads,1); + sumChannelsNHWCKernel<<>>(in,out,xySize,cSize,scaleSum); +} + +void customCudaValueHeadPoolNHWC(const float* in, float* out, int nSize, int xySize, int cSize, const float* maskSum) { + if(nSize > 65536) + throw std::runtime_error("customCudaValueHeadPoolNHWC: nSize too large"); + if(cSize > 65536) + throw std::runtime_error("customCudaValueHeadPoolNHWC: cSize too large"); + + //Use up to two warps worth of threads along the channel dimension, which is the + //most compact + int cThreads = 1; + while(cThreads < 64 && cThreads < cSize/2) + cThreads *= 2; + int cBlocks = (cSize + cThreads - 1) / cThreads; + + //Distribute the extra threads to perform parallel reduction along the xy dimension. + int xyThreads = (targetNumThreads < cThreads) ? 1 : (targetNumThreads / cThreads); + + //We need one shared memory spot per thread + int sharedMemSize = sizeof(float) * cThreads * xyThreads; + + dim3 grid(cBlocks,1,nSize); + dim3 threads(cThreads,xyThreads,1); + valueHeadPoolChannelsNHWCKernel<<>>(in,out,nSize,xySize,cSize,maskSum); +} + +void customCudaPoolRowsGPoolNHWC(const float* in, float* out, int nSize, int xySize, int cSize, const float* mask, const float* maskSum) { + if(nSize > 65536) + throw std::runtime_error("customCudaPoolRowsGPoolNHWC: nSize too large"); + if(cSize > 65536) + throw std::runtime_error("customCudaPoolRowsGPoolNHWC: cSize too large"); + + //Use up to two warps worth of threads along the channel dimension, which is the + //most compact + int cThreads = 1; + while(cThreads < 64 && cThreads < cSize/2) + cThreads *= 2; + int cBlocks = (cSize + cThreads - 1) / cThreads; + + //Distribute the extra threads to perform parallel reduction along the xy dimension. + int xyThreads = (targetNumThreads < cThreads) ? 1 : (targetNumThreads / cThreads); + + //We need one shared memory spot per thread, and then we double it because we need both sum and max. + //We also make sure it's a power of two to address any alignment concerns. + int sharedMemElts = 128; + while(sharedMemElts < cThreads * xyThreads) + sharedMemElts *= 2; + int sharedMemSize = sizeof(float) * sharedMemElts * 2; + + dim3 grid(cBlocks,1,nSize); + dim3 threads(cThreads,xyThreads,1); + if(mask != NULL) + gPoolChannelsNHWCMaskKernel<<>>(in,out,xySize,cSize,mask,maskSum,sharedMemElts); + else + gPoolChannelsNHWCKernel<<>>(in,out,xySize,cSize,maskSum,sharedMemElts); +} + +//-------------------------------------------------------------------------------------------------------------- + +__global__ +void gPoolChannelsNHWCHalfKernel(const half* in, half* out, int xySize, int cSize, const float* maskSum, int sharedMemElts) +{ +#ifdef HIP_SUPPORTS_FP16 + extern __shared__ float poolNHWCShared[]; + float* sumShared = (float*)poolNHWCShared; + float* maxShared = (float*)poolNHWCShared + sharedMemElts; + + int cId = threadIdx.x; + int cBlockDim = blockDim.x; + int xyId = threadIdx.y; + int xyBlockDim = blockDim.y; + + int cIdx = blockIdx.x * cBlockDim + cId; + int nIdx = blockIdx.z; + int sharedIdx = cId + cBlockDim * xyId; + int xycSize = xySize*cSize; + + sumShared[sharedIdx] = 0; + maxShared[sharedIdx] = -1.0f; + + if(cIdx < cSize) { + int xyIdx = xyId; + while(xyIdx < xySize) { + float a = __half2float(in[cIdx + xyIdx * cSize + nIdx * xycSize]); + sumShared[sharedIdx] += a; + maxShared[sharedIdx] = fmaxf(maxShared[sharedIdx], a); + xyIdx += xyBlockDim; + } + } + __syncthreads(); + + for(int s = xyBlockDim>>1; s > 0; s >>= 1) { + if(xyId < s) { + sumShared[sharedIdx] += sumShared[sharedIdx + cBlockDim * s]; + maxShared[sharedIdx] = fmaxf(maxShared[sharedIdx],maxShared[sharedIdx + cBlockDim * s]); + } + __syncthreads(); + } + if(xyId == 0 && cIdx < cSize) { + float sum = sumShared[sharedIdx]; + float div = maskSum[nIdx]; + float sqrtdiv = sqrt(div); + float mean = sum/div; + + out[cIdx + nIdx * (cSize*3)] = __float2half(mean); + out[cIdx + nIdx * (cSize*3) + cSize] = __float2half(mean * (sqrtdiv - 14.0f) * 0.1f); + out[cIdx + nIdx * (cSize*3) + cSize*2] = __float2half(maxShared[sharedIdx]); + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void gPoolChannelsNHWCHalfMaskKernel(const half* in, half* out, int xySize, int cSize, const half* mask, const float* maskSum, int sharedMemElts) +{ +#ifdef HIP_SUPPORTS_FP16 + extern __shared__ float poolNHWCShared[]; + float* sumShared = (float*)poolNHWCShared; + float* maxShared = (float*)poolNHWCShared + sharedMemElts; + + int cId = threadIdx.x; + int cBlockDim = blockDim.x; + int xyId = threadIdx.y; + int xyBlockDim = blockDim.y; + + int cIdx = blockIdx.x * cBlockDim + cId; + int nIdx = blockIdx.z; + int sharedIdx = cId + cBlockDim * xyId; + int xycSize = xySize*cSize; + + sumShared[sharedIdx] = 0; + maxShared[sharedIdx] = -1.0f; + + if(cIdx < cSize) { + int xyIdx = xyId; + while(xyIdx < xySize) { + float a = __half2float(in[cIdx + xyIdx * cSize + nIdx * xycSize]); + sumShared[sharedIdx] += a; + // Init to -1.0 above and + mask - 1.0 is because it will effectively make all padded space into -1.0 + // which is lower than the lowest value that any current activation function will produce. + // so the max over all valid spaces will the same as the mask over all spaces including padding + maxShared[sharedIdx] = fmaxf(maxShared[sharedIdx], a + (__half2float(mask[xyIdx + nIdx * xySize]) - 1.0f)); + xyIdx += xyBlockDim; + } + } + __syncthreads(); + + for(int s = xyBlockDim>>1; s > 0; s >>= 1) { + if(xyId < s) { + sumShared[sharedIdx] += sumShared[sharedIdx + cBlockDim * s]; + maxShared[sharedIdx] = fmaxf(maxShared[sharedIdx],maxShared[sharedIdx + cBlockDim * s]); + } + __syncthreads(); + } + if(xyId == 0 && cIdx < cSize) { + float sum = sumShared[sharedIdx]; + float div = maskSum[nIdx]; + float sqrtdiv = sqrt(div); + float mean = sum/div; + + out[cIdx + nIdx * (cSize*3)] = __float2half(mean); + out[cIdx + nIdx * (cSize*3) + cSize] = __float2half(mean * (sqrtdiv - 14.0f) * 0.1f); + out[cIdx + nIdx * (cSize*3) + cSize*2] = __float2half(maxShared[sharedIdx]); + } +#else + //Do nothing, FP16 not supported +#endif +} + +void customCudaPoolRowsGPoolNHWC(const half* in, half* out, int nSize, int xySize, int cSize, const half* mask, const float* maskSum) { + if(nSize > 65536) + throw std::runtime_error("customCudaPoolRowsGPoolNHWC: nSize too large"); + if(cSize > 65536) + throw std::runtime_error("customCudaPoolRowsGPoolNHWC: cSize too large"); + + //Use up to two warps worth of threads along the channel dimension, which is the + //most compact + int cThreads = 1; + while(cThreads < 64 && cThreads < cSize/2) + cThreads *= 2; + int cBlocks = (cSize + cThreads - 1) / cThreads; + + //Distribute the extra threads to perform parallel reduction along the xy dimension. + int xyThreads = (targetNumThreads < cThreads) ? 1 : (targetNumThreads / cThreads); + + //We need one shared memory spot per thread, and then we double it because we need both sum and max. + //We also make sure it's a power of two to address any alignment concerns. + int sharedMemElts = 128; + while(sharedMemElts < cThreads * xyThreads) + sharedMemElts *= 2; + int sharedMemSize = sizeof(float) * sharedMemElts * 2; + + dim3 grid(cBlocks,1,nSize); + dim3 threads(cThreads,xyThreads,1); + if(mask != NULL) + gPoolChannelsNHWCHalfMaskKernel<<>>(in,out,xySize,cSize,mask,maskSum,sharedMemElts); + else + gPoolChannelsNHWCHalfKernel<<>>(in,out,xySize,cSize,maskSum,sharedMemElts); +} + + +//-------------------------------------------------------------------------------------------------------------- + +__global__ +void copyToHalfKernel(const float *in, half* out, int n) +{ + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if(idx < n) { + out[idx] = __float2half(in[idx]); + } +} +__global__ +void copyFromHalfKernel(const half *in, float* out, int n) +{ + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if(idx < n) { + out[idx] = __half2float(in[idx]); + } +} + +void customCudaCopyToHalf(const float* in, half* out, int n) { + int blockSize = targetNumThreads; + int numBlocks = (n+blockSize-1)/blockSize; + copyToHalfKernel<<>>(in,out,n); +} +void customCudaCopyFromHalf(const half* in, float* out, int n) { + int blockSize = targetNumThreads; + int numBlocks = (n+blockSize-1)/blockSize; + copyFromHalfKernel<<>>(in,out,n); +} + +//-------------------------------------------------------------------------------------------------------------- + + +//-------------------------------------------------------------------------------------------------------------- +// Element-wise tensor add: buf[i] += toAdd[i], for float and half + +__global__ +void addTensorsInplaceKernel(float *buf, const float* toAdd, int nSize) +{ + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if(idx < nSize) { + buf[idx] += toAdd[idx]; + } +} +void customCudaAddTensorsInplace(float* buf, const float* toAdd, int nSize) { + int blockSize = targetNumThreads; + int numBlocks = (nSize+blockSize-1)/blockSize; + addTensorsInplaceKernel<<>>(buf,toAdd,nSize); +} + +__global__ +void addTensorsInplaceHalfKernel(half *buf, const half* toAdd, int nSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if(idx < nSize) { + buf[idx] = __hadd(buf[idx],toAdd[idx]); + } +#else + //Do nothing, FP16 not supported +#endif +} +void customCudaAddTensorsInplace(half* buf, const half* toAdd, int nSize) { + int blockSize = targetNumThreads; + int numBlocks = (nSize+blockSize-1)/blockSize; + addTensorsInplaceHalfKernel<<>>(buf,toAdd,nSize); +} + +//-------------------------------------------------------------------------------------------------------------- + +__global__ +void addTensorInplaceHalfKernel(half *buf, const half* biases, int nSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if(idx < nSize) { + buf[idx] = __hadd(buf[idx],biases[idx]); + } +#else + //Do nothing, FP16 not supported +#endif +} +void customCudaAddTensorInplace(half* buf, const half* biases, int nSize) { + int blockSize = targetNumThreads; + int numBlocks = (nSize+blockSize-1)/blockSize; + addTensorInplaceHalfKernel<<>>(buf,biases,nSize); +} + +//-------------------------------------------------------------------------------------------------------------- + + +__global__ +void addCBiasInplaceNCKernel(float *buf, const float* biases, int nSize, int cSize) +{ + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int nIdx = blockIdx.y * blockDim.y + threadIdx.y; + if(cIdx < cSize && nIdx < nSize) { + int idx = nIdx * cSize + cIdx; + buf[idx] = buf[idx] + biases[cIdx]; + } +} +__global__ +void addCBiasInplaceNCHalfKernel(half *buf, const half* biases, int nSize, int cSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int nIdx = blockIdx.y * blockDim.y + threadIdx.y; + if(cIdx < cSize && nIdx < nSize) { + int idx = nIdx * cSize + cIdx; + buf[idx] = __hadd(buf[idx],biases[cIdx]); + } +#else + //Do nothing, FP16 not supported +#endif +} + +__global__ +void addCBiasInplaceNCKernelRelu(float *buf, const float* biases, int nSize, int cSize) +{ + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int nIdx = blockIdx.y * blockDim.y + threadIdx.y; + if(cIdx < cSize && nIdx < nSize) { + int idx = nIdx * cSize + cIdx; + buf[idx] = fmaxf(buf[idx] + biases[cIdx],0.0f); + } +} +__global__ +void addCBiasInplaceNCHalfKernelRelu(half *buf, const half* biases, int nSize, int cSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int nIdx = blockIdx.y * blockDim.y + threadIdx.y; + if(cIdx < cSize && nIdx < nSize) { + int idx = nIdx * cSize + cIdx; + const half halfzero = __float2half(0.0f); + half a = __hadd(buf[idx],biases[cIdx]); + buf[idx] = __hgt(a,halfzero) ? a : halfzero; + } +#else + //Do nothing, FP16 not supported +#endif +} + +__global__ +void addCBiasInplaceNCKernelMish(float *buf, const float* biases, int nSize, int cSize) +{ + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int nIdx = blockIdx.y * blockDim.y + threadIdx.y; + if(cIdx < cSize && nIdx < nSize) { + int idx = nIdx * cSize + cIdx; + buf[idx] = mishf(buf[idx] + biases[cIdx]); + } +} +__global__ +void addCBiasInplaceNCHalfKernelMish(half *buf, const half* biases, int nSize, int cSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int nIdx = blockIdx.y * blockDim.y + threadIdx.y; + if(cIdx < cSize && nIdx < nSize) { + int idx = nIdx * cSize + cIdx; + half a = __hadd(buf[idx],biases[cIdx]); + buf[idx] = mishh(a); + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void addCBiasInplaceNCKernelMishScale8(float *buf, const float* biases, int nSize, int cSize) +{ + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int nIdx = blockIdx.y * blockDim.y + threadIdx.y; + if(cIdx < cSize && nIdx < nSize) { + int idx = nIdx * cSize + cIdx; + buf[idx] = mishf_scale8(buf[idx] + biases[cIdx]); + } +} +__global__ +void addCBiasInplaceNCHalfKernelMishScale8(half *buf, const half* biases, int nSize, int cSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int nIdx = blockIdx.y * blockDim.y + threadIdx.y; + if(cIdx < cSize && nIdx < nSize) { + int idx = nIdx * cSize + cIdx; + half a = __hadd(buf[idx],biases[cIdx]); + buf[idx] = mishh_scale8(a); + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void addCBiasInplaceNCKernelSilu(float *buf, const float* biases, int nSize, int cSize) +{ + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int nIdx = blockIdx.y * blockDim.y + threadIdx.y; + if(cIdx < cSize && nIdx < nSize) { + int idx = nIdx * cSize + cIdx; + buf[idx] = siluf(buf[idx] + biases[cIdx]); + } +} +__global__ +void addCBiasInplaceNCHalfKernelSilu(half *buf, const half* biases, int nSize, int cSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int nIdx = blockIdx.y * blockDim.y + threadIdx.y; + if(cIdx < cSize && nIdx < nSize) { + int idx = nIdx * cSize + cIdx; + half a = __hadd(buf[idx],biases[cIdx]); + buf[idx] = siluh(a); + } +#else + //Do nothing, FP16 not supported +#endif +} + +void sharedAddCBiasInplaceNC(void* buf, const void* biases, int nSize, int cSize, bool isHalf, int activation) { + int cThreads; + int cBlocks; + int nThreads; + int nBlocks; + splitThreadsAcrossDim01(cSize, nSize, cThreads, cBlocks, nThreads, nBlocks); + + if(nBlocks > 65536) + throw std::runtime_error("customCudaAddCBiasInplaceNC: nSize too large given cSize"); + + dim3 grid(cBlocks,nBlocks,1); + dim3 threads(cThreads,nThreads,1); + + if(activation == ACTIVATION_IDENTITY) { + if(isHalf) + addCBiasInplaceNCHalfKernel<<>>((half*)buf,(const half*)biases,nSize,cSize); + else + addCBiasInplaceNCKernel<<>>((float*)buf,(const float*)biases,nSize,cSize); + } + else if(activation == ACTIVATION_RELU) { + if(isHalf) + addCBiasInplaceNCHalfKernelRelu<<>>((half*)buf,(const half*)biases,nSize,cSize); + else + addCBiasInplaceNCKernelRelu<<>>((float*)buf,(const float*)biases,nSize,cSize); + } + else if(activation == ACTIVATION_MISH) { + if(isHalf) + addCBiasInplaceNCHalfKernelMish<<>>((half*)buf,(const half*)biases,nSize,cSize); + else + addCBiasInplaceNCKernelMish<<>>((float*)buf,(const float*)biases,nSize,cSize); + } + else if(activation == ACTIVATION_MISH_SCALE8) { + if(isHalf) + addCBiasInplaceNCHalfKernelMishScale8<<>>((half*)buf,(const half*)biases,nSize,cSize); + else + addCBiasInplaceNCKernelMishScale8<<>>((float*)buf,(const float*)biases,nSize,cSize); + } + else if(activation == ACTIVATION_SILU) { + if(isHalf) + addCBiasInplaceNCHalfKernelSilu<<>>((half*)buf,(const half*)biases,nSize,cSize); + else + addCBiasInplaceNCKernelSilu<<>>((float*)buf,(const float*)biases,nSize,cSize); + } + else { + throw std::runtime_error("customCudaAddCBiasInplaceNC: unsupported activation"); + } +} + +void customCudaAddCBiasInplaceNC(float* buf, const float* biases, int nSize, int cSize, int activation) { + sharedAddCBiasInplaceNC(buf,biases,nSize,cSize,false,activation); +} +void customCudaAddCBiasInplaceNC(half* buf, const half* biases, int nSize, int cSize, int activation) { + sharedAddCBiasInplaceNC(buf,biases,nSize,cSize,true,activation); +} + +//-------------------------------------------------------------------------------------------------------------- + +__global__ +void addNCBiasInplaceNCHWKernel(float *buf, const float* biases, int cSize, int sSize) +{ + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int ncIdx = nIdx * cSize + cIdx; + int idx = ncIdx * sSize + sIdx; + buf[idx] = buf[idx] + biases[ncIdx]; + } +} +__global__ +void addNCBiasInplaceNCHWHalfKernel(half *buf, const half* biases, int cSize, int sSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int ncIdx = nIdx * cSize + cIdx; + int idx = ncIdx * sSize + sIdx; + buf[idx] = __hadd(buf[idx],biases[ncIdx]); + } +#else + //Do nothing, FP16 not supported +#endif +} + +void sharedAddNCBiasInplaceNCHW(void *buf, const void* biases, int nSize, int cSize, int xySize, bool isHalf) { + if(nSize > 65536) + throw std::runtime_error("customCudaAddNCBiasInplaceNCHW: nSize too large"); + if(cSize > 65536) + throw std::runtime_error("customCudaAddNCBiasInplaceNCHW: cSize too large"); + + int sSize = xySize; + int sThreads; + int sBlocks; + int cThreads; + int cBlocks; + splitThreadsAcrossDim01(sSize, cSize, sThreads, sBlocks, cThreads, cBlocks); + + dim3 grid(sBlocks,cBlocks,nSize); + dim3 threads(sThreads,cThreads,1); + if(isHalf) + addNCBiasInplaceNCHWHalfKernel<<>>((half*)buf,(const half*)biases,cSize,sSize); + else + addNCBiasInplaceNCHWKernel<<>>((float*)buf,(const float*)biases,cSize,sSize); +} + +void customCudaAddNCBiasInplaceNCHW(float *buf, const float* biases, int nSize, int cSize, int xySize) { + sharedAddNCBiasInplaceNCHW(buf,biases,nSize,cSize,xySize,false); +} +void customCudaAddNCBiasInplaceNCHW(half *buf, const half* biases, int nSize, int cSize, int xySize) { + sharedAddNCBiasInplaceNCHW(buf,biases,nSize,cSize,xySize,true); +} + +//-------------------------------------------------------------------------------------------------------------- + +__global__ +void addNCBiasInplaceNHWCKernel(float *buf, const float* biases, int sSize, int cSize) +{ + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int ncIdx = nIdx * cSize + cIdx; + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + buf[idx] = buf[idx] + biases[ncIdx]; + } +} +__global__ +void addNCBiasInplaceNHWCHalfKernel(half *buf, const half* biases, int sSize, int cSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int ncIdx = nIdx * cSize + cIdx; + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + buf[idx] = __hadd(buf[idx],biases[ncIdx]); + } +#else + //Do nothing, FP16 not supported +#endif +} + +void sharedAddNCBiasInplaceNHWC(void *buf, const void* biases, int nSize, int xySize, int cSize, bool isHalf) { + if(nSize > 65536) + throw std::runtime_error("customCudaAddNCBiasInplaceNHWC: nSize too large"); + if(xySize > 65536) + throw std::runtime_error("customCudaAddNCBiasInplaceNHWC: xySize too large"); + + int sSize = xySize; + int cThreads; + int cBlocks; + int sThreads; + int sBlocks; + splitThreadsAcrossDim01(cSize, sSize, cThreads, cBlocks, sThreads, sBlocks); + + dim3 grid(cBlocks,sBlocks,nSize); + dim3 threads(cThreads,sThreads,1); + if(isHalf) + addNCBiasInplaceNHWCHalfKernel<<>>((half*)buf,(const half*)biases,sSize,cSize); + else + addNCBiasInplaceNHWCKernel<<>>((float*)buf,(const float*)biases,sSize,cSize); +} + +void customCudaAddNCBiasInplaceNHWC(float *buf, const float* biases, int nSize, int xySize, int cSize) { + sharedAddNCBiasInplaceNHWC(buf,biases,nSize,xySize,cSize,false); +} +void customCudaAddNCBiasInplaceNHWC(half *buf, const half* biases, int nSize, int xySize, int cSize) { + sharedAddNCBiasInplaceNHWC(buf,biases,nSize,xySize,cSize,true); +} + +//-------------------------------------------------------------------------------------------------------------- + +__global__ +void applyCScaleBiasNCHWKernel(const float *in, float* out, const float* scale, const float* biases, int cSize, int sSize) +{ + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + out[idx] = in[idx] * scale[cIdx] + biases[cIdx]; + } +} +__global__ +void applyCScaleBiasNCHWReluKernel(const float *in, float* out, const float* scale, const float* biases, int cSize, int sSize) +{ + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + out[idx] = fmaxf(in[idx] * scale[cIdx] + biases[cIdx],0.0f); + } +} +__global__ +void applyCScaleBiasNCHWMishKernel(const float *in, float* out, const float* scale, const float* biases, int cSize, int sSize) +{ + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + out[idx] = mishf(in[idx] * scale[cIdx] + biases[cIdx]); + } +} +__global__ +void applyCScaleBiasNCHWMishScale8Kernel(const float *in, float* out, const float* scale, const float* biases, int cSize, int sSize) +{ + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + out[idx] = mishf_scale8(in[idx] * scale[cIdx] + biases[cIdx]); + } +} +__global__ +void applyCScaleBiasNCHWSiluKernel(const float *in, float* out, const float* scale, const float* biases, int cSize, int sSize) +{ + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + out[idx] = siluf(in[idx] * scale[cIdx] + biases[cIdx]); + } +} +__global__ +void applyCScaleBiasNCHWMaskKernel(const float *in, float* out, const float* scale, const float* biases, const float* mask, int cSize, int sSize) +{ + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + out[idx] = (in[idx] * scale[cIdx] + biases[cIdx]) * mask[nIdx*sSize+sIdx]; + } +} +__global__ +void applyCScaleBiasNCHWReluMaskKernel(const float *in, float* out, const float* scale, const float* biases, const float* mask, int cSize, int sSize) +{ + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + out[idx] = fmaxf(in[idx] * scale[cIdx] + biases[cIdx],0.0f) * mask[nIdx*sSize+sIdx]; + } +} +__global__ +void applyCScaleBiasNCHWMishMaskKernel(const float *in, float* out, const float* scale, const float* biases, const float* mask, int cSize, int sSize) +{ + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + out[idx] = mishf(in[idx] * scale[cIdx] + biases[cIdx]) * mask[nIdx*sSize+sIdx]; + } +} +__global__ +void applyCScaleBiasNCHWMishScale8MaskKernel(const float *in, float* out, const float* scale, const float* biases, const float* mask, int cSize, int sSize) +{ + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + out[idx] = mishf_scale8(in[idx] * scale[cIdx] + biases[cIdx]) * mask[nIdx*sSize+sIdx]; + } +} +__global__ +void applyCScaleBiasNCHWSiluMaskKernel(const float *in, float* out, const float* scale, const float* biases, const float* mask, int cSize, int sSize) +{ + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + out[idx] = siluf(in[idx] * scale[cIdx] + biases[cIdx]) * mask[nIdx*sSize+sIdx]; + } +} +__global__ +void applyCScaleBiasNCHWHalfKernel(const half *in, half* out, const half* scale, const half* biases, int cSize, int sSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + out[idx] = __hfma(in[idx],scale[cIdx],biases[cIdx]); + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void applyCScaleBiasNCHWReluHalfKernel(const half *in, half* out, const half* scale, const half* biases, int cSize, int sSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + half a = __hfma(in[idx],scale[cIdx],biases[cIdx]); + const half halfzero = __float2half(0.0f); + out[idx] = __hgt(a,halfzero) ? a : halfzero; + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void applyCScaleBiasNCHWMishHalfKernel(const half *in, half* out, const half* scale, const half* biases, int cSize, int sSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + half a = __hfma(in[idx],scale[cIdx],biases[cIdx]); + out[idx] = mishh(a); + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void applyCScaleBiasNCHWMishScale8HalfKernel(const half *in, half* out, const half* scale, const half* biases, int cSize, int sSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + half a = __hfma(in[idx],scale[cIdx],biases[cIdx]); + out[idx] = mishh_scale8(a); + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void applyCScaleBiasNCHWSiluHalfKernel(const half *in, half* out, const half* scale, const half* biases, int cSize, int sSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + half a = __hfma(in[idx],scale[cIdx],biases[cIdx]); + out[idx] = siluh(a); + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void applyCScaleBiasNCHWMaskHalfKernel(const half *in, half* out, const half* scale, const half* biases, const half* mask, int cSize, int sSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + out[idx] = __hmul(__hfma(in[idx],scale[cIdx],biases[cIdx]),mask[nIdx*sSize+sIdx]); + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void applyCScaleBiasNCHWReluMaskHalfKernel(const half *in, half* out, const half* scale, const half* biases, const half* mask, int cSize, int sSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + half a = __hmul(__hfma(in[idx],scale[cIdx],biases[cIdx]),mask[nIdx*sSize+sIdx]); + const half halfzero = __float2half(0.0f); + out[idx] = __hgt(a,halfzero) ? a : halfzero; + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void applyCScaleBiasNCHWMishMaskHalfKernel(const half *in, half* out, const half* scale, const half* biases, const half* mask, int cSize, int sSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + half a = __hmul(__hfma(in[idx],scale[cIdx],biases[cIdx]),mask[nIdx*sSize+sIdx]); + out[idx] = mishh(a); + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void applyCScaleBiasNCHWMishScale8MaskHalfKernel(const half *in, half* out, const half* scale, const half* biases, const half* mask, int cSize, int sSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + half a = __hmul(__hfma(in[idx],scale[cIdx],biases[cIdx]),mask[nIdx*sSize+sIdx]); + out[idx] = mishh_scale8(a); + } +#else + //Do nothing, FP16 not supported +#endif +} + +__global__ +void applyCScaleBiasNCHWSiluMaskHalfKernel(const half *in, half* out, const half* scale, const half* biases, const half* mask, int cSize, int sSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int sIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * cSize + cIdx) * sSize + sIdx; + half a = __hmul(__hfma(in[idx],scale[cIdx],biases[cIdx]),mask[nIdx*sSize+sIdx]); + out[idx] = siluh(a); + } +#else + //Do nothing, FP16 not supported +#endif +} + +void sharedApplyCScaleBiasNCHW(const void* in, void* out, const void* scale, const void* biases, const void* mask, int nSize, int cSize, int xySize, bool isHalf, int activation) { + if(nSize > 65536) + throw std::runtime_error("customCudaApplyCScaleBiasNCHW: nSize too large"); + if(cSize > 65536) + throw std::runtime_error("customCudaApplyCScaleBiasNCHW: cSize too large"); + + int sSize = xySize; + int sThreads; + int sBlocks; + int cThreads; + int cBlocks; + splitThreadsAcrossDim01(sSize, cSize, sThreads, sBlocks, cThreads, cBlocks); + + dim3 grid(sBlocks,cBlocks,nSize); + dim3 threads(sThreads,cThreads,1); + if(mask == NULL) { + if(activation == ACTIVATION_IDENTITY) { + if(isHalf) + applyCScaleBiasNCHWHalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,cSize,sSize); + else + applyCScaleBiasNCHWKernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,cSize,sSize); + } + else if(activation == ACTIVATION_RELU) { + if(isHalf) + applyCScaleBiasNCHWReluHalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,cSize,sSize); + else + applyCScaleBiasNCHWReluKernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,cSize,sSize); + } + else if(activation == ACTIVATION_MISH) { + if(isHalf) + applyCScaleBiasNCHWMishHalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,cSize,sSize); + else + applyCScaleBiasNCHWMishKernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,cSize,sSize); + } + else if(activation == ACTIVATION_MISH_SCALE8) { + if(isHalf) + applyCScaleBiasNCHWMishScale8HalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,cSize,sSize); + else + applyCScaleBiasNCHWMishScale8Kernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,cSize,sSize); + } + else if(activation == ACTIVATION_SILU) { + if(isHalf) + applyCScaleBiasNCHWSiluHalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,cSize,sSize); + else + applyCScaleBiasNCHWSiluKernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,cSize,sSize); + } + else { + throw std::runtime_error("customCudaApplyCScaleBiasNCHW: unsupported activation"); + } + } + else { + if(activation == ACTIVATION_IDENTITY) { + if(isHalf) + applyCScaleBiasNCHWMaskHalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,(const half*)mask,cSize,sSize); + else + applyCScaleBiasNCHWMaskKernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,(const float*)mask,cSize,sSize); + } + else if(activation == ACTIVATION_RELU) { + if(isHalf) + applyCScaleBiasNCHWReluMaskHalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,(const half*)mask,cSize,sSize); + else + applyCScaleBiasNCHWReluMaskKernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,(const float*)mask,cSize,sSize); + } + else if(activation == ACTIVATION_MISH) { + if(isHalf) + applyCScaleBiasNCHWMishMaskHalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,(const half*)mask,cSize,sSize); + else + applyCScaleBiasNCHWMishMaskKernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,(const float*)mask,cSize,sSize); + } + else if(activation == ACTIVATION_MISH_SCALE8) { + if(isHalf) + applyCScaleBiasNCHWMishScale8MaskHalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,(const half*)mask,cSize,sSize); + else + applyCScaleBiasNCHWMishScale8MaskKernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,(const float*)mask,cSize,sSize); + } + else if(activation == ACTIVATION_SILU) { + if(isHalf) + applyCScaleBiasNCHWSiluMaskHalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,(const half*)mask,cSize,sSize); + else + applyCScaleBiasNCHWSiluMaskKernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,(const float*)mask,cSize,sSize); + } + else { + throw std::runtime_error("customCudaApplyCScaleBiasNCHW: unsupported activation"); + } + } +} + +void customCudaApplyCScaleBiasNCHW(const float* in, float* out, const float* scale, const float* biases, const float* mask, int nSize, int cSize, int xySize, int activation) { + sharedApplyCScaleBiasNCHW(in,out,scale,biases,mask,nSize,cSize,xySize,false,activation); +} +void customCudaApplyCScaleBiasNCHW(const half* in, half* out, const half* scale, const half* biases, const half* mask, int nSize, int cSize, int xySize, int activation) { + sharedApplyCScaleBiasNCHW(in,out,scale,biases,mask,nSize,cSize,xySize,true,activation); +} + + +//-------------------------------------------------------------------------------------------------------------- + +__global__ +void applyCScaleBiasNHWCKernel(const float* in, float* out, const float* scale, const float* biases, int sSize, int cSize) +{ + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + out[idx] = in[idx] * scale[cIdx] + biases[cIdx]; + } +} +__global__ +void applyCScaleBiasNHWCReluKernel(const float* in, float* out, const float* scale, const float* biases, int sSize, int cSize) +{ + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + out[idx] = fmaxf(in[idx] * scale[cIdx] + biases[cIdx],0.0f); + } +} +__global__ +void applyCScaleBiasNHWCMishKernel(const float* in, float* out, const float* scale, const float* biases, int sSize, int cSize) +{ + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + out[idx] = mishf(in[idx] * scale[cIdx] + biases[cIdx]); + } +} +__global__ +void applyCScaleBiasNHWCMishScale8Kernel(const float* in, float* out, const float* scale, const float* biases, int sSize, int cSize) +{ + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + out[idx] = mishf_scale8(in[idx] * scale[cIdx] + biases[cIdx]); + } +} +__global__ +void applyCScaleBiasNHWCSiluKernel(const float* in, float* out, const float* scale, const float* biases, int sSize, int cSize) +{ + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + out[idx] = siluf(in[idx] * scale[cIdx] + biases[cIdx]); + } +} +__global__ +void applyCScaleBiasNHWCMaskKernel(const float* in, float* out, const float* scale, const float* biases, const float* mask, int sSize, int cSize) +{ + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + out[idx] = (in[idx] * scale[cIdx] + biases[cIdx]) * mask[nIdx*sSize+sIdx]; + } +} +__global__ +void applyCScaleBiasNHWCReluMaskKernel(const float* in, float* out, const float* scale, const float* biases, const float* mask, int sSize, int cSize) +{ + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + out[idx] = fmaxf(in[idx] * scale[cIdx] + biases[cIdx],0.0f) * mask[nIdx*sSize+sIdx]; + } +} +__global__ +void applyCScaleBiasNHWCMishMaskKernel(const float* in, float* out, const float* scale, const float* biases, const float* mask, int sSize, int cSize) +{ + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + out[idx] = mishf(in[idx] * scale[cIdx] + biases[cIdx]) * mask[nIdx*sSize+sIdx]; + } +} +__global__ +void applyCScaleBiasNHWCMishScale8MaskKernel(const float* in, float* out, const float* scale, const float* biases, const float* mask, int sSize, int cSize) +{ + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + out[idx] = mishf_scale8(in[idx] * scale[cIdx] + biases[cIdx]) * mask[nIdx*sSize+sIdx]; + } +} +__global__ +void applyCScaleBiasNHWCSiluMaskKernel(const float* in, float* out, const float* scale, const float* biases, const float* mask, int sSize, int cSize) +{ + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + out[idx] = siluf(in[idx] * scale[cIdx] + biases[cIdx]) * mask[nIdx*sSize+sIdx]; + } +} +__global__ +void applyCScaleBiasNHWCHalfKernel(const half* in, half* out, const half* scale, const half* biases, int sSize, int cSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + out[idx] = __hfma(in[idx],scale[cIdx],biases[cIdx]); + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void applyCScaleBiasNHWCReluHalfKernel(const half* in, half* out, const half* scale, const half* biases, int sSize, int cSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + half a = __hfma(in[idx],scale[cIdx],biases[cIdx]); + const half halfzero = __float2half(0.0f); + out[idx] = __hgt(a,halfzero) ? a : halfzero; + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void applyCScaleBiasNHWCMishHalfKernel(const half* in, half* out, const half* scale, const half* biases, int sSize, int cSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + half a = __hfma(in[idx],scale[cIdx],biases[cIdx]); + out[idx] = mishh(a); + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void applyCScaleBiasNHWCMishScale8HalfKernel(const half* in, half* out, const half* scale, const half* biases, int sSize, int cSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + half a = __hfma(in[idx],scale[cIdx],biases[cIdx]); + out[idx] = mishh_scale8(a); + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void applyCScaleBiasNHWCSiluHalfKernel(const half* in, half* out, const half* scale, const half* biases, int sSize, int cSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + half a = __hfma(in[idx],scale[cIdx],biases[cIdx]); + out[idx] = siluh(a); + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void applyCScaleBiasNHWCMaskHalfKernel(const half* in, half* out, const half* scale, const half* biases, const half* mask, int sSize, int cSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + out[idx] = __hmul(__hfma(in[idx],scale[cIdx],biases[cIdx]),mask[nIdx*sSize+sIdx]); + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void applyCScaleBiasNHWCReluMaskHalfKernel(const half* in, half* out, const half* scale, const half* biases, const half* mask, int sSize, int cSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + half a = __hmul(__hfma(in[idx],scale[cIdx],biases[cIdx]),mask[nIdx*sSize+sIdx]); + const half halfzero = __float2half(0.0f); + out[idx] = __hgt(a,halfzero) ? a : halfzero; + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void applyCScaleBiasNHWCMishMaskHalfKernel(const half* in, half* out, const half* scale, const half* biases, const half* mask, int sSize, int cSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + half a = __hmul(__hfma(in[idx],scale[cIdx],biases[cIdx]),mask[nIdx*sSize+sIdx]); + out[idx] = mishh(a); + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void applyCScaleBiasNHWCMishScale8MaskHalfKernel(const half* in, half* out, const half* scale, const half* biases, const half* mask, int sSize, int cSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + half a = __hmul(__hfma(in[idx],scale[cIdx],biases[cIdx]),mask[nIdx*sSize+sIdx]); + out[idx] = mishh_scale8(a); + } +#else + //Do nothing, FP16 not supported +#endif +} +__global__ +void applyCScaleBiasNHWCSiluMaskHalfKernel(const half* in, half* out, const half* scale, const half* biases, const half* mask, int sSize, int cSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int sIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx < cSize && sIdx < sSize) { + int idx = (nIdx * sSize + sIdx) * cSize + cIdx; + half a = __hmul(__hfma(in[idx],scale[cIdx],biases[cIdx]),mask[nIdx*sSize+sIdx]); + out[idx] = siluh(a); + } +#else + //Do nothing, FP16 not supported +#endif +} + +void sharedApplyCScaleBiasNHWC(const void* in, void* out, const void* scale, const void* biases, const void* mask, int nSize, int xySize, int cSize, bool isHalf, int activation) { + if(nSize > 65536) + throw std::runtime_error("customCudaApplyCScaleBiasNHWC: nSize too large"); + if(xySize > 65536) + throw std::runtime_error("customCudaApplyCScaleBiasNHWC: xySize too large"); + + int sSize = xySize; + int cThreads; + int cBlocks; + int sThreads; + int sBlocks; + splitThreadsAcrossDim01(cSize, sSize, cThreads, cBlocks, sThreads, sBlocks); + + dim3 grid(cBlocks,sBlocks,nSize); + dim3 threads(cThreads,sThreads,1); + if(mask == NULL) { + if(activation == ACTIVATION_IDENTITY) { + if(isHalf) + applyCScaleBiasNHWCHalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,sSize,cSize); + else + applyCScaleBiasNHWCKernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,sSize,cSize); + } + else if(activation == ACTIVATION_RELU) { + if(isHalf) + applyCScaleBiasNHWCReluHalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,sSize,cSize); + else + applyCScaleBiasNHWCReluKernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,sSize,cSize); + } + else if(activation == ACTIVATION_MISH) { + if(isHalf) + applyCScaleBiasNHWCMishHalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,sSize,cSize); + else + applyCScaleBiasNHWCMishKernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,sSize,cSize); + } + else if(activation == ACTIVATION_MISH_SCALE8) { + if(isHalf) + applyCScaleBiasNHWCMishScale8HalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,sSize,cSize); + else + applyCScaleBiasNHWCMishScale8Kernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,sSize,cSize); + } + else if(activation == ACTIVATION_SILU) { + if(isHalf) + applyCScaleBiasNHWCSiluHalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,sSize,cSize); + else + applyCScaleBiasNHWCSiluKernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,sSize,cSize); + } + else { + throw std::runtime_error("customCudaApplyCScaleBiasNHWC: unsupported activation"); + } + } + else { + if(activation == ACTIVATION_IDENTITY) { + if(isHalf) + applyCScaleBiasNHWCMaskHalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,(const half*)mask,sSize,cSize); + else + applyCScaleBiasNHWCMaskKernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,(const float*)mask,sSize,cSize); + } + else if(activation == ACTIVATION_RELU) { + if(isHalf) + applyCScaleBiasNHWCReluMaskHalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,(const half*)mask,sSize,cSize); + else + applyCScaleBiasNHWCReluMaskKernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,(const float*)mask,sSize,cSize); + } + else if(activation == ACTIVATION_MISH) { + if(isHalf) + applyCScaleBiasNHWCMishMaskHalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,(const half*)mask,sSize,cSize); + else + applyCScaleBiasNHWCMishMaskKernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,(const float*)mask,sSize,cSize); + } + else if(activation == ACTIVATION_MISH_SCALE8) { + if(isHalf) + applyCScaleBiasNHWCMishScale8MaskHalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,(const half*)mask,sSize,cSize); + else + applyCScaleBiasNHWCMishScale8MaskKernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,(const float*)mask,sSize,cSize); + } + else if(activation == ACTIVATION_SILU) { + if(isHalf) + applyCScaleBiasNHWCSiluMaskHalfKernel<<>>((const half*)in,(half*)out,(const half*)scale,(const half*)biases,(const half*)mask,sSize,cSize); + else + applyCScaleBiasNHWCSiluMaskKernel<<>>((const float*)in,(float*)out,(const float*)scale,(const float*)biases,(const float*)mask,sSize,cSize); + } + else { + throw std::runtime_error("customCudaApplyCScaleBiasNHWC: unsupported activation"); + } + } +} + +void customCudaApplyCScaleBiasNHWC(const float* in, float* out, const float* scale, const float* biases, const float* mask, int nSize, int xySize, int cSize, int activation) { + sharedApplyCScaleBiasNHWC(in,out,scale,biases,mask,nSize,xySize,cSize,false,activation); +} +void customCudaApplyCScaleBiasNHWC(const half* in, half* out, const half* scale, const half* biases, const half* mask, int nSize, int xySize, int cSize, int activation) { + sharedApplyCScaleBiasNHWC(in,out,scale,biases,mask,nSize,xySize,cSize,true,activation); +} + +//============================================================================================== +// Transformer support kernels +//============================================================================================== + +//-------------------------------------------------------------------------------------------------------------- +// RoPE: apply rotary position embeddings in-place. +// buf: [totalDim, seqLen*batchSize] column-major (totalDim = numBufHeads*qHeadDim, fast-moving). +// Each thread handles one (pair, xy, n, h) combination. + +__global__ +void applyRoPEKernel( + float* buf, const float* cosTable, const float* sinTable, + int batchSize, int seqLen, int numBufHeads, int numKVHeads, int qHeadDim, int totalDim, int numPairs, int learnableRope +) { + int xy = blockIdx.x; + int n = blockIdx.y; + int hp = threadIdx.x; + int totalHP = numBufHeads * numPairs; + if(xy >= seqLen || n >= batchSize || hp >= totalHP) + return; + + int h = hp / numPairs; + int pairIdx = hp % numPairs; + int c0 = h * qHeadDim + 2 * pairIdx; + int c1 = c0 + 1; + size_t col = (size_t)n * seqLen + xy; + size_t idx0 = c0 + col * totalDim; + size_t idx1 = c1 + col * totalDim; + + int tableIdx; + if(learnableRope) { + int kvh = h * numKVHeads / numBufHeads; + tableIdx = (kvh * numPairs + pairIdx) * seqLen + xy; + } else { + tableIdx = pairIdx * seqLen + xy; + } + + float cosVal = cosTable[tableIdx]; + float sinVal = sinTable[tableIdx]; + float x0 = buf[idx0]; + float x1 = buf[idx1]; + buf[idx0] = x0 * cosVal - x1 * sinVal; + buf[idx1] = x0 * sinVal + x1 * cosVal; +} + +__global__ +void applyRoPEHalfKernel( + half* buf, const half* cosTable, const half* sinTable, + int batchSize, int seqLen, int numBufHeads, int numKVHeads, int qHeadDim, int totalDim, int numPairs, int learnableRope +) { +#ifdef HIP_SUPPORTS_FP16 + int xy = blockIdx.x; + int n = blockIdx.y; + int hp = threadIdx.x; + int totalHP = numBufHeads * numPairs; + if(xy >= seqLen || n >= batchSize || hp >= totalHP) + return; + + int h = hp / numPairs; + int pairIdx = hp % numPairs; + int c0 = h * qHeadDim + 2 * pairIdx; + int c1 = c0 + 1; + size_t col = (size_t)n * seqLen + xy; + size_t idx0 = c0 + col * totalDim; + size_t idx1 = c1 + col * totalDim; + + int tableIdx; + if(learnableRope) { + int kvh = h * numKVHeads / numBufHeads; + tableIdx = (kvh * numPairs + pairIdx) * seqLen + xy; + } else { + tableIdx = pairIdx * seqLen + xy; + } + + float cosVal = __half2float(cosTable[tableIdx]); + float sinVal = __half2float(sinTable[tableIdx]); + float x0 = __half2float(buf[idx0]); + float x1 = __half2float(buf[idx1]); + buf[idx0] = __float2half(x0 * cosVal - x1 * sinVal); + buf[idx1] = __float2half(x0 * sinVal + x1 * cosVal); +#else + //Do nothing, FP16 not supported +#endif +} + +// Learnable RoPE, table-free variant: recompute cos/sin in-kernel from the per-head frequencies +// instead of reading a precomputed [numKVHeads, numPairs, seqLen] cos/sin table. +// freqs layout: (numKVHeads, numPairs, 2) flattened; [...,0]=freqX (width/x), [...,1]=freqY (height/y). +__global__ +void applyRoPELearnableRecomputeKernel( + float* buf, const float* freqs, + int batchSize, int seqLen, int numBufHeads, int numKVHeads, int qHeadDim, int totalDim, int numPairs, int nnXLen +) { + int xy = blockIdx.x; + int n = blockIdx.y; + int hp = threadIdx.x; + int totalHP = numBufHeads * numPairs; + if(xy >= seqLen || n >= batchSize || hp >= totalHP) + return; + + int h = hp / numPairs; + int pairIdx = hp % numPairs; + int c0 = h * qHeadDim + 2 * pairIdx; + int c1 = c0 + 1; + size_t col = (size_t)n * seqLen + xy; + size_t idx0 = c0 + col * totalDim; + size_t idx1 = c1 + col * totalDim; + + int kvh = h * numKVHeads / numBufHeads; + int x = xy % nnXLen; + int y = xy / nnXLen; + float freqX = freqs[(kvh * numPairs + pairIdx) * 2 + 0]; + float freqY = freqs[(kvh * numPairs + pairIdx) * 2 + 1]; + float angle = (float)x * freqX + (float)y * freqY; + float sinVal, cosVal; + sincosf(angle, &sinVal, &cosVal); + + float x0 = buf[idx0]; + float x1 = buf[idx1]; + buf[idx0] = x0 * cosVal - x1 * sinVal; + buf[idx1] = x0 * sinVal + x1 * cosVal; +} + +__global__ +void applyRoPELearnableRecomputeHalfKernel( + half* buf, const float* freqs, + int batchSize, int seqLen, int numBufHeads, int numKVHeads, int qHeadDim, int totalDim, int numPairs, int nnXLen +) { +#ifdef HIP_SUPPORTS_FP16 + int xy = blockIdx.x; + int n = blockIdx.y; + int hp = threadIdx.x; + int totalHP = numBufHeads * numPairs; + if(xy >= seqLen || n >= batchSize || hp >= totalHP) + return; + + int h = hp / numPairs; + int pairIdx = hp % numPairs; + int c0 = h * qHeadDim + 2 * pairIdx; + int c1 = c0 + 1; + size_t col = (size_t)n * seqLen + xy; + size_t idx0 = c0 + col * totalDim; + size_t idx1 = c1 + col * totalDim; + + int kvh = h * numKVHeads / numBufHeads; + int x = xy % nnXLen; + int y = xy / nnXLen; + float freqX = freqs[(kvh * numPairs + pairIdx) * 2 + 0]; + float freqY = freqs[(kvh * numPairs + pairIdx) * 2 + 1]; + float angle = (float)x * freqX + (float)y * freqY; + float sinVal, cosVal; + sincosf(angle, &sinVal, &cosVal); + + float x0 = __half2float(buf[idx0]); + float x1 = __half2float(buf[idx1]); + buf[idx0] = __float2half(x0 * cosVal - x1 * sinVal); + buf[idx1] = __float2half(x0 * sinVal + x1 * cosVal); +#else + //Do nothing, FP16 not supported +#endif +} + +// One block per (xy, n). threadIdx.x = h*numPairs + pairIdx covers all channel pairs for the position. +void customCudaApplyRoPE( + float* buf, const float* cosTable, const float* sinTable, + int batchSize, int seqLen, int numBufHeads, int numKVHeads, int qHeadDim, int numPairs, bool learnableRope +) { + int totalDim = numBufHeads * qHeadDim; + int totalHP = numBufHeads * numPairs; + if(totalHP > 1024) + throw std::runtime_error("customCudaApplyRoPE: numHeads*qHeadDim/2 (" + std::to_string(totalHP) + ") exceeds the 1024 threads/block limit"); + int threads = ((totalHP + 31) / 32) * 32; + dim3 blocks(seqLen, batchSize, 1); + applyRoPEKernel<<>>( + buf, cosTable, sinTable, batchSize, seqLen, numBufHeads, numKVHeads, qHeadDim, totalDim, numPairs, learnableRope ? 1 : 0 + ); +} +void customCudaApplyRoPE( + half* buf, const half* cosTable, const half* sinTable, + int batchSize, int seqLen, int numBufHeads, int numKVHeads, int qHeadDim, int numPairs, bool learnableRope +) { + int totalDim = numBufHeads * qHeadDim; + int totalHP = numBufHeads * numPairs; + if(totalHP > 1024) + throw std::runtime_error("customCudaApplyRoPE: numHeads*qHeadDim/2 (" + std::to_string(totalHP) + ") exceeds the 1024 threads/block limit"); + int threads = ((totalHP + 31) / 32) * 32; + dim3 blocks(seqLen, batchSize, 1); + applyRoPEHalfKernel<<>>( + buf, cosTable, sinTable, batchSize, seqLen, numBufHeads, numKVHeads, qHeadDim, totalDim, numPairs, learnableRope ? 1 : 0 + ); +} + +void customCudaApplyRoPELearnableRecompute( + float* buf, const float* freqs, + int batchSize, int seqLen, int numBufHeads, int numKVHeads, int qHeadDim, int numPairs, int nnXLen +) { + int totalDim = numBufHeads * qHeadDim; + int totalHP = numBufHeads * numPairs; + if(totalHP > 1024) + throw std::runtime_error("customCudaApplyRoPELearnableRecompute: numHeads*qHeadDim/2 (" + std::to_string(totalHP) + ") exceeds the 1024 threads/block limit"); + int threads = ((totalHP + 31) / 32) * 32; + dim3 blocks(seqLen, batchSize, 1); + applyRoPELearnableRecomputeKernel<<>>( + buf, freqs, batchSize, seqLen, numBufHeads, numKVHeads, qHeadDim, totalDim, numPairs, nnXLen + ); +} +void customCudaApplyRoPELearnableRecompute( + half* buf, const float* freqs, + int batchSize, int seqLen, int numBufHeads, int numKVHeads, int qHeadDim, int numPairs, int nnXLen +) { + int totalDim = numBufHeads * qHeadDim; + int totalHP = numBufHeads * numPairs; + if(totalHP > 1024) + throw std::runtime_error("customCudaApplyRoPELearnableRecompute: numHeads*qHeadDim/2 (" + std::to_string(totalHP) + ") exceeds the 1024 threads/block limit"); + int threads = ((totalHP + 31) / 32) * 32; + dim3 blocks(seqLen, batchSize, 1); + applyRoPELearnableRecomputeHalfKernel<<>>( + buf, freqs, batchSize, seqLen, numBufHeads, numKVHeads, qHeadDim, totalDim, numPairs, nnXLen + ); +} + +//-------------------------------------------------------------------------------------------------------------- +// FlashAttention-style scaled dot product attention with online softmax (tiled). +// Grid: (numQGroups, batchSize * numHeads), block: BLOCK_Q threads. +// Layout: BSHD row-major. Templated on qHeadDim/vHeadDim so inner loops unroll. + +template +__device__ __forceinline__ +void flashAttentionTiledImpl( + const T* Q, const T* K, const T* V, const T* mask, T* output, + int seqLen, int numHeads, int numKVHeads, float scale +) { + const int tid = threadIdx.x; + const int qBlockStart = blockIdx.x * (BLOCK_Q * Q_PER_THREAD); + const int bh = blockIdx.y; + const int n = bh / numHeads; + const int h = bh % numHeads; + const int kvh = h * numKVHeads / numHeads; + + const int qTotalDim = numHeads * qHeadDim; + const int kTotalDim = numKVHeads * qHeadDim; + const int vTotalDim = numKVHeads * vHeadDim; + const int oTotalDim = numHeads * vHeadDim; + + constexpr int K_TILE_STRIDE = qHeadDim; + constexpr int V_TILE_STRIDE = vHeadDim; + __shared__ float kTile[BLOCK_KV * K_TILE_STRIDE]; + __shared__ float vTile[BLOCK_KV * V_TILE_STRIDE]; + __shared__ float kMaskTile[BLOCK_KV]; + + float qReg[Q_PER_THREAD * qHeadDim]; + float qMask[Q_PER_THREAD]; + float runningMax[Q_PER_THREAD]; + float runningSum[Q_PER_THREAD]; + float acc[Q_PER_THREAD * vHeadDim]; + + // Load Q for the Q_PER_THREAD positions this thread owns. + #pragma unroll + for(int qi = 0; qi < Q_PER_THREAD; qi++) { + int qPos = qBlockStart + qi * BLOCK_Q + tid; + qMask[qi] = 0.0f; + if(qPos < seqLen) { + if(mask != NULL) { + qMask[qi] = (float)mask[n * seqLen + qPos]; + } else { + qMask[qi] = 1.0f; + } + if(qMask[qi] != 0.0f) { + const T* qPtr = Q + ((size_t)n * seqLen + qPos) * qTotalDim + h * qHeadDim; + #pragma unroll + for(int d = 0; d < qHeadDim; d++) qReg[qi * qHeadDim + d] = (float)qPtr[d]; + } + } + runningMax[qi] = -1e30f; + runningSum[qi] = 0.0f; + #pragma unroll + for(int d = 0; d < vHeadDim; d++) acc[qi * vHeadDim + d] = 0.0f; + } + + // Iterate over K/V in BLOCK_KV-row tiles. + for(int kvStart = 0; kvStart < seqLen; kvStart += BLOCK_KV) { + // Cooperatively load K tile: BLOCK_KV rows of qHeadDim values (stride K_TILE_STRIDE). + #pragma unroll + for(int t = tid; t < BLOCK_KV * qHeadDim; t += BLOCK_Q) { + int tileKPos = t / qHeadDim; + int tileD = t % qHeadDim; + int globalKPos = kvStart + tileKPos; + float v = 0.0f; + if(globalKPos < seqLen) { + const T* kPtr = K + ((size_t)n * seqLen + globalKPos) * kTotalDim + kvh * qHeadDim; + v = (float)kPtr[tileD]; + } + kTile[tileKPos * K_TILE_STRIDE + tileD] = v; + } + // Cooperatively load V tile: BLOCK_KV rows of vHeadDim values (stride V_TILE_STRIDE). + #pragma unroll + for(int t = tid; t < BLOCK_KV * vHeadDim; t += BLOCK_Q) { + int tileKPos = t / vHeadDim; + int tileD = t % vHeadDim; + int globalKPos = kvStart + tileKPos; + float v = 0.0f; + if(globalKPos < seqLen) { + const T* vPtr = V + ((size_t)n * seqLen + globalKPos) * vTotalDim + kvh * vHeadDim; + v = (float)vPtr[tileD]; + } + vTile[tileKPos * V_TILE_STRIDE + tileD] = v; + } + // Cooperatively load mask tile. + for(int t = tid; t < BLOCK_KV; t += BLOCK_Q) { + int globalKPos = kvStart + t; + float m = 0.0f; + if(globalKPos < seqLen) { + m = (mask != NULL) ? (float)mask[n * seqLen + globalKPos] : 1.0f; + } + kMaskTile[t] = m; + } + __syncthreads(); + + int kvEnd = min(BLOCK_KV, seqLen - kvStart); + + // Each thread updates its Q_PER_THREAD queries against the shared K/V tile. + #pragma unroll + for(int qi = 0; qi < Q_PER_THREAD; qi++) { + int qPos = qBlockStart + qi * BLOCK_Q + tid; + if(qPos >= seqLen || qMask[qi] == 0.0f) continue; + + for(int tk = 0; tk < kvEnd; tk++) { + if(kMaskTile[tk] == 0.0f) continue; + + float dot = 0.0f; + #pragma unroll + for(int d = 0; d < qHeadDim; d++) { + dot += qReg[qi * qHeadDim + d] * kTile[tk * K_TILE_STRIDE + d]; + } + dot *= scale; + + float newMax = fmaxf(runningMax[qi], dot); + float expOldMax = expf(runningMax[qi] - newMax); + float expCur = expf(dot - newMax); + + #pragma unroll + for(int d = 0; d < vHeadDim; d++) { + acc[qi * vHeadDim + d] = acc[qi * vHeadDim + d] * expOldMax + expCur * vTile[tk * V_TILE_STRIDE + d]; + } + runningSum[qi] = runningSum[qi] * expOldMax + expCur; + runningMax[qi] = newMax; + } + } + __syncthreads(); + } + + // Write outputs. + #pragma unroll + for(int qi = 0; qi < Q_PER_THREAD; qi++) { + int qPos = qBlockStart + qi * BLOCK_Q + tid; + if(qPos >= seqLen) continue; + T* outRow = output + ((size_t)n * seqLen + qPos) * oTotalDim + h * vHeadDim; + if(qMask[qi] == 0.0f) { + #pragma unroll + for(int d = 0; d < vHeadDim; d++) outRow[d] = (T)0.0f; + } else { + float invSum = (runningSum[qi] > 0.0f) ? (1.0f / runningSum[qi]) : 0.0f; + #pragma unroll + for(int d = 0; d < vHeadDim; d++) outRow[d] = (T)(acc[qi * vHeadDim + d] * invSum); + } + } +} + +template +__global__ +void flashAttentionKernelFloat( + const float* Q, const float* K, const float* V, const float* mask, float* output, + int seqLen, int numHeads, int numKVHeads, float scale +) { + flashAttentionTiledImpl( + Q, K, V, mask, output, seqLen, numHeads, numKVHeads, scale); +} + +template +__global__ +void flashAttentionKernelHalf( + const half* Q, const half* K, const half* V, const half* mask, half* output, + int seqLen, int numHeads, int numKVHeads, float scale +) { +#ifdef HIP_SUPPORTS_FP16 + flashAttentionTiledImpl( + Q, K, V, mask, output, seqLen, numHeads, numKVHeads, scale); +#endif +} + +#define FA_LAUNCH_FLOAT(QD, VD, BQ, BKV, QPT) \ + do { \ + int totalQPerBlock = (BQ) * (QPT); \ + dim3 grid((seqLen + totalQPerBlock - 1) / totalQPerBlock, batchSize * numHeads); \ + flashAttentionKernelFloat<(QD), (VD), (BQ), (BKV), (QPT)><<>>( \ + Q, K, V, mask, output, seqLen, numHeads, numKVHeads, scale); \ + } while(0) + +#define FA_LAUNCH_HALF(QD, VD, BQ, BKV, QPT) \ + do { \ + int totalQPerBlock = (BQ) * (QPT); \ + dim3 grid((seqLen + totalQPerBlock - 1) / totalQPerBlock, batchSize * numHeads); \ + flashAttentionKernelHalf<(QD), (VD), (BQ), (BKV), (QPT)><<>>( \ + Q, K, V, mask, output, seqLen, numHeads, numKVHeads, scale); \ + } while(0) + +// Only (qHeadDim,vHeadDim) pairs known to be used by KataGo's transformer models so far are +// instantiated; add more shapes as needed (mirrors the CUDA backend's supported-shape list). +void customCudaFlashAttention( + const float* Q, const float* K, const float* V, const float* mask, float* output, + int batchSize, int seqLen, int numHeads, int numKVHeads, int qHeadDim, int vHeadDim +) { + if(batchSize * numHeads > 65536) + throw std::runtime_error("customCudaFlashAttention: batchSize * numHeads too large"); + float scale = 1.0f / sqrtf((float)qHeadDim); + if(qHeadDim == 32 && vHeadDim == 32) FA_LAUNCH_FLOAT(32, 32, 128, 32, 1); + else if(qHeadDim == 32 && vHeadDim == 16) FA_LAUNCH_FLOAT(32, 16, 128, 32, 1); + else if(qHeadDim == 64 && vHeadDim == 64) FA_LAUNCH_FLOAT(64, 64, 128, 32, 1); + else if(qHeadDim == 64 && vHeadDim == 32) FA_LAUNCH_FLOAT(64, 32, 128, 32, 1); + else if(qHeadDim == 32 && vHeadDim == 64) FA_LAUNCH_FLOAT(32, 64, 128, 32, 1); + else throw std::runtime_error("customCudaFlashAttention: unsupported (qHeadDim,vHeadDim) combination"); +} +void customCudaFlashAttention( + const half* Q, const half* K, const half* V, const half* mask, half* output, + int batchSize, int seqLen, int numHeads, int numKVHeads, int qHeadDim, int vHeadDim +) { + if(batchSize * numHeads > 65536) + throw std::runtime_error("customCudaFlashAttention: batchSize * numHeads too large"); + float scale = 1.0f / sqrtf((float)qHeadDim); + if(qHeadDim == 32 && vHeadDim == 32) FA_LAUNCH_HALF(32, 32, 128, 32, 1); + else if(qHeadDim == 32 && vHeadDim == 16) FA_LAUNCH_HALF(32, 16, 128, 32, 1); + else if(qHeadDim == 64 && vHeadDim == 64) FA_LAUNCH_HALF(64, 64, 128, 32, 1); + else if(qHeadDim == 64 && vHeadDim == 32) FA_LAUNCH_HALF(64, 32, 128, 32, 1); + else if(qHeadDim == 32 && vHeadDim == 64) FA_LAUNCH_HALF(32, 64, 128, 32, 1); + else throw std::runtime_error("customCudaFlashAttention: unsupported (qHeadDim,vHeadDim) combination"); +} + +#undef FA_LAUNCH_FLOAT +#undef FA_LAUNCH_HALF + +//-------------------------------------------------------------------------------------------------------------- +// Convert mask [batchSize, seqLen] (0/1) into a fully-materialized additive attention bias of shape +// [batchSize, seqLen, seqLen]: bias[b,q,k] = (mask[b,k] != 0 ? 0 : -3e4). +// Note: the q dim is fully replicated since the mask only depends on k. +// +// The constant must be a large finite negative to avoid any chance of misbehavior in the fused +// attention softmax, and it must fit in fp16 (max ~65504) since the bias tensor's dtype must match +// Q/K/V's. We use -3e4, the largest round value that leaves fp16 headroom for the model's own +// logits on top. Measured logit magnitudes on real models as of mid-2026 are < ~500. + +__global__ +void maskToAttnBiasFullKernel(const float* mask, float* outBias, int seqLen) { + int k = blockIdx.x * blockDim.x + threadIdx.x; + int q = blockIdx.y; + int b = blockIdx.z; + if(k >= seqLen) + return; + float m = mask[b * seqLen + k]; + outBias[((size_t)b * seqLen + q) * seqLen + k] = (m != 0.0f) ? 0.0f : -3e4f; +} + +__global__ +void maskToAttnBiasFullHalfKernel(const half* mask, half* outBias, int seqLen) { +#ifdef HIP_SUPPORTS_FP16 + int k = blockIdx.x * blockDim.x + threadIdx.x; + int q = blockIdx.y; + int b = blockIdx.z; + if(k >= seqLen) + return; + float m = __half2float(mask[b * seqLen + k]); + outBias[((size_t)b * seqLen + q) * seqLen + k] = __float2half((m != 0.0f) ? 0.0f : -3e4f); +#endif +} + +void customCudaMaskToAttnBiasFull(const float* mask, float* outBias, int batchSize, int seqLen) { + if(batchSize <= 0 || seqLen <= 0) + return; + int threads = 128; + dim3 blocks((seqLen + threads - 1) / threads, seqLen, batchSize); + maskToAttnBiasFullKernel<<>>(mask, outBias, seqLen); +} +void customCudaMaskToAttnBiasFull(const half* mask, half* outBias, int batchSize, int seqLen) { + if(batchSize <= 0 || seqLen <= 0) + return; + int threads = 128; + dim3 blocks((seqLen + threads - 1) / threads, seqLen, batchSize); + maskToAttnBiasFullHalfKernel<<>>(mask, outBias, seqLen); +} + +//-------------------------------------------------------------------------------------------------------------- +// SwiGLU: out[i] = SiLU(a[i]) * b[i] + +__global__ +void swiGLUKernel(const float* a, const float* b, float* out, int size) +{ + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if(idx < size) { + out[idx] = siluf(a[idx]) * b[idx]; + } +} + +__global__ +void swiGLUHalfKernel(const half* a, const half* b, half* out, int size) +{ +#ifdef HIP_SUPPORTS_FP16 + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if(idx < size) { + float av = __half2float(a[idx]); + float bv = __half2float(b[idx]); + out[idx] = __float2half(siluf(av) * bv); + } +#else + //Do nothing, FP16 not supported +#endif +} + +void customCudaSwiGLU(const float* a, const float* b, float* out, int size) { + if(size <= 0) + return; + int threads = targetNumThreads; + int blocks = (size + threads - 1) / threads; + swiGLUKernel<<>>(a, b, out, size); +} +void customCudaSwiGLU(const half* a, const half* b, half* out, int size) { + if(size <= 0) + return; + int threads = targetNumThreads; + int blocks = (size + threads - 1) / threads; + swiGLUHalfKernel<<>>(a, b, out, size); +} + +//-------------------------------------------------------------------------------------------------------------- +// Masked residual add: trunk[i] += residual[i] * mask[spatial_idx] +// NCHW: trunk/residual [n, c, xy], mask [n, xy] +// NHWC: trunk/residual [n, xy, c], mask [n, xy] + +__global__ +void maskedResidualAddNCHWKernel(float* trunk, const float* residual, const float* mask, int cSize, int xySize) +{ + int xyIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(xyIdx >= xySize || cIdx >= cSize) + return; + int idx = (nIdx * cSize + cIdx) * xySize + xyIdx; + float m = (mask != NULL) ? mask[nIdx * xySize + xyIdx] : 1.0f; + trunk[idx] += residual[idx] * m; +} + +__global__ +void maskedResidualAddNCHWHalfKernel(half* trunk, const half* residual, const half* mask, int cSize, int xySize) +{ +#ifdef HIP_SUPPORTS_FP16 + int xyIdx = blockIdx.x * blockDim.x + threadIdx.x; + int cIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(xyIdx >= xySize || cIdx >= cSize) + return; + int idx = (nIdx * cSize + cIdx) * xySize + xyIdx; + float m = (mask != NULL) ? __half2float(mask[nIdx * xySize + xyIdx]) : 1.0f; + trunk[idx] = __float2half(__half2float(trunk[idx]) + __half2float(residual[idx]) * m); +#else + //Do nothing, FP16 not supported +#endif +} + +void customCudaMaskedResidualAddNCHW(float* trunk, const float* residual, const float* mask, int nSize, int cSize, int xySize) { + if(nSize > 65536) + throw std::runtime_error("customCudaMaskedResidualAddNCHW: nSize too large"); + checkBufferIndexFitsInt(nSize, cSize, xySize, "customCudaMaskedResidualAddNCHW"); + int xyThreads, xyBlocks, cThreads, cBlocks; + splitThreadsAcrossDim01(xySize, cSize, xyThreads, xyBlocks, cThreads, cBlocks); + dim3 grid(xyBlocks, cBlocks, nSize); + dim3 threads(xyThreads, cThreads, 1); + maskedResidualAddNCHWKernel<<>>(trunk, residual, mask, cSize, xySize); +} +void customCudaMaskedResidualAddNCHW(half* trunk, const half* residual, const half* mask, int nSize, int cSize, int xySize) { + if(nSize > 65536) + throw std::runtime_error("customCudaMaskedResidualAddNCHW: nSize too large"); + checkBufferIndexFitsInt(nSize, cSize, xySize, "customCudaMaskedResidualAddNCHW"); + int xyThreads, xyBlocks, cThreads, cBlocks; + splitThreadsAcrossDim01(xySize, cSize, xyThreads, xyBlocks, cThreads, cBlocks); + dim3 grid(xyBlocks, cBlocks, nSize); + dim3 threads(xyThreads, cThreads, 1); + maskedResidualAddNCHWHalfKernel<<>>(trunk, residual, mask, cSize, xySize); +} + +__global__ +void maskedResidualAddNHWCKernel(float* trunk, const float* residual, const float* mask, int xySize, int cSize) +{ + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int xyIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx >= cSize || xyIdx >= xySize) + return; + int idx = (nIdx * xySize + xyIdx) * cSize + cIdx; + float m = (mask != NULL) ? mask[nIdx * xySize + xyIdx] : 1.0f; + trunk[idx] += residual[idx] * m; +} + +__global__ +void maskedResidualAddNHWCHalfKernel(half* trunk, const half* residual, const half* mask, int xySize, int cSize) +{ +#ifdef HIP_SUPPORTS_FP16 + int cIdx = blockIdx.x * blockDim.x + threadIdx.x; + int xyIdx = blockIdx.y * blockDim.y + threadIdx.y; + int nIdx = blockIdx.z; + if(cIdx >= cSize || xyIdx >= xySize) + return; + int idx = (nIdx * xySize + xyIdx) * cSize + cIdx; + float m = (mask != NULL) ? __half2float(mask[nIdx * xySize + xyIdx]) : 1.0f; + trunk[idx] = __float2half(__half2float(trunk[idx]) + __half2float(residual[idx]) * m); +#else + //Do nothing, FP16 not supported +#endif +} + +void customCudaMaskedResidualAddNHWC(float* trunk, const float* residual, const float* mask, int nSize, int xySize, int cSize) { + if(nSize > 65536) + throw std::runtime_error("customCudaMaskedResidualAddNHWC: nSize too large"); + checkBufferIndexFitsInt(nSize, xySize, cSize, "customCudaMaskedResidualAddNHWC"); + int cThreads, cBlocks, xyThreads, xyBlocks; + splitThreadsAcrossDim01(cSize, xySize, cThreads, cBlocks, xyThreads, xyBlocks); + dim3 grid(cBlocks, xyBlocks, nSize); + dim3 threads(cThreads, xyThreads, 1); + maskedResidualAddNHWCKernel<<>>(trunk, residual, mask, xySize, cSize); +} +void customCudaMaskedResidualAddNHWC(half* trunk, const half* residual, const half* mask, int nSize, int xySize, int cSize) { + if(nSize > 65536) + throw std::runtime_error("customCudaMaskedResidualAddNHWC: nSize too large"); + checkBufferIndexFitsInt(nSize, xySize, cSize, "customCudaMaskedResidualAddNHWC"); + int cThreads, cBlocks, xyThreads, xyBlocks; + splitThreadsAcrossDim01(cSize, xySize, cThreads, cBlocks, xyThreads, xyBlocks); + dim3 grid(cBlocks, xyBlocks, nSize); + dim3 threads(cThreads, xyThreads, 1); + maskedResidualAddNHWCHalfKernel<<>>(trunk, residual, mask, xySize, cSize); +} + +//-------------------------------------------------------------------------------------------------------------- +// RMSNorm with gamma/beta/activation (for trunk tip and transformer pre-norm, non-spatial mode). +// NHWC: input/output [n, xy, c], gamma/beta [c], mask [n, xy] +// Each block handles one (n, xy) position. Scalar (non-vectorized) kernels only, for portability +// across AMD wavefront sizes (no warp-shuffle assumptions). + +__global__ +void rmsNormGammaBetaNHWCKernel( + const float* in, float* out, const float* gamma, const float* beta, const float* mask, + int nSize, int xySize, int cSize, float epsilon, int activation +) { + extern __shared__ float rmsShared[]; + int pos = blockIdx.x; // n * xySize + xy + int tid = threadIdx.x; + int n = pos / xySize; + int xy = pos % xySize; + if(n >= nSize) + return; + + float maskVal = (mask != NULL) ? mask[n * xySize + xy] : 1.0f; + + const float* inRow = in + (size_t)pos * cSize; + + float acc = 0.0f; + for(int c = tid; c < cSize; c += blockDim.x) { + float val = inRow[c] * maskVal; + acc += val * val; + } + rmsShared[tid] = acc; + __syncthreads(); + for(int s = blockDim.x / 2; s > 0; s >>= 1) { + if(tid < s) rmsShared[tid] += rmsShared[tid + s]; + __syncthreads(); + } + float rms = rsqrtf(rmsShared[0] / (float)cSize + epsilon); + + float* outRow = out + (size_t)pos * cSize; + for(int c = tid; c < cSize; c += blockDim.x) { + float val = inRow[c] * maskVal * rms * gamma[c] + beta[c]; + if(activation == ACTIVATION_RELU) val = fmaxf(val, 0.0f); + else if(activation == ACTIVATION_MISH) val = mishf(val); + else if(activation == ACTIVATION_SILU) val = siluf(val); + val *= maskVal; + outRow[c] = val; + } +} + +__global__ +void rmsNormGammaBetaNHWCHalfKernel( + const half* in, half* out, const half* gamma, const half* beta, const half* mask, + int nSize, int xySize, int cSize, float epsilon, int activation +) { +#ifdef HIP_SUPPORTS_FP16 + extern __shared__ float rmsShared[]; + int pos = blockIdx.x; + int tid = threadIdx.x; + int n = pos / xySize; + int xy = pos % xySize; + if(n >= nSize) + return; + + float maskVal = (mask != NULL) ? __half2float(mask[n * xySize + xy]) : 1.0f; + + const half* inRow = in + (size_t)pos * cSize; + + float acc = 0.0f; + for(int c = tid; c < cSize; c += blockDim.x) { + float val = __half2float(inRow[c]) * maskVal; + acc += val * val; + } + rmsShared[tid] = acc; + __syncthreads(); + for(int s = blockDim.x / 2; s > 0; s >>= 1) { + if(tid < s) rmsShared[tid] += rmsShared[tid + s]; + __syncthreads(); + } + float rms = rsqrtf(rmsShared[0] / (float)cSize + epsilon); + + half* outRow = out + (size_t)pos * cSize; + for(int c = tid; c < cSize; c += blockDim.x) { + float val = __half2float(inRow[c]) * maskVal * rms * __half2float(gamma[c]) + __half2float(beta[c]); + if(activation == ACTIVATION_RELU) val = fmaxf(val, 0.0f); + else if(activation == ACTIVATION_MISH) val = mishf(val); + else if(activation == ACTIVATION_SILU) val = siluf(val); + val *= maskVal; + outRow[c] = __float2half(val); + } +#else + //Do nothing, FP16 not supported +#endif +} + +void customCudaRMSNormGammaBetaNHWC( + const float* in, float* out, const float* gamma, const float* beta, const float* mask, + int nSize, int xySize, int cSize, float epsilon, int activation +) { + checkBufferIndexFitsInt(nSize, xySize, cSize, "customCudaRMSNormGammaBetaNHWC"); + int totalPositions = nSize * xySize; + if(totalPositions <= 0) + return; + int threads = 1; + while(threads < cSize && threads < targetNumThreads) threads *= 2; + int sharedMem = threads * sizeof(float); + rmsNormGammaBetaNHWCKernel<<>>( + in, out, gamma, beta, mask, nSize, xySize, cSize, epsilon, activation); +} +void customCudaRMSNormGammaBetaNHWC( + const half* in, half* out, const half* gamma, const half* beta, const half* mask, + int nSize, int xySize, int cSize, float epsilon, int activation +) { + checkBufferIndexFitsInt(nSize, xySize, cSize, "customCudaRMSNormGammaBetaNHWC"); + int totalPositions = nSize * xySize; + if(totalPositions <= 0) + return; + int threads = 1; + while(threads < cSize && threads < targetNumThreads) threads *= 2; + int sharedMem = threads * sizeof(float); + rmsNormGammaBetaNHWCHalfKernel<<>>( + in, out, gamma, beta, mask, nSize, xySize, cSize, epsilon, activation); +} + +// NCHW variant: input/output [n, c, xy], gamma/beta [c], mask [n, xy] +__global__ +void rmsNormGammaBetaNCHWKernel( + const float* in, float* out, const float* gamma, const float* beta, const float* mask, + int nSize, int cSize, int xySize, float epsilon, int activation +) { + extern __shared__ float rmsShared[]; + int pos = blockIdx.x; // n * xySize + xy + int tid = threadIdx.x; + int n = pos / xySize; + int xy = pos % xySize; + if(n >= nSize) + return; + + float maskVal = (mask != NULL) ? mask[n * xySize + xy] : 1.0f; + + float acc = 0.0f; + for(int c = tid; c < cSize; c += blockDim.x) { + float val = in[(n * cSize + c) * xySize + xy] * maskVal; + acc += val * val; + } + rmsShared[tid] = acc; + __syncthreads(); + for(int s = blockDim.x / 2; s > 0; s >>= 1) { + if(tid < s) rmsShared[tid] += rmsShared[tid + s]; + __syncthreads(); + } + float rms = rsqrtf(rmsShared[0] / (float)cSize + epsilon); + + for(int c = tid; c < cSize; c += blockDim.x) { + float val = in[(n * cSize + c) * xySize + xy] * maskVal * rms * gamma[c] + beta[c]; + if(activation == ACTIVATION_RELU) val = fmaxf(val, 0.0f); + else if(activation == ACTIVATION_MISH) val = mishf(val); + else if(activation == ACTIVATION_SILU) val = siluf(val); + val *= maskVal; + out[(n * cSize + c) * xySize + xy] = val; + } +} + +__global__ +void rmsNormGammaBetaNCHWHalfKernel( + const half* in, half* out, const half* gamma, const half* beta, const half* mask, + int nSize, int cSize, int xySize, float epsilon, int activation +) { +#ifdef HIP_SUPPORTS_FP16 + extern __shared__ float rmsShared[]; + int pos = blockIdx.x; + int tid = threadIdx.x; + int n = pos / xySize; + int xy = pos % xySize; + if(n >= nSize) + return; + + float maskVal = (mask != NULL) ? __half2float(mask[n * xySize + xy]) : 1.0f; + + float acc = 0.0f; + for(int c = tid; c < cSize; c += blockDim.x) { + float val = __half2float(in[(n * cSize + c) * xySize + xy]) * maskVal; + acc += val * val; + } + rmsShared[tid] = acc; + __syncthreads(); + for(int s = blockDim.x / 2; s > 0; s >>= 1) { + if(tid < s) rmsShared[tid] += rmsShared[tid + s]; + __syncthreads(); + } + float rms = rsqrtf(rmsShared[0] / (float)cSize + epsilon); + + for(int c = tid; c < cSize; c += blockDim.x) { + float val = __half2float(in[(n * cSize + c) * xySize + xy]) * maskVal * rms * __half2float(gamma[c]) + __half2float(beta[c]); + if(activation == ACTIVATION_RELU) val = fmaxf(val, 0.0f); + else if(activation == ACTIVATION_MISH) val = mishf(val); + else if(activation == ACTIVATION_SILU) val = siluf(val); + val *= maskVal; + out[(n * cSize + c) * xySize + xy] = __float2half(val); + } +#else + //Do nothing, FP16 not supported +#endif +} + +void customCudaRMSNormGammaBetaNCHW( + const float* in, float* out, const float* gamma, const float* beta, const float* mask, + int nSize, int cSize, int xySize, float epsilon, int activation +) { + checkBufferIndexFitsInt(nSize, cSize, xySize, "customCudaRMSNormGammaBetaNCHW"); + int totalPositions = nSize * xySize; + if(totalPositions <= 0) + return; + int threads = 1; + while(threads < cSize && threads < targetNumThreads) threads *= 2; + int sharedMem = threads * sizeof(float); + rmsNormGammaBetaNCHWKernel<<>>( + in, out, gamma, beta, mask, nSize, cSize, xySize, epsilon, activation); +} +void customCudaRMSNormGammaBetaNCHW( + const half* in, half* out, const half* gamma, const half* beta, const half* mask, + int nSize, int cSize, int xySize, float epsilon, int activation +) { + checkBufferIndexFitsInt(nSize, cSize, xySize, "customCudaRMSNormGammaBetaNCHW"); + int totalPositions = nSize * xySize; + if(totalPositions <= 0) + return; + int threads = 1; + while(threads < cSize && threads < targetNumThreads) threads *= 2; + int sharedMem = threads * sizeof(float); + rmsNormGammaBetaNCHWHalfKernel<<>>( + in, out, gamma, beta, mask, nSize, cSize, xySize, epsilon, activation); +} + +//-------------------------------------------------------------------------------------------------------------- +// Spatial RMSNorm: normalize over all C*H*W per batch element. +// Three-pass, deterministic: +// Pass 1 (SumSq): grid (numBlocksPerBatch, nSize). Many blocks per batch element grid-stride over +// the flat C*xy range, reduce in-block, write one partial per block into partialBuf. +// Pass 2 (Reduce): grid (nSize). One block per batch element sums its numBlocksPerBatch partials +// (fixed order) into sumSqBuf[n]. +// Pass 3 (Apply): grid (numApplyBlocks, nSize). Normalize + activation + remask. + +static const int SPATIAL_RMSNORM_BLOCKS_PER_BATCH = 8; +static_assert(CUDA_SPATIAL_RMSNORM_SUMSQ_STRIDE == SPATIAL_RMSNORM_BLOCKS_PER_BATCH + 1, + "CUDA_SPATIAL_RMSNORM_SUMSQ_STRIDE must equal SPATIAL_RMSNORM_BLOCKS_PER_BATCH + 1"); + +static int spatialRMSNormBlocksPerBatch(int totalElems) { + int maxUseful = (totalElems + targetNumThreads - 1) / targetNumThreads; + if(maxUseful < 1) maxUseful = 1; + int b = SPATIAL_RMSNORM_BLOCKS_PER_BATCH; + if(b > maxUseful) b = maxUseful; + return b; +} + +template +__global__ +void spatialRMSNormSumSqKernel( + const float* in, const float* mask, float* partialBuf, + int totalElems, int cSize, int xySize, int numBlocksPerBatch, int partialStride +) { + extern __shared__ float srmsShared[]; + int n = blockIdx.y; + int blk = blockIdx.x; + int tid = threadIdx.x; + + const float* inRow = in + (size_t)n * totalElems; + + float acc = 0.0f; + for(int i = blk * blockDim.x + tid; i < totalElems; i += blockDim.x * numBlocksPerBatch) { + int xy = IS_NHWC ? (i / cSize) : (i % xySize); + float m = (mask != NULL) ? mask[n * xySize + xy] : 1.0f; + float val = inRow[i] * m; + acc += val * val; + } + srmsShared[tid] = acc; + __syncthreads(); + for(int s = blockDim.x / 2; s > 0; s >>= 1) { + if(tid < s) srmsShared[tid] += srmsShared[tid + s]; + __syncthreads(); + } + if(tid == 0) partialBuf[n * partialStride + blk] = srmsShared[0]; +} + +template +__global__ +void spatialRMSNormSumSqHalfKernel( + const half* in, const half* mask, float* partialBuf, + int totalElems, int cSize, int xySize, int numBlocksPerBatch, int partialStride +) { +#ifdef HIP_SUPPORTS_FP16 + extern __shared__ float srmsShared[]; + int n = blockIdx.y; + int blk = blockIdx.x; + int tid = threadIdx.x; + + const half* inRow = in + (size_t)n * totalElems; + + float acc = 0.0f; + for(int i = blk * blockDim.x + tid; i < totalElems; i += blockDim.x * numBlocksPerBatch) { + int xy = IS_NHWC ? (i / cSize) : (i % xySize); + float m = (mask != NULL) ? __half2float(mask[n * xySize + xy]) : 1.0f; + float val = __half2float(inRow[i]) * m; + acc += val * val; + } + srmsShared[tid] = acc; + __syncthreads(); + for(int s = blockDim.x / 2; s > 0; s >>= 1) { + if(tid < s) srmsShared[tid] += srmsShared[tid + s]; + __syncthreads(); + } + if(tid == 0) partialBuf[n * partialStride + blk] = srmsShared[0]; +#else + //Do nothing, FP16 not supported +#endif +} + +__global__ +void spatialRMSNormReduceKernel( + const float* partialBuf, float* sumSqBuf, int numBlocksPerBatch, int partialStride +) { + int n = blockIdx.x; + if(threadIdx.x != 0) + return; + float total = 0.0f; + const float* row = partialBuf + (size_t)n * partialStride; + for(int b = 0; b < numBlocksPerBatch; b++) + total += row[b]; + sumSqBuf[n * partialStride + numBlocksPerBatch] = total; +} + +__global__ +void spatialRMSNormApplyNHWCKernel( + const float* in, float* out, const float* gamma, const float* beta, const float* mask, + const float* maskSum, const float* sumSqBuf, + int totalElems, int cSize, int xySize, float epsilon, int activation, int numBlocksPerBatch, int partialStride +) { + int n = blockIdx.y; + float mSum = maskSum[n]; + float totalSize = mSum * (float)cSize; + float rms = rsqrtf(sumSqBuf[n * partialStride + numBlocksPerBatch] / totalSize + epsilon); + + const float* inRow = in + (size_t)n * totalElems; + float* outRow = out + (size_t)n * totalElems; + + for(int i = blockIdx.x * blockDim.x + threadIdx.x; i < totalElems; i += blockDim.x * gridDim.x) { + int xy = i / cSize; + int c = i - xy * cSize; + float maskVal = (mask != NULL) ? mask[n * xySize + xy] : 1.0f; + float val = inRow[i] * maskVal * rms * gamma[c] + beta[c]; + if(activation == ACTIVATION_RELU) val = fmaxf(val, 0.0f); + else if(activation == ACTIVATION_MISH) val = mishf(val); + else if(activation == ACTIVATION_SILU) val = siluf(val); + val *= maskVal; + outRow[i] = val; + } +} + +__global__ +void spatialRMSNormApplyNHWCHalfKernel( + const half* in, half* out, const half* gamma, const half* beta, const half* mask, + const float* maskSum, const float* sumSqBuf, + int totalElems, int cSize, int xySize, float epsilon, int activation, int numBlocksPerBatch, int partialStride +) { +#ifdef HIP_SUPPORTS_FP16 + int n = blockIdx.y; + float mSum = maskSum[n]; + float totalSize = mSum * (float)cSize; + float rms = rsqrtf(sumSqBuf[n * partialStride + numBlocksPerBatch] / totalSize + epsilon); + + const half* inRow = in + (size_t)n * totalElems; + half* outRow = out + (size_t)n * totalElems; + + for(int i = blockIdx.x * blockDim.x + threadIdx.x; i < totalElems; i += blockDim.x * gridDim.x) { + int xy = i / cSize; + int c = i - xy * cSize; + float maskVal = (mask != NULL) ? __half2float(mask[n * xySize + xy]) : 1.0f; + float val = __half2float(inRow[i]) * maskVal * rms * __half2float(gamma[c]) + __half2float(beta[c]); + if(activation == ACTIVATION_RELU) val = fmaxf(val, 0.0f); + else if(activation == ACTIVATION_MISH) val = mishf(val); + else if(activation == ACTIVATION_SILU) val = siluf(val); + val *= maskVal; + outRow[i] = __float2half(val); + } +#else + //Do nothing, FP16 not supported +#endif +} + +__global__ +void spatialRMSNormApplyNCHWKernel( + const float* in, float* out, const float* gamma, const float* beta, const float* mask, + const float* maskSum, const float* sumSqBuf, + int totalElems, int cSize, int xySize, float epsilon, int activation, int numBlocksPerBatch, int partialStride +) { + int n = blockIdx.y; + float mSum = maskSum[n]; + float totalSize = mSum * (float)cSize; + float rms = rsqrtf(sumSqBuf[n * partialStride + numBlocksPerBatch] / totalSize + epsilon); + + const float* inRow = in + (size_t)n * totalElems; + float* outRow = out + (size_t)n * totalElems; + + for(int i = blockIdx.x * blockDim.x + threadIdx.x; i < totalElems; i += blockDim.x * gridDim.x) { + int c = i / xySize; + int xy = i - c * xySize; + float maskVal = (mask != NULL) ? mask[n * xySize + xy] : 1.0f; + float val = inRow[i] * maskVal * rms * gamma[c] + beta[c]; + if(activation == ACTIVATION_RELU) val = fmaxf(val, 0.0f); + else if(activation == ACTIVATION_MISH) val = mishf(val); + else if(activation == ACTIVATION_SILU) val = siluf(val); + val *= maskVal; + outRow[i] = val; + } +} + +__global__ +void spatialRMSNormApplyNCHWHalfKernel( + const half* in, half* out, const half* gamma, const half* beta, const half* mask, + const float* maskSum, const float* sumSqBuf, + int totalElems, int cSize, int xySize, float epsilon, int activation, int numBlocksPerBatch, int partialStride +) { +#ifdef HIP_SUPPORTS_FP16 + int n = blockIdx.y; + float mSum = maskSum[n]; + float totalSize = mSum * (float)cSize; + float rms = rsqrtf(sumSqBuf[n * partialStride + numBlocksPerBatch] / totalSize + epsilon); + + const half* inRow = in + (size_t)n * totalElems; + half* outRow = out + (size_t)n * totalElems; + + for(int i = blockIdx.x * blockDim.x + threadIdx.x; i < totalElems; i += blockDim.x * gridDim.x) { + int c = i / xySize; + int xy = i - c * xySize; + float maskVal = (mask != NULL) ? __half2float(mask[n * xySize + xy]) : 1.0f; + float val = __half2float(inRow[i]) * maskVal * rms * __half2float(gamma[c]) + __half2float(beta[c]); + if(activation == ACTIVATION_RELU) val = fmaxf(val, 0.0f); + else if(activation == ACTIVATION_MISH) val = mishf(val); + else if(activation == ACTIVATION_SILU) val = siluf(val); + val *= maskVal; + outRow[i] = __float2half(val); + } +#else + //Do nothing, FP16 not supported +#endif +} + +static int spatialRMSNormApplyBlocks(int totalElems, int threads) { + int blocks = (totalElems + threads - 1) / threads; + if(blocks < 1) blocks = 1; + if(blocks > 256) blocks = 256; + return blocks; +} + +void customCudaSpatialRMSNormNHWC( + const float* in, float* out, const float* gamma, const float* beta, const float* mask, const float* maskSum, + int nSize, int xySize, int cSize, float epsilon, int activation, float* sumSqBuf +) { + if(nSize <= 0) + return; + if(nSize > 65536) + throw std::runtime_error("customCudaSpatialRMSNormNHWC: nSize too large"); + checkBufferIndexFitsInt(nSize, xySize, cSize, "customCudaSpatialRMSNormNHWC"); + int totalElems = xySize * cSize; + int numBlocksPerBatch = spatialRMSNormBlocksPerBatch(totalElems); + int partialStride = SPATIAL_RMSNORM_BLOCKS_PER_BATCH + 1; + + int threads1 = targetNumThreads; + int sharedMem1 = threads1 * sizeof(float); + dim3 grid1(numBlocksPerBatch, nSize); + spatialRMSNormSumSqKernel<<>>( + in, mask, sumSqBuf, totalElems, cSize, xySize, numBlocksPerBatch, partialStride); + + spatialRMSNormReduceKernel<<>>(sumSqBuf, sumSqBuf, numBlocksPerBatch, partialStride); + + int threads2 = targetNumThreads; + int applyBlocks = spatialRMSNormApplyBlocks(totalElems, threads2); + dim3 grid2(applyBlocks, nSize); + spatialRMSNormApplyNHWCKernel<<>>( + in, out, gamma, beta, mask, maskSum, sumSqBuf, totalElems, cSize, xySize, epsilon, activation, numBlocksPerBatch, partialStride); +} +void customCudaSpatialRMSNormNHWC( + const half* in, half* out, const half* gamma, const half* beta, const half* mask, const float* maskSum, + int nSize, int xySize, int cSize, float epsilon, int activation, float* sumSqBuf +) { + if(nSize <= 0) + return; + if(nSize > 65536) + throw std::runtime_error("customCudaSpatialRMSNormNHWC: nSize too large"); + checkBufferIndexFitsInt(nSize, xySize, cSize, "customCudaSpatialRMSNormNHWC"); + int totalElems = xySize * cSize; + int numBlocksPerBatch = spatialRMSNormBlocksPerBatch(totalElems); + int partialStride = SPATIAL_RMSNORM_BLOCKS_PER_BATCH + 1; + + int threads1 = targetNumThreads; + int sharedMem1 = threads1 * sizeof(float); + dim3 grid1(numBlocksPerBatch, nSize); + spatialRMSNormSumSqHalfKernel<<>>( + in, mask, sumSqBuf, totalElems, cSize, xySize, numBlocksPerBatch, partialStride); + + spatialRMSNormReduceKernel<<>>(sumSqBuf, sumSqBuf, numBlocksPerBatch, partialStride); + + int threads2 = targetNumThreads; + int applyBlocks = spatialRMSNormApplyBlocks(totalElems, threads2); + dim3 grid2(applyBlocks, nSize); + spatialRMSNormApplyNHWCHalfKernel<<>>( + in, out, gamma, beta, mask, maskSum, sumSqBuf, totalElems, cSize, xySize, epsilon, activation, numBlocksPerBatch, partialStride); +} + +void customCudaSpatialRMSNormNCHW( + const float* in, float* out, const float* gamma, const float* beta, const float* mask, const float* maskSum, + int nSize, int cSize, int xySize, float epsilon, int activation, float* sumSqBuf +) { + if(nSize <= 0) + return; + if(nSize > 65536) + throw std::runtime_error("customCudaSpatialRMSNormNCHW: nSize too large"); + checkBufferIndexFitsInt(nSize, cSize, xySize, "customCudaSpatialRMSNormNCHW"); + int totalElems = cSize * xySize; + int numBlocksPerBatch = spatialRMSNormBlocksPerBatch(totalElems); + int partialStride = SPATIAL_RMSNORM_BLOCKS_PER_BATCH + 1; + + int threads1 = targetNumThreads; + int sharedMem1 = threads1 * sizeof(float); + dim3 grid1(numBlocksPerBatch, nSize); + spatialRMSNormSumSqKernel<<>>( + in, mask, sumSqBuf, totalElems, cSize, xySize, numBlocksPerBatch, partialStride); + + spatialRMSNormReduceKernel<<>>(sumSqBuf, sumSqBuf, numBlocksPerBatch, partialStride); + + int threads2 = targetNumThreads; + int applyBlocks = spatialRMSNormApplyBlocks(totalElems, threads2); + dim3 grid2(applyBlocks, nSize); + spatialRMSNormApplyNCHWKernel<<>>( + in, out, gamma, beta, mask, maskSum, sumSqBuf, totalElems, cSize, xySize, epsilon, activation, numBlocksPerBatch, partialStride); +} +void customCudaSpatialRMSNormNCHW( + const half* in, half* out, const half* gamma, const half* beta, const half* mask, const float* maskSum, + int nSize, int cSize, int xySize, float epsilon, int activation, float* sumSqBuf +) { + if(nSize <= 0) + return; + if(nSize > 65536) + throw std::runtime_error("customCudaSpatialRMSNormNCHW: nSize too large"); + checkBufferIndexFitsInt(nSize, cSize, xySize, "customCudaSpatialRMSNormNCHW"); + int totalElems = cSize * xySize; + int numBlocksPerBatch = spatialRMSNormBlocksPerBatch(totalElems); + int partialStride = SPATIAL_RMSNORM_BLOCKS_PER_BATCH + 1; + + int threads1 = targetNumThreads; + int sharedMem1 = threads1 * sizeof(float); + dim3 grid1(numBlocksPerBatch, nSize); + spatialRMSNormSumSqHalfKernel<<>>( + in, mask, sumSqBuf, totalElems, cSize, xySize, numBlocksPerBatch, partialStride); + + spatialRMSNormReduceKernel<<>>(sumSqBuf, sumSqBuf, numBlocksPerBatch, partialStride); + + int threads2 = targetNumThreads; + int applyBlocks = spatialRMSNormApplyBlocks(totalElems, threads2); + dim3 grid2(applyBlocks, nSize); + spatialRMSNormApplyNCHWHalfKernel<<>>( + in, out, gamma, beta, mask, maskSum, sumSqBuf, totalElems, cSize, xySize, epsilon, activation, numBlocksPerBatch, partialStride); +} diff --git a/cpp/neuralnet/rocmincludes.h b/cpp/neuralnet/rocmincludes.h new file mode 100644 index 0000000000..8b494a37e1 --- /dev/null +++ b/cpp/neuralnet/rocmincludes.h @@ -0,0 +1,15 @@ +#ifndef NEURALNET_ROCMINCLUDES_H +#define NEURALNET_ROCMINCLUDES_H + +//Ensure that CUDA_API_PER_THREAD_DEFAULT_STREAM is always defined +//before any cuda headers are included so that we get the desired threading behavior for CUDA. + +#define CUDA_API_PER_THREAD_DEFAULT_STREAM +#include +#include + +#include +#include + + +#endif //NEURALNET_ROCMINCLUDES_H diff --git a/cpp/neuralnet/rocmutils.cpp b/cpp/neuralnet/rocmutils.cpp new file mode 100644 index 0000000000..752298b7ff --- /dev/null +++ b/cpp/neuralnet/rocmutils.cpp @@ -0,0 +1,170 @@ +#include "../neuralnet/rocmutils.h" + +#include +#include "../neuralnet/rocmerrorcheck.h" +#include "../neuralnet/rocmincludes.h" +#include "../neuralnet/rocmhelpers.h" + +#include "../external/half-2.2.0/include/half.hpp" + +//------------------------ +#include "../core/using.h" +//------------------------ + +using half_t = half_float::half; + +void CudaUtils::mallocOnDevice(const string& name, int numWeights, void*& deviceBuf, bool useFP16) { + if(useFP16) { + size_t halfBytes = numWeights * sizeof(half_t); + CUDA_ERR(name.c_str(),hipMalloc(&deviceBuf, halfBytes)); + } + else { + size_t floatBytes = numWeights * sizeof(float); + CUDA_ERR(name.c_str(),hipMalloc(&deviceBuf, floatBytes)); + } +} + +void CudaUtils::mallocAndCopyToDevice(const string& name, const vector& weights, void*& deviceBuf, bool useFP16) { + size_t numWeights = weights.size(); + if(useFP16) { + size_t halfBytes = numWeights * sizeof(half_t); + vector weightsHalf(weights.size()); + for(size_t i = 0; i(weights[i]); + CUDA_ERR(name.c_str(),hipMalloc(&deviceBuf, halfBytes)); + CUDA_ERR(name.c_str(),hipMemcpy(deviceBuf, weightsHalf.data(), halfBytes, hipMemcpyHostToDevice)); + } + else { + size_t floatBytes = numWeights * sizeof(float); + CUDA_ERR(name.c_str(),hipMalloc(&deviceBuf, floatBytes)); + CUDA_ERR(name.c_str(),hipMemcpy(deviceBuf, weights.data(), floatBytes, hipMemcpyHostToDevice)); + } +} + +void CudaUtils::mallocAndCopyToDevice(const string& name, const float* weights, int numWeights, void*& deviceBuf, bool useFP16) { + if(useFP16) { + size_t halfBytes = numWeights * sizeof(half_t); + vector weightsHalf(numWeights); + for(int i = 0; i(weights[i]); + CUDA_ERR(name.c_str(),hipMalloc(&deviceBuf, halfBytes)); + CUDA_ERR(name.c_str(),hipMemcpy(deviceBuf, weightsHalf.data(), halfBytes, hipMemcpyHostToDevice)); + } + else { + size_t floatBytes = numWeights * sizeof(float); + CUDA_ERR(name.c_str(),hipMalloc(&deviceBuf, floatBytes)); + CUDA_ERR(name.c_str(),hipMemcpy(deviceBuf, weights, floatBytes, hipMemcpyHostToDevice)); + } +} + +//Only use in testing, allocates an intermediate buffer in the case of FP16 which will be very slow. +void CudaUtils::expensiveCopyFromDevice(const string& name, float* weights, int numWeights, const void* deviceBuf, bool useFP16) { + if(useFP16) { + vector weightsHalf(numWeights); + size_t halfBytes = numWeights * sizeof(half_t); + CUDA_ERR(name.c_str(),hipMemcpy(weightsHalf.data(), deviceBuf, halfBytes, hipMemcpyDeviceToHost)); + for(int i = 0; i values(batchSize * cSize); + expensiveCopyFromDevice(name, values.data(), values.size(), deviceBuf, useFP16); + cout << "=========================================================" << endl; + cout << "TENSOR" << endl; + cout << name << endl; + cout << std::setprecision(8); + int i = 0; + for(int n = 0; n values(batchSize * cSize * xSize * ySize); + expensiveCopyFromDevice(name, values.data(), values.size(), deviceBuf, useFP16); + cout << "=========================================================" << endl; + cout << "TENSOR" << endl; + cout << name << endl; + cout << std::setprecision(8); + int i = 0; + double total1 = 0; + double total2 = 0; + double total3 = 0; + for(int n = 0; n= (int64_t)1 << 31) + throw StringError("Batch size too large, resulting GPU buffers might exceed 2^31 entries which is not currently supported"); +} + +void CudaUtils::hostMallocZeroOneBufs(void*& zeroBuf, void*& oneBuf, bool useFP16) { + if(!useFP16) { + zeroBuf = malloc(sizeof(float)); + oneBuf = malloc(sizeof(float)); + *((float*)zeroBuf) = 0.0f; + *((float*)oneBuf) = 1.0f; + } + else { + //Convert to FP16 on the device, then copy back so we have it in host memory + float zero = 0.0f; + float one = 1.0f; + void* zeroTmp; + void* oneTmp; + mallocAndCopyToDevice("Buffers",&zero,1,zeroTmp,useFP16); + mallocAndCopyToDevice("Buffers",&one,1,oneTmp,useFP16); + zeroBuf = malloc(sizeof(half_t)); + oneBuf = malloc(sizeof(half_t)); + CUDA_ERR("Buffers",hipMemcpy(zeroBuf,zeroTmp,sizeof(half_t),hipMemcpyDeviceToHost)); + CUDA_ERR("Buffers",hipMemcpy(oneBuf,oneTmp,sizeof(half_t),hipMemcpyDeviceToHost)); + hipFree(zeroTmp); + hipFree(oneTmp); + } +} diff --git a/cpp/neuralnet/rocmutils.h b/cpp/neuralnet/rocmutils.h new file mode 100644 index 0000000000..d438684026 --- /dev/null +++ b/cpp/neuralnet/rocmutils.h @@ -0,0 +1,21 @@ +#ifndef NEURALNET_ROCMUTILS_H +#define NEURALNET_ROCMUTILS_H + +#include "../core/global.h" + +namespace CudaUtils { + void mallocOnDevice(const std::string& name, int numWeights, void*& deviceBuf, bool useFP16); + void mallocAndCopyToDevice(const std::string& name, const std::vector& weights, void*& deviceBuf, bool useFP16); + void mallocAndCopyToDevice(const std::string& name, const float* weights, int numWeights, void*& deviceBuf, bool useFP16); + + //Only use in testing, allocates an intermediate buffer in the case of FP16 which will be very slow. + void expensiveCopyFromDevice(const std::string& name, float* weights, int numWeights, const void* deviceBuf, bool useFP16); + + void debugPrint2D(const std::string& name, const void* deviceBuf, int batchSize, int cSize, bool useFP16); + void debugPrint4D(const std::string& name, const void* deviceBuf, int batchSize, int cSize, int xSize, int ySize, bool useNHWC, bool useFP16); + + void checkBufferSize(int batchSize, int xSize, int ySize, int channels); + void hostMallocZeroOneBufs(void*& zeroBuf, void*& oneBuf, bool useFP16); +} + +#endif // NEURALNET_ROCMUTILS_H diff --git a/cpp/program/gtpconfig.cpp b/cpp/program/gtpconfig.cpp index 3fa8651e5b..64c890b6c2 100644 --- a/cpp/program/gtpconfig.cpp +++ b/cpp/program/gtpconfig.cpp @@ -537,6 +537,9 @@ string GTPConfig::makeConfig( #endif #ifdef USE_OPENCL_BACKEND replacement += "openclDeviceToUseThread" + Global::intToString(i) + " = " + Global::intToString(deviceIdxs[i]) + "\n"; +#endif +#ifdef USE_ROCM_BACKEND + replacement += "rocmDeviceToUseThread" + Global::intToString(i) + " = " + Global::intToString(deviceIdxs[i]) + "\n"; #endif } replace("$$MULTIPLE_GPUS", replacement); diff --git a/cpp/program/setup.cpp b/cpp/program/setup.cpp index bf7950315e..d40f06c614 100644 --- a/cpp/program/setup.cpp +++ b/cpp/program/setup.cpp @@ -20,6 +20,7 @@ std::vector Setup::getBackendPrefixes() { prefixes.push_back("trt"); prefixes.push_back("metal"); prefixes.push_back("opencl"); + prefixes.push_back("rocm"); prefixes.push_back("eigen"); prefixes.push_back("dummybackend"); return prefixes; @@ -87,6 +88,8 @@ vector Setup::initializeNNEvaluators( string backendPrefix = "metal"; #elif defined(USE_OPENCL_BACKEND) string backendPrefix = "opencl"; + #elif defined(USE_ROCM_BACKEND) + string backendPrefix = "rocm"; #elif defined(USE_EIGEN_BACKEND) string backendPrefix = "eigen"; #else @@ -142,7 +145,7 @@ vector Setup::initializeNNEvaluators( requireExactNNLen = cfg.getBool("requireMaxBoardSize"); } - bool inputsUseNHWC = backendPrefix == "opencl" || backendPrefix == "trt" || backendPrefix == "metal" ? false : true; + bool inputsUseNHWC = backendPrefix == "opencl" || backendPrefix == "trt" || backendPrefix == "metal" || backendPrefix == "rocm" || backendPrefix == "mgx" ? false : true; if(cfg.contains(backendPrefix+"InputsUseNHWC"+idxStr)) inputsUseNHWC = cfg.getBool(backendPrefix+"InputsUseNHWC"+idxStr); else if(cfg.contains("inputsUseNHWC"+idxStr))