diff --git a/AGENTS.md b/AGENTS.md index 1a25eb63..3d2f13e1 100644 --- a/AGENTS.md +++ b/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 ` 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 `/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++ (`_cpp` target, `ParArch::CPU`) and as CUDA (`_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_&` 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_&`) 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 diff --git a/cmake/libs/hip/archs.cmake b/cmake/libs/hip/archs.cmake index 61e0e422..de5a8725 100644 --- a/cmake/libs/hip/archs.cmake +++ b/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() diff --git a/cmake/libs/hip/hip.cmake b/cmake/libs/hip/hip.cmake index df3ce8a1..7031ca7c 100644 --- a/cmake/libs/hip/hip.cmake +++ b/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) diff --git a/cmake/libs/hip/hip_init.cmake b/cmake/libs/hip/hip_init.cmake index bde73688..6da57f34 100644 --- a/cmake/libs/hip/hip_init.cmake +++ b/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= (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) diff --git a/cmake/tests/discover_tests.cmake b/cmake/tests/discover_tests.cmake index 748d43a6..48a44a0b 100644 --- a/cmake/tests/discover_tests.cmake +++ b/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() diff --git a/include/fused_kernel/algorithms/attention/attention_wave.h b/include/fused_kernel/algorithms/attention/attention_wave.h new file mode 100644 index 00000000..11227e1d --- /dev/null +++ b/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 + +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 +__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 diff --git a/include/fused_kernel/algorithms/attention/flash_attention.h b/include/fused_kernel/algorithms/attention/flash_attention.h index 2768ae66..5dd14e8a 100644 --- a/include/fused_kernel/algorithms/attention/flash_attention.h +++ b/include/fused_kernel/algorithms/attention/flash_attention.h @@ -47,8 +47,8 @@ * Mapping (SM 12x friendly - no TMEM/tcgen05, per the fa-5090 analysis, * correctness-first SIMT baseline before mma.sync tiling): * grid.y = batch*heads; grid.x = ceil(seq_q / WARPS_PER_BLOCK) - * 1 warp per query row; q and o live in registers spread across lanes - * (HEAD_DIM/32 each); dot products warp-reduce via __shfl_xor_sync; + * 1 warp (wave32 on HIP) per query row; q and o live in registers spread + * across lanes (HEAD_DIM/32 each); dot products warp-reduce via attnShflXor; * K/V tiles staged cooperatively in shared memory THROUGH the prologue * IOps (tiles hold post-prologue fp32, so the prologue runs exactly * once per element). @@ -59,6 +59,7 @@ * thread.z = batch*head plane */ +#include #include #include @@ -146,7 +147,7 @@ inline auto makeInt8KVRead(const int8_t* data, const float* scales, return Int8TokenDequantRead::build(Int8TokenDequantReadParams{ ptr, scales }); } -#if defined(__NVCC__) +#if defined(__NVCC__) || defined(__HIPCC__) /* The DPP. QIOp/KIOp/VIOp are INSTANTIABLE Read or ReadBack IOps (possibly * fused with compute continuations): the algorithm reads each element of @@ -162,12 +163,12 @@ struct FlashAttentionDPP { public: FK_STATIC_STRUCT(FlashAttentionDPP, SelfType) - static_assert(HEAD_DIM % 32 == 0, "HEAD_DIM must be a multiple of 32"); + static_assert(HEAD_DIM % ATTN_WAVE_SIZE == 0, "HEAD_DIM must be a multiple of 32"); static_assert(isAnyReadType, "Q prologue must be a Read or ReadBack IOp"); static_assert(isAnyReadType, "K prologue must be a Read or ReadBack IOp"); static_assert(isAnyReadType, "V prologue must be a Read or ReadBack IOp"); - static constexpr int ELEMS_PER_LANE = HEAD_DIM / 32; - static constexpr int THREADS = WARPS_PER_BLOCK * 32; + static constexpr int ELEMS_PER_LANE = HEAD_DIM / ATTN_WAVE_SIZE; + static constexpr int THREADS = WARPS_PER_BLOCK * ATTN_WAVE_SIZE; struct Params { QIOp q; // prologue Read/ReadBack IOp for Q @@ -193,8 +194,9 @@ struct FlashAttentionDPP { __shared__ float kTile[BLOCK_N][HEAD_DIM]; __shared__ float vTile[BLOCK_N][HEAD_DIM]; - const int lane = threadIdx.x & 31; - const int warp = threadIdx.x >> 5; + attnRequireWave32(); + const int lane = threadIdx.x % ATTN_WAVE_SIZE; + const int warp = threadIdx.x / ATTN_WAVE_SIZE; const int qIdx = blockIdx.x * WARPS_PER_BLOCK + warp; const int bh = blockIdx.y; const bool active = qIdx < p.seq_q; @@ -204,7 +206,7 @@ struct FlashAttentionDPP { #pragma unroll for (int e = 0; e < ELEMS_PER_LANE; ++e) { // Q PROLOGUE: read through the IOp, in-register, at load time. - qReg[e] = active ? readElem(p.q, lane + 32 * e, qIdx, bh) : 0.f; + qReg[e] = active ? readElem(p.q, lane + ATTN_WAVE_SIZE * e, qIdx, bh) : 0.f; oAcc[e] = 0.f; } float m = -FLT_MAX; @@ -237,11 +239,11 @@ struct FlashAttentionDPP { float partial = 0.f; #pragma unroll for (int e = 0; e < ELEMS_PER_LANE; ++e) { - partial += qReg[e] * kTile[j][lane + 32 * e]; + partial += qReg[e] * kTile[j][lane + ATTN_WAVE_SIZE * e]; } #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - partial += __shfl_xor_sync(0xffffffffu, partial, offset); + for (int offset = ATTN_WAVE_SIZE / 2; offset > 0; offset >>= 1) { + partial += attnShflXor(partial, offset); } const float s = partial * p.scale; @@ -251,7 +253,7 @@ struct FlashAttentionDPP { l = l * corr + pj; #pragma unroll for (int e = 0; e < ELEMS_PER_LANE; ++e) { - oAcc[e] = oAcc[e] * corr + pj * vTile[j][lane + 32 * e]; + oAcc[e] = oAcc[e] * corr + pj * vTile[j][lane + ATTN_WAVE_SIZE * e]; } m = mNew; } @@ -265,7 +267,7 @@ struct FlashAttentionDPP { // EPILOGUE FUSION: the FKL IOp chain runs in-register on the // normalized output, then ONE global write. const float r = (oAcc[e] * invL) | p.epilogue; - p.o[oRow + lane + 32 * e] = low_precision::attnFromF32(r); + p.o[oRow + lane + ATTN_WAVE_SIZE * e] = low_precision::attnFromF32(r); } } } @@ -285,7 +287,7 @@ template & stream, + const bool causal, Stream_& stream, const float scaleOverride = -1.f, const EpilogueIOp& epilogue = {}) { using DPP = FlashAttentionDPP; @@ -294,9 +296,15 @@ inline void executeFlashAttention( const typename DPP::Params params{ q, k, v, o, seqQ, seqK, scale, causal, epilogue }; const dim3 grid((seqQ + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK, batchHeads, 1); const dim3 block(DPP::THREADS, 1, 1); +#if defined(__HIPCC__) + launchFlashAttentionDPP_Kernel + <<>>(params); + gpuErrchk(hipGetLastError()); +#else launchFlashAttentionDPP_Kernel <<>>(params); gpuErrchk(cudaGetLastError()); +#endif } /* Pointer convenience API (kept for compatibility): builds the canonical @@ -310,7 +318,7 @@ inline void executeFlashAttention( const std::conditional_t* k, const std::conditional_t* v, T* o, const int batchHeads, const int seqQ, const int seqK, - const bool causal, Stream_& stream, + const bool causal, Stream_& stream, const float* kScale = nullptr, const float* vScale = nullptr, const float scaleOverride = -1.f, const EpilogueIOp& epilogue = {}) { @@ -330,7 +338,7 @@ inline void executeFlashAttention( } } -#endif // defined(__NVCC__) +#endif // defined(__NVCC__) || defined(__HIPCC__) // ---- host-side KV cache compression helper (reference packing) ----------- // Per-token symmetric int8: scale[t] = max|row|/127; q(x) = round(x/scale). @@ -361,11 +369,42 @@ inline void quantizeKVCacheHost(const T* dense, int8_t* q8, float* scales, // but the e4m3 grid is non-uniform (more precision near 0), which suits // attention tails. Storage: data fp8 e4m3, one fp32 scale per token // (scale = max|row| / 448, e4m3 max normal = 448). -#if __has_include() +// The format is OCP e4m3 (bias 7, max 448) on both backends: on HIP that is +// __hip_fp8_e4m3, not the FNUZ __hip_fp8_e4m3_fnuz of MI300 (gfx942). +#if defined(__HIPCC__) && __has_include() +#include +#define FK_HAS_FP8 1 + +namespace fk { +using attn_fp8_e4m3 = __hip_fp8_e4m3; +using attn_fp8_storage_t = __hip_fp8_storage_t; +} // namespace fk +#elif __has_include() #include #define FK_HAS_FP8 1 namespace fk { +using attn_fp8_e4m3 = __nv_fp8_e4m3; +using attn_fp8_storage_t = __nv_fp8_storage_t; +} // namespace fk +#endif + +#ifdef FK_HAS_FP8 +namespace fk { + +// One raw e4m3 byte (int8_t storage) to float. +__host__ __device__ __forceinline__ float attnFp8ToF32(const int8_t raw) { + attn_fp8_e4m3 f8; + f8.__x = static_cast(raw); + return static_cast(f8); +} + +// True when attn_fp8_e4m3 decodes as OCP e4m3: 0x7E = 448 (max normal), +// 0x38 = 1, 0x80 = -0 (FNUZ would give 0x80 = NaN, 0x7E = 240). +inline bool attnFp8IsOcpE4M3() { + const float zero = attnFp8ToF32(static_cast(0x80)); + return attnFp8ToF32(0x7E) == 448.f && attnFp8ToF32(0x38) == 1.f && zero == 0.f && std::signbit(zero); +} struct Fp8TokenDequantReadParams { RawPtr data; // raw e4m3 bytes (int8_t storage) @@ -388,14 +427,7 @@ struct Fp8TokenDequantRead { const int8_t raw = *PtrAccessor::cr_point(thread, params.data); const int seq = static_cast(params.data.dims.height); const float sc = params.scales[(long)thread.z * seq + thread.y]; -#if defined(__CUDA_ARCH__) - const __nv_fp8_e4m3* f8 = reinterpret_cast(&raw); - return static_cast(*f8) * sc; -#else - __nv_fp8_e4m3 f8; - f8.__x = static_cast<__nv_fp8_storage_t>(raw); - return static_cast(f8) * sc; -#endif + return attnFp8ToF32(raw) * sc; } FK_HOST_DEVICE_FUSE uint num_elems_x(const Point thread, const OperationDataType& opData) { @@ -443,7 +475,7 @@ inline void quantizeKVCacheFp8Host(const T* dense, void* f8out, float* scales, scales[t] = sc; for (int d = 0; d < headDim; ++d) { const float x = low_precision::attnToF32(dense[(long)t * headDim + d]) / sc; - const __nv_fp8_e4m3 f8(x); + const attn_fp8_e4m3 f8(x); out[(long)t * headDim + d] = static_cast(f8.__x); } } @@ -451,6 +483,90 @@ inline void quantizeKVCacheFp8Host(const T* dense, void* f8out, float* scales, } // namespace fk -#endif // __has_include() +#endif // FK_HAS_FP8 + +// ============================ bf16 READ, IOp TRAITS ========================== +// The raw bf16 prologue and the traits the DPPs use to pick a schedule +// (FlashAttentionMmaDPP) or a load width (FlashDecodeDPP) for a prologue type. +#if defined(__NVCC__) || defined(__HIPCC__) +#if defined(__HIPCC__) +#include +#else +#include +#endif + +namespace fk { + +#if defined(__HIPCC__) +using attn_bf16 = __hip_bfloat16; +using attn_bf162 = __hip_bfloat162; +#else +using attn_bf16 = __nv_bfloat16; +using attn_bf162 = __nv_bfloat162; +#endif + +// ---- bf16 identity Read IOp (the "raw" prologue) --------------------------- +struct Bf16AttentionRead { +private: + using Parent = ReadOperation, + float, TF::DISABLED, Bf16AttentionRead>; + using SelfType = Bf16AttentionRead; +public: + FK_STATIC_STRUCT(Bf16AttentionRead, SelfType) + DECLARE_READ_PARENT + + FK_HOST_DEVICE_FUSE float exec(const Point thread, const ParamsType& params) { +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) + return __bfloat162float(*PtrAccessor::cr_point(thread, params)); +#else + return 0.f; // host never executes this +#endif + } + FK_HOST_DEVICE_FUSE uint num_elems_x(const Point thread, const OperationDataType& opData) { + return opData.params.dims.width; + } + FK_HOST_DEVICE_FUSE uint num_elems_y(const Point thread, const OperationDataType& opData) { + return opData.params.dims.height; + } + FK_HOST_DEVICE_FUSE uint num_elems_z(const Point thread, const OperationDataType& opData) { + return opData.params.dims.planes; + } + FK_HOST_DEVICE_FUSE uint pitch(const Point thread, const OperationDataType& opData) { + return opData.params.dims.pitch; + } + FK_HOST_DEVICE_FUSE ActiveThreads getActiveThreads(const OperationDataType& opData) { + return { num_elems_x(Point{0,0,0}, opData), + num_elems_y(Point{0,0,0}, opData), + num_elems_z(Point{0,0,0}, opData) }; + } +}; + +// bf16 overload of the canonical prologue builder. +inline auto makeAttentionRead(const attn_bf16* data, const int batchHeads, + const int seq, const int headDim) { + const RawPtr ptr{ + const_cast(data), + PtrDims(static_cast(headDim), static_cast(seq), + static_cast(batchHeads), 1, + static_cast(headDim * sizeof(attn_bf16))) }; + return Bf16AttentionRead::build(ptr); +} + +template +constexpr bool isRawBf16Read = std::is_same_v; + +#ifdef FK_HAS_FP8 +template +constexpr bool isFp8KVRead = std::is_same_v; +#else +template +constexpr bool isFp8KVRead = false; +#endif +template +constexpr bool isInt8KVRead = std::is_same_v; + +} // namespace fk + +#endif // defined(__NVCC__) || defined(__HIPCC__) #endif // FK_ATTENTION_FLASH_ATTENTION_H diff --git a/include/fused_kernel/algorithms/attention/flash_attention_mma.h b/include/fused_kernel/algorithms/attention/flash_attention_mma.h index 03b8213f..0e20eb22 100644 --- a/include/fused_kernel/algorithms/attention/flash_attention_mma.h +++ b/include/fused_kernel/algorithms/attention/flash_attention_mma.h @@ -52,65 +52,8 @@ namespace fk { -// ---- bf16 identity Read IOp (the "raw" prologue) --------------------------- -struct Bf16AttentionRead { -private: - using Parent = ReadOperation<__nv_bfloat16, RawPtr, - float, TF::DISABLED, Bf16AttentionRead>; - using SelfType = Bf16AttentionRead; -public: - FK_STATIC_STRUCT(Bf16AttentionRead, SelfType) - DECLARE_READ_PARENT - - FK_HOST_DEVICE_FUSE float exec(const Point thread, const ParamsType& params) { -#if defined(__CUDA_ARCH__) - return __bfloat162float(*PtrAccessor::cr_point(thread, params)); -#else - return 0.f; // host never executes this -#endif - } - FK_HOST_DEVICE_FUSE uint num_elems_x(const Point thread, const OperationDataType& opData) { - return opData.params.dims.width; - } - FK_HOST_DEVICE_FUSE uint num_elems_y(const Point thread, const OperationDataType& opData) { - return opData.params.dims.height; - } - FK_HOST_DEVICE_FUSE uint num_elems_z(const Point thread, const OperationDataType& opData) { - return opData.params.dims.planes; - } - FK_HOST_DEVICE_FUSE uint pitch(const Point thread, const OperationDataType& opData) { - return opData.params.dims.pitch; - } - FK_HOST_DEVICE_FUSE ActiveThreads getActiveThreads(const OperationDataType& opData) { - return { num_elems_x(Point{0,0,0}, opData), - num_elems_y(Point{0,0,0}, opData), - num_elems_z(Point{0,0,0}, opData) }; - } -}; - -// bf16 overload of the canonical prologue builder. -inline auto makeAttentionRead(const __nv_bfloat16* data, const int batchHeads, - const int seq, const int headDim) { - const RawPtr ptr{ - const_cast<__nv_bfloat16*>(data), - PtrDims(static_cast(headDim), static_cast(seq), - static_cast(batchHeads), 1, - static_cast(headDim * sizeof(__nv_bfloat16))) }; - return Bf16AttentionRead::build(ptr); -} - -template -constexpr bool isRawBf16Read = std::is_same_v; - -#ifdef FK_HAS_FP8 -template -constexpr bool isFp8KVRead = std::is_same_v; -#else -template -constexpr bool isFp8KVRead = false; -#endif -template -constexpr bool isInt8KVRead = std::is_same_v; +// Bf16AttentionRead, its makeAttentionRead overload and the isRawBf16Read, +// isFp8KVRead and isInt8KVRead traits live in flash_attention.h. // ---- FLEX ATTENTION: score mods (FlexAttention-style score_mod) ------------ // A score mod is a tiny device functor applied to each attention score diff --git a/include/fused_kernel/algorithms/attention/flash_decode.h b/include/fused_kernel/algorithms/attention/flash_decode.h index 6c913df6..10f9debb 100644 --- a/include/fused_kernel/algorithms/attention/flash_decode.h +++ b/include/fused_kernel/algorithms/attention/flash_decode.h @@ -35,10 +35,10 @@ * This pairs with FA4-sm12x's fp8 decode direction (Dao-AILab PR #2634) * but keeps FKL's universal prologue: ANY Read/ReadBack IOp works. */ -#include -#include // IOp traits + bf16 read +#include // IOp traits + bf16 read +#include -#if defined(__NVCC__) +#if defined(__NVCC__) || defined(__HIPCC__) namespace fk { @@ -50,9 +50,9 @@ struct FlashDecodeDPP { public: FK_STATIC_STRUCT(FlashDecodeDPP, SelfType) - static_assert(HEAD_DIM % 32 == 0, "HEAD_DIM must be a multiple of 32"); - static constexpr int THREADS = NUM_WARPS * 32; - static constexpr int ELEMS = HEAD_DIM / 32; // per-lane o/q elements + static_assert(HEAD_DIM % ATTN_WAVE_SIZE == 0, "HEAD_DIM must be a multiple of 32"); + static constexpr int THREADS = NUM_WARPS * ATTN_WAVE_SIZE; + static constexpr int ELEMS = HEAD_DIM / ATTN_WAVE_SIZE; // per-lane o/q elements struct Params { QIOp q; KIOp k; VIOp v; @@ -66,29 +66,31 @@ struct FlashDecodeDPP { template FK_DEVICE_FUSE float readElem(const IOp& iop, const int x, const int y, const int z) { - return attnToF32(IOp::Operation::exec(Point{ x, y, z }, iop)); + return low_precision::attnToF32(IOp::Operation::exec(Point{ x, y, z }, iop)); } /* Vectorized row loader: lane owns elements [ELEMS*lane, ELEMS*lane+ELEMS). Raw bf16 rows load as bfloat162 pairs (4B), fp8/int8 rows as packed bytes (ELEMS bytes in one load) — decode is bandwidth-bound, wide - loads are the whole game. Generic IOps fall back to per-element. */ + loads are the whole game. Generic IOps fall back to per-element, and + so do the widths the packed loads don't cover: bf16 pairs need an even + ELEMS, the byte loads 2 or 4 bytes (HEAD_DIM 64 or 128). */ template FK_DEVICE_FUSE void loadRow(const IOp& iop, const int lane, const int t, const int bh, const int seqLen, float (&out)[ELEMS]) { - if constexpr (isRawBf16Read) { + if constexpr (isRawBf16Read && ELEMS % 2 == 0) { const auto& prm = iop.params; - const __nv_bfloat16* base = prm.data + const attn_bf16* base = prm.data + ((long)bh * seqLen + t) * HEAD_DIM + ELEMS * lane; #pragma unroll for (int e = 0; e < ELEMS; e += 2) { - const __nv_bfloat162 pair = - *reinterpret_cast(base + e); + const attn_bf162 pair = + *reinterpret_cast(base + e); out[e] = __bfloat162float(pair.x); out[e + 1] = __bfloat162float(pair.y); } - } else if constexpr (isFp8KVRead || isInt8KVRead) { + } else if constexpr ((isFp8KVRead || isInt8KVRead) && (ELEMS == 2 || ELEMS == 4)) { const auto& prm = iop.params; const float sc = prm.scales[(long)bh * seqLen + t]; const int8_t* base = reinterpret_cast(prm.data.data) @@ -102,9 +104,7 @@ struct FlashDecodeDPP { const int8_t b = static_cast((packed >> (8 * e)) & 0xFF); if constexpr (isFp8KVRead) { #ifdef FK_HAS_FP8 - __nv_fp8_e4m3 f8; - f8.__x = static_cast<__nv_fp8_storage_t>(b); - out[e] = static_cast(f8) * sc; + out[e] = attnFp8ToF32(b) * sc; #endif } else { out[e] = static_cast(b) * sc; @@ -118,8 +118,9 @@ struct FlashDecodeDPP { } static __device__ void exec(const Params& p) { - const int lane = threadIdx.x & 31; - const int warp = threadIdx.x >> 5; + attnRequireWave32(); + const int lane = threadIdx.x % ATTN_WAVE_SIZE; + const int warp = threadIdx.x / ATTN_WAVE_SIZE; const int split = blockIdx.x; const int bh = blockIdx.y; @@ -159,8 +160,8 @@ struct FlashDecodeDPP { #pragma unroll for (int e = 0; e < ELEMS; ++e) partialDot += qReg[e] * kReg[e]; #pragma unroll - for (int off = 16; off > 0; off >>= 1) - partialDot += __shfl_xor_sync(0xFFFFFFFFu, partialDot, off); + for (int off = ATTN_WAVE_SIZE / 2; off > 0; off >>= 1) + partialDot += attnShflXor(partialDot, off); const float s = partialDot * p.scale; const float mNew = fmaxf(m, s); @@ -310,7 +311,7 @@ inline void executeFlashDecode( float* workspace /* flashDecodeWorkspaceFloats(); zeroed at alloc; nullptr ok if splits==1 */, const int batchHeads, const int seqK, - Stream_& stream, + Stream_& stream, const float scaleOverride = -1.f, const int splitsOverride = 0) { using DPP = FlashDecodeDPP; const float scale = scaleOverride > 0.f ? scaleOverride @@ -328,13 +329,19 @@ inline void executeFlashDecode( // single kernel: the split merge is FUSED (last-CTA-arrives pattern) — // profiling showed the separate 1.4us merge kernel + launch gap was pure // overhead at decode latencies. +#if defined(__HIPCC__) + launchFlashDecodeDPP_Kernel + <<>>(params); + gpuErrchk(hipGetLastError()); +#else launchFlashDecodeDPP_Kernel <<>>(params); gpuErrchk(cudaGetLastError()); +#endif } } // namespace fk -#endif // defined(__NVCC__) +#endif // defined(__NVCC__) || defined(__HIPCC__) #endif // FK_ATTENTION_FLASH_DECODE_H diff --git a/include/fused_kernel/algorithms/attention/softmax.h b/include/fused_kernel/algorithms/attention/softmax.h index 949af316..b2b13e0a 100644 --- a/include/fused_kernel/algorithms/attention/softmax.h +++ b/include/fused_kernel/algorithms/attention/softmax.h @@ -44,6 +44,8 @@ #include #if defined(__NVCC__) #include +#elif defined(__HIPCC__) +#include #endif namespace fk { @@ -53,8 +55,11 @@ namespace fk { #define FK_COOP_DEVICE_FUSE static __device__ __forceinline__ void namespace low_precision { + // Host and device: the host KV-cache packing helpers (quantizeKVCacheHost...) + // call them too, which nvcc accepted for a constexpr __device__ function + // and HIP's clang does not. template - FK_DEVICE_CNST float attnToF32(const T &v) { + FK_HOST_DEVICE_CNST float attnToF32(const T &v) { if constexpr (std::is_same_v) { return __half2float(v); } else { @@ -63,7 +68,7 @@ namespace fk { } template - FK_DEVICE_CNST T attnFromF32(const float &v) { + FK_HOST_DEVICE_CNST T attnFromF32(const float &v) { if constexpr (std::is_same_v) { return __float2half(v); } else { @@ -171,12 +176,17 @@ __global__ void launchDPP_Kernel(const __grid_constant__ DPPDetails details, /* IOp-first API: input is the prologue Read/ReadBack IOp. */ template inline void executeSoftmax(const InIOp& input, const OutIOp& output, - Stream_& stream) { + Stream_& stream) { const SoftmaxDPPDetails details{}; const dim3 grid(InIOp::Operation::num_elems_y(Point{0,0,0}, input), 1, 1); const dim3 block(BLOCK_SIZE, 1, 1); +#if defined(__HIPCC__) + launchDPP_Kernel<<>>(details, input, NullType{}, output); + gpuErrchk(hipGetLastError()); +#else launchDPP_Kernel<<>>(details, input, NullType{}, output); gpuErrchk(cudaGetLastError()); +#endif } } // namespace fk diff --git a/utests/algorithm/attention/attention_test_common.h b/utests/algorithm/attention/attention_test_common.h new file mode 100644 index 00000000..1524f8aa --- /dev/null +++ b/utests/algorithm/attention/attention_test_common.h @@ -0,0 +1,112 @@ +/* Copyright 2026 Oscar Amoros Huguet + Copyright 2026 Johnny Nunez + + 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_UTESTS_ATTENTION_TEST_COMMON_H +#define FK_UTESTS_ATTENTION_TEST_COMMON_H + +/* Shared by the attention utests (the "_common" in the name keeps test + * discovery from turning it into a test): the double-precision CPU oracle, + * pass/fail reporting, and Ptr1D buffers moved through their pinned mirror. + * Include it after the attention header under test. */ + +#include +#include + +#include +#include +#include +#include + +namespace fk { + +// double-precision CPU oracle: O = softmax(scale * Q K^T [causal]) V +inline void cpuAttention(const std::vector& q, const std::vector& k, + const std::vector& v, std::vector& o, + const int bh, const int seqQ, const int seqK, const int d, + const double scale, const bool causal) { + o.assign((size_t)bh * seqQ * d, 0.0); + std::vector s(seqK); + for (int b = 0; b < bh; ++b) { + const double* Q = q.data() + (size_t)b * seqQ * d; + const double* K = k.data() + (size_t)b * seqK * d; + const double* V = v.data() + (size_t)b * seqK * d; + double* O = o.data() + (size_t)b * seqQ * d; + for (int i = 0; i < seqQ; ++i) { + const int kEnd = causal ? std::min(seqK, i + 1) : seqK; + double m = -1e300; + for (int j = 0; j < kEnd; ++j) { + double dot = 0.0; + for (int c = 0; c < d; ++c) dot += Q[(size_t)i*d+c] * K[(size_t)j*d+c]; + s[j] = dot * scale; + m = std::max(m, s[j]); + } + double l = 0.0; + for (int j = 0; j < kEnd; ++j) { s[j] = std::exp(s[j] - m); l += s[j]; } + for (int j = 0; j < kEnd; ++j) { + const double pj = s[j] / l; + for (int c = 0; c < d; ++c) O[(size_t)i*d+c] += pj * V[(size_t)j*d+c]; + } + } + } +} + +inline int failures = 0; + +inline void report(const char* name, const double maxErr, const double tol) { + if (maxErr > tol) { + std::cout << "FAIL " << name << ": maxErr=" << maxErr << " tol=" << tol << std::endl; + ++failures; + } else { + std::cout << "Running test " << name << ": Success!! (maxErr=" << maxErr << ")" << std::endl; + } +} + +template +inline double maxAbsErr(const std::vector& got, const std::vector& ref) { + double maxErr = 0.0; + for (size_t i = 0; i < ref.size(); ++i) { + maxErr = std::max(maxErr, std::abs(static_cast(got[i]) - ref[i])); + } + return maxErr; +} + +// Device buffer holding a copy of host: written to the pinned mirror, then +// uploaded on the stream. +template +inline Ptr1D uploadBuffer(const std::vector& host, Stream& stream) { + Ptr1D buffer(static_cast(host.size())); + std::copy(host.begin(), host.end(), buffer.ptrPinned().data); + buffer.upload(stream); + return buffer; +} + +// Device buffer of n elements, zeroed (through the pinned mirror). +template +inline Ptr1D zeroBuffer(const size_t n, Stream& stream) { + return uploadBuffer(std::vector(n, T{}), stream); +} + +// Waits for the stream and returns the contents of the device buffer. +template +inline std::vector downloadBuffer(Ptr1D& buffer, Stream& stream) { + buffer.download(stream); + stream.sync(); + const T* host = buffer.ptrPinned().data; + return std::vector(host, host + buffer.dims().width); +} + +} // namespace fk + +#endif // FK_UTESTS_ATTENTION_TEST_COMMON_H diff --git a/utests/algorithm/attention/utest_attention_wave.h b/utests/algorithm/attention/utest_attention_wave.h new file mode 100644 index 00000000..1d4b584f --- /dev/null +++ b/utests/algorithm/attention/utest_attention_wave.h @@ -0,0 +1,85 @@ +/* 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. */ + +/* attention_wave.h: attnShflXor exchanges values inside each 32-lane wave + * (never across waves) and attnRequireWave32 lets a wave32 kernel run. Two + * waves per block: each lane holds its thread index, a butterfly sum must give + * the sum of its own wave, and a single exchange must return lane ^ mask. */ + +#define __ONLY_GPU__ // Generated as _cu and _hip, not as _cpp: attention_wave.h has no CPU code +#include + +#include +#include +#include + +#include + +namespace fk { + +constexpr int WAVE_TEST_THREADS = 2 * ATTN_WAVE_SIZE; +constexpr int WAVE_TEST_MASK = 5; + +__global__ void attnWaveTestKernel(const RawPtr sums, const RawPtr exchanged) { + attnRequireWave32(); + const int tid = static_cast(threadIdx.x); + float sum = static_cast(tid); + for (int offset = ATTN_WAVE_SIZE / 2; offset > 0; offset >>= 1) { + sum += attnShflXor(sum, offset); + } + sums.data[tid] = sum; + exchanged.data[tid] = attnShflXor(tid, WAVE_TEST_MASK); +} + +int launch_impl() { + Ptr1D sums(WAVE_TEST_THREADS); + Ptr1D exchanged(WAVE_TEST_THREADS); + Stream stream; +#if defined(__HIPCC__) + attnWaveTestKernel<<<1, WAVE_TEST_THREADS, 0, stream.getHIPStream()>>>(sums.ptr(), exchanged.ptr()); + gpuErrchk(hipGetLastError()); +#else + attnWaveTestKernel<<<1, WAVE_TEST_THREADS, 0, stream.getCUDAStream()>>>(sums.ptr(), exchanged.ptr()); + gpuErrchk(cudaGetLastError()); +#endif + sums.download(stream); + exchanged.download(stream); + stream.sync(); + + int failures = 0; + for (int tid = 0; tid < WAVE_TEST_THREADS; ++tid) { + const int first = tid - tid % ATTN_WAVE_SIZE; + const int expectedIntSum = ATTN_WAVE_SIZE * first + ATTN_WAVE_SIZE * (ATTN_WAVE_SIZE - 1) / 2; + const float expectedSum = static_cast(expectedIntSum); + const int expectedExchange = tid ^ WAVE_TEST_MASK; + const float gotSum = sums.ptrPinned().data[tid]; + const int gotExchange = exchanged.ptrPinned().data[tid]; + if (gotSum != expectedSum || gotExchange != expectedExchange) { + std::cout << "FAIL attnShflXor lane " << tid << ": sum " << gotSum << " (expected " << expectedSum + << "), exchange " << gotExchange << " (expected " << expectedExchange << ")" << std::endl; + ++failures; + } + } + if (failures == 0) { + std::cout << "Running test attnShflXor butterfly sum and exchange, 2 waves: Success!!" << std::endl; + return 0; + } + return -1; +} + +} // namespace fk + +int launch() { + return fk::launch_impl(); +} diff --git a/utests/algorithm/attention/utest_flash_attention.h b/utests/algorithm/attention/utest_flash_attention.h index 4524c3e4..1c581aca 100644 --- a/utests/algorithm/attention/utest_flash_attention.h +++ b/utests/algorithm/attention/utest_flash_attention.h @@ -13,11 +13,12 @@ See the License for the specific language governing permissions and limitations under the License. */ -#define __ONLY_CU__ // This file is only generated and compiled with nvcc, not with the host compiler +#define __ONLY_GPU__ // Generated as _cu and _hip, not as _cpp: the attention DPPs are GPU-only #include #include #include +#include #include #include @@ -26,48 +27,6 @@ namespace fk { -// double-precision CPU oracle: O = softmax(scale * Q K^T [causal]) V -static void cpuAttention(const std::vector& q, const std::vector& k, - const std::vector& v, std::vector& o, - const int bh, const int seqQ, const int seqK, const int d, - const double scale, const bool causal) { - o.assign((size_t)bh * seqQ * d, 0.0); - std::vector s(seqK); - for (int b = 0; b < bh; ++b) { - const double* Q = q.data() + (size_t)b * seqQ * d; - const double* K = k.data() + (size_t)b * seqK * d; - const double* V = v.data() + (size_t)b * seqK * d; - double* O = o.data() + (size_t)b * seqQ * d; - for (int i = 0; i < seqQ; ++i) { - const int kEnd = causal ? std::min(seqK, i + 1) : seqK; - double m = -1e300; - for (int j = 0; j < kEnd; ++j) { - double dot = 0.0; - for (int c = 0; c < d; ++c) dot += Q[(size_t)i*d+c] * K[(size_t)j*d+c]; - s[j] = dot * scale; - m = std::max(m, s[j]); - } - double l = 0.0; - for (int j = 0; j < kEnd; ++j) { s[j] = std::exp(s[j] - m); l += s[j]; } - for (int j = 0; j < kEnd; ++j) { - const double pj = s[j] / l; - for (int c = 0; c < d; ++c) O[(size_t)i*d+c] += pj * V[(size_t)j*d+c]; - } - } - } -} - -static int failures = 0; - -static void report(const char* name, const double maxErr, const double tol) { - if (maxErr > tol) { - std::cout << "FAIL " << name << ": maxErr=" << maxErr << " tol=" << tol << std::endl; - ++failures; - } else { - std::cout << "Running test " << name << ": Success!! (maxErr=" << maxErr << ")" << std::endl; - } -} - template static void testDense(const char* name, const int bh, const int seqQ, const int seqK, const bool causal, const double tol, const unsigned seed) { @@ -82,26 +41,15 @@ static void testDense(const char* name, const int bh, const int seqQ, const int cpuAttention(dq, dk, dv, dref, bh, seqQ, seqK, HEAD_DIM, 1.0 / std::sqrt((double)HEAD_DIM), causal); - float *q, *k, *v, *o; - gpuErrchk(cudaMalloc(&q, nQ * sizeof(float))); - gpuErrchk(cudaMalloc(&k, nK * sizeof(float))); - gpuErrchk(cudaMalloc(&v, nK * sizeof(float))); - gpuErrchk(cudaMalloc(&o, nQ * sizeof(float))); - gpuErrchk(cudaMemcpy(q, hq.data(), nQ * sizeof(float), cudaMemcpyHostToDevice)); - gpuErrchk(cudaMemcpy(k, hk.data(), nK * sizeof(float), cudaMemcpyHostToDevice)); - gpuErrchk(cudaMemcpy(v, hv.data(), nK * sizeof(float), cudaMemcpyHostToDevice)); - Stream stream; - executeFlashAttention(q, k, v, o, bh, seqQ, seqK, causal, stream); - stream.sync(); - - std::vector got(nQ); - gpuErrchk(cudaMemcpy(got.data(), o, nQ * sizeof(float), cudaMemcpyDeviceToHost)); - cudaFree(q); cudaFree(k); cudaFree(v); cudaFree(o); - - double maxErr = 0.0; - for (size_t i = 0; i < nQ; ++i) maxErr = std::max(maxErr, std::abs((double)got[i] - dref[i])); - report(name, maxErr, tol); + auto q = uploadBuffer(hq, stream); + auto k = uploadBuffer(hk, stream); + auto v = uploadBuffer(hv, stream); + auto o = zeroBuffer(nQ, stream); + executeFlashAttention(q.ptr().data, k.ptr().data, v.ptr().data, o.ptr().data, bh, seqQ, + seqK, causal, stream); + const std::vector got = downloadBuffer(o, stream); + report(name, maxAbsErr(got, dref), tol); } template @@ -131,33 +79,77 @@ static void testInt8KV(const char* name, const int bh, const int seqQ, const int cpuAttention(dq, dkD, dvD, dref, bh, seqQ, seqK, HEAD_DIM, 1.0 / std::sqrt((double)HEAD_DIM), causal); - float *q, *o, *dkS, *dvS; - int8_t *dk8, *dv8; - gpuErrchk(cudaMalloc(&q, nQ * sizeof(float))); - gpuErrchk(cudaMalloc(&o, nQ * sizeof(float))); - gpuErrchk(cudaMalloc(&dk8, nK)); - gpuErrchk(cudaMalloc(&dv8, nK)); - gpuErrchk(cudaMalloc(&dkS, nTok * sizeof(float))); - gpuErrchk(cudaMalloc(&dvS, nTok * sizeof(float))); - gpuErrchk(cudaMemcpy(q, hq.data(), nQ * sizeof(float), cudaMemcpyHostToDevice)); - gpuErrchk(cudaMemcpy(dk8, k8.data(), nK, cudaMemcpyHostToDevice)); - gpuErrchk(cudaMemcpy(dv8, v8.data(), nK, cudaMemcpyHostToDevice)); - gpuErrchk(cudaMemcpy(dkS, kSc.data(), nTok * sizeof(float), cudaMemcpyHostToDevice)); - gpuErrchk(cudaMemcpy(dvS, vSc.data(), nTok * sizeof(float), cudaMemcpyHostToDevice)); - Stream stream; + auto q = uploadBuffer(hq, stream); + auto dk8 = uploadBuffer(k8, stream); + auto dv8 = uploadBuffer(v8, stream); + auto dkS = uploadBuffer(kSc, stream); + auto dvS = uploadBuffer(vSc, stream); + auto o = zeroBuffer(nQ, stream); executeFlashAttention( - q, dk8, dv8, o, bh, seqQ, seqK, causal, stream, dkS, dvS); - stream.sync(); + q.ptr().data, dk8.ptr().data, dv8.ptr().data, o.ptr().data, bh, seqQ, seqK, causal, stream, + dkS.ptr().data, dvS.ptr().data); + const std::vector got = downloadBuffer(o, stream); + report(name, maxAbsErr(got, dref), tol); +} - std::vector got(nQ); - gpuErrchk(cudaMemcpy(got.data(), o, nQ * sizeof(float), cudaMemcpyDeviceToHost)); - cudaFree(q); cudaFree(o); cudaFree(dk8); cudaFree(dv8); cudaFree(dkS); cudaFree(dvS); +#ifdef FK_HAS_FP8 +// The oracle decodes with attnFp8ToF32 on the host, so first make sure it is +// OCP e4m3 and not another 8-bit float (FNUZ, e5m2): a wrong format would +// otherwise fool the kernel and the oracle alike. +static void testFp8Format() { + if (attnFp8IsOcpE4M3()) { + std::cout << "Running test attn_fp8_e4m3 decodes OCP e4m3 (0x7E=448, 0x38=1, 0x80=-0): Success!!" + << std::endl; + } else { + std::cout << "FAIL attn_fp8_e4m3 is not OCP e4m3: 0x7E=" << attnFp8ToF32(0x7E) + << " 0x38=" << attnFp8ToF32(0x38) << " 0x80=" << attnFp8ToF32(static_cast(0x80)) + << std::endl; + ++failures; + } +} - double maxErr = 0.0; - for (size_t i = 0; i < nQ; ++i) maxErr = std::max(maxErr, std::abs((double)got[i] - dref[i])); - report(name, maxErr, tol); +template +static void testFp8KV(const char* name, const int bh, const int seqQ, const int seqK, + const bool causal, const double tol, const unsigned seed) { + std::mt19937 rng(seed); + std::uniform_real_distribution dist(-1.f, 1.f); + const size_t nQ = (size_t)bh * seqQ * HEAD_DIM, nK = (size_t)bh * seqK * HEAD_DIM; + const size_t nTok = (size_t)bh * seqK; + std::vector hq(nQ), hk(nK), hv(nK); + for (size_t i = 0; i < nQ; ++i) hq[i] = dist(rng); + for (size_t i = 0; i < nK; ++i) hk[i] = dist(rng); + for (size_t i = 0; i < nK; ++i) hv[i] = dist(rng); + + // host-side compression (per-token fp8 e4m3) + std::vector k8(nK), v8(nK); + std::vector kSc(nTok), vSc(nTok); + quantizeKVCacheFp8Host(hk.data(), k8.data(), kSc.data(), (int)nTok, HEAD_DIM); + quantizeKVCacheFp8Host(hv.data(), v8.data(), vSc.data(), (int)nTok, HEAD_DIM); + + // the oracle sees the DEQUANTIZED values, as in testInt8KV + std::vector dq(nQ), dkD(nK), dvD(nK), dref; + for (size_t i = 0; i < nQ; ++i) dq[i] = hq[i]; + for (size_t i = 0; i < nK; ++i) dkD[i] = (double)attnFp8ToF32(k8[i]) * kSc[i / HEAD_DIM]; + for (size_t i = 0; i < nK; ++i) dvD[i] = (double)attnFp8ToF32(v8[i]) * vSc[i / HEAD_DIM]; + cpuAttention(dq, dkD, dvD, dref, bh, seqQ, seqK, HEAD_DIM, + 1.0 / std::sqrt((double)HEAD_DIM), causal); + + Stream stream; + auto q = uploadBuffer(hq, stream); + auto dk8 = uploadBuffer(k8, stream); + auto dv8 = uploadBuffer(v8, stream); + auto dkS = uploadBuffer(kSc, stream); + auto dvS = uploadBuffer(vSc, stream); + auto o = zeroBuffer(nQ, stream); + const auto qIOp = makeAttentionRead(q.ptr().data, bh, seqQ, HEAD_DIM); + const auto kIOp = makeFp8KVRead(dk8.ptr().data, dkS.ptr().data, bh, seqK, HEAD_DIM); + const auto vIOp = makeFp8KVRead(dv8.ptr().data, dvS.ptr().data, bh, seqK, HEAD_DIM); + executeFlashAttention(qIOp, kIOp, vIOp, o.ptr().data, bh, seqQ, seqK, causal, stream); + const std::vector got = downloadBuffer(o, stream); + report(name, maxAbsErr(got, dref), tol); } +#endif // FK_HAS_FP8 static void testFusedEpilogue() { // attention output | Mul(2) | Add(0.5) fused in-register: compare against @@ -170,27 +162,20 @@ static void testFusedEpilogue() { for (auto* vec : { &hq, &hk, &hv }) for (auto& x : *vec) x = dist(rng); - float *q, *k, *v, *o1, *o2; - gpuErrchk(cudaMalloc(&q, n * sizeof(float))); - gpuErrchk(cudaMalloc(&k, n * sizeof(float))); - gpuErrchk(cudaMalloc(&v, n * sizeof(float))); - gpuErrchk(cudaMalloc(&o1, n * sizeof(float))); - gpuErrchk(cudaMalloc(&o2, n * sizeof(float))); - gpuErrchk(cudaMemcpy(q, hq.data(), n * sizeof(float), cudaMemcpyHostToDevice)); - gpuErrchk(cudaMemcpy(k, hk.data(), n * sizeof(float), cudaMemcpyHostToDevice)); - gpuErrchk(cudaMemcpy(v, hv.data(), n * sizeof(float), cudaMemcpyHostToDevice)); - Stream stream; - executeFlashAttention(q, k, v, o1, BH, SQ, SK, false, stream); + auto q = uploadBuffer(hq, stream); + auto k = uploadBuffer(hk, stream); + auto v = uploadBuffer(hv, stream); + auto o1 = zeroBuffer(n, stream); + auto o2 = zeroBuffer(n, stream); + executeFlashAttention(q.ptr().data, k.ptr().data, v.ptr().data, o1.ptr().data, BH, SQ, SK, + false, stream); const auto epilogue = Mul::build(2.f).then(Add::build(0.5f)); executeFlashAttention( - q, k, v, o2, BH, SQ, SK, false, stream, nullptr, nullptr, -1.f, epilogue); - stream.sync(); - - std::vector g1(n), g2(n); - gpuErrchk(cudaMemcpy(g1.data(), o1, n * sizeof(float), cudaMemcpyDeviceToHost)); - gpuErrchk(cudaMemcpy(g2.data(), o2, n * sizeof(float), cudaMemcpyDeviceToHost)); - cudaFree(q); cudaFree(k); cudaFree(v); cudaFree(o1); cudaFree(o2); + q.ptr().data, k.ptr().data, v.ptr().data, o2.ptr().data, BH, SQ, SK, false, stream, nullptr, nullptr, + -1.f, epilogue); + const std::vector g1 = downloadBuffer(o1, stream); + const std::vector g2 = downloadBuffer(o2, stream); double maxErr = 0.0; for (size_t i = 0; i < n; ++i) @@ -213,16 +198,15 @@ static void testFusedPrologue() { for (auto& x : hk) x = dist(rng); for (auto& x : hv) x = dist(rng); - float *q, *k, *v, *o; - gpuErrchk(cudaMalloc(&q, nQ * sizeof(float))); - gpuErrchk(cudaMalloc(&k, nK * sizeof(float))); - gpuErrchk(cudaMalloc(&v, nK * sizeof(float))); - gpuErrchk(cudaMalloc(&o, nQ * sizeof(float))); - gpuErrchk(cudaMemcpy(q, hq.data(), nQ * sizeof(float), cudaMemcpyHostToDevice)); - gpuErrchk(cudaMemcpy(k, hk.data(), nK * sizeof(float), cudaMemcpyHostToDevice)); - gpuErrchk(cudaMemcpy(v, hv.data(), nK * sizeof(float), cudaMemcpyHostToDevice)); - Stream stream; + auto qBuf = uploadBuffer(hq, stream); + auto kBuf = uploadBuffer(hk, stream); + auto vBuf = uploadBuffer(hv, stream); + auto oBuf = zeroBuffer(nQ, stream); + const float* q = qBuf.ptr().data; + const float* k = kBuf.ptr().data; + const float* v = vBuf.ptr().data; + float* o = oBuf.ptr().data; const double scl = 1.0 / std::sqrt((double)HEAD_DIM); std::vector dq(nQ), dk(nK), dv(nK), ref; for (size_t i = 0; i < nK; ++i) dk[i] = hk[i]; @@ -236,9 +220,7 @@ static void testFusedPrologue() { const auto vIOp = makeAttentionRead(v, BH, SK, HEAD_DIM); executeFlashAttention(qIOp, kIOp, vIOp, o, BH, SQ, SK, false, stream); - stream.sync(); - std::vector got(nQ); - gpuErrchk(cudaMemcpy(got.data(), o, nQ * sizeof(float), cudaMemcpyDeviceToHost)); + const std::vector got = downloadBuffer(oBuf, stream); for (size_t i = 0; i < nQ; ++i) dq[i] = 2.0 * hq[i]; // host-applied cpuAttention(dq, dk, dv, ref, BH, SQ, SK, HEAD_DIM, scl, false); double maxErr = 0.0; @@ -256,9 +238,7 @@ static void testFusedPrologue() { .then(Add::build(1.f)); executeFlashAttention(qIOp, kIOp, vIOp, o, BH, SQ, SK, false, stream); - stream.sync(); - std::vector got(nQ); - gpuErrchk(cudaMemcpy(got.data(), o, nQ * sizeof(float), cudaMemcpyDeviceToHost)); + const std::vector got = downloadBuffer(oBuf, stream); for (size_t i = 0; i < nQ; ++i) dq[i] = hq[i]; cpuAttention(dq, dk, dv, ref, BH, SQ, SK, HEAD_DIM, scl, false); double maxErr = 0.0; @@ -267,7 +247,6 @@ static void testFusedPrologue() { report("FlashAttention V-prologue ReadIOp.then(Mul(3)).then(Add(1))", maxErr, 5e-6); } - cudaFree(q); cudaFree(k); cudaFree(v); cudaFree(o); } int launch_impl() { @@ -277,6 +256,12 @@ int launch_impl() { testDense<128>("FA dense d128 b1 s64", 1, 64, 64, false, 5e-6, 4); testInt8KV<64>("FA int8-KV d64 b2 s64 causal", 2, 64, 64, true, 5e-6, 5); testInt8KV<64>("FA int8-KV d64 ragged s50/s100", 2, 50, 100, false, 5e-6, 6); +#ifdef FK_HAS_FP8 + testFp8Format(); + testFp8KV<64>("FA fp8-KV d64 b2 s64 causal", 2, 64, 64, true, 5e-6, 7); + testFp8KV<64>("FA fp8-KV d64 ragged s50/s100", 2, 50, 100, false, 5e-6, 8); + testFp8KV<128>("FA fp8-KV d128 b1 s64", 1, 64, 64, false, 5e-6, 9); +#endif testFusedEpilogue(); testFusedPrologue(); if (failures == 0) { return 0; } diff --git a/utests/algorithm/attention/utest_flash_attention_mma.h b/utests/algorithm/attention/utest_flash_attention_mma.h new file mode 100644 index 00000000..709ab680 --- /dev/null +++ b/utests/algorithm/attention/utest_flash_attention_mma.h @@ -0,0 +1,94 @@ +/* Copyright 2026 Oscar Amoros Huguet + Copyright 2026 Johnny Nunez + + 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. */ + +/* FlashAttentionMmaDPP (mma.sync.m16n8k16 bf16) against the double-precision + * oracle, through executeFlashAttentionMma with raw bf16 Q/K/V (the cp.async + * schedule): dense, causal, ragged sequence lengths, and a small grid that + * takes the split-KV forward and its combine kernel. The oracle sees the + * bf16-rounded inputs. The FP8 QK^T path (FK_ENABLE_FP8_QK, sm_120a) is not + * tested here. + * + * The kernel's inline PTX (mma.sync bf16, ldmatrix, cp.async) needs sm_80; + * on older GPUs the test prints SKIPPED and passes. */ + +#define __ONLY_CU__ // This file is only generated and compiled with nvcc, not with the host compiler +#include + +#include +#include + +#include + +namespace fk { + +// P is rounded to bf16 (8 significant bits) for the PV mma: each output is a +// convex combination of V values in [-1, 1], so the rounding error is below +// 2^-8 = 3.9e-3. The softmax uses ex2.approx; 1e-2 leaves margin for both. +constexpr double MMA_BF16_TOL = 1e-2; + +template +static void testMmaBf16(const char* name, const int bh, const int seqQ, const int seqK, const bool causal, + const unsigned seed) { + std::mt19937 rng(seed); + std::uniform_real_distribution dist(-1.f, 1.f); + const size_t nQ = (size_t)bh * seqQ * HEAD_DIM, nK = (size_t)bh * seqK * HEAD_DIM; + std::vector<__nv_bfloat16> bq(nQ), bk(nK), bv(nK); + std::vector dq(nQ), dk(nK), dv(nK), ref; + for (size_t i = 0; i < nQ; ++i) { bq[i] = __float2bfloat16(dist(rng)); dq[i] = __bfloat162float(bq[i]); } + for (size_t i = 0; i < nK; ++i) { bk[i] = __float2bfloat16(dist(rng)); dk[i] = __bfloat162float(bk[i]); } + for (size_t i = 0; i < nK; ++i) { bv[i] = __float2bfloat16(dist(rng)); dv[i] = __bfloat162float(bv[i]); } + cpuAttention(dq, dk, dv, ref, bh, seqQ, seqK, HEAD_DIM, 1.0 / std::sqrt((double)HEAD_DIM), causal); + + Stream stream; + auto q = uploadBuffer(bq, stream); + auto k = uploadBuffer(bk, stream); + auto v = uploadBuffer(bv, stream); + auto o = zeroBuffer(nQ, stream); + const auto qIOp = makeAttentionRead(static_cast(q.ptr().data), bh, seqQ, HEAD_DIM); + const auto kIOp = makeAttentionRead(static_cast(k.ptr().data), bh, seqK, HEAD_DIM); + const auto vIOp = makeAttentionRead(static_cast(v.ptr().data), bh, seqK, HEAD_DIM); + static_assert(isRawBf16Read>, "bf16 K/V must take the raw (cp.async) schedule"); + executeFlashAttentionMma(qIOp, kIOp, vIOp, o.ptr().data, bh, seqQ, seqK, causal, stream); + const std::vector got = downloadBuffer(o, stream); + report(name, maxAbsErr(got, ref), MMA_BF16_TOL); +} + +int launch_impl() { + int device = 0; + cudaDeviceProp prop{}; + gpuErrchk(cudaGetDevice(&device)); + gpuErrchk(cudaGetDeviceProperties(&prop, device)); + if (prop.major < 8) { + std::cout << "SKIPPED (FlashAttentionMmaDPP needs sm_80, this GPU is sm_" << prop.major << prop.minor << ")" + << std::endl; + return 0; + } + testMmaBf16<64>("FA-MMA bf16 d64 bh2 s256", 2, 256, 256, false, 1); + testMmaBf16<64>("FA-MMA bf16 d64 bh2 s256 causal", 2, 256, 256, true, 2); + testMmaBf16<64>("FA-MMA bf16 d64 ragged s100/s200", 2, 100, 200, false, 3); + testMmaBf16<128>("FA-MMA bf16 d128 bh2 s128 causal", 2, 128, 128, true, 4); + testMmaBf16<128>("FA-MMA bf16 d128 ragged s70/s150 causal", 2, 70, 150, true, 5); + // 16 q-blocks and 16 KV tiles per row: under one wave, so split-KV + combine + testMmaBf16<128>("FA-MMA bf16 d128 bh1 s1024 (split-KV)", 1, 1024, 1024, false, 6); + if (failures == 0) { return 0; } + std::cout << failures << " flash attention mma test(s) FAILED" << std::endl; + return -1; +} + +} // namespace fk + +int launch() { + return fk::launch_impl(); +} diff --git a/utests/algorithm/attention/utest_flash_decode.h b/utests/algorithm/attention/utest_flash_decode.h new file mode 100644 index 00000000..26f81d42 --- /dev/null +++ b/utests/algorithm/attention/utest_flash_decode.h @@ -0,0 +1,179 @@ +/* Copyright 2026 Oscar Amoros Huguet + Copyright 2026 Johnny Nunez + + 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. */ + +/* FlashDecodeDPP (seq_q == 1, split-KV) against the double-precision oracle, + * through executeFlashDecode, for every K/V read path of loadRow: generic + * (fp32 PerThreadRead), raw bf16 (vectorized bfloat162 loads), and the + * int8 / fp8 per-token caches (packed byte loads). The quantized cases give + * the oracle the dequantized values, as utest_flash_attention does. + * Forced split counts larger than seqK / 4 leave some splits empty, which + * must still arrive at the fused last-block merge; every case runs twice on + * the same workspace to check that the arrival counters reset themselves. */ + +#define __ONLY_GPU__ // Generated as _cu and _hip, not as _cpp: the attention DPPs are GPU-only +#include + +#include +#include + +#include + +namespace fk { + +// Decode accumulates in fp32 like FlashAttentionDPP, but its softmax uses +// __expf (fast exp, a few ulp) instead of expf, hence a looser bound than the +// 5e-6 of utest_flash_attention. +constexpr double DECODE_TOL = 1e-5; + +enum class DecodeKV { FP32, BF16, INT8, FP8 }; + +template +static void runDecode(const char* name, const QIOp& q, const KIOp& k, const VIOp& v, + const std::vector& ref, const int bh, const int seqK, + const int splitsOverride, Stream& stream) { + const int splits = splitsOverride > 0 ? splitsOverride : flashDecodeSplits(bh, seqK); + auto workspace = zeroBuffer(flashDecodeWorkspaceFloats(bh, splits, HEAD_DIM), stream); + auto o = zeroBuffer((size_t)bh * HEAD_DIM, stream); + for (int run = 0; run < 2; ++run) { + executeFlashDecode(q, k, v, o.ptr().data, workspace.ptr().data, bh, seqK, stream, + -1.f, splitsOverride); + const std::vector got = downloadBuffer(o, stream); + const std::string runName = std::string(name) + " splits=" + std::to_string(splits) + + (run == 0 ? "" : " (2nd call, same workspace)"); + report(runName.c_str(), maxAbsErr(got, ref), DECODE_TOL); + } +} + +template +static void testDecode(const char* name, const int bh, const int seqK, const int splitsOverride, + const unsigned seed) { + std::mt19937 rng(seed); + std::uniform_real_distribution dist(-1.f, 1.f); + const size_t nQ = (size_t)bh * HEAD_DIM, nK = (size_t)bh * seqK * HEAD_DIM; + const size_t nTok = (size_t)bh * seqK; + std::vector hq(nQ), hk(nK), hv(nK); + for (auto& x : hq) x = dist(rng); + for (auto& x : hk) x = dist(rng); + for (auto& x : hv) x = dist(rng); + + std::vector dq(nQ), dk(nK), dv(nK), ref; + Stream stream; + const double scale = 1.0 / std::sqrt((double)HEAD_DIM); + if constexpr (KV == DecodeKV::FP32) { + for (size_t i = 0; i < nQ; ++i) dq[i] = hq[i]; + for (size_t i = 0; i < nK; ++i) { dk[i] = hk[i]; dv[i] = hv[i]; } + cpuAttention(dq, dk, dv, ref, bh, 1, seqK, HEAD_DIM, scale, false); + auto q = uploadBuffer(hq, stream); + auto kd = uploadBuffer(hk, stream); + auto vd = uploadBuffer(hv, stream); + runDecode(name, makeAttentionRead(q.ptr().data, bh, 1, HEAD_DIM), + makeAttentionRead(kd.ptr().data, bh, seqK, HEAD_DIM), + makeAttentionRead(vd.ptr().data, bh, seqK, HEAD_DIM), ref, bh, seqK, + splitsOverride, stream); + } else if constexpr (KV == DecodeKV::BF16) { + // Q, K and V are all bf16 (the raw prologue); the oracle sees the + // bf16-rounded values + std::vector bq(nQ), bk(nK), bv(nK); + for (size_t i = 0; i < nQ; ++i) { bq[i] = __float2bfloat16(hq[i]); dq[i] = __bfloat162float(bq[i]); } + for (size_t i = 0; i < nK; ++i) { + bk[i] = __float2bfloat16(hk[i]); dk[i] = __bfloat162float(bk[i]); + bv[i] = __float2bfloat16(hv[i]); dv[i] = __bfloat162float(bv[i]); + } + cpuAttention(dq, dk, dv, ref, bh, 1, seqK, HEAD_DIM, scale, false); + auto q = uploadBuffer(bq, stream); + auto kd = uploadBuffer(bk, stream); + auto vd = uploadBuffer(bv, stream); + const auto kIOp = makeAttentionRead(static_cast(kd.ptr().data), bh, seqK, HEAD_DIM); + static_assert(isRawBf16Read>, "bf16 K must take the raw (vectorized) path"); + runDecode(name, makeAttentionRead(static_cast(q.ptr().data), bh, 1, HEAD_DIM), + kIOp, + makeAttentionRead(static_cast(vd.ptr().data), bh, seqK, HEAD_DIM), + ref, bh, seqK, splitsOverride, stream); + } else { + std::vector k8(nK), v8(nK); + std::vector kSc(nTok), vSc(nTok); + for (size_t i = 0; i < nQ; ++i) dq[i] = hq[i]; + if constexpr (KV == DecodeKV::INT8) { + quantizeKVCacheHost(hk.data(), k8.data(), kSc.data(), (int)nTok, HEAD_DIM); + quantizeKVCacheHost(hv.data(), v8.data(), vSc.data(), (int)nTok, HEAD_DIM); + for (size_t i = 0; i < nK; ++i) { + dk[i] = (double)k8[i] * kSc[i / HEAD_DIM]; + dv[i] = (double)v8[i] * vSc[i / HEAD_DIM]; + } + } else { +#ifdef FK_HAS_FP8 + quantizeKVCacheFp8Host(hk.data(), k8.data(), kSc.data(), (int)nTok, HEAD_DIM); + quantizeKVCacheFp8Host(hv.data(), v8.data(), vSc.data(), (int)nTok, HEAD_DIM); + for (size_t i = 0; i < nK; ++i) { + dk[i] = (double)attnFp8ToF32(k8[i]) * kSc[i / HEAD_DIM]; + dv[i] = (double)attnFp8ToF32(v8[i]) * vSc[i / HEAD_DIM]; + } +#endif + } + cpuAttention(dq, dk, dv, ref, bh, 1, seqK, HEAD_DIM, scale, false); + auto q = uploadBuffer(hq, stream); + auto kd = uploadBuffer(k8, stream); + auto vd = uploadBuffer(v8, stream); + auto ks = uploadBuffer(kSc, stream); + auto vs = uploadBuffer(vSc, stream); + const auto qIOp = makeAttentionRead(q.ptr().data, bh, 1, HEAD_DIM); + if constexpr (KV == DecodeKV::INT8) { + runDecode(name, qIOp, makeInt8KVRead(kd.ptr().data, ks.ptr().data, bh, seqK, HEAD_DIM), + makeInt8KVRead(vd.ptr().data, vs.ptr().data, bh, seqK, HEAD_DIM), ref, bh, seqK, + splitsOverride, stream); + } else { +#ifdef FK_HAS_FP8 + runDecode(name, qIOp, makeFp8KVRead(kd.ptr().data, ks.ptr().data, bh, seqK, HEAD_DIM), + makeFp8KVRead(vd.ptr().data, vs.ptr().data, bh, seqK, HEAD_DIM), ref, bh, seqK, + splitsOverride, stream); +#endif + } + } +} + +int launch_impl() { + // splits=0: flashDecodeSplits picks them (4 for bh4 s1000, 3 for bh2 s513) + testDecode<64, DecodeKV::FP32>("FlashDecode fp32 d64 bh4 s1000", 4, 1000, 0, 1); + testDecode<32, DecodeKV::FP32>("FlashDecode fp32 d32 bh2 s300", 2, 300, 1, 2); + testDecode<128, DecodeKV::FP32>("FlashDecode fp32 d128 bh2 s513", 2, 513, 0, 3); + testDecode<64, DecodeKV::FP32>("FlashDecode fp32 d64 bh2 s100 empty splits", 2, 100, 30, 4); + testDecode<64, DecodeKV::BF16>("FlashDecode bf16 d64 bh4 s1000", 4, 1000, 0, 5); + testDecode<128, DecodeKV::BF16>("FlashDecode bf16 d128 bh2 s100 empty splits", 2, 100, 30, 6); + testDecode<64, DecodeKV::INT8>("FlashDecode int8-KV d64 bh4 s1000", 4, 1000, 0, 7); + testDecode<128, DecodeKV::INT8>("FlashDecode int8-KV d128 bh2 s100 empty splits", 2, 100, 30, 8); +#ifdef FK_HAS_FP8 + testDecode<64, DecodeKV::FP8>("FlashDecode fp8-KV d64 bh4 s1000", 4, 1000, 0, 9); + testDecode<128, DecodeKV::FP8>("FlashDecode fp8-KV d128 bh2 s100 empty splits", 2, 100, 30, 10); +#endif + // d32 (1 element per lane) and d256 (8): the packed loads don't cover + // these widths, loadRow falls back to reading each element + testDecode<32, DecodeKV::BF16>("FlashDecode bf16 d32 bh2 s300", 2, 300, 0, 11); + testDecode<256, DecodeKV::BF16>("FlashDecode bf16 d256 bh2 s300", 2, 300, 0, 12); + testDecode<32, DecodeKV::INT8>("FlashDecode int8-KV d32 bh2 s300", 2, 300, 0, 13); + testDecode<256, DecodeKV::INT8>("FlashDecode int8-KV d256 bh2 s300", 2, 300, 0, 14); +#ifdef FK_HAS_FP8 + testDecode<32, DecodeKV::FP8>("FlashDecode fp8-KV d32 bh2 s300", 2, 300, 0, 15); + testDecode<256, DecodeKV::FP8>("FlashDecode fp8-KV d256 bh2 s300", 2, 300, 0, 16); +#endif + if (failures == 0) { return 0; } + std::cout << failures << " flash decode test(s) FAILED" << std::endl; + return -1; +} + +} // namespace fk + +int launch() { + return fk::launch_impl(); +} diff --git a/utests/algorithm/attention/utest_softmax.h b/utests/algorithm/attention/utest_softmax.h index 0500d797..a6317597 100644 --- a/utests/algorithm/attention/utest_softmax.h +++ b/utests/algorithm/attention/utest_softmax.h @@ -12,12 +12,13 @@ See the License for the specific language governing permissions and limitations under the License. */ -#define __ONLY_CU__ // This file is only generated and compiled with nvcc, not with the host compiler +#define __ONLY_GPU__ // Generated as _cu and _hip, not as _cpp: the attention DPPs are GPU-only #include #include #include +#include #include #include #include @@ -27,31 +28,50 @@ namespace fk { static int failures = 0; +// Row y of the pinned mirror of a pitched Ptr2D. +static float* pinnedRow(const Ptr2D& image, const int y) { + const RawPtr pinned = image.ptrPinned(); + return reinterpret_cast(reinterpret_cast(pinned.data) + (size_t)y * pinned.dims.pitch); +} + +// Uploads a width x height host image (rows packed) into the pitched Ptr2D. +static void uploadRows(Ptr2D& image, const std::vector& host, const int width, const int height, + Stream& stream) { + for (int y = 0; y < height; ++y) { + std::copy(host.begin() + (size_t)y * width, host.begin() + (size_t)(y + 1) * width, pinnedRow(image, y)); + } + image.upload(stream); +} + +// Waits for the stream and returns the image with its rows packed. +static std::vector downloadRows(Ptr2D& image, const int width, const int height, Stream& stream) { + image.download(stream); + stream.sync(); + std::vector host((size_t)width * height); + for (int y = 0; y < height; ++y) { + const float* row = pinnedRow(image, y); + std::copy(row, row + width, host.begin() + (size_t)y * width); + } + return host; +} + static void runCase(const char* name, const int width, const int height, const float lo, const float hi, const double tol, const unsigned seed) { std::mt19937 rng(seed); std::uniform_real_distribution dist(lo, hi); - Ptr2D input(width, height, 0, MemType::Device); - Ptr2D output(width, height, 0, MemType::Device); + Ptr2D input(width, height); + Ptr2D output(width, height); std::vector host((size_t)width * height); for (auto& v : host) v = dist(rng); - gpuErrchk(cudaMemcpy2D(input.ptr().data, input.ptr().dims.pitch, - host.data(), width * sizeof(float), - width * sizeof(float), height, cudaMemcpyHostToDevice)); - Stream stream; + uploadRows(input, host, width, height, stream); const auto inIOp = PerThreadRead::build(input.ptr()); const auto outIOp = PerThreadWrite::build(output.ptr()); executeSoftmax<256>(inIOp, outIOp, stream); - stream.sync(); - - std::vector got((size_t)width * height); - gpuErrchk(cudaMemcpy2D(got.data(), width * sizeof(float), - output.ptr().data, output.ptr().dims.pitch, - width * sizeof(float), height, cudaMemcpyDeviceToHost)); + const std::vector got = downloadRows(output, width, height, stream); double maxErr = 0.0; for (int y = 0; y < height; ++y) { @@ -85,27 +105,19 @@ static void runPrologueCase(const char* name, const int width, const int height, std::mt19937 rng(seed); std::uniform_real_distribution dist(-4.f, 4.f); - Ptr2D input(width, height, 0, MemType::Device); - Ptr2D output(width, height, 0, MemType::Device); + Ptr2D input(width, height); + Ptr2D output(width, height); std::vector host((size_t)width * height); for (auto& v : host) v = dist(rng); - gpuErrchk(cudaMemcpy2D(input.ptr().data, input.ptr().dims.pitch, - host.data(), width * sizeof(float), - width * sizeof(float), height, cudaMemcpyHostToDevice)); - Stream stream; + uploadRows(input, host, width, height, stream); const auto inIOp = PerThreadRead::build(input.ptr()) .then(Mul::build(2.f)) .then(Add::build(1.f)); const auto outIOp = PerThreadWrite::build(output.ptr()); executeSoftmax<256>(inIOp, outIOp, stream); - stream.sync(); - - std::vector got((size_t)width * height); - gpuErrchk(cudaMemcpy2D(got.data(), width * sizeof(float), - output.ptr().data, output.ptr().dims.pitch, - width * sizeof(float), height, cudaMemcpyDeviceToHost)); + const std::vector got = downloadRows(output, width, height, stream); double maxErr = 0.0; for (int y = 0; y < height; ++y) {