diff --git a/Compiling.md b/Compiling.md index 744251ccef..1676d8a20c 100644 --- a/Compiling.md +++ b/Compiling.md @@ -85,13 +85,19 @@ As also mentioned in the instructions below but repeated here for visibility, if generic `/opt/rocm/lib` path is included as a fallback after the pinned path, so the binary can also run on a machine with a different (soname-compatible) ROCm version installed. * 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 (compiled kernels in `~/.cache/miopen/`, the tuning find-db 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 (CK, packaged as - `composablekernel-dev` / `amdrocm-ck*`) is also installed, the build additionally enables a fused-attention fast path + * **Transformer/attention models (model version 17+):** supported on all architectures via + built-in kernels. On RDNA2 (gfx103x) GPUs the built-in attention is a dedicated dot2 kernel + by default (on a gfx1030/RX 6800 XT, transformer nets ran ~1.5x faster end to end with it + than with the generic kernel); other GPUs with `v_dot2_f32_f16` can opt into it with + `rocmUseDot2Attention = true`, untested there. If a matching version of AMD's Composable + Kernel (CK, packaged as `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 + RDNA1/RDNA2 always use the built-in kernels, 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. + See `cpp/external/composable_kernel_fmha/README.md` for details. For RDNA2 (gfx103x) + targets, an installed CK also enables its DPP GEMM kernels, which the backend uses for the + FP16 transformer projections and 1x1 convs where they time faster than hipBLAS. ## Windows * TLDR: diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 32852705a1..859e6ed366 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -899,10 +899,12 @@ elseif(USE_BACKEND STREQUAL "ROCM") # rocmbackend.cpp, rocmutils.cpp, and rocmhelpers.hip are thin wrappers around # cudaandrocmbackend.inc, cudaandrocmutils.inc, and cudaandrocmhelpers.inc, which are shared # with the CUDA backend (tracked via compiler depfiles, so they do not appear here). + # rocmkernels.hip holds the ROCm-only kernels with no CUDA counterpart. set(NEURALNET_BACKEND_SOURCES neuralnet/rocmbackend.cpp neuralnet/rocmutils.cpp neuralnet/rocmhelpers.hip + neuralnet/rocmkernels.hip ) elseif(USE_BACKEND STREQUAL "") @@ -1624,6 +1626,98 @@ elseif(USE_BACKEND STREQUAL "ROCM") endif() target_compile_definitions(katago PRIVATE KATAGO_ROCM_HAS_CK_FMHA=0) endif() + + # ---------- Optional CK DPP GEMM (RDNA2) ---------- + # Composable Kernel's DPP GEMM kernels beat rocBLAS on many transformer projection GEMM shapes on + # gfx103x (the backend picks per shape at runtime, see rocmckgemm.h). They only exist for gfx103x, + # and CK's headers are slow to compile, so rocmckgemm.hip is built as a separate object for just + # the gfx103x archs of the build, after a one-time configure probe (cached) that it compiles + # against the installed CK headers. + find_path(KATAGO_CK_INCLUDE_DIR + NAMES ck/tensor_operation/gpu/device/impl/device_gemm_dpp.hpp + HINTS ${CMAKE_PREFIX_PATH} ENV HIP_PATH ENV ROCM_PATH + PATH_SUFFIXES include + ) + set(KATAGO_CK_GEMM_ARCHS "") + foreach(_arch ${CMAKE_HIP_ARCHITECTURES}) + string(REGEX REPLACE ":.*$" "" _arch_base "${_arch}") + if(_arch_base MATCHES "^gfx103[0-6]$") + list(APPEND KATAGO_CK_GEMM_ARCHS "${_arch}") + endif() + endforeach() + set(_katago_ck_gemm_enabled FALSE) + if(KATAGO_CK_GEMM_ARCHS AND NOT KATAGO_CK_INCLUDE_DIR) + message(STATUS "CK headers (ck/tensor_operation/gpu/device/impl/device_gemm_dpp.hpp) not found; ROCm backend will use hipBLAS for all GEMMs on ${KATAGO_CK_GEMM_ARCHS}") + elseif(KATAGO_CK_GEMM_ARCHS) + # As for the CK FMHA probe above: the signature covers everything the verdict depends on, and + # the probe compiles with the per-configuration flags the real build will use. + file(SHA1 "${CMAKE_SOURCE_DIR}/neuralnet/rocmckgemm.hip" _ck_gemm_src_hash) + file(SHA1 "${KATAGO_CK_INCLUDE_DIR}/ck/tensor_operation/gpu/device/impl/device_gemm_dpp.hpp" _ck_gemm_hdr_hash) + set(_ck_gemm_probe_sig "${KATAGO_CK_INCLUDE_DIR}|${CMAKE_HIP_COMPILER}|${CMAKE_HIP_COMPILER_VERSION}|${CMAKE_BUILD_TYPE}|${CMAKE_HIP_FLAGS}|${KATAGO_CK_GEMM_ARCHS}|${_ck_gemm_src_hash}|${_ck_gemm_hdr_hash}") + if(NOT DEFINED KATAGO_CK_GEMM_COMPILE_OK OR NOT "${KATAGO_CK_GEMM_PROBED_SIG}" STREQUAL "${_ck_gemm_probe_sig}") + message(STATUS "Probing whether the CK DPP GEMM kernels compile for ${KATAGO_CK_GEMM_ARCHS} (may take a little while)...") + set(_ck_gemm_probe_dir "${CMAKE_BINARY_DIR}/_katago_ck_gemm_probe") + file(MAKE_DIRECTORY "${_ck_gemm_probe_dir}") + separate_arguments(_ck_gemm_probe_flags NATIVE_COMMAND "${CMAKE_HIP_FLAGS}") + string(TOUPPER "${CMAKE_BUILD_TYPE}" _ck_gemm_probe_build_type) + if(NOT _ck_gemm_probe_build_type) + set(_ck_gemm_probe_build_type "RELEASE") + endif() + if(CMAKE_HIP_FLAGS_${_ck_gemm_probe_build_type}) + separate_arguments(_ck_gemm_probe_config_flags NATIVE_COMMAND "${CMAKE_HIP_FLAGS_${_ck_gemm_probe_build_type}}") + list(APPEND _ck_gemm_probe_flags ${_ck_gemm_probe_config_flags}) + else() + list(APPEND _ck_gemm_probe_flags "-O3") + endif() + if(KATAGO_HIP_SYSTEM_INCLUDE_DIR) + list(APPEND _ck_gemm_probe_flags "-isystem${KATAGO_HIP_SYSTEM_INCLUDE_DIR}") + endif() + set(_ck_gemm_probe_arch_flags "") + foreach(_arch ${KATAGO_CK_GEMM_ARCHS}) + list(APPEND _ck_gemm_probe_arch_flags "--offload-arch=${_arch}") + endforeach() + # The arch list macro only feeds rocmCkGemmSupportsArch, so any value compiles the same. + execute_process( + COMMAND "${CMAKE_HIP_COMPILER}" -x hip ${_ck_gemm_probe_arch_flags} -std=c++17 -c + "-isystem${KATAGO_CK_INCLUDE_DIR}" -DCK_ENABLE_FP16 "-DKATAGO_ROCM_CK_GEMM_ARCHS=\"probe\"" + ${_ck_gemm_probe_flags} + "${CMAKE_SOURCE_DIR}/neuralnet/rocmckgemm.hip" -o "${_ck_gemm_probe_dir}/probe.o" + RESULT_VARIABLE _ck_gemm_rc + OUTPUT_QUIET + ERROR_VARIABLE _ck_gemm_err) + if(_ck_gemm_rc EQUAL 0) + set(KATAGO_CK_GEMM_COMPILE_OK TRUE CACHE INTERNAL "") + else() + file(WRITE "${_ck_gemm_probe_dir}/probe_error.log" "${_ck_gemm_err}") + set(KATAGO_CK_GEMM_COMPILE_OK FALSE CACHE INTERNAL "") + endif() + set(KATAGO_CK_GEMM_PROBED_SIG "${_ck_gemm_probe_sig}" CACHE INTERNAL "") + endif() + if(KATAGO_CK_GEMM_COMPILE_OK) + set(_katago_ck_gemm_enabled TRUE) + else() + message(STATUS "The CK DPP GEMM kernels failed to compile against ${KATAGO_CK_INCLUDE_DIR} (see ${CMAKE_BINARY_DIR}/_katago_ck_gemm_probe/probe_error.log); ROCm backend will use hipBLAS for all GEMMs") + endif() + endif() + if(_katago_ck_gemm_enabled) + message(STATUS "Enabling CK DPP GEMM kernels for ${KATAGO_CK_GEMM_ARCHS}") + add_library(katago_ck_gemm OBJECT neuralnet/rocmckgemm.hip) + set_target_properties(katago_ck_gemm PROPERTIES HIP_ARCHITECTURES "${KATAGO_CK_GEMM_ARCHS}") + target_include_directories(katago_ck_gemm SYSTEM PRIVATE ${KATAGO_CK_INCLUDE_DIR}) + string(REPLACE ";" "\\;" _ck_gemm_archs_escaped "${KATAGO_CK_GEMM_ARCHS}") + target_compile_definitions(katago_ck_gemm PRIVATE CK_ENABLE_FP16 "KATAGO_ROCM_CK_GEMM_ARCHS=\"${_ck_gemm_archs_escaped}\"") + # Same stale-/usr/include/hip workaround as for katago_ck_fmha_kernels above. + if(KATAGO_HIP_SYSTEM_INCLUDE_DIR) + target_compile_options(katago_ck_gemm PRIVATE "-isystem${KATAGO_HIP_SYSTEM_INCLUDE_DIR}") + endif() + target_link_libraries(katago katago_ck_gemm) + target_compile_definitions(katago PRIVATE KATAGO_ROCM_HAS_CK_GEMM=1) + else() + if(NOT KATAGO_CK_GEMM_ARCHS) + message(STATUS "No gfx103x arch targeted; ROCm backend will use hipBLAS for all GEMMs") + endif() + target_compile_definitions(katago PRIVATE KATAGO_ROCM_HAS_CK_GEMM=0) + endif() # On Windows plain clang++ (not hipcc) is the CXX compiler, so rocmbackend.cpp - which contains # the device-code-including shared backend - must get real HIP compilation via a source property, # CK or no CK. (The CK generated files get the same property inside the CK branch above.) diff --git a/cpp/configs/analysis_example.cfg b/cpp/configs/analysis_example.cfg index cc09490414..8057337c4c 100644 --- a/cpp/configs/analysis_example.cfg +++ b/cpp/configs/analysis_example.cfg @@ -292,6 +292,8 @@ nnRandomize = true # rocmUseFP16 = auto # rocmUseNHWC = auto # Uses NHWC tensor layout. Default: auto - NHWC with FP16 on GPUs where NHWC convolutions are faster (CDNA), else NCHW; transformers always NHWC. # rocmUse1x1Matmul = auto # Whether 1x1 NHWC convs run as a hipBLAS GEMM instead of a MIOpen conv. auto = GEMM when possible (MIOpen's NHWC conv are slow on some GPUs) +# rocmUseDot2Attention = auto # Whether FP16 attention uses the built-in dot2 attention kernels when the CK fused-attention path does not run. auto = on for RDNA2 (gfx103x), where they were measured; true = on for any GPU with v_dot2_f32_f16 (untested there) +# rocmUseCkGemm = auto # Whether FP16 transformer projections and 1x1 convs may run on CK's DPP GEMM kernels instead of hipBLAS, picked per shape by timing both. auto = on where built (RDNA2/gfx103x targets with CK installed) # ONNX Runtime backend settings-------------------------------------- diff --git a/cpp/configs/gtp_example.cfg b/cpp/configs/gtp_example.cfg index 2d9610c7bc..44f0271f87 100644 --- a/cpp/configs/gtp_example.cfg +++ b/cpp/configs/gtp_example.cfg @@ -626,6 +626,16 @@ searchFactorWhenWinningThreshold = 0.95 # Whether 1x1 NHWC convs run as a hipBLAS GEMM instead of a MIOpen conv. # auto = GEMM when possible (MIOpen's NHWC conv are slow on some GPUs) # rocmUse1x1Matmul = auto +# Whether FP16 attention (transformer models) uses the built-in dot2 attention kernels when +# the CK fused-attention path does not run (e.g. always on RDNA2). +# auto = on for RDNA2 (gfx103x), where they were measured; true = on for any GPU with +# v_dot2_f32_f16 (everything but gfx900-class GCN and gfx1010/gfx1013), untested there. +# rocmUseDot2Attention = auto +# Whether FP16 transformer projections and 1x1 convs may run on Composable Kernel's DPP GEMM +# kernels instead of hipBLAS, picked per GEMM shape by timing both during the first +# evaluations. Only built for RDNA2 (gfx103x) targets when CK is installed. +# auto = on where built +# rocmUseCkGemm = auto # ------------------------------ # OpenCL GPU settings diff --git a/cpp/external/composable_kernel_fmha/README.md b/cpp/external/composable_kernel_fmha/README.md index e622daf2c6..bbf6fbbf6f 100644 --- a/cpp/external/composable_kernel_fmha/README.md +++ b/cpp/external/composable_kernel_fmha/README.md @@ -58,8 +58,9 @@ regenerating from a different tag, verify against the installed - 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.) + for gfx10/RDNA2 - gfx1030/1031/1032 always use the built-in attention kernels, the dot2 ones of + `cpp/neuralnet/rocmkernels.h` where they apply. gfx125 doesn't exist as a target in this codegen + version either.) ## Regenerating diff --git a/cpp/neuralnet/cudaandrocmbackend.inc b/cpp/neuralnet/cudaandrocmbackend.inc index cd19a7e7da..3702924fed 100644 --- a/cpp/neuralnet/cudaandrocmbackend.inc +++ b/cpp/neuralnet/cudaandrocmbackend.inc @@ -48,6 +48,7 @@ #error "Define only one of KATAGO_GPU_CUDA and KATAGO_GPU_HIP (the gates below are if/else)" #endif +#include #include #include #include @@ -442,8 +443,14 @@ enum class SharedWeightsKind : int { CONV_1X1_MATMUL = 2, QKV_PACKED = 3, FFN_PACKED = 4, + // Out-major ([outC][inC]) linear1 then gate weights: for the CUTLASS fused FFN on CUDA, the CK + // GEMM path on ROCm. FFN_FUSED_PACKED = 5, QKV_COMBINED = 6, + // Out-major ([outC][inC]) variants used by the ROCm CK GEMM path, see rocmTokenGemmCkOrHipblas. + MATMUL_OUT_MAJOR = 7, + CONV_1X1_OUT_MAJOR = 8, + QKV_COMBINED_OUT_MAJOR = 9, }; //Upload a weight buffer to the device, taking ownership of the allocation even if the upload @@ -644,6 +651,38 @@ struct CudaHandles { // Whether 1x1 NHWC convs run as a hipBLAS GEMM instead of a MIOpen conv, already resolved // against the GPU architecture, so Auto here means on. Set before Model construction. enabled_t use1x1MatmulMode; + // Whether FP16 attention may use the dot2 kernels in rocmkernels.h: the arch has + // v_dot2_f32_f16, and rocmUseDot2Attention is true, or auto on an arch it is measured on (see + // rocmDot2AttentionAutoArch). Set before Model construction, which picks the QKV layout by it. + bool dot2AttentionEnabled; + bool loggedUsingDot2Attention; + // hipDeviceProp_t::multiProcessorCount and warpSize, which size the dot2 attention launch. + int numMultiprocessors; + int waveSize; + // Whether FP16 token GEMMs (transformer projections, NHWC 1x1 convs) may run on the CK DPP GEMM + // kernels (rocmckgemm.h): built for this arch and not disabled by config. Set before Model + // construction, since the layers that use it upload their weights out-major for it. + bool ckGemmEnabled; + // Per GEMM problem {M, N, K, lda, ldb, ldc, accumulate}: the CK instance to run (or -1 for + // hipBLAS), chosen from timings of the real calls, see rocmTokenGemmCkOrHipblas. Only touched by + // this handle's server thread. + struct CkGemmChoice { + // hipBLAS (-1) and the CK instances supporting the problem. + std::vector candidates; + // While undecided, the start and stop events of each timed call, the calls rotating through + // the candidates. + std::vector sampleEvents; + bool decided = false; + int instance = -1; + CkGemmChoice() = default; + CkGemmChoice(const CkGemmChoice&) = delete; + CkGemmChoice& operator=(const CkGemmChoice&) = delete; + ~CkGemmChoice() { + for(hipEvent_t e : sampleEvents) + (void)hipEventDestroy(e); + } + }; + std::map, CkGemmChoice> ckGemmChoices; CudaHandles(int major, int minor, Logger* logger_, bool ckFmhaArchSupported_, bool disableFusedAttention_) : majorComputeCapability(major), @@ -658,7 +697,12 @@ struct CudaHandles { loggedUsingCkFmha(false), loggedCkFmhaDisabled(false), loggedUsing1x1Matmul(false), - use1x1MatmulMode(enabled_t::Auto) + use1x1MatmulMode(enabled_t::Auto), + dot2AttentionEnabled(false), + loggedUsingDot2Attention(false), + numMultiprocessors(0), + waveSize(0), + ckGemmEnabled(false) { CUBLAS_ERR("CudaHandles",cublasSetStream(cublas.handle,stream.stream)); CUDNN_ERR("CudaHandles",cudnnSetStream(cudnn.handle,stream.stream)); @@ -676,6 +720,131 @@ struct CudaHandles { }; #endif +#ifdef KATAGO_GPU_HIP +//FP16 GEMM over tokens on the CK DPP kernels (see rocmckgemm.h): +// out[t][o] (row stride outStride) (+)= sum_i in[t][i] (row stride inStride) * weights[o][i] +//with weights out-major, [outC][inC] row-major, and FP32 accumulation, adding into out if +//accumulate. runHipblas must compute the same through hipBLAS (from the layer's in-major +//weights). Which of the two (and which CK instance) is fastest per problem shape is decided by +//timing the real calls: the calls of each shape rotate through the candidates, each timed with +//events, until every candidate has CK_GEMM_SAMPLES timings and the last one has completed (polled, +//never waited for); the candidate with the best timing runs from then on. Timing in place rather +//than in isolation matters: back-to-back runs on cache-hot operands rank the candidates +//differently than the model's interleaved calls do. All candidates compute the same product with +//FP32 accumulation, so until the choice settles, results differ from the steady state only by +//summation order (and, when accumulating, CK's FP16 add, see rocmCkGemmRun). The choice is per +//handle, so it is not bit-reproducible across runs. +template +static void rocmTokenGemmCkOrHipblas( + CudaHandles* cudaHandles, + const string& name, + const void* weights, + int outChannels, + int inChannels, + const void* inputBuf, + int inStride, + void* outputBuf, + int outStride, + int tokens, + bool accumulate, + RunHipblas runHipblas +) { +#if KATAGO_ROCM_HAS_CK_GEMM + constexpr int CK_GEMM_SAMPLES = 4; + auto runCandidate = [&](int instance) { + if(instance < 0) { + runHipblas(); + return; + } + rocmCkGemmRun( + instance, accumulate, + (const half*)inputBuf, inStride, (const half*)weights, inChannels, (half*)outputBuf, outStride, + tokens, outChannels, inChannels, cudaHandles->stream); + CUDA_ERR(name.c_str(), hipPeekAtLastError()); + }; + const std::array key = {tokens, outChannels, inChannels, inStride, inChannels, outStride, accumulate ? 1 : 0}; + CudaHandles::CkGemmChoice& choice = cudaHandles->ckGemmChoices[key]; + if(!choice.decided) { + if(choice.candidates.empty()) { + choice.candidates.push_back(-1); + for(int i = 0; i < rocmCkGemmNumInstances(); i++) { + if(rocmCkGemmSupports(i, tokens, outChannels, inChannels, inStride, inChannels, outStride)) + choice.candidates.push_back(i); + } + } + const size_t numTaken = choice.sampleEvents.size() / 2; + bool samplesReady = false; + if(numTaken >= choice.candidates.size() * CK_GEMM_SAMPLES) { + const hipError_t status = hipEventQuery(choice.sampleEvents.back()); + if(status != hipErrorNotReady) + CUDA_ERR(name.c_str(), status); + samplesReady = status == hipSuccess; + } + if(!samplesReady) { + const int candidate = choice.candidates[numTaken % choice.candidates.size()]; + hipEvent_t start, stop; + CUDA_ERR(name.c_str(), hipEventCreate(&start)); + choice.sampleEvents.push_back(start); + CUDA_ERR(name.c_str(), hipEventCreate(&stop)); + choice.sampleEvents.push_back(stop); + CUDA_ERR(name.c_str(), hipEventRecord(start, cudaHandles->stream)); + runCandidate(candidate); + CUDA_ERR(name.c_str(), hipEventRecord(stop, cudaHandles->stream)); + return; + } + // The candidate with the best (fastest) sample wins, ties going to hipBLAS, the first + // candidate. Every sample has completed: the stream ran them in order. + vector bestMs(choice.candidates.size(), 1e30f); + for(size_t s = 0; s < numTaken; s++) { + float ms = 0.0f; + CUDA_ERR(name.c_str(), hipEventElapsedTime(&ms, choice.sampleEvents[2 * s], choice.sampleEvents[2 * s + 1])); + const size_t j = s % choice.candidates.size(); + bestMs[j] = std::min(bestMs[j], ms); + } + size_t best = 0; + for(size_t j = 1; j < choice.candidates.size(); j++) { + if(bestMs[j] < bestMs[best]) + best = j; + } + choice.instance = choice.candidates[best]; + choice.decided = true; + for(hipEvent_t e : choice.sampleEvents) + CUDA_ERR(name.c_str(), hipEventDestroy(e)); + choice.sampleEvents.clear(); + } + runCandidate(choice.instance); +#else + // Never called without CK: the out-major weights only exist when cudaHandles->ckGemmEnabled. + (void)cudaHandles; (void)name; (void)weights; (void)outChannels; (void)inChannels; + (void)inputBuf; (void)inStride; (void)outputBuf; (void)outStride; (void)tokens; (void)accumulate; + runHipblas(); +#endif +} + +//Transposes in-major [inC][outC] weights (the layout of MatMulLayerDesc) to out-major [outC][inC]. +static vector transposeInMajorWeights(const vector& w, int inChannels, int outChannels) { + vector t((size_t)inChannels * outChannels); + for(int i = 0; i < inChannels; i++) + for(int o = 0; o < outChannels; o++) + t[(size_t)o * inChannels + i] = w[(size_t)i * outChannels + o]; + return t; +} + +//Whether TransformerAttentionBlock::apply attempts the CK fused attention path for FP16 attention. +//Over a handle's lifetime this can only turn from true to false (fusedAttentionRuntimeDisabled), +//so a false at model construction holds for every apply. +static bool rocmCkFmhaWillBeTried(const CudaHandles* cudaHandles) { +#if KATAGO_ROCM_HAS_CK_FMHA && KATAGO_ROCM_CK_FMHA_ARCH_OK + return + !cudaHandles->disableFusedAttention && cudaHandles->ckFmhaArchSupported && + !cudaHandles->fusedAttentionRuntimeDisabled; +#else + (void)cudaHandles; + return false; +#endif +} +#endif + //--------------------------------------------------------------------------------- //Destructor helpers for ByBatchSize. Descriptor-style entries are opaque pointers and may be @@ -835,6 +1004,7 @@ struct CudnnManager { struct ScratchBuffers { + const int maxBatchSize; const size_t batchXYFloatBytes; const size_t batchFloatBytes; const size_t batchXYBytes; @@ -858,7 +1028,8 @@ struct ScratchBuffers { ScratchBuffers& operator=(const ScratchBuffers&) = delete; ScratchBuffers(int maxBatchSize, int nnXLen, int nnYLen, bool useFP16) - : batchXYFloatBytes((size_t)maxBatchSize * nnXLen * nnYLen * sizeof(float)), + : maxBatchSize(maxBatchSize), + 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))) @@ -1243,6 +1414,10 @@ struct ConvLayer { bool use1x1Matmul; int matmulSpatialSize; SharedDeviceBuf matmulWeightBuf; + // With use1x1Matmul in FP16 when the CK GEMM kernels are enabled, also the desc weights + // unchanged, [outC, inC] (out-major), for the CK kernels via rocmTokenGemmCkOrHipblas, which + // falls back to hipBLAS on matmulWeightBuf where that is faster. + SharedDeviceBuf matmulWeightOutMajorBuf; // Set when MIOpen's only NHWC solution for this conv is its naive fallback. The conv is then // re-expressed as NCHW at runtime, transposing the input and output around it. bool useNchwFallback; @@ -1313,6 +1488,12 @@ struct ConvLayer { name + ":matmulW", [&](void*& b){ CudaUtils::mallocAndCopyToDevice(name + ":matmulW",wT,b,useFP16); }, desc, SharedWeightsKind::CONV_1X1_MATMUL, useFP16, cudaHandles->shareModelWeights); + if(useFP16 && cudaHandles->ckGemmEnabled && inChannels % 8 == 0 && outChannels % 8 == 0) { + matmulWeightOutMajorBuf = mallocAndCopyToDeviceMaybeShared( + name + ":matmulWOutMajor", + [&](void*& b){ CudaUtils::mallocAndCopyToDevice(name + ":matmulWOutMajor",desc->weights,b,useFP16); }, + desc, SharedWeightsKind::CONV_1X1_OUT_MAJOR, useFP16, cudaHandles->shareModelWeights); + } if(!cudaHandles->loggedUsing1x1Matmul) { cudaHandles->loggedUsing1x1Matmul = true; if(cudaHandles->logger != NULL) @@ -1565,36 +1746,44 @@ struct ConvLayer { &beta, (float*)outputBuf, outChannels)); } else { + auto runHipblas = [&]() { #if defined(hipblasVersionMajor) && hipblasVersionMajor >= 2 - // FP32 compute for accuracy, as in MatMulLayer. - const float alpha = 1.0f; - const float beta = accumulate ? 1.0f : 0.0f; - CUBLAS_ERR(name.c_str(),hipblasGemmEx( - cudaHandles->cublas, - HIPBLAS_OP_N, - HIPBLAS_OP_N, - outChannels, - tokens, - inChannels, - &alpha, - matmulWeightBuf.buf, HIP_R_16F, outChannels, - inputBuf, HIP_R_16F, inChannels, - &beta, - outputBuf, HIP_R_16F, outChannels, - HIPBLAS_COMPUTE_32F, - HIPBLAS_GEMM_DEFAULT - )); + // FP32 compute for accuracy, as in MatMulLayer. + const float alpha = 1.0f; + const float beta = accumulate ? 1.0f : 0.0f; + CUBLAS_ERR(name.c_str(),hipblasGemmEx( + cudaHandles->cublas, + HIPBLAS_OP_N, + HIPBLAS_OP_N, + outChannels, + tokens, + inChannels, + &alpha, + matmulWeightBuf.buf, HIP_R_16F, outChannels, + inputBuf, HIP_R_16F, inChannels, + &beta, + outputBuf, HIP_R_16F, outChannels, + HIPBLAS_COMPUTE_32F, + HIPBLAS_GEMM_DEFAULT + )); #else - // hipBLAS older than 2.0, below the minimum Compiling.md documents. - const cublas_half_t* alpha = (const cublas_half_t*)scratch->oneBuf; - const cublas_half_t* betaHalf = (const cublas_half_t*)(accumulate ? scratch->oneBuf : scratch->zeroBuf); - CUBLAS_ERR(name.c_str(),cublasHgemm( - cudaHandles->cublas, CUBLAS_OP_N, CUBLAS_OP_N, - outChannels, tokens, inChannels, - alpha, (const cublas_half_t*)matmulWeightBuf.buf, outChannels, - (const cublas_half_t*)inputBuf, inChannels, - betaHalf, (cublas_half_t*)outputBuf, outChannels)); + // hipBLAS older than 2.0, below the minimum Compiling.md documents. + const cublas_half_t* alpha = (const cublas_half_t*)scratch->oneBuf; + const cublas_half_t* betaHalf = (const cublas_half_t*)(accumulate ? scratch->oneBuf : scratch->zeroBuf); + CUBLAS_ERR(name.c_str(),cublasHgemm( + cudaHandles->cublas, CUBLAS_OP_N, CUBLAS_OP_N, + outChannels, tokens, inChannels, + alpha, (const cublas_half_t*)matmulWeightBuf.buf, outChannels, + (const cublas_half_t*)inputBuf, inChannels, + betaHalf, (cublas_half_t*)outputBuf, outChannels)); #endif + }; + if(matmulWeightOutMajorBuf.buf != nullptr) + rocmTokenGemmCkOrHipblas( + cudaHandles, name, matmulWeightOutMajorBuf.buf, outChannels, inChannels, + inputBuf, inChannels, outputBuf, outChannels, tokens, accumulate, runHipblas); + else + runHipblas(); } return; } @@ -1774,15 +1963,23 @@ struct MatMulLayer { const int outChannels; const bool usingFP16; SharedDeviceBuf matBuf; + // ROCm only: an additional out-major ([outC][inC]) copy of the weights for the CK GEMM kernels + // via rocmTokenGemmCkOrHipblas. Uploaded for layers that multiply over all board positions + // (tokenGemm) when the CK GEMM kernels are enabled. matBuf stays for hipBLAS, which wins on + // some shapes. + SharedDeviceBuf matBufOutMajor; MatMulLayer() = delete; MatMulLayer(const MatMulLayer&) = delete; MatMulLayer& operator=(const MatMulLayer&) = delete; + // tokenGemm: this layer's apply() always multiplies over every board position of the batch + // (transformer projections), so it is worth the CK GEMM path (ROCm only). MatMulLayer( CudaHandles* cudaHandles, const MatMulLayerDesc* desc, - bool useFP16 + bool useFP16, + bool tokenGemm = false ) : name(desc->name), inChannels(desc->inChannels), @@ -1795,6 +1992,16 @@ struct MatMulLayer { name, [&](void*& b){ CudaUtils::mallocAndCopyToDevice(name,desc->weights,b,useFP16); }, desc, SharedWeightsKind::MATMUL, useFP16, cudaHandles->shareModelWeights); +#ifdef KATAGO_GPU_HIP + if(tokenGemm && useFP16 && cudaHandles->ckGemmEnabled && inChannels % 8 == 0 && outChannels % 8 == 0) { + matBufOutMajor = mallocAndCopyToDeviceMaybeShared( + name + ":outMajor", + [&](void*& b){ CudaUtils::mallocAndCopyToDevice(name + ":outMajor",transposeInMajorWeights(desc->weights,inChannels,outChannels),b,useFP16); }, + desc, SharedWeightsKind::MATMUL_OUT_MAJOR, useFP16, cudaHandles->shareModelWeights); + } +#else + (void)tokenGemm; +#endif } } @@ -1808,6 +2015,7 @@ struct MatMulLayer { //If accumulate is true, adds the matmul result into outputBuf (GEMM beta=1) instead of //overwriting it, letting a residual add fuse into the GEMM epilogue. + //inputLd is the row stride of inputBuf in elements, 0 for packed rows (inChannels). void apply( CudaHandles* cudaHandles, ScratchBuffers* scratch, @@ -1816,11 +2024,14 @@ struct MatMulLayer { void* outputBuf, void* workspaceBuf, size_t workspaceBytes, - bool accumulate = false + bool accumulate = false, + int inputLd = 0 ) const { (void)workspaceBuf; (void)workspaceBytes; assert(inChannels > 0 && outChannels > 0); + if(inputLd == 0) + inputLd = inChannels; if(!usingFP16) { const float alpha = 1.0f; @@ -1834,7 +2045,7 @@ struct MatMulLayer { inChannels, &alpha, (const float*)matBuf.buf,outChannels, - (const float*)inputBuf,inChannels, + (const float*)inputBuf,inputLd, &beta, (float*)outputBuf,outChannels )); @@ -1846,23 +2057,31 @@ struct MatMulLayer { //FP16 in/out and FP32 compute fixes that and costs nothing in speed. alpha/beta must match //the compute type, so they are plain host floats here rather than the half zeroBuf/oneBuf //the CUDA path uses. - const float alpha = 1.0f; - const float beta = accumulate ? 1.0f : 0.0f; - CUBLAS_ERR(name.c_str(),hipblasGemmEx( - cudaHandles->cublas, - HIPBLAS_OP_N, - HIPBLAS_OP_N, - outChannels, - batchSize, - inChannels, - &alpha, - matBuf.buf, HIP_R_16F, outChannels, - inputBuf, HIP_R_16F, inChannels, - &beta, - outputBuf, HIP_R_16F, outChannels, - HIPBLAS_COMPUTE_32F, - HIPBLAS_GEMM_DEFAULT - )); + auto runHipblas = [&]() { + const float alpha = 1.0f; + const float beta = accumulate ? 1.0f : 0.0f; + CUBLAS_ERR(name.c_str(),hipblasGemmEx( + cudaHandles->cublas, + HIPBLAS_OP_N, + HIPBLAS_OP_N, + outChannels, + batchSize, + inChannels, + &alpha, + matBuf.buf, HIP_R_16F, outChannels, + inputBuf, HIP_R_16F, inputLd, + &beta, + outputBuf, HIP_R_16F, outChannels, + HIPBLAS_COMPUTE_32F, + HIPBLAS_GEMM_DEFAULT + )); + }; + if(matBufOutMajor.buf != nullptr) + rocmTokenGemmCkOrHipblas( + cudaHandles, name, matBufOutMajor.buf, outChannels, inChannels, + inputBuf, inputLd, outputBuf, outChannels, batchSize, accumulate, runHipblas); + else + runHipblas(); #else //CUDA, or hipBLAS older than 2.0 (below ROCm 6.0, itself below the minimum Compiling.md //documents). There hipblasGemmEx has no spelling that takes hipblasComputeType_t at all, so @@ -1878,7 +2097,7 @@ struct MatMulLayer { inChannels, alpha, (const cublas_half_t*)matBuf.buf,outChannels, - (const cublas_half_t*)inputBuf,inChannels, + (const cublas_half_t*)inputBuf,inputLd, beta, (cublas_half_t*)outputBuf,outChannels )); @@ -1892,9 +2111,11 @@ struct MatMulLayer { //--------------------------------------------------------------------------------- //Run `count` GEMMs that share one input (B stride 0) over weight matrices packed contiguously -//in device memory, writing the outputs at a fixed stride. Used to batch the Q/K/V projections -//and the FFN linear+gate projections of transformer blocks into single GEMM calls. -//Precision behavior matches MatMulLayer::apply. +//in device memory, writing the outputs at a fixed stride, each output row at outputLd elements. +//Used to batch the Q/K/V projections and the FFN linear+gate projections of transformer blocks +//into single GEMM calls: outputs one after another (outputStrideElts a whole output, outputLd = +//outChannels), or interleaved per row (outputStrideElts = outChannels, outputLd = count * +//outChannels). Precision behavior matches MatMulLayer::apply. static void applySharedInputStridedMatMuls( CudaHandles* cudaHandles, ScratchBuffers* scratch, @@ -1907,6 +2128,7 @@ static void applySharedInputStridedMatMuls( const void* inputBuf, void* outputBuf, long long outputStrideElts, + int outputLd, int count ) { const long long weightStrideElts = (long long)inChannels * outChannels; @@ -1924,7 +2146,7 @@ static void applySharedInputStridedMatMuls( (const float*)packedWeights, outChannels, weightStrideElts, (const float*)inputBuf, inChannels, 0LL, &beta, - (float*)outputBuf, outChannels, outputStrideElts, + (float*)outputBuf, outputLd, outputStrideElts, count )); } @@ -1944,7 +2166,7 @@ static void applySharedInputStridedMatMuls( packedWeights, HIP_R_16F, outChannels, weightStrideElts, inputBuf, HIP_R_16F, inChannels, 0LL, &beta, - outputBuf, HIP_R_16F, outChannels, outputStrideElts, + outputBuf, HIP_R_16F, outputLd, outputStrideElts, count, HIPBLAS_COMPUTE_32F, HIPBLAS_GEMM_DEFAULT @@ -1963,7 +2185,7 @@ static void applySharedInputStridedMatMuls( (const cublas_half_t*)packedWeights, outChannels, weightStrideElts, (const cublas_half_t*)inputBuf, inChannels, 0LL, beta, - (cublas_half_t*)outputBuf, outChannels, outputStrideElts, + (cublas_half_t*)outputBuf, outputLd, outputStrideElts, count )); #endif @@ -2624,13 +2846,16 @@ struct TransformerAttentionBlock { // three projections run as a single strided-batched GEMM sharing the input. Otherwise the // three separate projection layers below are constructed and applied individually. const bool sameQKVShapes; - // When the mma attention path will consume Q/K/V (FP16, mma-supported shapes), the three - // projections instead run as ONE combined GEMM writing an interleaved [tokens, q|k|v] - // buffer that downstream kernels read via strides. qkvPackedWeights then holds the - // combined interleaved layout (SharedWeightsKind::QKV_COMBINED) instead of the - // strided-batched slices. + // When the mma attention path (CUDA) or the dot2 attention path (ROCm, with the CK GEMM + // kernels enabled) will consume Q/K/V, the three projections instead run as ONE combined GEMM + // writing an interleaved [tokens, q|k|v] buffer that downstream kernels read via strides. + // qkvPackedWeights then holds the combined interleaved layout (SharedWeightsKind::QKV_COMBINED) + // instead of the strided-batched slices. const bool useCombinedQKV; SharedDeviceBuf qkvPackedWeights; + // ROCm only, with useCombinedQKV: the combined weights out-major ([q|k|v][inC]) for the CK GEMM + // kernels, see rocmTokenGemmCkOrHipblas. + SharedDeviceBuf qkvCombinedOutMajorWeights; std::unique_ptr qProj; std::unique_ptr kProj; std::unique_ptr vProj; @@ -2644,13 +2869,12 @@ struct TransformerAttentionBlock { int ropeNumKVHeads; static bool shouldCombineQKV(CudaHandles* cudaHandles, const TransformerAttentionDesc* desc, bool useFP16) { -#if defined(DEBUG_INTERMEDIATE_VALUES) || !defined(KATAGO_GPU_CUDA) +#if defined(DEBUG_INTERMEDIATE_VALUES) // The intermediate-value debug prints assume packed per-tensor layouts, so keep them // meaningful by never using the interleaved combined layout in debug-print builds. - // HIP has no combined-QKV path at all. (void)cudaHandles; (void)desc; (void)useFP16; return false; -#else +#elif defined(KATAGO_GPU_CUDA) return useFP16 && cudaHandles->mmaAttentionEnabled && @@ -2658,6 +2882,19 @@ struct TransformerAttentionBlock { desc->qProj.inChannels == desc->kProj.inChannels && desc->qProj.inChannels == desc->vProj.inChannels && customCudaFlashAttentionMmaSupportsShape(desc->numHeads, desc->numKVHeads, desc->qHeadDim, desc->vHeadDim); +#else + // ROCm: for the CK GEMM kernels, which run one combined GEMM faster than hipBLAS runs the + // projections at small batch sizes. Only when dot2 attention is the sole consumer of Q/K/V: + // it reads them through strides and applies RoPE itself, while the CK fused attention path, + // which would otherwise be tried first, and Step 3's separate RoPE kernels do not. + return + useFP16 && + cudaHandles->ckGemmEnabled && + cudaHandles->dot2AttentionEnabled && + !rocmCkFmhaWillBeTried(cudaHandles) && + desc->qProj.inChannels == desc->kProj.inChannels && + desc->qProj.inChannels == desc->vProj.inChannels && + rocmDot2AttentionSupportsShape(desc->numHeads, desc->numKVHeads, desc->qHeadDim, desc->vHeadDim); #endif } @@ -2686,7 +2923,7 @@ struct TransformerAttentionBlock { usingFP16(useFP16), usingNHWC(useNHWC), preLN(cudaHandles, &desc->preLN, useFP16), - outProj(cudaHandles, &desc->outProj, useFP16), + outProj(cudaHandles, &desc->outProj, useFP16, true), sameQKVShapes( desc->qProj.inChannels == desc->kProj.inChannels && desc->qProj.inChannels == desc->vProj.inChannels && @@ -2744,6 +2981,12 @@ struct TransformerAttentionBlock { name + ":qkvCombined", [&](void*& b){ CudaUtils::mallocAndCopyToDevice(name + ":qkvCombined", packed, b, useFP16); }, desc, SharedWeightsKind::QKV_COMBINED, useFP16, cudaHandles->shareModelWeights); +#ifdef KATAGO_GPU_HIP + qkvCombinedOutMajorWeights = mallocAndCopyToDeviceMaybeShared( + name + ":qkvCombinedOutMajor", + [&](void*& b){ CudaUtils::mallocAndCopyToDevice(name + ":qkvCombinedOutMajor", transposeInMajorWeights(packed, inChannels, (int)combinedN), b, useFP16); }, + desc, SharedWeightsKind::QKV_COMBINED_OUT_MAJOR, useFP16, cudaHandles->shareModelWeights); +#endif } else if(sameQKVShapes) { vector packed; @@ -2823,6 +3066,21 @@ struct TransformerAttentionBlock { CudaUtils::debugPrint3D(KATAGO_GPU_BACKEND_NAME " Attn RMSNorm out", trunkScratchBuf, batchSize, inChannels, seqLen, usingNHWC, usingFP16, maskBuf); #endif +#ifdef KATAGO_GPU_HIP + // The dot2 attention path (rocmkernels.h) compacts and re-lays-out Q/K/V in a prep kernel that + // can apply RoPE on the way, replacing Step 3's separate RoPE launch. That only happens when + // nothing else will consume Q/K first, i.e. when the CK fused path will not be attempted. + const bool useDot2Attention = + usingFP16 && cudaHandles->dot2AttentionEnabled && + rocmDot2AttentionSupportsShape(numHeads, numKVHeads, qHeadDim, vHeadDim); + const bool ckWillBeTried = usingFP16 && rocmCkFmhaWillBeTried(cudaHandles); + const bool fuseRopeIntoDot2Prep = useRope && useDot2Attention && !ckWillBeTried; + // The combined layout is chosen at construction only for dot2 attention as the sole consumer. + testAssert(!useCombinedQKV || (useDot2Attention && !ckWillBeTried)); +#else + const bool fuseRopeIntoDot2Prep = false; +#endif + // Step 2: Q/K/V projections // trunkScratchBuf is [N, XY, C] NHWC = [C, N*seqLen] column-major. // MatMulLayer expects input as [inChannels, batchSize], which matches. @@ -2852,11 +3110,20 @@ struct TransformerAttentionBlock { // The mma kernel wants 16-byte-aligned slice pointers, so the interior offsets must be // multiples of 8 halfs. True for every real head configuration, so assert. testAssert(qTotalDim % 8 == 0 && kTotalDim % 8 == 0); - applySharedInputStridedMatMuls( - cudaHandles, scratch, usingFP16, name, - combinedN, matBatchSize, inChannels, - qkvPackedWeights.buf, trunkScratchBuf, qkvBuf.buf, 0, 1 - ); + auto runBlas = [&]() { + applySharedInputStridedMatMuls( + cudaHandles, scratch, usingFP16, name, + combinedN, matBatchSize, inChannels, + qkvPackedWeights.buf, trunkScratchBuf, qkvBuf.buf, 0, combinedN, 1 + ); + }; +#ifdef KATAGO_GPU_HIP + rocmTokenGemmCkOrHipblas( + cudaHandles, name, qkvCombinedOutMajorWeights.buf, combinedN, inChannels, + trunkScratchBuf, inChannels, qkvBuf.buf, combinedN, matBatchSize, false, runBlas); +#else + runBlas(); +#endif } else { qPtr = qkvBuf.buf; @@ -2876,7 +3143,7 @@ struct TransformerAttentionBlock { applySharedInputStridedMatMuls( cudaHandles, scratch, usingFP16, name, qTotalDim, matBatchSize, inChannels, - qkvPackedWeights.buf, trunkScratchBuf, qPtr, strideElts, 3 + qkvPackedWeights.buf, trunkScratchBuf, qPtr, strideElts, qTotalDim, 3 ); } else { @@ -2892,7 +3159,7 @@ struct TransformerAttentionBlock { // Step 3: Apply RoPE to Q and K // Q is [qTotalDim, seqLen*batchSize] column-major = [batchSize*seqLen, qTotalDim] row-major - if(useRope) { + if(useRope && !fuseRopeIntoDot2Prep) { // The fused Q+K kernels assume each KV head maps to a whole number of query heads. // Models always satisfy this, but fall back to separate launches if not. bool fuseQK = (numHeads % numKVHeads == 0); @@ -3070,8 +3337,7 @@ struct TransformerAttentionBlock { bool usedSDPA = false; #if KATAGO_ROCM_HAS_CK_FMHA && KATAGO_ROCM_CK_FMHA_ARCH_OK - if(usingFP16 && !cudaHandles->disableFusedAttention && - cudaHandles->ckFmhaArchSupported && !cudaHandles->fusedAttentionRuntimeDisabled) { + if(ckWillBeTried) { bool hasMask = (maskBuf != NULL); // CK's fused path takes a pre-materialized additive bias rather than a raw per-position @@ -3181,7 +3447,7 @@ struct TransformerAttentionBlock { cudaHandles->logger->write( "ROCm backend: CK fused attention unavailable" + (ckFailure.empty() ? string(" (unsupported shape/config)") : (": " + ckFailure)) + - "; using the plain attention kernel instead"); + (useDot2Attention ? "; using the dot2 attention kernel instead" : "; using the plain attention kernel instead")); } } } @@ -3191,10 +3457,50 @@ struct TransformerAttentionBlock { !cudaHandles->loggedCkFmhaDisabled) { cudaHandles->loggedCkFmhaDisabled = true; if(cudaHandles->logger != NULL) - cudaHandles->logger->write( - "ROCm backend: rocmDisableFusedAttention is set, using the plain attention kernel instead of the CK fused attention path that would otherwise have been used"); + cudaHandles->logger->write(string( + "ROCm backend: rocmDisableFusedAttention is set, using the ") + + (useDot2Attention ? "dot2" : "plain") + + " attention kernel instead of the CK fused attention path that would otherwise have been used"); } #endif + // dot2 attention (see useDot2Attention above), when the CK path was not taken or failed. + if(!usedSDPA && useDot2Attention) { + RocmAttnRope rope; + if(fuseRopeIntoDot2Prep) { + if(learnableRope) { + rope.mode = RocmAttnRope::RECOMPUTE; + rope.freqs = (const float*)ropeFreqsBuf.buf; + rope.nnXLen = nnXLen; + } + else { + rope.mode = RocmAttnRope::TABLE; + rope.cosTable = (const half*)ropeCosTable.buf; + rope.sinTable = (const half*)ropeSinTable.buf; + rope.learnableTable = false; + } + } + // Sized by max batch size like every transient buffer here, see the comment on qkvBuf. + const size_t kvBytes = rocmDot2AttentionKVBytes(scratch->maxBatchSize, seqLen, numKVHeads, qHeadDim); + SizedBuf kcBuf(scratch->allocator, kvBytes); + SizedBuf vcBuf(scratch->allocator, kvBytes); + SizedBuf permBuf(scratch->allocator, rocmDot2AttentionPermBytes(scratch->maxBatchSize, seqLen)); + rocmDot2AttentionPrep( + (half*)qPtr, (const half*)kPtr, (const half*)vPtr, (const half*)maskBuf, + (half*)kcBuf.buf, (half*)vcBuf.buf, (short*)permBuf.buf, + batchSize, seqLen, numHeads, numKVHeads, qHeadDim, qStrideElts, kvStrideElts, rope, cudaHandles->stream); + rocmDot2Attention( + (const half*)qPtr, (const half*)kcBuf.buf, (const half*)vcBuf.buf, (const short*)permBuf.buf, + (half*)attnOutBuf.buf, batchSize, seqLen, numHeads, numKVHeads, qHeadDim, qStrideElts, + cudaHandles->numMultiprocessors, cudaHandles->waveSize, cudaHandles->stream); + CUDA_ERR(name.c_str(), hipPeekAtLastError()); + // Not SDPA proper, but it is what keeps the plain fallback below from running. + usedSDPA = true; + if(!cudaHandles->loggedUsingDot2Attention) { + cudaHandles->loggedUsingDot2Attention = true; + if(cudaHandles->logger != NULL) + cudaHandles->logger->write("ROCm backend: using dot2 attention"); + } + } #endif if(!usedSDPA) { @@ -3280,6 +3586,12 @@ struct TransformerFFNBlock { #if defined(KATAGO_GPU_CUDA) && defined(USE_CUTLASS_FUSED_FFN) bool usingFusedFFN = false; #endif + // ROCm only: additionally the out-major layout (as for the CUTLASS kernel) when the CK GEMM + // kernels are enabled for FP16. linear1 and the gate then write one buffer interleaved per + // token, [tokens, linear1 | gate], through either one CK GEMM or the strided-batched hipBLAS + // GEMM, whichever times faster (see rocmTokenGemmCkOrHipblas), and linear2 reads the SwiGLU + // result in place at that row stride. + SharedDeviceBuf ffnOutMajorWeights; TransformerFFNBlock() = delete; TransformerFFNBlock(const TransformerFFNBlock&) = delete; @@ -3302,7 +3614,7 @@ struct TransformerFFNBlock { usingFP16(useFP16), usingNHWC(useNHWC), preLN(cudaHandles, &desc->preLN, useFP16), - linear2(cudaHandles, &desc->linear2, useFP16) + linear2(cudaHandles, &desc->linear2, useFP16, true) { if(!useSwiGLU) { throw StringError("Non-SwiGLU transformer FFN is not yet supported in " KATAGO_GPU_BACKEND_NAME " backend"); @@ -3359,6 +3671,17 @@ struct TransformerFFNBlock { name + ":ffnPacked", [&](void*& b){ CudaUtils::mallocAndCopyToDevice(name + ":ffnPacked", packed, b, useFP16); }, desc, SharedWeightsKind::FFN_PACKED, useFP16, cudaHandles->shareModelWeights); +#if defined(KATAGO_GPU_HIP) && !defined(DEBUG_INTERMEDIATE_VALUES) + if(useFP16 && cudaHandles->ckGemmEnabled && ffnChannels % 8 == 0 && numChannels % 8 == 0) { + vector outMajor = transposeInMajorWeights(desc->linear1.weights, numChannels, ffnChannels); + const vector gateOutMajor = transposeInMajorWeights(desc->linearGate.weights, numChannels, ffnChannels); + outMajor.insert(outMajor.end(), gateOutMajor.begin(), gateOutMajor.end()); + ffnOutMajorWeights = mallocAndCopyToDeviceMaybeShared( + name + ":ffnOutMajor", + [&](void*& b){ CudaUtils::mallocAndCopyToDevice(name + ":ffnOutMajor", outMajor, b, useFP16); }, + desc, SharedWeightsKind::FFN_FUSED_PACKED, useFP16, cudaHandles->shareModelWeights); + } +#endif } } @@ -3424,6 +3747,29 @@ struct TransformerFFNBlock { CUDA_ERR(name.c_str(), cudaPeekAtLastError()); usedFusedFFN = true; } +#endif + // Row stride (elements) of the SwiGLU result in ffnBuf, which linear2 reads. + int swigluRowStride = ffnChannels; +#ifdef KATAGO_GPU_HIP + if(ffnOutMajorWeights.buf != nullptr) { + // See ffnOutMajorWeights: [tokens, linear1 | gate], SwiGLU in place into the linear1 half. + swigluRowStride = 2 * ffnChannels; + auto runHipblas = [&]() { + applySharedInputStridedMatMuls( + cudaHandles, scratch, usingFP16, name, + ffnChannels, matBatchSize, numChannels, + ffnPackedWeights.buf, trunkScratchBuf, ffnBuf.buf, ffnChannels, swigluRowStride, 2 + ); + }; + rocmTokenGemmCkOrHipblas( + cudaHandles, name, ffnOutMajorWeights.buf, swigluRowStride, numChannels, + trunkScratchBuf, numChannels, ffnBuf.buf, swigluRowStride, matBatchSize, false, runHipblas); + customCudaSwiGLURows( + (const half*)ffnBuf.buf, (const half*)ffnBuf.buf + ffnChannels, (half*)ffnBuf.buf, + matBatchSize, ffnChannels, swigluRowStride, cudaHandles->stream); + CUDA_ERR(name.c_str(), cudaPeekAtLastError()); + } + else #endif if(!usedFusedFFN) { { @@ -3432,7 +3778,7 @@ struct TransformerFFNBlock { applySharedInputStridedMatMuls( cudaHandles, scratch, usingFP16, name, ffnChannels, matBatchSize, numChannels, - ffnPackedWeights.buf, trunkScratchBuf, ffnBuf.buf, strideElts, 2 + ffnPackedWeights.buf, trunkScratchBuf, ffnBuf.buf, strideElts, ffnChannels, 2 ); } @@ -3456,11 +3802,11 @@ struct TransformerFFNBlock { // Same masking logic as the attention block's output projection: fuse the residual into the // GEMM only when there is no mask. if(maskBuf == NULL) { - linear2.apply(cudaHandles, scratch, matBatchSize, ffnBuf.buf, trunkBuf, workspaceBuf, workspaceBytes, true); + linear2.apply(cudaHandles, scratch, matBatchSize, ffnBuf.buf, trunkBuf, workspaceBuf, workspaceBytes, true, swigluRowStride); CUDA_ERR(name.c_str(), cudaPeekAtLastError()); } else { - linear2.apply(cudaHandles, scratch, matBatchSize, ffnBuf.buf, trunkScratchBuf, workspaceBuf, workspaceBytes); + linear2.apply(cudaHandles, scratch, matBatchSize, ffnBuf.buf, trunkScratchBuf, workspaceBuf, workspaceBytes, false, swigluRowStride); if(pendingTrunkDelta != NULL) { // Defer the add so the next block can fuse it into its pre-norm. @@ -4835,6 +5181,11 @@ struct ComputeContext { bool fuseResidualNorm; // See CudaHandles::use1x1MatmulMode. enabled_t use1x1MatmulMode; + // See CudaHandles::dot2AttentionEnabled. Auto = on for the archs the kernels were measured on + // (rocmDot2AttentionAutoArch), True = on wherever the arch has v_dot2_f32_f16. + enabled_t useDot2AttentionMode; + // See CudaHandles::ckGemmEnabled. Auto = on where this build has CK GEMM kernels for the arch. + enabled_t useCkGemmMode; }; #endif @@ -4917,6 +5268,10 @@ ComputeContext* NeuralNet::createComputeContext( context->fuseResidualNorm = cfg.contains("rocmFuseResidualNorm") ? cfg.getBool("rocmFuseResidualNorm") : true; context->use1x1MatmulMode = use1x1MatmulMode; + context->useDot2AttentionMode = + cfg.contains("rocmUseDot2Attention") ? cfg.getEnabled("rocmUseDot2Attention") : enabled_t::Auto; + context->useCkGemmMode = + cfg.contains("rocmUseCkGemm") ? cfg.getEnabled("rocmUseCkGemm") : enabled_t::Auto; return context; } #endif @@ -5031,6 +5386,11 @@ struct ComputeHandle { int minorComputeCapability, bool ckFmhaArchSupported, bool preferMiopen1x1Conv, + bool archHasDot2, + bool archDot2AttentionAuto, + bool archHasCkGemm, + int numMultiprocessors, + int waveSize, int maxBatchSize, bool requireExactNNLen_, bool inputsUseNHWC_, @@ -5057,6 +5417,21 @@ struct ComputeHandle { cudaHandles->use1x1MatmulMode = (context->use1x1MatmulMode == enabled_t::Auto && preferMiopen1x1Conv) ? enabled_t::False : context->use1x1MatmulMode; + if(context->useDot2AttentionMode == enabled_t::True && !archHasDot2) + throw StringError("rocmUseDot2Attention = true, but this GPU architecture has no v_dot2_f32_f16"); + cudaHandles->dot2AttentionEnabled = + archHasDot2 && + (context->useDot2AttentionMode == enabled_t::True || + (context->useDot2AttentionMode == enabled_t::Auto && archDot2AttentionAuto)); + cudaHandles->numMultiprocessors = numMultiprocessors; + cudaHandles->waveSize = waveSize; + if(context->useCkGemmMode == enabled_t::True && !archHasCkGemm) + throw StringError("rocmUseCkGemm = true, but this build has no CK GEMM kernels for this GPU architecture (only gfx103x is supported)"); + // Only the FP16 GEMMs have CK kernels; must be set before building the model, whose layers + // choose their weight layout by it. + cudaHandles->ckGemmEnabled = useFP16 && archHasCkGemm && context->useCkGemmMode != enabled_t::False; + if(cudaHandles->ckGemmEnabled && logger != NULL) + logger->write("ROCm backend: CK DPP GEMM kernels enabled for FP16 token GEMMs, used per shape where faster than hipBLAS"); model = std::make_unique( cudaHandles.get(), &(loadedModel->modelDesc), maxBatchSize, nnXLen, nnYLen, inputsUseNHWC, useFP16, useNHWC @@ -5246,11 +5621,18 @@ ComputeHandle* NeuralNet::createComputeHandle( ); } +#if KATAGO_ROCM_HAS_CK_GEMM + const bool archHasCkGemm = rocmCkGemmSupportsArch(prop.gcnArchName); +#else + const bool archHasCkGemm = false; +#endif // Reusing the NHWC predicate: the archs that prefer MIOpen's NHWC 1x1 conv over a GEMM are the // same ones that prefer NHWC at all. ComputeHandle* gpuHandle = new ComputeHandle( context,loadedModel,logger,prop.major,prop.minor,isCkFmhaSupportedArch(prop), - isNhwcFp16PreferredArch(prop), + isNhwcFp16PreferredArch(prop), rocmArchHasDot2(prop.gcnArchName), rocmDot2AttentionAutoArch(prop.gcnArchName), + archHasCkGemm, + prop.multiProcessorCount, prop.warpSize, maxBatchSize,requireExactNNLen,inputsUseNHWC,useFP16,useNHWC ); return gpuHandle; diff --git a/cpp/neuralnet/cudaandrocmhelpers.h b/cpp/neuralnet/cudaandrocmhelpers.h index 85c992f2ba..f61732db81 100644 --- a/cpp/neuralnet/cudaandrocmhelpers.h +++ b/cpp/neuralnet/cudaandrocmhelpers.h @@ -151,6 +151,11 @@ bool customCudaFlashAttentionMmaSupported(); //SwiGLU: out[i] = SiLU(a[i]) * b[i] void customCudaSwiGLU(const float* a, const float* b, float* out, int size, cudaStream_t stream); void customCudaSwiGLU(const half* a, const half* b, half* out, int size, cudaStream_t stream); +//Row-strided SwiGLU: out[r][c] = SiLU(a[r][c]) * b[r][c] for r < rows, c < cols, every operand at +//row stride rowStride (elements); out may alias a. cols and rowStride must be multiples of 8 and +//the pointers 16-byte aligned. Used by the ROCm backend's interleaved FFN layout; the body is only +//compiled on ROCm, but the declaration is harmless on CUDA. +void customCudaSwiGLURows(const half* a, const half* b, half* out, int rows, int cols, int rowStride, cudaStream_t stream); //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 diff --git a/cpp/neuralnet/cudaandrocmhelpers.inc b/cpp/neuralnet/cudaandrocmhelpers.inc index 76c5d2277a..ea9d80b4c5 100644 --- a/cpp/neuralnet/cudaandrocmhelpers.inc +++ b/cpp/neuralnet/cudaandrocmhelpers.inc @@ -1125,28 +1125,43 @@ void customCudaAddTensorInplace(float* buf, const float* biases, int nSize, cuda // convolution when the two differ. Only the initial convolution's input ever needs this. // // out[n][b][a] = in[n][a][b]. For NCHW->NHWC call with (a=C, b=HW), for NHWC->NCHW with (a=HW, b=C). -// Reads are coalesced over b, in's contiguous axis. The tensor is small, being the raw input -// planes, so the strided writes are not worth tiling for. +// Besides the input planes, this also transposes the output of a conv re-expressed as NCHW (see +// ConvLayer::useNchwFallback), e.g. a 512-channel trunk, so it is tiled through shared memory to +// keep both the reads and the writes coalesced: one 32x32 tile per block, 32x8 threads. template __global__ void batchedTranspose2DKernel(const T* in, T* out, int aSize, int bSize) { - int b = blockIdx.x * blockDim.x + threadIdx.x; - int a = blockIdx.y; - int n = blockIdx.z; - if(b < bSize) - out[((size_t)n * bSize + b) * aSize + a] = in[((size_t)n * aSize + a) * bSize + b]; + __shared__ T tile[32][33]; + const int b0 = blockIdx.x * 32; + const int a0 = blockIdx.y * 32; + const int n = blockIdx.z; + const T* inN = in + (size_t)n * aSize * bSize; + T* outN = out + (size_t)n * aSize * bSize; + for(int i = threadIdx.y; i < 32; i += 8) { + int a = a0 + i; + int b = b0 + threadIdx.x; + if(a < aSize && b < bSize) + tile[i][threadIdx.x] = inN[(size_t)a * bSize + b]; + } + __syncthreads(); + for(int i = threadIdx.y; i < 32; i += 8) { + int b = b0 + i; + int a = a0 + threadIdx.x; + if(a < aSize && b < bSize) + outN[(size_t)b * aSize + a] = tile[threadIdx.x][i]; + } } template static void applyBatchedTranspose2D(const T* in, T* out, int nSize, int aSize, int bSize, const char* name, cudaStream_t stream) { checkBufferIndexFitsInt(nSize, aSize, bSize, name); - if(nSize > 65535 || aSize > 65535) + if(nSize > 65535 || (aSize + 31) / 32 > 65535) throw std::runtime_error(std::string(name) + ": dimension exceeds the grid y/z size limit"); - int blockSize = targetNumThreads; - dim3 grid((bSize + blockSize - 1) / blockSize, aSize, nSize); - batchedTranspose2DKernel<<>>(in, out, aSize, bSize); + dim3 grid((bSize + 31) / 32, (aSize + 31) / 32, nSize); + dim3 threads(32, 8); + batchedTranspose2DKernel<<>>(in, out, aSize, bSize); } void customCudaCopyNCHWtoNHWC(const float* in, float* out, int nSize, int cSize, int xySize, cudaStream_t stream) { @@ -2131,6 +2146,90 @@ void applyCScaleBiasNHWCSiluMaskHalfKernel(const half* in, half* out, const half #endif } +#ifdef KATAGO_GPU_HIP +// ROCm only so far: these vectorized FP16 kernels (here and the masked residual add, wave-per-row +// RMSNorm, and SwiGLU ones below) were measured on ROCm (gfx1030) only, so the CUDA backend keeps +// the kernels they bypass. +// +// Vectorized FP16 NHWC variant of the kernels above: one thread per 8 consecutive channels of one +// position, with 16-byte loads and stores, used whenever cSize is a multiple of 8 and the buffers +// are 16-byte aligned. The scalar kernels move 2 bytes per thread and reach only a fraction of +// memory bandwidth, which made them a visible cost next to the GEMMs. Per-element arithmetic is +// exactly that of the scalar half kernels (half fma, then half mask multiply, then the activation), +// so results are bit-identical. +#ifdef KATAGO_GPU_SUPPORTS_FP16 +template +__device__ __forceinline__ half applyActivationHalf(half a) { + if(ACTIVATION == ACTIVATION_RELU) { + const half halfzero = __float2half(0.0f); + return __hgt(a,halfzero) ? a : halfzero; + } + else if(ACTIVATION == ACTIVATION_MISH) + return mishh(a); + else if(ACTIVATION == ACTIVATION_MISH_SCALE8) + return mishh_scale8(a); + else if(ACTIVATION == ACTIVATION_SILU) + return siluh(a); + else + return a; +} +#endif + +template +__global__ +void applyCScaleBiasNHWCHalfVec8Kernel(const half* in, half* out, const half* scale, const half* biases, const half* mask, int numPositions, int cSize) +{ +#ifdef KATAGO_GPU_SUPPORTS_FP16 + const int c8Size = cSize / 8; + const int t = blockIdx.x * blockDim.x + threadIdx.x; + if(t >= numPositions * c8Size) + return; + const int pos = t / c8Size; + const int c8 = t - pos * c8Size; + uint4 v = ((const uint4*)in)[t]; + const uint4 s = ((const uint4*)scale)[c8]; + const uint4 b = ((const uint4*)biases)[c8]; + half* hv = (half*)&v; + const half* hs = (const half*)&s; + const half* hb = (const half*)&b; + const half m = HAS_MASK ? mask[pos] : __float2half(1.0f); + #pragma unroll + for(int k = 0; k < 8; k++) { + half a = __hfma(hv[k],hs[k],hb[k]); + if(HAS_MASK) + a = __hmul(a,m); + hv[k] = applyActivationHalf(a); + } + ((uint4*)out)[t] = v; +#else + //Do nothing, FP16 not supported +#endif +} + +template +static void launchApplyCScaleBiasNHWCHalfVec8(const half* in, half* out, const half* scale, const half* biases, const half* mask, int numPositions, int cSize, int activation, cudaStream_t stream) { + const int total = numPositions * (cSize / 8); + const int threads = 256; + const int blocks = (total + threads - 1) / threads; + if(activation == ACTIVATION_IDENTITY) + applyCScaleBiasNHWCHalfVec8Kernel<<>>(in,out,scale,biases,mask,numPositions,cSize); + else if(activation == ACTIVATION_RELU) + applyCScaleBiasNHWCHalfVec8Kernel<<>>(in,out,scale,biases,mask,numPositions,cSize); + else if(activation == ACTIVATION_MISH) + applyCScaleBiasNHWCHalfVec8Kernel<<>>(in,out,scale,biases,mask,numPositions,cSize); + else if(activation == ACTIVATION_SILU) + applyCScaleBiasNHWCHalfVec8Kernel<<>>(in,out,scale,biases,mask,numPositions,cSize); + else if(activation == ACTIVATION_MISH_SCALE8) + applyCScaleBiasNHWCHalfVec8Kernel<<>>(in,out,scale,biases,mask,numPositions,cSize); + else + throw std::runtime_error("customCudaApplyCScaleBiasNHWC: unsupported activation"); +} + +static bool isAligned16(const void* p) { + return ((uintptr_t)p & 15) == 0; +} +#endif // KATAGO_GPU_HIP + static 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, cudaStream_t stream) { if(nSize > 65535) throw std::runtime_error("customCudaApplyCScaleBiasNHWC: nSize too large"); @@ -2138,6 +2237,16 @@ static void sharedApplyCScaleBiasNHWC(const void* in, void* out, const void* sca throw std::runtime_error("customCudaApplyCScaleBiasNHWC: xySize too large"); checkBufferIndexFitsInt(nSize, xySize, cSize, "customCudaApplyCScaleBiasNHWC"); +#ifdef KATAGO_GPU_HIP + if(isHalf && cSize % 8 == 0 && isAligned16(in) && isAligned16(out) && isAligned16(scale) && isAligned16(biases)) { + if(mask == NULL) + launchApplyCScaleBiasNHWCHalfVec8((const half*)in,(half*)out,(const half*)scale,(const half*)biases,NULL,nSize*xySize,cSize,activation,stream); + else + launchApplyCScaleBiasNHWCHalfVec8((const half*)in,(half*)out,(const half*)scale,(const half*)biases,(const half*)mask,nSize*xySize,cSize,activation,stream); + return; + } +#endif + int sSize = xySize; int cThreads; int cBlocks; @@ -3120,9 +3229,52 @@ void customCudaSwiGLU(const float* a, const float* b, float* out, int size, cuda int blocks = (size + threads - 1) / threads; swiGLUKernel<<>>(a, b, out, size); } +#ifdef KATAGO_GPU_HIP +// Vectorized variant over rows (ROCm only, see applyCScaleBiasNHWCHalfVec8Kernel): out[r][c] = +// SiLU(a[r][c]) * b[r][c] for c < cols, every operand at row stride rowStride (elements), one thread +// per 8 consecutive elements with 16-byte accesses. Same per-element arithmetic as the kernel above. +__global__ +void swiGLURowsHalfVec8Kernel(const half* a, const half* b, half* out, int rows, int cols, int rowStride) +{ +#ifdef KATAGO_GPU_SUPPORTS_FP16 + const int c8Size = cols / 8; + const int t = blockIdx.x * blockDim.x + threadIdx.x; + if(t >= rows * c8Size) + return; + const int r = t / c8Size; + const int c = (t - r * c8Size) * 8; + const size_t offset = (size_t)r * rowStride + c; + uint4 av = *(const uint4*)(a + offset); + const uint4 bv = *(const uint4*)(b + offset); + half* ha = (half*)&av; + const half* hb = (const half*)&bv; + #pragma unroll + for(int e = 0; e < 8; e++) + ha[e] = __float2half(siluf(__half2float(ha[e])) * __half2float(hb[e])); + *(uint4*)(out + offset) = av; +#else + (void)a; (void)b; (void)out; (void)rows; (void)cols; (void)rowStride; +#endif +} + +static void launchSwiGLURowsHalfVec8(const half* a, const half* b, half* out, int rows, int cols, int rowStride, cudaStream_t stream) { + const int total = rows * (cols / 8); + const int threads = 256; + swiGLURowsHalfVec8Kernel<<<(total + threads - 1) / threads, threads, 0, stream>>>(a, b, out, rows, cols, rowStride); +} +#endif // KATAGO_GPU_HIP + void customCudaSwiGLU(const half* a, const half* b, half* out, int size, cudaStream_t stream) { if(size <= 0) return; +#ifdef KATAGO_GPU_HIP + // Contiguous operands are a single row for the vectorized kernel, about 1.5x faster at the + // transformer FFN sizes of small batches. + if(size % 8 == 0 && isAligned16(a) && isAligned16(b) && isAligned16(out)) { + launchSwiGLURowsHalfVec8(a, b, out, 1, size, size, stream); + return; + } +#endif constexpr int ELTS_PER_THREAD = 4; // half2 pairs per thread int threads = 256; int pairCount = size >> 1; @@ -3130,6 +3282,17 @@ void customCudaSwiGLU(const half* a, const half* b, half* out, int size, cudaStr if(blocks < 1) blocks = 1; // ensure the odd-size tail still gets a block swiGLUHalfStrideKernel<<>>(a, b, out, size); } +#ifdef KATAGO_GPU_HIP +void customCudaSwiGLURows(const half* a, const half* b, half* out, int rows, int cols, int rowStride, cudaStream_t stream) { + if(rows <= 0 || cols <= 0) + return; + if(cols % 8 != 0 || rowStride % 8 != 0 || !isAligned16(a) || !isAligned16(b) || !isAligned16(out)) + throw std::runtime_error("customCudaSwiGLURows: cols and rowStride must be multiples of 8, pointers 16-byte aligned"); + if((long long)rows * (cols / 8) >= (1LL << 31)) + throw std::runtime_error("customCudaSwiGLURows: too many elements"); + launchSwiGLURowsHalfVec8(a, b, out, rows, cols, rowStride, stream); +} +#endif //-------------------------------------------------------------------------------------------------------------- // Masked residual add: trunk[i] += residual[i] * mask[spatial_idx] @@ -3227,10 +3390,44 @@ void customCudaMaskedResidualAddNHWC(float* trunk, const float* residual, const dim3 threads(cThreads, xyThreads, 1); maskedResidualAddNHWCKernel<<>>(trunk, residual, mask, xySize, cSize); } +#ifdef KATAGO_GPU_HIP +// Vectorized variant (ROCm only, see applyCScaleBiasNHWCHalfVec8Kernel): one thread per 8 +// consecutive channels (16-byte accesses), same arithmetic. +__global__ +void maskedResidualAddNHWCHalfVec8Kernel(half* trunk, const half* residual, const half* mask, int numPositions, int cSize) +{ +#ifdef KATAGO_GPU_SUPPORTS_FP16 + const int c8Size = cSize / 8; + const int t = blockIdx.x * blockDim.x + threadIdx.x; + if(t >= numPositions * c8Size) + return; + const int pos = t / c8Size; + const float m = (mask != NULL) ? __half2float(mask[pos]) : 1.0f; + uint4 tv = ((const uint4*)trunk)[t]; + const uint4 rv = ((const uint4*)residual)[t]; + half* ht = (half*)&tv; + const half* hr = (const half*)&rv; + #pragma unroll + for(int e = 0; e < 8; e++) + ht[e] = __float2half(__half2float(ht[e]) + __half2float(hr[e]) * m); + ((uint4*)trunk)[t] = tv; +#else + //Do nothing, FP16 not supported +#endif +} +#endif + void customCudaMaskedResidualAddNHWC(half* trunk, const half* residual, const half* mask, int nSize, int xySize, int cSize, cudaStream_t stream) { if(nSize > 65535) throw std::runtime_error("customCudaMaskedResidualAddNHWC: nSize too large"); checkBufferIndexFitsInt(nSize, xySize, cSize, "customCudaMaskedResidualAddNHWC"); +#ifdef KATAGO_GPU_HIP + if(cSize % 8 == 0 && isAligned16(trunk) && isAligned16(residual)) { + const int total = nSize * xySize * (cSize / 8); + maskedResidualAddNHWCHalfVec8Kernel<<<(total + 255) / 256, 256, 0, stream>>>(trunk, residual, mask, nSize * xySize, cSize); + return; + } +#endif int cThreads, cBlocks, xyThreads, xyBlocks; splitThreadsAcrossDim01(cSize, xySize, cThreads, cBlocks, xyThreads, xyBlocks); dim3 grid(cBlocks, xyBlocks, nSize); @@ -3646,6 +3843,113 @@ void residualAddRmsNormGammaBetaNHWCHalfVecKernel( #endif } +#ifdef KATAGO_GPU_HIP +// Wave-per-row variants (ROCm only, see applyCScaleBiasNHWCHalfVec8Kernel) of the two vectorized +// kernels above (RMSNorm, optionally fused with the masked residual add), for cSize a multiple of 8 +// and at most 8 * MAX_CHUNKS * 32. One warp / wavefront handles one position with 16-byte accesses +// and a pure shuffle reduction, instead of one block per position with a two-level shared memory +// reduction: at the trunk widths of transformer blocks (256-768 channels) the per-block overhead +// dominated those kernels. Positions are walked with a grid-stride loop over warps, so the launch +// does not depend on the warp width. Arithmetic per element matches the kernels above; only the +// order of the float sum of squares differs. +template +__global__ +void rmsNormGammaBetaNHWCHalfWaveRowKernel( + half* trunk, const half* in, half* out, const half* gamma, const half* beta, const half* mask, + int totalPositions, int cSize, float epsilon, int activation +) { +#ifdef KATAGO_GPU_SUPPORTS_FP16 + const int lanes = KATAGO_WARP_LANES; + const int lane = threadIdx.x % lanes; + const int wavesPerBlock = blockDim.x / lanes; + const int numChunks = cSize / 8; + for(int pos = blockIdx.x * wavesPerBlock + threadIdx.x / lanes; pos < totalPositions; pos += gridDim.x * wavesPerBlock) { + const float maskVal = (mask != NULL) ? __half2float(mask[pos]) : 1.0f; + const uint4* inRow = (const uint4*)(in + (size_t)pos * cSize); + uint4* trunkRow = RESIDUAL ? (uint4*)(trunk + (size_t)pos * cSize) : NULL; + float vals[MAX_CHUNKS][8]; + float acc = 0.0f; + #pragma unroll + for(int j = 0; j < MAX_CHUNKS; j++) { + const int k = lane + j * lanes; + if(k < numChunks) { + uint4 v = inRow[k]; + const half* hv = (const half*)&v; + if(RESIDUAL) { + // in is the delta: trunk += delta * mask, rounded through half, then normalize trunk. + uint4 t = trunkRow[k]; + half* ht = (half*)&t; + #pragma unroll + for(int e = 0; e < 8; e++) + ht[e] = __float2half(__half2float(ht[e]) + __half2float(hv[e]) * maskVal); + trunkRow[k] = t; + #pragma unroll + for(int e = 0; e < 8; e++) + vals[j][e] = __half2float(ht[e]) * maskVal; + } + else { + #pragma unroll + for(int e = 0; e < 8; e++) + vals[j][e] = __half2float(hv[e]) * maskVal; + } + #pragma unroll + for(int e = 0; e < 8; e++) + acc += vals[j][e] * vals[j][e]; + } + } + acc = warpReduceSumFloat(acc); + const float rms = rsqrtf(acc / (float)cSize + epsilon); + #pragma unroll + for(int j = 0; j < MAX_CHUNKS; j++) { + const int k = lane + j * lanes; + if(k < numChunks) { + const uint4 g = ((const uint4*)gamma)[k]; + const uint4 b = ((const uint4*)beta)[k]; + const half* hg = (const half*)&g; + const half* hb = (const half*)&b; + uint4 o; + half* ho = (half*)&o; + #pragma unroll + for(int e = 0; e < 8; e++) { + float x = vals[j][e] * rms * __half2float(hg[e]) + __half2float(hb[e]); + if(activation == ACTIVATION_RELU) + x = fmaxf(x, 0.0f); + else if(activation == ACTIVATION_MISH) + x = mishf(x); + else if(activation == ACTIVATION_SILU) + x = siluf(x); + ho[e] = __float2half(x * maskVal); + } + ((uint4*)(out + (size_t)pos * cSize))[k] = o; + } + } + } +#else + (void)trunk; (void)in; (void)out; (void)gamma; (void)beta; (void)mask; + (void)totalPositions; (void)cSize; (void)epsilon; (void)activation; +#endif +} + +// Launches the wave-per-row kernel if the shape and alignment allow it, returning whether it did. +template +static bool tryLaunchRmsNormWaveRow( + half* trunk, const half* in, half* out, const half* gamma, const half* beta, const half* mask, + int totalPositions, int cSize, float epsilon, int activation, cudaStream_t stream +) { + constexpr int MAX_CHUNKS = 4; + if(cSize % 8 != 0 || cSize / 8 > MAX_CHUNKS * 32) + return false; + if(!isAligned16(in) || !isAligned16(out) || !isAligned16(gamma) || !isAligned16(beta) || (RESIDUAL && !isAligned16(trunk))) + return false; + // 8 positions per 256-thread block at a warp width of 32 (64-wide wavefronts loop twice). + const int threads = 256; + const int blocks = (totalPositions + 7) / 8; + rmsNormGammaBetaNHWCHalfWaveRowKernel<<>>( + trunk, in, out, gamma, beta, mask, totalPositions, cSize, epsilon, activation); + return true; +} +#endif // KATAGO_GPU_HIP + void customCudaMaskedResidualAddRMSNormNHWC( float* trunk, const float* delta, float* out, const float* gamma, const float* beta, const float* mask, int nSize, int xySize, int cSize, float epsilon, int activation, cudaStream_t stream @@ -3668,6 +3972,10 @@ void customCudaMaskedResidualAddRMSNormNHWC( int totalPositions = nSize * xySize; if(totalPositions <= 0) return; +#ifdef KATAGO_GPU_HIP + if(tryLaunchRmsNormWaveRow(trunk, delta, out, gamma, beta, mask, totalPositions, cSize, epsilon, activation, stream)) + return; +#endif // Same vectorized-path eligibility as customCudaRMSNormGammaBetaNHWC. if(cSize % 2 == 0) { int halfPairs = cSize / 2; @@ -3716,6 +4024,10 @@ void customCudaRMSNormGammaBetaNHWC( int totalPositions = nSize * xySize; if(totalPositions <= 0) return; +#ifdef KATAGO_GPU_HIP + if(tryLaunchRmsNormWaveRow(NULL, in, out, gamma, beta, mask, totalPositions, cSize, epsilon, activation, stream)) + return; +#endif // Vectorized path: cSize must be even (use half2). Pick ELTS_PER_THREAD so we end up with a // thread count that is a whole number of warps and <= 512. // For cSize=384 -> 192 half2 pairs -> 192 threads, 1 elt/thread. diff --git a/cpp/neuralnet/rocmbackend.cpp b/cpp/neuralnet/rocmbackend.cpp index e66403cb5d..4c014bf2a4 100644 --- a/cpp/neuralnet/rocmbackend.cpp +++ b/cpp/neuralnet/rocmbackend.cpp @@ -32,6 +32,10 @@ #endif #include "../neuralnet/rocmhelpers.h" +#include "../neuralnet/rocmkernels.h" +#if KATAGO_ROCM_HAS_CK_GEMM + #include "../neuralnet/rocmckgemm.h" +#endif #include "../neuralnet/rocmutils.h" #include "../neuralnet/rocmcudanames.h" diff --git a/cpp/neuralnet/rocmckgemm.h b/cpp/neuralnet/rocmckgemm.h new file mode 100644 index 0000000000..1f5f777097 --- /dev/null +++ b/cpp/neuralnet/rocmckgemm.h @@ -0,0 +1,32 @@ +// FP16 GEMM through Composable Kernel's DPP GEMM kernels (rocmckgemm.hip), which on RDNA2 +// (gfx103x, no matrix cores) run many transformer projection shapes faster than the rocBLAS / +// Tensile kernels hipBLAS picks. Compiled only for the gfx103x archs of the build (see the CK GEMM +// section of CMakeLists.txt), so callers must check rocmCkGemmSupportsArch first. + +#ifndef NEURALNET_ROCMCKGEMM_H_ +#define NEURALNET_ROCMCKGEMM_H_ + +#include "../neuralnet/rocmincludes.h" + +// Whether this build has CK GEMM kernels for the device architecture (hipDeviceProp_t::gcnArchName). +bool rocmCkGemmSupportsArch(const char* gcnArchName); + +// out[t][o] = sum_i in[t][i] * w[o][i] for t < tokens, o < outChannels, i < inChannels, with in +// (NHWC activations), w (weights, out-major) and out (NHWC) row-major at row strides inStride, +// wStride, outStride (elements), and FP32 accumulation. +// +// The kernel instances, which perform differently per problem shape, so callers time them. +int rocmCkGemmNumInstances(); +// Whether instance (in [0, rocmCkGemmNumInstances())) supports this problem. +bool rocmCkGemmSupports(int instance, int tokens, int outChannels, int inChannels, int inStride, int wStride, int outStride); +// Runs instance on the problem, which it must support. With accumulate, adds the product into out +// instead, rounded to FP16 first and then added in FP16: the same rounding as the transformer +// residual adds on masked boards, where the projection is written to FP16 and added into the +// trunk by the next RMSNorm, but not hipBLAS's beta = 1, which adds in FP32 and rounds once. +// Pointers must be 16-byte aligned. +void rocmCkGemmRun( + int instance, bool accumulate, + const half* in, int inStride, const half* w, int wStride, half* out, int outStride, + int tokens, int outChannels, int inChannels, hipStream_t stream); + +#endif // NEURALNET_ROCMCKGEMM_H_ diff --git a/cpp/neuralnet/rocmckgemm.hip b/cpp/neuralnet/rocmckgemm.hip new file mode 100644 index 0000000000..f0bfc6f57b --- /dev/null +++ b/cpp/neuralnet/rocmckgemm.hip @@ -0,0 +1,154 @@ +// CK DPP GEMM instances, see rocmckgemm.h. Built as its own object only for the gfx103x archs of +// the build (CK's DPP GEMM targets RDNA2 and its headers are slow to compile), listed in +// KATAGO_ROCM_CK_GEMM_ARCHS by CMakeLists.txt. + +#include "../neuralnet/rocmckgemm.h" + +#include +#include + +#include "ck/ck.hpp" +#include "ck/tensor_operation/gpu/device/tensor_layout.hpp" +#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp" +#include "ck/tensor_operation/gpu/device/impl/device_gemm_dpp.hpp" +#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp" + +#ifndef KATAGO_ROCM_CK_GEMM_ARCHS +#error "KATAGO_ROCM_CK_GEMM_ARCHS must list the archs this file is compiled for, separated by ';'" +#endif + +namespace { + +using F16 = ck::half_t; +using F32 = float; +using Row = ck::tensor_layout::gemm::RowMajor; +using Col = ck::tensor_layout::gemm::ColumnMajor; +template using S = ck::Sequence; +using PassThrough = ck::tensor_operation::element_wise::PassThrough; +constexpr auto MNPadding = ck::tensor_operation::device::GemmSpecialization::MNPadding; + +// In CK's terms each instance computes c[m][n] = sum_k a[m][k] * b[n][k] with A = the weights +// ([outC][inC], row-major), B = the activations ([tokens][inC], CK's column-major B), and C = the +// output as column-major [outC][tokens], i.e. KataGo's row-major [tokens][outC]. With M along the +// output channels, each thread's 8 consecutive M values of C are contiguous in memory, so C is +// stored (and for Add, read) 16 bytes at a time, which measured 0-17% faster than the transposed +// formulation (M along tokens, C row-major), whose C accesses can only be scalar. +// +// Tile configurations: the subset of CK's device_gemm_dpp_f16_f16_f16_mk_nk_mn instance list that +// measured fastest on gfx1030 for 361 to 23104 tokens and 256 to 1536 channels, with MN padding so +// that the token count (batch * board area) need not be a tile multiple. Instantiated directly as +// the gridwise GEMM that DeviceGemmDpp wraps, because DeviceGemmDpp fixes the C memory operation +// to Set, while accumulating needs Add. +template +struct DppGemm { + static constexpr ck::index_t BlockSize = BS; + // As DeviceGemmDpp, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<4, BC1, 1>, + // S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, 4, 8> would instantiate it: C vectorized along M2 + // (dimension 4 of the per-thread C slice, 8 elements). + using Gridwise = ck::GridwiseGemm_ak0mak1_bk0nbk1_mn_dpp< + BS, F16, F32, F16, COp, Row, Col, Col, PassThrough, PassThrough, PassThrough, MNPadding, + MB, NB, KB, MD, ND, 8, 8, MW, NW, + S<4, AC1, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, false, true, + S<4, BC1, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, false, true, + S<0, 2, 4, 1, 3, 5>, 4, 8>; +}; + +template +struct Instances { + using I0 = DppGemm; + using I1 = DppGemm; + using I2 = DppGemm; + using I3 = DppGemm; +}; + +constexpr int NUM_INSTANCES = 4; + +template +bool runInstance( + bool checkOnly, const half* in, int inStride, const half* w, int wStride, half* out, int outStride, + int tokens, int outChannels, int inChannels, hipStream_t stream +) { + using Gridwise = typename Inst::Gridwise; + const typename Gridwise::Argument arg( + reinterpret_cast(w), reinterpret_cast(in), reinterpret_cast(out), + outChannels, tokens, inChannels, wStride, inStride, outStride); + if(!Gridwise::CheckValidity(arg)) + return false; + if(checkOnly) + return true; + // As DeviceGemmDpp's invoker launches it. + const auto [gdx, gdy, gdz] = Gridwise::CalculateGridSize(outChannels, tokens); + if(Gridwise::CalculateHasMainKBlockLoop(inChannels)) + ck::kernel_gemm_dpp<<>>(arg); + else + ck::kernel_gemm_dpp<<>>(arg); + return true; +} + +template +bool runIndexed( + int instance, bool checkOnly, const half* in, int inStride, const half* w, int wStride, half* out, int outStride, + int tokens, int outChannels, int inChannels, hipStream_t stream +) { + using Is = Instances; + switch(instance) { + case 0: return runInstance(checkOnly, in, inStride, w, wStride, out, outStride, tokens, outChannels, inChannels, stream); + case 1: return runInstance(checkOnly, in, inStride, w, wStride, out, outStride, tokens, outChannels, inChannels, stream); + case 2: return runInstance(checkOnly, in, inStride, w, wStride, out, outStride, tokens, outChannels, inChannels, stream); + case 3: return runInstance(checkOnly, in, inStride, w, wStride, out, outStride, tokens, outChannels, inChannels, stream); + default: return false; + } +} + +} // namespace + +bool rocmCkGemmSupportsArch(const char* gcnArchName) { + std::string arch(gcnArchName); + size_t colonPos = arch.find(':'); + if(colonPos != std::string::npos) + arch = arch.substr(0, colonPos); + std::string list = KATAGO_ROCM_CK_GEMM_ARCHS; + size_t start = 0; + while(start <= list.size()) { + size_t end = list.find(';', start); + if(end == std::string::npos) + end = list.size(); + std::string entry = list.substr(start, end - start); + size_t entryColon = entry.find(':'); + if(entryColon != std::string::npos) + entry = entry.substr(0, entryColon); + if(!entry.empty() && entry == arch) + return true; + start = end + 1; + } + return false; +} + +int rocmCkGemmNumInstances() { + return NUM_INSTANCES; +} + +bool rocmCkGemmSupports(int instance, int tokens, int outChannels, int inChannels, int inStride, int wStride, int outStride) { + // The vectorized loads and stores need 8-element (16-byte) aligned rows, which CK's shape check + // does not cover for C. Set and Add instances share their shape constraints. + if(outChannels % 8 != 0 || inStride % 8 != 0 || wStride % 8 != 0 || outStride % 8 != 0) + return false; + return runIndexed( + instance, true, nullptr, inStride, nullptr, wStride, nullptr, outStride, tokens, outChannels, inChannels, nullptr); +} + +void rocmCkGemmRun( + int instance, bool accumulate, + const half* in, int inStride, const half* w, int wStride, half* out, int outStride, + int tokens, int outChannels, int inChannels, hipStream_t stream +) { + const bool ran = accumulate + ? runIndexed( + instance, false, in, inStride, w, wStride, out, outStride, tokens, outChannels, inChannels, stream) + : runIndexed( + instance, false, in, inStride, w, wStride, out, outStride, tokens, outChannels, inChannels, stream); + if(!ran) + throw std::runtime_error("rocmCkGemmRun: instance does not support this problem"); +} diff --git a/cpp/neuralnet/rocmkernels.h b/cpp/neuralnet/rocmkernels.h new file mode 100644 index 0000000000..5158a64c9a --- /dev/null +++ b/cpp/neuralnet/rocmkernels.h @@ -0,0 +1,68 @@ +// Declarations for ROCm-only GPU kernels (rocmkernels.hip). Unlike the kernels shared with the +// CUDA backend in cudaandrocmhelpers.inc, these are written directly against AMDGPU features - +// v_dot2_f32_f16 and SGPR (scalar-cache) broadcast operands - to make the RDNA archs that have +// no matrix cores (and no CK FMHA kernels) fast, so they have no CUDA counterpart. + +#ifndef NEURALNET_ROCMKERNELS_H_ +#define NEURALNET_ROCMKERNELS_H_ + +#include "../neuralnet/rocmincludes.h" + +// Whether the device architecture has v_dot2_f32_f16, which every kernel here is built on. Takes +// hipDeviceProp_t::gcnArchName (e.g. "gfx1030" or "gfx90a:sramecc+:xnack-"). The archs without it +// (GCN gfx900-class and RDNA1 gfx1010/gfx1013) get empty kernel bodies, so callers must check this. +bool rocmArchHasDot2(const char* gcnArchName); +// Whether rocmUseDot2Attention = auto enables the dot2 attention kernels on this arch: only RDNA2 +// (gfx1030-gfx1036), where they were developed and measured (on gfx1030). Other archs with +// v_dot2_f32_f16 run them only with rocmUseDot2Attention = true. +bool rocmDot2AttentionAutoArch(const char* gcnArchName); + +//============================================================================================== +// dot2 scaled dot product attention +//============================================================================================== + +// How the attention prep kernel applies rotary position embeddings to Q and K, mirroring the +// customCudaApplyRoPE* launchers in cudaandrocmhelpers.h. +struct RocmAttnRope { + enum Mode { NONE = 0, TABLE = 1, RECOMPUTE = 2 }; + Mode mode = NONE; + // TABLE: cos/sin tables, (numPairs, seqLen) if !learnableTable, else (numKVHeads, numPairs, seqLen). + const half* cosTable = nullptr; + const half* sinTable = nullptr; + bool learnableTable = false; + // RECOMPUTE: per-head frequencies (numKVHeads, numPairs, 2), cos/sin recomputed from board x/y. + const float* freqs = nullptr; + int nnXLen = 0; +}; + +// Whether the dot2 attention kernels support this head configuration. +bool rocmDot2AttentionSupportsShape(int numHeads, int numKVHeads, int qHeadDim, int vHeadDim); + +// Scratch sizes for the prep kernel's outputs: the compacted K rows (also the size of the +// key-pair-interleaved V rows), and the per-batch-element position permutation. +size_t rocmDot2AttentionKVBytes(int batchSize, int seqLen, int numKVHeads, int headDim); +size_t rocmDot2AttentionPermBytes(int batchSize, int seqLen); + +// Prep: applies RoPE to Q in place (on-board positions only, the only ones attention reads), +// writes RoPE'd K rows compacted over on-board positions into kc and the V rows compacted and +// interleaved in key pairs into vc, and the permutation of positions (on-board first) plus the +// on-board count into perm. Q/K/V are BSHD with per-token row strides qStride/kvStride (halfs). +// mask (can be null) is [batchSize, seqLen]. +void rocmDot2AttentionPrep( + half* q, const half* k, const half* v, const half* mask, + half* kc, half* vc, short* perm, + int batchSize, int seqLen, int numHeads, int numKVHeads, int headDim, int qStride, int kvStride, + const RocmAttnRope& rope, hipStream_t stream +); + +// Attention over the prep kernel's outputs. output is BSHD with row stride numHeads*headDim, and +// off-board query rows are written as zeros. Same semantics as customCudaFlashAttention. +// numMultiprocessors and waveSize (hipDeviceProp_t::multiProcessorCount and warpSize) decide +// whether small batches split the keys over more waves. +void rocmDot2Attention( + const half* q, const half* kc, const half* vc, const short* perm, half* output, + int batchSize, int seqLen, int numHeads, int numKVHeads, int headDim, int qStride, + int numMultiprocessors, int waveSize, hipStream_t stream +); + +#endif // NEURALNET_ROCMKERNELS_H_ diff --git a/cpp/neuralnet/rocmkernels.hip b/cpp/neuralnet/rocmkernels.hip new file mode 100644 index 0000000000..4fc6882d38 --- /dev/null +++ b/cpp/neuralnet/rocmkernels.hip @@ -0,0 +1,536 @@ +// ROCm-only GPU kernels, see rocmkernels.h. + +#include "../neuralnet/rocmkernels.h" + +#include +#include + +// v_dot2_f32_f16 (__builtin_amdgcn_fdot2) does not exist on these archs. Their device passes +// compile the kernels below to empty bodies, and rocmArchHasDot2 keeps the host from launching +// them there. Keep the two lists in sync. +#if defined(__gfx900__) || defined(__gfx902__) || defined(__gfx909__) || defined(__gfx90c__) \ + || defined(__gfx1010__) || defined(__gfx1013__) || defined(__gfx9_generic__) || defined(__gfx10_1_generic__) + #define KATAGO_ROCM_DEVICE_HAS_DOT2 0 +#else + #define KATAGO_ROCM_DEVICE_HAS_DOT2 1 +#endif + +static std::string rkArchBaseName(const char* gcnArchName) { + std::string arch(gcnArchName); + size_t colonPos = arch.find(':'); + if(colonPos != std::string::npos) + arch = arch.substr(0, colonPos); + return arch; +} + +bool rocmArchHasDot2(const char* gcnArchName) { + const std::string arch = rkArchBaseName(gcnArchName); + return !( + arch == "gfx900" || arch == "gfx902" || arch == "gfx909" || arch == "gfx90c" || + arch == "gfx1010" || arch == "gfx1013" || arch == "gfx9-generic" || arch == "gfx10-1-generic" + ); +} + +bool rocmDot2AttentionAutoArch(const char* gcnArchName) { + const std::string arch = rkArchBaseName(gcnArchName); + return arch.size() == 7 && arch.compare(0, 6, "gfx103") == 0 && arch[6] >= '0' && arch[6] <= '6'; +} + +typedef _Float16 rk_h2 __attribute__((ext_vector_type(2))); + +static __device__ __forceinline__ rk_h2 rkU32ToH2(uint32_t u) { return __builtin_bit_cast(rk_h2, u); } +static __device__ __forceinline__ uint32_t rkPackF32ToH2(float a, float b) { + return __builtin_bit_cast(uint32_t, __builtin_amdgcn_cvt_pkrtz(a, b)); +} +static __device__ __forceinline__ float rkDot2(rk_h2 a, rk_h2 b, float c) { +#if KATAGO_ROCM_DEVICE_HAS_DOT2 + return __builtin_amdgcn_fdot2(a, b, c, false); +#else + (void)a; (void)b; + return c; +#endif +} + +//============================================================================================== +// dot2 scaled dot product attention +//============================================================================================== +// +// RDNA2 has no matrix cores, so attention runs on v_dot2_f32_f16 (two FP16 products accumulated +// in FP32 per lane per cycle). Each lane owns whole query rows, which keeps the online softmax +// entirely lane-local, so every K and V value is a broadcast operand shared by the 32 lanes of a +// wave. Broadcasting through LDS or vector memory costs 32x the bytes in data-path bandwidth and +// tops out well below the dot2 rate, so K and V are instead read through the scalar cache into +// SGPRs and used directly as the scalar operand of dot2, which costs no vector bandwidth at all. +// +// That requires K and V rows contiguous per (batch element, KV head) in the order the keys are +// visited, and V interleaved in key pairs so that one dot2 accumulates two keys into one output +// channel. A prep kernel produces that layout, compacting over on-board positions so that masked +// keys (small boards in a larger buffer) cost nothing, and applies RoPE on the way, replacing the +// separate RoPE kernel. + +// Kc/Vc rows per (batch element, KV head): seqLen rounded up to whole key pairs. +static __host__ __device__ inline int rkAttnPaddedSeq(int seqLen) { return (seqLen + 1) & ~1; } + +bool rocmDot2AttentionSupportsShape(int numHeads, int numKVHeads, int qHeadDim, int vHeadDim) { + return + qHeadDim == vHeadDim && (qHeadDim == 32 || qHeadDim == 64) && + numKVHeads > 0 && numHeads % numKVHeads == 0; +} + +size_t rocmDot2AttentionKVBytes(int batchSize, int seqLen, int numKVHeads, int headDim) { + return (size_t)batchSize * numKVHeads * rkAttnPaddedSeq(seqLen) * headDim * sizeof(half); +} +size_t rocmDot2AttentionPermBytes(int batchSize, int seqLen) { + // Per batch element: seqLen positions followed by the on-board count, padded to 4 bytes. + return (size_t)batchSize * ((seqLen + 2) & ~1) * sizeof(short); +} +static __host__ __device__ inline int rkPermStride(int seqLen) { return (seqLen + 2) & ~1; } + +// Block-wide ordered partition of the positions of batch element n: on-board positions in order +// into perm[0..count), off-board ones into perm[count..seqLen) (in reverse order, which is +// irrelevant to every consumer). Returns count. Requires blockDim.x == NT. +template +static __device__ __forceinline__ int rkPartitionPositions(const half* __restrict__ mask, int n, int seqLen, short* perm) { + __shared__ int waveCounts[NT / 32 + 1]; + const int tid = threadIdx.x; + const int lane = __lane_id(); + const int wave = tid / warpSize; + const int numWaves = NT / warpSize; + int validBase = 0; + int invalidBase = 0; + for(int p0 = 0; p0 < seqLen; p0 += NT) { + int pos = p0 + tid; + bool inRange = pos < seqLen; + bool valid = inRange && (mask == NULL || (float)mask[(size_t)n * seqLen + pos] != 0.0f); + unsigned long long below = (1ull << lane) - 1ull; + unsigned long long vb = __ballot(valid); + unsigned long long ib = __ballot(inRange && !valid); + if(lane == 0) + waveCounts[wave] = __popcll(vb) | (__popcll(ib) << 16); + __syncthreads(); + int validOff = validBase, invalidOff = invalidBase; + for(int w = 0; w < numWaves; w++) { + int c = waveCounts[w]; + if(w < wave) { validOff += c & 0xffff; invalidOff += c >> 16; } + validBase += c & 0xffff; + invalidBase += c >> 16; + } + if(valid) + perm[validOff + __popcll(vb & below)] = (short)pos; + else if(inRange) + perm[seqLen - 1 - (invalidOff + __popcll(ib & below))] = (short)pos; + __syncthreads(); + } + return validBase; +} + +// Grid (numKVHeads, batchSize), block NT. See rocmDot2AttentionPrep. +template +__global__ __launch_bounds__(NT) +void rkAttnPrepKernel( + half* __restrict__ Q, const half* __restrict__ K, const half* __restrict__ V, const half* __restrict__ mask, + half* __restrict__ Kc, uint32_t* __restrict__ Vc, short* __restrict__ permOut, + int seqLen, int numHeads, int numKVHeads, int qStride, int kvStride, float kScale, + const half* __restrict__ cosTable, const half* __restrict__ sinTable, int learnableTable, + const float* __restrict__ freqs, int nnXLen +) { +#if KATAGO_ROCM_DEVICE_HAS_DOT2 + constexpr int DV4 = D / 8; // 16-byte chunks per row + constexpr int NUM_PAIRS = D / 2; + extern __shared__ short perm[]; + const int kvh = blockIdx.x; + const int n = blockIdx.y; + const int count = rkPartitionPositions(mask, n, seqLen, perm); + if(kvh == 0) { + short* dst = permOut + (size_t)n * rkPermStride(seqLen); + for(int i = threadIdx.x; i < seqLen; i += NT) + dst[i] = perm[i]; + if(threadIdx.x == 0) + dst[seqLen] = (short)count; + } + if(count == 0) + return; + + const int SP = rkAttnPaddedSeq(seqLen); + const size_t rowBase = ((size_t)n * numKVHeads + kvh) * SP; + const int countP = (count + 1) & ~1; + const int groupSize = numHeads / numKVHeads; + + // K rows (rotated and pre-scaled by kScale, folded into the one rounding to half that the + // rotation needs anyway) and, for real slots, the rotation of the Q rows of this KV head's + // group. The slot past an odd count repeats the last key, so it cannot change the softmax max, + // and the attention kernel gives it zero weight. + for(int t = threadIdx.x; t < countP * DV4; t += NT) { + const int j = t / DV4; + const int c = t - j * DV4; + const bool realSlot = j < count; + const int pos = perm[realSlot ? j : count - 1]; + float cosV[4] = {1.0f, 1.0f, 1.0f, 1.0f}; + float sinV[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + if(ROPE != 0) { + #pragma unroll + for(int p = 0; p < 4; p++) { + const int pairIdx = c * 4 + p; + if(ROPE == 1) { + int tableIdx = learnableTable ? (kvh * NUM_PAIRS + pairIdx) * seqLen + pos : pairIdx * seqLen + pos; + cosV[p] = __half2float(cosTable[tableIdx]); + sinV[p] = __half2float(sinTable[tableIdx]); + } + else { + int x = pos % nnXLen; + int y = pos / nnXLen; + float angle = (float)x * freqs[(kvh * NUM_PAIRS + pairIdx) * 2 + 0] + (float)y * freqs[(kvh * NUM_PAIRS + pairIdx) * 2 + 1]; + sincosf(angle, &sinV[p], &cosV[p]); + } + } + } + // Same arithmetic as the customCudaApplyRoPE* half kernels: rotate in float, round once. + auto rotate = [&](uint4 in, float scale) { + uint32_t w[4] = {in.x, in.y, in.z, in.w}; + #pragma unroll + for(int p = 0; p < 4; p++) { + rk_h2 x = rkU32ToH2(w[p]); + float a = (float)x[0], b = (float)x[1]; + rk_h2 r; + r[0] = (_Float16)((a * cosV[p] - b * sinV[p]) * scale); + r[1] = (_Float16)((a * sinV[p] + b * cosV[p]) * scale); + w[p] = __builtin_bit_cast(uint32_t, r); + } + return make_uint4(w[0], w[1], w[2], w[3]); + }; + const uint4 kv = *(const uint4*)(K + ((size_t)n * seqLen + pos) * kvStride + kvh * D + c * 8); + *(uint4*)(Kc + (rowBase + j) * D + c * 8) = rotate(kv, kScale); + if(ROPE != 0 && realSlot) { + for(int g = 0; g < groupSize; g++) { + uint4* qp = (uint4*)(Q + ((size_t)n * seqLen + pos) * qStride + (kvh * groupSize + g) * D + c * 8); + *qp = rotate(*qp, 1.0f); + } + } + } + + // V rows, interleaved in key pairs: word d of pair row jp holds (V[key 2jp][d], V[key 2jp+1][d]). + // The slot past an odd count gets a zero V row. + for(int t = threadIdx.x; t < (countP / 2) * DV4; t += NT) { + const int jp = t / DV4; + const int c = t - jp * DV4; + const int key0 = 2 * jp; + uint4 a = *(const uint4*)(V + ((size_t)n * seqLen + perm[key0]) * kvStride + kvh * D + c * 8); + uint4 b = make_uint4(0, 0, 0, 0); + if(key0 + 1 < count) + b = *(const uint4*)(V + ((size_t)n * seqLen + perm[key0 + 1]) * kvStride + kvh * D + c * 8); + uint4 lo, hi; + lo.x = (a.x & 0xffffu) | (b.x << 16); lo.y = (a.x >> 16) | (b.x & 0xffff0000u); + lo.z = (a.y & 0xffffu) | (b.y << 16); lo.w = (a.y >> 16) | (b.y & 0xffff0000u); + hi.x = (a.z & 0xffffu) | (b.z << 16); hi.y = (a.z >> 16) | (b.z & 0xffff0000u); + hi.z = (a.w & 0xffffu) | (b.w << 16); hi.w = (a.w >> 16) | (b.w & 0xffff0000u); + uint4* dst = (uint4*)(Vc + (rowBase / 2 + jp) * D + c * 8); + dst[0] = lo; + dst[1] = hi; + } +#endif +} + +// Grid (ceil(seqLen / (NT*TM)), batchSize*numHeads), block NT*KS. Each lane owns TM query slots of +// the prep kernel's permutation (on-board positions first). Per key pair it computes the two +// scores of each of its rows against K rows held in SGPRs, updates a lazily-rescaled online +// softmax, and accumulates both keys into every output channel with one dot2 per channel against +// the interleaved V pair row, also in SGPRs. +// With KS > 1, the KS groups of NT threads cover the same query slots but split the key pairs into +// contiguous chunks, then merge their partial softmax states through LDS. Small batches launch too +// few waves to hide the latency of this loop, and the split multiplies them. +template +__global__ __launch_bounds__(NT * KS) +void rkDot2AttnKernel( + const half* __restrict__ Q, const half* __restrict__ Kc, const uint32_t* __restrict__ Vc, + const short* __restrict__ permIn, half* __restrict__ out, + int seqLen, int numHeads, int numKVHeads, int qStride +) { +#if KATAGO_ROCM_DEVICE_HAS_DOT2 + constexpr int DH2 = D / 2; + constexpr int DV4 = D / 8; + // NT is a multiple of the wave size, so split is wave-uniform; readfirstlane tells the compiler, + // which keeps the K/V loads below on the scalar unit. + const int split = __builtin_amdgcn_readfirstlane(threadIdx.x / NT); + const int tid = threadIdx.x - split * NT; + const int bh = blockIdx.y; + const int n = bh / numHeads; + const int h = bh - n * numHeads; + const int kvh = h / (numHeads / numKVHeads); + const int oStride = numHeads * D; + const short* perm = permIn + (size_t)n * rkPermStride(seqLen); + const int count = perm[seqLen]; + + const int slotBase = blockIdx.x * (NT * TM); + int qPos[TM]; + #pragma unroll + for(int i = 0; i < TM; i++) { + const int slot = slotBase + i * NT + tid; + qPos[i] = slot < seqLen ? (int)perm[slot] : -1; + if(slot >= count && qPos[i] >= 0) { + // Off-board query: zero output row. + if(split == 0) { + uint4* o = (uint4*)(out + ((size_t)n * seqLen + qPos[i]) * oStride + h * D); + #pragma unroll + for(int j = 0; j < DV4; j++) o[j] = make_uint4(0, 0, 0, 0); + } + qPos[i] = -1; + } + } + // Block-uniform, so no thread skips the merge barrier below. + if(slotBase >= count) + return; + + // Q rows. The prep kernel pre-scaled K by softmax scale * log2(e), so dot products with these + // are directly exp2 arguments. + rk_h2 q[TM][DH2]; + #pragma unroll + for(int i = 0; i < TM; i++) { + const int p = qPos[i] < 0 ? 0 : qPos[i]; + const uint4* src = (const uint4*)(Q + ((size_t)n * seqLen + p) * qStride + h * D); + #pragma unroll + for(int j = 0; j < DV4; j++) { + uint4 v = src[j]; + q[i][j * 4 + 0] = rkU32ToH2(v.x); + q[i][j * 4 + 1] = rkU32ToH2(v.y); + q[i][j * 4 + 2] = rkU32ToH2(v.z); + q[i][j * 4 + 3] = rkU32ToH2(v.w); + } + } + + // Online softmax with a lazily moved reference max: scores are kept relative to -negM, and the + // reference only moves (rescaling acc and l) when the running max leaves [-8, 8], so p <= 256 + // always fits FP16 with full relative precision, and rescaling is rare after the first pairs. + float negM[TM], mrun[TM], l[TM]; + float acc[TM][D]; + #pragma unroll + for(int i = 0; i < TM; i++) { + negM[i] = 0.0f; + mrun[i] = -INFINITY; + l[i] = 0.0f; + #pragma unroll + for(int d = 0; d < D; d++) acc[i][d] = 0.0f; + } + + const int SP = rkAttnPaddedSeq(seqLen); + const uint32_t* __restrict__ kBase = (const uint32_t*)(Kc + ((size_t)n * numKVHeads + kvh) * SP * D); + const uint32_t* __restrict__ vBase = Vc + ((size_t)n * numKVHeads + kvh) * (SP / 2) * D; + const int numPairs = (count + 1) >> 1; + // This group's chunk of the key pairs. Split 0's is never empty, since count >= 1 here. + const int chunk = (numPairs + KS - 1) / KS; + const int jpBegin = min(split * chunk, numPairs); + const int jpEnd = min(jpBegin + chunk, numPairs); + + // K rows of the current key pair. Loaded through the scalar cache (wave-uniform addresses into + // read-only memory) and refilled with the next pair's rows right after use, so that their load + // latency overlaps the softmax and PV work. + uint32_t kr[2 * DH2]; + { + const uint32_t* k0 = kBase + (size_t)(2 * min(jpBegin, numPairs - 1)) * DH2; + #pragma unroll + for(int dd = 0; dd < 2 * DH2; dd++) kr[dd] = k0[dd]; + } + for(int jp = jpBegin; jp < jpEnd; jp++) { + uint32_t vr[D]; + const uint32_t* vp = vBase + (size_t)jp * D; + #pragma unroll + for(int d = 0; d < D; d++) vr[d] = vp[d]; + + float s0[TM], s1[TM]; + #pragma unroll + for(int i = 0; i < TM; i++) { s0[i] = negM[i]; s1[i] = negM[i]; } + #pragma unroll + for(int dd = 0; dd < DH2; dd++) { + #pragma unroll + for(int i = 0; i < TM; i++) { + s0[i] = rkDot2(rkU32ToH2(kr[dd]), q[i][dd], s0[i]); + s1[i] = rkDot2(rkU32ToH2(kr[DH2 + dd]), q[i][dd], s1[i]); + } + } + { + const uint32_t* kn = kBase + (size_t)(2 * min(jp + 1, jpEnd - 1)) * DH2; + #pragma unroll + for(int dd = 0; dd < 2 * DH2; dd++) kr[dd] = kn[dd]; + } + + bool needRescale = false; + float mr[TM]; + #pragma unroll + for(int i = 0; i < TM; i++) { + mr[i] = fmaxf(mrun[i], fmaxf(s0[i], s1[i])); + needRescale |= (mr[i] > 8.0f || mr[i] < -8.0f); + } + // Wave-uniform branch: a divergent one makes the compiler shuffle the whole accumulator + // array around the join point, which costs far more than rescaling a few extra rows by 1. + if(__any(needRescale)) { + #pragma unroll + for(int i = 0; i < TM; i++) { + float shift = (mr[i] > 8.0f || mr[i] < -8.0f) ? mr[i] : 0.0f; + // The reference only moves down (shift < 0) before the row has accumulated anything + // (running max -inf), when l and acc are still 0, so alpha would scale nothing, but + // exp2(-shift) can overflow to inf there and turn the zeros into NaN: clamp it to 1. + float alpha = __builtin_amdgcn_exp2f(-fmaxf(shift, 0.0f)); + negM[i] -= shift; + l[i] *= alpha; + #pragma unroll + for(int d = 0; d < D; d++) acc[i][d] *= alpha; + s0[i] -= shift; + s1[i] -= shift; + mr[i] -= shift; + } + } + // The padded slot past an odd count has a zero V row; this also drops it from the denominator. + const rk_h2 ones = rkU32ToH2((2 * jp + 1 < count) ? 0x3C003C00u : 0x00003C00u); + #pragma unroll + for(int i = 0; i < TM; i++) { + mrun[i] = mr[i]; + rk_h2 ph = rkU32ToH2(rkPackF32ToH2(__builtin_amdgcn_exp2f(s0[i]), __builtin_amdgcn_exp2f(s1[i]))); + l[i] = rkDot2(ones, ph, l[i]); + #pragma unroll + for(int d = 0; d < D; d++) + acc[i][d] = rkDot2(rkU32ToH2(vr[d]), ph, acc[i][d]); + } + } + + if constexpr(KS > 1) { + // Splits 1..KS-1 hand their partial states to split 0 one at a time through one LDS buffer, + // lane-minor so that the accesses are conflict-free, keeping the LDS per block (and so the + // occupancy) independent of KS. + __shared__ float part[TM][D + 2][NT]; + // Split 0's weights are relative to the reference max -negM; incoming ones are rebased onto + // the larger of the two references. + for(int s = 1; s < KS; s++) { + if(split == s) { + #pragma unroll + for(int i = 0; i < TM; i++) { + // An empty chunk contributes nothing, whatever its (initial) reference. + part[i][0][tid] = l[i] > 0.0f ? -negM[i] : -INFINITY; + part[i][1][tid] = l[i]; + #pragma unroll + for(int d = 0; d < D; d++) part[i][2 + d][tid] = acc[i][d]; + } + } + __syncthreads(); + if(split == 0) { + #pragma unroll + for(int i = 0; i < TM; i++) { + const float rs = part[i][0][tid]; + const float r = fmaxf(-negM[i], rs); + const float f = __builtin_amdgcn_exp2f(-negM[i] - r); + const float fs = __builtin_amdgcn_exp2f(rs - r); + negM[i] = -r; + l[i] = l[i] * f + part[i][1][tid] * fs; + #pragma unroll + for(int d0 = 0; d0 < D; d0 += 16) { + // Keeps the compiler from hoisting all the LDS reads at once, whose registers would + // cut the occupancy of the whole kernel. + __builtin_amdgcn_sched_barrier(0); + #pragma unroll + for(int d = d0; d < d0 + 16; d++) acc[i][d] = acc[i][d] * f + part[i][2 + d][tid] * fs; + } + } + } + __syncthreads(); + } + if(split > 0) + return; + } + + #pragma unroll + for(int i = 0; i < TM; i++) { + if(qPos[i] < 0) + continue; + const float inv = 1.0f / l[i]; + uint4* o = (uint4*)(out + ((size_t)n * seqLen + qPos[i]) * oStride + h * D); + #pragma unroll + for(int j = 0; j < DV4; j++) { + uint4 v; + v.x = rkPackF32ToH2(acc[i][j * 8 + 0] * inv, acc[i][j * 8 + 1] * inv); + v.y = rkPackF32ToH2(acc[i][j * 8 + 2] * inv, acc[i][j * 8 + 3] * inv); + v.z = rkPackF32ToH2(acc[i][j * 8 + 4] * inv, acc[i][j * 8 + 5] * inv); + v.w = rkPackF32ToH2(acc[i][j * 8 + 6] * inv, acc[i][j * 8 + 7] * inv); + o[j] = v; + } + } +#endif +} + +void rocmDot2AttentionPrep( + half* q, const half* k, const half* v, const half* mask, + half* kc, half* vc, short* perm, + int batchSize, int seqLen, int numHeads, int numKVHeads, int headDim, int qStride, int kvStride, + const RocmAttnRope& rope, hipStream_t stream +) { + if(!rocmDot2AttentionSupportsShape(numHeads, numKVHeads, headDim, headDim)) + throw std::runtime_error("rocmDot2AttentionPrep: unsupported head configuration"); + if(seqLen >= 32767 || batchSize > 65535) + throw std::runtime_error("rocmDot2AttentionPrep: seqLen or batchSize too large"); + if(qStride % 8 != 0 || kvStride % 8 != 0) + throw std::runtime_error("rocmDot2AttentionPrep: row strides must be multiples of 8 halfs"); + constexpr int NT = 256; + dim3 grid(numKVHeads, batchSize); + size_t shm = ((size_t)seqLen * sizeof(short) + 15) / 16 * 16; + uint32_t* vcw = (uint32_t*)vc; + // Softmax scale and the log2(e) of computing exp as exp2, folded into K. + const float kScale = 1.4426950408889634f / sqrtf((float)headDim); +#define RK_PREP_LAUNCH(DIM, MODE) \ + rkAttnPrepKernel<<>>( \ + q, k, v, mask, kc, vcw, perm, seqLen, numHeads, numKVHeads, qStride, kvStride, kScale, \ + rope.cosTable, rope.sinTable, rope.learnableTable ? 1 : 0, rope.freqs, rope.nnXLen) + const int mode = (int)rope.mode; + if(headDim == 32) { + if(mode == RocmAttnRope::NONE) RK_PREP_LAUNCH(32, 0); + else if(mode == RocmAttnRope::TABLE) RK_PREP_LAUNCH(32, 1); + else RK_PREP_LAUNCH(32, 2); + } + else { + if(mode == RocmAttnRope::NONE) RK_PREP_LAUNCH(64, 0); + else if(mode == RocmAttnRope::TABLE) RK_PREP_LAUNCH(64, 1); + else RK_PREP_LAUNCH(64, 2); + } +#undef RK_PREP_LAUNCH +} + +// Launches rkDot2AttnKernel with KS key splits, or none. +template +static void rkLaunchDot2Attention( + bool splitKeys, const half* q, const half* kc, const uint32_t* vc, const short* perm, half* output, + int batchSize, int seqLen, int numHeads, int numKVHeads, int qStride, hipStream_t stream +) { + dim3 grid((seqLen + NT * TM - 1) / (NT * TM), batchSize * numHeads); + if(splitKeys) + rkDot2AttnKernel<<>>( + q, kc, vc, perm, output, seqLen, numHeads, numKVHeads, qStride); + else + rkDot2AttnKernel<<>>( + q, kc, vc, perm, output, seqLen, numHeads, numKVHeads, qStride); +} + +void rocmDot2Attention( + const half* q, const half* kc, const half* vc, const short* perm, half* output, + int batchSize, int seqLen, int numHeads, int numKVHeads, int headDim, int qStride, + int numMultiprocessors, int waveSize, hipStream_t stream +) { + if(!rocmDot2AttentionSupportsShape(numHeads, numKVHeads, headDim, headDim)) + throw std::runtime_error("rocmDot2Attention: unsupported head configuration"); + if((long long)batchSize * numHeads > 65535) + throw std::runtime_error("rocmDot2Attention: batchSize * numHeads too large"); + const uint32_t* vcw = (const uint32_t*)vc; + // Measured on gfx1030: head dim 32 prefers two query rows per lane (more scalar-operand reuse at + // a still high occupancy), head dim 64 one row (its accumulators are twice as large). + constexpr int TM32 = 2, NT32 = 64, TM64 = 1, NT64 = 128; + const int TM = headDim == 32 ? TM32 : TM64; + const int NT = headDim == 32 ? NT32 : NT64; + // Unsplit waves, against the ~8 per SIMD it takes to hide the latency of the per-key-pair + // dependency chain, with 4 SIMDs per multiprocessor (a CU on GCN/CDNA, a WGP on RDNA). Also + // measured on gfx1030: head dim 32 gains from 4 key splits up to about twice that many waves + // (2.7x faster at batch size 1, 1.25x at 8, even at 48) and loses beyond; head dim 64, whose + // merge costs it occupancy, only from 2 splits below half. + const long long waves = (long long)((seqLen + NT * TM - 1) / (NT * TM)) * batchSize * numHeads * (NT / waveSize); + const long long targetWaves = 32LL * numMultiprocessors; + if(headDim == 32) + rkLaunchDot2Attention<32, TM32, NT32, 4>( + waves <= 2 * targetWaves, q, kc, vcw, perm, output, batchSize, seqLen, numHeads, numKVHeads, qStride, stream); + else + rkLaunchDot2Attention<64, TM64, NT64, 2>( + 2 * waves <= targetWaves, q, kc, vcw, perm, output, batchSize, seqLen, numHeads, numKVHeads, qStride, stream); +}