RDNA4 phase 3: scalar FlashAttention, FP8/int8 KV, decode and softmax on HIP

AGENTS.md

@@ -29,7 +29,7 @@ ctest -R '^utest_softmax_cu$'                    # or via ctest, from build/
 ctest -R '_cpp$'                                 # CPU-backend tests only (no GPU needed at runtime)
 ```
 
-Key CMake options: `ENABLE_CPU` (ON), `ENABLE_CPU_MULTITHREAD` (ON if the `thread-pool` submodule is checked out, OFF otherwise; forcing it ON without the submodule is a `FATAL_ERROR`; when ON the FKL target defines `FK_CPU_MULTITHREAD_ENABLED`, adds `thread-pool/include` (installed as `include/fused_kernel/thread_pool`) and links `Threads::Threads`), `ENABLE_CUDA` (ON if nvcc found), `ENABLE_HIP` (ON if amdclang is found), `BUILD_TEST` (ON), `BUILD_UTEST` (ON), `ENABLE_BENCHMARK` (OFF), `CUDA_ARCH` ("native", passed verbatim to `CUDA_ARCHITECTURES` — no arch filtering exists), gfx1200 by default for HIP, `ARCH_FLAGS` (CPU SIMD), `ENABLE_NVTX`, `ENABLE_DEBUG`, `TEMPLATE_DEPTH` (1000).
+Key CMake options: `ENABLE_CPU` (ON), `ENABLE_CPU_MULTITHREAD` (ON if the `thread-pool` submodule is checked out, OFF otherwise; forcing it ON without the submodule is a `FATAL_ERROR`; when ON the FKL target defines `FK_CPU_MULTITHREAD_ENABLED`, adds `thread-pool/include` (installed as `include/fused_kernel/thread_pool`) and links `Threads::Threads`), `ENABLE_CUDA` (ON if nvcc found), `ENABLE_HIP` (ON if amdclang is found), `BUILD_TEST` (ON), `BUILD_UTEST` (ON), `ENABLE_BENCHMARK` (OFF), `CUDA_ARCH` ("native", passed verbatim to `CUDA_ARCHITECTURES` — no arch filtering exists), HIP: the build machine's GPU by default (`CMAKE_HIP_ARCHITECTURES` as detected by CMake), or `HIP_ARCH` (e.g. `gfx1200;gfx1201`, RX 9060 and RX 9070, which can't load each other's code objects), which seeds `CMAKE_HIP_ARCHITECTURES` on the first configure of a build directory; an explicit `CMAKE_HIP_ARCHITECTURES` wins, and `hip::device` follows it through `GPU_TARGETS`, `ARCH_FLAGS` (CPU SIMD), `ENABLE_NVTX`, `ENABLE_DEBUG`, `TEMPLATE_DEPTH` (1000).
 
 `cmake --install build --prefix <dir>` installs the headers (`include/fused_kernel`, plus `include/fused_kernel/thread_pool` with the backend) and the `FKL` package (`share/cmake/FKL`): consumers use `find_package(FKL)` and link `FKL::FKL`, as with `add_subdirectory`. `FKL::FKL` requires C++20 with the `cxx_std_20` compile feature, which CMake also applies to CUDA and HIP sources that don't set `CUDA_STANDARD`/`HIP_STANDARD` (never add `cuda_std_20`/`hip_std_20`: with `add_subdirectory` from a C++-only project they fail with "No known features for CUDA compiler"); the project using FKL must enable CXX, and both the root `CMakeLists.txt` (as a subproject) and the package configuration stop with a clear message otherwise. Without `-DCMAKE_INSTALL_PREFIX` the install goes to `<build>/install`. `lib/export/` is not the installed package: it is the package that projects with FKL as a git submodule (cvGPUSpeedup, FastNPP) find without building FKL, and configure keeps its version file up to date; don't remove it while they use it. The package configuration template is `cmake/FKLConfig.cmake.in`, which `find_dependency(Threads)` when the backend is enabled; the version file accepts the same minor version before 1.0 (`find_package(FKL 0.2)` finds only 0.2.x) and the same major version from 1.0 on.
 
@@ -39,7 +39,7 @@ There is no lint/format gate. Format manually with `clang-format -i` using the r
 
 - Tests are header files auto-discovered by CMake (`cmake/tests/discover_tests.cmake`, GLOB_RECURSE with CONFIGURE_DEPENDS) inside immediate **subdirectories** of `tests/` and `utests/`. Top-level files (`tests/main.h`, `tests/operation_test_utils.h`) are never tests. Paths containing `_common` are excluded (shared helpers). Adding a new `.h` test requires no CMake edits.
 - Every test header defines `int launch()` returning 0 for pass. A generated launcher TU compiles the header twice: as C++ (`<name>_cpp` target, `ParArch::CPU`) and as CUDA (`<name>_cu`, `ParArch::GPU_NVIDIA`). The same source targets both backends because `defaultParArch` switches on `__NVCC__`.
-- Backend suppression is a raw **substring** search: any occurrence of `ONLY_CU` / `ONLY_CPU` anywhere in the file (even in a comment) suppresses the `_cpp` / `_cu` target. Repo convention: `#define __ONLY_CU__`.
+- Backend suppression is a raw **substring** search: any occurrence of `ONLY_CU` / `ONLY_CPU` / `ONLY_GPU` anywhere in the file (even in a comment) suppresses targets: `ONLY_CU` keeps only `_cu`, `ONLY_CPU` keeps only `_cpp`, `ONLY_GPU` keeps `_cu` and `_hip` (no `_cpp`). A HIP target (`_hip`) is generated only when neither `ONLY_CU` nor `ONLY_CPU` appears. Repo convention: `#define __ONLY_CU__` / `#define __ONLY_GPU__`.
 - Test harness helpers (`START_ADDING_TESTS`, `STOP_ADDING_TESTS`, `RUN_ALL_TESTS`, `TestCaseBuilder`, the `testCases` map) live in `tests/operation_test_utils.h`; `tests/main.h` only declares `launch()`.
 - Every public alias and `build()` overload needs a utest that instantiates it — template code only breaks on instantiation (see issue #244: shipped aliases that never compiled).
 - Test targets link `FKL::FKL`, so with `ENABLE_CPU_MULTITHREAD=ON` both `_cpp` and `_cu` get `FK_CPU_MULTITHREAD_ENABLED`. CPU_MULTITHREAD coverage goes in each DPP family's existing test under `#ifdef FK_CPU_MULTITHREAD_ENABLED` (it must also compile under nvcc): run every case through the `Stream_<ParArch::CPU_MULTITHREAD>&` entry point with host memory, `sync()`, and require the output to match `ParArch::CPU`, plus the oracle: exactly (`memcmp`/`==`) for integer outputs and float pipelines without multiply-add, within a tight tolerance (e.g. 1e-5 relative) for float pipelines that multiply and add (fused multiply-add contraction). A test that means "GPU" must say `PA == ParArch::GPU_NVIDIA || PA == ParArch::GPU_AMD`, not `PA != ParArch::CPU`. The backend itself is tested by `utests/core/execution_model/utest_cpu_multithread.h` (`_cpp` only).
@@ -85,7 +85,7 @@ Pattern (see `Mul` in `algorithms/basic_ops/arithmetic.h`, `Crop` in `algorithms
 
 ### attention/ — GPU-only exception to the rules
 
-`algorithms/attention/` (softmax, flash_attention, flash_attention_mma, flash_decode) contains self-contained cooperative DPPs (whole kernels: online softmax, FlashAttention-2 SIMT and mma.sync tensor-core variants, split-KV FlashDecoding), not composable Operations. The DPP structs are inside `#if defined(__NVCC__)` and remain **GPU-only: no `ParArch::CPU` nor `ParArch::CPU_MULTITHREAD` implementation**. This is a documented exception, decided by the repo owner, to the DPP skill's rule that every DPP has both; it is not a precedent for other DPPs. The directory also holds composable CPU-compiling pieces outside the guard (e.g. `Int8TokenDequantRead`, a normal Read Operation usable in any FKL pipeline, plus host packing helpers). Nothing in `attention/` is included by any umbrella header. Cooperative exec bodies can't be constexpr (shared memory + barriers), so softmax and flash_attention use `FK_COOP_DEVICE_FUSE` while the mma and decode variants hand-write `static __device__ void exec`. Fusion happens via Read-IOp prologues (dequantization runs in-register at load time) and IOp-chain epilogues. `DivergentBatchTransformDPP` (Divergent HF) is not an exception: it has `ParArch::CPU` and `ParArch::CPU_MULTITHREAD` DPPs and executors.
+`algorithms/attention/` (softmax, flash_attention, flash_attention_mma, flash_decode) contains self-contained cooperative DPPs (whole kernels: online softmax, FlashAttention-2 SIMT and mma.sync tensor-core variants, split-KV FlashDecoding), not composable Operations. The DPP structs are inside `#if defined(__NVCC__) || defined(__HIPCC__)` (softmax, flash_attention, flash_decode; their entry points take `Stream_<defaultParArch>&`) or `#if defined(__NVCC__)` (flash_attention_mma: PTX `mma.sync`, CUDA only), and remain **GPU-only: no `ParArch::CPU` nor `ParArch::CPU_MULTITHREAD` implementation**, so their tests use `__ONLY_GPU__`. The HIP DPPs need wave32 (`attention_wave.h`: `ATTN_WAVE_SIZE`, `attnShflXor`, `attnRequireWave32`). FP8 is OCP e4m3 on both backends (`attn_fp8_e4m3` = `__nv_fp8_e4m3` / `__hip_fp8_e4m3`, decoded by `attnFp8ToF32`; `FK_HAS_FP8` when the header exists), bf16 is `attn_bf16`. This is a documented exception, decided by the repo owner, to the DPP skill's rule that every DPP has both; it is not a precedent for other DPPs. The directory also holds composable CPU-compiling pieces outside the guard (e.g. `Int8TokenDequantRead`, a normal Read Operation usable in any FKL pipeline, plus host packing helpers). Nothing in `attention/` is included by any umbrella header. Cooperative exec bodies can't be constexpr (shared memory + barriers), so softmax and flash_attention use `FK_COOP_DEVICE_FUSE` while the mma and decode variants hand-write `static __device__ void exec`. Fusion happens via Read-IOp prologues (dequantization runs in-register at load time) and IOp-chain epilogues. `DivergentBatchTransformDPP` (Divergent HF) is not an exception: it has `ParArch::CPU` and `ParArch::CPU_MULTITHREAD` DPPs and executors.
 
 ### Decoding common compile failures
 

cmake/libs/hip/archs.cmake

@@ -1,9 +1,4 @@
-set(HIP_ARCH "gfx1200" CACHE STRING "HIP architecture(s) to build")
-
+# HIP_ARCH and the CMAKE_HIP_ARCHITECTURES it seeds are set in hip_init.cmake, before enable_language(HIP).
 function(set_target_hip_arch_flags TARGET_NAME)
-    if(DEFINED CMAKE_HIP_ARCHITECTURES)
-        set_target_properties(${TARGET_NAME} PROPERTIES HIP_ARCHITECTURES "${CMAKE_HIP_ARCHITECTURES}")
-    else()
-        set_target_properties(${TARGET_NAME} PROPERTIES HIP_ARCHITECTURES "${HIP_ARCH}")
-    endif()
+    set_target_properties(${TARGET_NAME} PROPERTIES HIP_ARCHITECTURES "${CMAKE_HIP_ARCHITECTURES}")
 endfunction()

cmake/libs/hip/hip.cmake

@@ -1,6 +1,12 @@
 include(cmake/libs/hip/debug.cmake)
 include(cmake/libs/hip/target_generation.cmake)
 
+# hip::device compiles the C++ sources of a target with --offload-arch=${GPU_TARGETS}, and hip-config sets GPU_TARGETS
+# to the GPU of this machine when it is empty: make it follow CMAKE_HIP_ARCHITECTURES, as the HIP sources do.
+# hip-config reads it through the cache variable GPU_BUILD_TARGETS, which it only creates, so drop that variable to
+# update it on a reconfigure.
+set(GPU_TARGETS "${CMAKE_HIP_ARCHITECTURES}")
+unset(GPU_BUILD_TARGETS CACHE)
 find_package(HIP CONFIG REQUIRED)
 option(ENABLE_HIP_LINE_INFO "Enable line info for HIP kernels compilation" ON)
 

cmake/libs/hip/hip_init.cmake

@@ -9,6 +9,14 @@ if (NOT CMAKE_HIP_COMPILER)
 endif()
 
 set(CMAKE_CXX_SCAN_FOR_MODULES OFF) #clang with HIP does not come with  clang-scan-deps
+
+# By default enable_language(HIP) sets CMAKE_HIP_ARCHITECTURES to the GPU of this machine (or the compiler default).
+# -DHIP_ARCH=<list> (e.g. "gfx1200;gfx1201": Navi 44 and Navi 48 can't load each other's code objects) seeds it instead
+# on the first configure of a build directory; an explicit -DCMAKE_HIP_ARCHITECTURES wins over HIP_ARCH.
+
+if(DEFINED HIP_ARCH AND NOT DEFINED CACHE{CMAKE_HIP_ARCHITECTURES})
+  set(CMAKE_HIP_ARCHITECTURES "${HIP_ARCH}" CACHE STRING "HIP architectures")
+endif()
 enable_language(HIP)
 
 include(cmake/libs/hip/hip.cmake)

cmake/tests/discover_tests.cmake

@@ -15,12 +15,13 @@ function (discover_tests DIR)
 
         string(FIND "${TEST_SOURCE_CONTENTS}" "ONLY_CU"  POS_ONLY_CU)
         string(FIND "${TEST_SOURCE_CONTENTS}" "ONLY_CPU"  POS_ONLY_CPU)
+        string(FIND "${TEST_SOURCE_CONTENTS}" "ONLY_GPU"  POS_ONLY_GPU) # no _cpp target, but _cu and _hip ones
         set (TEST_SOURCE1 "${TEST_SOURCE}")
         string(REPLACE ${CMAKE_SOURCE_DIR} "" TEST_SOURCE1 "${TEST_SOURCE1}") #make the path relative to the current directory)
 
         cmake_path(GET TEST_SOURCE1 RELATIVE_PART DIR_RELATIVE_PATH) 
 
-        if (${POS_ONLY_CU} EQUAL -1) #if the source file does not contain "__ONLY_CU__"    
+        if (${POS_ONLY_CU} EQUAL -1 AND ${POS_ONLY_GPU} EQUAL -1) #if the source file contains neither "__ONLY_CU__" nor "__ONLY_GPU__"
             if (${ENABLE_CPU})                                    
                 add_generated_test("${TARGET_NAME}" "${TEST_SOURCE}" "cpp" "${DIR_RELATIVE_PATH}")                
              endif()

include/fused_kernel/algorithms/attention/attention_wave.h

@@ -0,0 +1,69 @@
+/* Copyright 2026 Oscar Amoros Huguet
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License. */
+
+#ifndef FK_ATTENTION_ATTENTION_WAVE_H
+#define FK_ATTENTION_ATTENTION_WAVE_H
+
+/* Wave (warp) size of the attention DPPs: they map one row, or one fragment
+ * row group, onto exactly 32 lanes (NVIDIA warps, and RDNA waves compiled as
+ * wave32, the HIP default on gfx12). Wave64 is not supported.
+ *
+ * ROCm's clang no longer exposes the wave size at compile time
+ * (__AMDGCN_WAVEFRONT_SIZE__ was removed, and neither warpSize nor
+ * __builtin_amdgcn_wavefrontsize() is a constant expression), so:
+ *  - with a compiler that still defines __AMDGCN_WAVEFRONT_SIZE__, a wave64
+ *    device compilation fails here with #error;
+ *  - otherwise the kernels call attnRequireWave32(), which the optimizer folds
+ *    away in a wave32 build and turns into a trap in a wave64 one. */
+
+#include <fused_kernel/core/utils/utils.h>
+
+namespace fk {
+
+constexpr int ATTN_WAVE_SIZE = 32;
+
+#if defined(__HIPCC__) && defined(__HIP_DEVICE_COMPILE__) && defined(__AMDGCN_WAVEFRONT_SIZE__)
+#if __AMDGCN_WAVEFRONT_SIZE__ != 32
+#error "The attention DPPs need wave32: don't compile them with -mwavefrontsize64"
+#endif
+#endif
+
+#if defined(__NVCC__) || defined(__HIPCC__)
+
+// Traps when the kernel runs as wave64. Compiles to nothing in a wave32 (and in
+// any CUDA) build.
+__device__ __forceinline__ void attnRequireWave32() {
+#if defined(__HIPCC__) && defined(__HIP_DEVICE_COMPILE__)
+    if (__builtin_amdgcn_wavefrontsize() != ATTN_WAVE_SIZE) {
+        __builtin_trap();
+    }
+#endif
+}
+
+// Butterfly exchange across the 32 lanes of the wave: returns the value of
+// lane (lane ^ laneMask).
+template <typename T>
+__device__ __forceinline__ T attnShflXor(const T value, const int laneMask) {
+#if defined(__HIPCC__)
+    return __shfl_xor(value, laneMask, ATTN_WAVE_SIZE);
+#else
+    return __shfl_xor_sync(0xffffffffu, value, laneMask);
+#endif
+}
+
+#endif // defined(__NVCC__) || defined(__HIPCC__)
+
+} // namespace fk
+
+#endif // FK_ATTENTION_ATTENTION_WAVE_H