diff --git a/.github/workflows/unsloth-sd-prebuilt.yml b/.github/workflows/unsloth-sd-prebuilt.yml index 308082a96a..4ffdbd95ae 100644 --- a/.github/workflows/unsloth-sd-prebuilt.yml +++ b/.github/workflows/unsloth-sd-prebuilt.yml @@ -425,7 +425,8 @@ jobs: # prefix (it ships as libcublas / libcublas-dev), so it has to go in the # other list or apt cannot find it. cudart-dev is what carries the headers; # cudart alone is the runtime and does not compile anything. - sub-packages: '["nvcc", "cudart", "cudart-dev", "thrust"]' + # nvrtc-dev: cudnn-frontend's headers include nvrtc.h (it dlopens libnvrtc itself, so nothing links it). + sub-packages: '["nvcc", "cudart", "cudart-dev", "thrust", "nvrtc-dev"]' non-cuda-sub-packages: '["libcublas", "libcublas-dev"]' # This leg rebuilt every object on every run: 3903 s of the 4121 s job on @@ -456,6 +457,29 @@ jobs: # keeps what it compiled. save: false + - name: cuDNN headers (build time only) + run: | + set -euo pipefail + # ggml-cuda's cuDNN attention path (GGML_CUDA_CUDNN) compiles against cudnn.h, but the + # library itself is opened with dlopen when the first attention op runs and is NOT + # bundled: a host without a loadable libcudnn.so.9 keeps the ggml attention kernels. + # Only the headers package is fetched (80 KB), pinned by version and checksum. + deb=libcudnn9-headers-cuda-12_9.27.0.42-1_amd64.deb + curl -fsSL --retry 3 -o "$RUNNER_TEMP/$deb" \ + "https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/$deb" + echo "68243ab04422da565e0bafec00f64ac08cfd82d805c8015c0c503af7ab4dd1d6 $RUNNER_TEMP/$deb" | sha256sum -c - + dpkg -x "$RUNNER_TEMP/$deb" "$RUNNER_TEMP/cudnn-deb" + # The package ships cudnn*_v9.h; the unversioned names normally come from the -dev + # package's alternatives, so link them here. + inc="$RUNNER_TEMP/cudnn-include" + mkdir -p "$inc" + for f in "$RUNNER_TEMP"/cudnn-deb/usr/include/x86_64-linux-gnu/cudnn*_v9.h; do + b="$(basename "$f")" + ln -s "$f" "$inc/${b%_v9.h}.h" + done + test -e "$inc/cudnn.h" + echo "CUDNN_INCLUDE_DIR=$inc" >> "$GITHUB_ENV" + - name: Build sd-cli + sd-server (CUDA) working-directory: src run: | @@ -474,6 +498,8 @@ jobs: -DSD_WEBP=OFF -DSD_WEBM=OFF \ -DGGML_NATIVE=OFF \ -DSD_CUDA=ON \ + -DGGML_CUDA_CUDNN=ON \ + -DGGML_CUDA_CUDNN_INCLUDE_DIR="$CUDNN_INCLUDE_DIR" \ -DCMAKE_CUDA_ARCHITECTURES="$CUDA_ARCHS" \ -DCMAKE_C_COMPILER_LAUNCHER=ccache \ -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \ @@ -502,6 +528,14 @@ jobs: patchelf --set-rpath '$ORIGIN' "$BIN/$exe" done ldd "$BIN/sd-cli" | sed -n '1,40p' + # cuDNN and NVRTC are opened at run time only; a link-time dependency would stop the binaries + # loading on every host without it. + for exe in sd-cli sd-server; do + if readelf -d "$BIN/$exe" | grep -qE 'NEEDED.*(cudnn|nvrtc)'; then + echo "ERROR: $exe links libcudnn or libnvrtc; they must stay run-time dlopens" >&2 + exit 1 + fi + done - name: Package bundle env: @@ -511,8 +545,16 @@ jobs: LABEL: Linux-Ubuntu-22.04-x86_64-cuda12 COMMIT: ${{ needs.resolve.outputs.commit }} SOURCE_REPO: ${{ github.repository }} - LICENSE_FILE: ${{ github.workspace }}/src/LICENSE - run: python3 tooling/scripts/unsloth/package_bundle.py + LICENSE_FILE: ${{ runner.temp }}/LICENSE + run: | + set -euo pipefail + # The binaries compile in cudnn-frontend's header-only templates, so its MIT notice ships too. + { + cat src/LICENSE + printf '\n\n---- cudnn-frontend (https://github.com/NVIDIA/cudnn-frontend) ----\n\n' + cat src/build/_deps/ggml_cudnn_frontend-src/LICENSE.txt + } > "$LICENSE_FILE" + python3 tooling/scripts/unsloth/package_bundle.py - name: Upload bundle uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 diff --git a/cmake/ggml.cmake b/cmake/ggml.cmake index 39d77eae35..82d4814933 100644 --- a/cmake/ggml.cmake +++ b/cmake/ggml.cmake @@ -21,6 +21,23 @@ endif() target_include_directories(${SD_LIB} PRIVATE "${sd_ggml_private_include}") set_property(TARGET ${SD_LIB} PROPERTY SD_GGML_PRIVATE_INCLUDE_DIR "${sd_ggml_private_include}") +# The fused DiT ops come from scripts/unsloth/ggml-patches/0004; an unpatched ggml keeps the unfused H3 graph. +set(sd_ggml_header "${sd_ggml_private_include}/../include/ggml.h") +if(EXISTS "${sd_ggml_header}") + set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS "${sd_ggml_header}") + file(STRINGS "${sd_ggml_header}" sd_ggml_h3_fused_ops REGEX "ggml_rope_pe_permute\\(") + file(STRINGS "${sd_ggml_header}" sd_ggml_h3_fused_qk_norm REGEX "ggml_rope_pe_permute_rms\\(") +endif() +if(sd_ggml_h3_fused_ops) + target_compile_definitions(${SD_LIB} PUBLIC SD_GGML_H3_FUSED_OPS) +else() + message(STATUS "ggml lacks the fused DiT ops patch: MiniMax-H3 uses the unfused graph") +endif() +# The fused q/k RMS norm comes from patch 0005. +if(sd_ggml_h3_fused_qk_norm) + target_compile_definitions(${SD_LIB} PUBLIC SD_GGML_H3_FUSED_QK_NORM) +endif() + if(SD_USE_UPSTREAM_GGML) target_compile_definitions(${SD_LIB} PUBLIC SD_USE_UPSTREAM_GGML) message(WARNING "Using upstream GGML: INT8 tensorwise/convrot is disabled and FP8 weights are converted to F16 at load time. Some operators may be unsupported and performance may be lower than with patched GGML.") diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index bcceb73c47..d50b9f8b18 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -52,6 +52,82 @@ which removes most of its attention cost without touching the text encoder. Set keep that attention by default (on Vulkan the flash attention decode was slower); `SD_H3_VAE_FLASH_ATTN=1` turns it on there. +On NVIDIA GPUs from Ampere on, the unmasked DiT and VAE attention runs a pipelined +long-sequence variant of the ggml-cuda flash attention kernel (up to 4x faster at H3 +shapes; the arithmetic per output element is unchanged). `GGML_CUDA_FA_LONGSEQ=0` restores the +stock kernel. `GGML_CUDA_FA_LONGSEQ_NCOLS=128` opts into a wider tile that is faster on A100 and L4 +but not bit-identical. + +Builds configured with `-DGGML_CUDA_CUDNN=ON` (the Linux CUDA prebuilt is) can run that unmasked +attention, and the `--sage-attn` attention, through cuDNN's fused attention. By default it replaces +the flash attention kernel on Ada (sm89) and B200-class (sm100) GPUs and the sage kernel on sm100 only, +where it measured faster; on an RTX 3090 (sm86) both existing kernels were faster, on an RTX 6000 Ada +the sage kernel was, and other architectures were not measured. +Only `cudnn.h` is needed at build time (`GGML_CUDA_CUDNN_INCLUDE_DIR`; the header-only +cudnn-frontend v1.26.0 is downloaded by CMake unless `GGML_CUDA_CUDNN_FRONTEND_DIR` points at a +checkout). `libcudnn.so.9` is opened when the first attention op runs and is not shipped: put a cuDNN 9 +built for the same CUDA major version as the binary (cuDNN 9 for CUDA 12 for the prebuilt) on the +library path or beside the binary, or point `GGML_CUDA_CUDNN_LIB` at it (cuDNN 9.27 needs nothing +else; older 9.x releases build these kernels with NVRTC and also need that CUDA major's `libnvrtc`). +Without a loadable cuDNN, on Turing and older, for attention calls under about a million scores +(where the extra conversions cost more than cuDNN saves) and for any shape cuDNN declines, the +existing kernels run. The first call of each attention shape builds a cuDNN plan (about 0.4 to 1.3 s +on a B200, logged once). cuDNN runs F16 with F32 accumulation: for the H3 DiT at 960x544x124 it +takes about 8 ms per attention call on a B200 against 28 ms for the sage kernel and 43 ms for the +ggml flash attention kernel, and it is closer to an exact (fp64) reference than either, so frames +and audio differ slightly from the previous kernels. `GGML_CUDA_CUDNN_ATTN=0` restores them and +`GGML_CUDA_CUDNN_ATTN=1` uses cuDNN for the flash attention op on any Ampere or newer GPU; +`GGML_CUDA_CUDNN_SAGE=0` / `=1` does the same for the sage op only, and +`GGML_CUDA_CUDNN_ATTN_BF16=1` runs cuDNN in BF16 (less accurate, same speed). + +The video VAE decodes one 16x16 latent tile per decoder graph by default. `SD_H3_VAE_TILE=N` +uses N x N latent tiles instead (20 decodes about 0.7 s faster at 960x544x124 on B200); the tile +seams move, so the frames differ from the default (36 dB PSNR at 20) and it is opt-in. +`SD_H3_VAE_TILE_BATCH=auto` puts several tiles into one graph, sized from free device memory (at +most 4 unless `SD_H3_VAE_TILE_BATCH_MAX` raises it), and `SD_H3_VAE_TILE_BATCH=N` forces N; the +batched projections can round differently from the per-tile decode on some GPUs, so it is +opt-in. The decoder weights stay on the device across temporal chunks +(`SD_H3_VAE_KEEP_RESIDENT=0` releases them after every chunk), and the decoder blocks use a +table-based rotary embedding and a fused SwiGLU (`SD_H3_VAE_GRAPH_OPT=0` restores the previous +graph). Each batched tile still goes through its own attention call. With one tile per graph and +flash attention, each decoder attention normalises q/k in place on the projection and one fused +RoPE op per tensor writes the head-major Q (F32) and K/V (F16) that the attention kernel reads, +replacing the table RoPE, chunk copies, permutes and casts; the frames are bit-identical +(`SD_H3_VAE_FUSED_QKV=0` restores the unfused graph). On CUDA the q/k RMS norm runs inside that +RoPE op with the same reduction as the standalone norm (`SD_H3_VAE_FUSED_QK_NORM=0` keeps it +separate). + +The host side of the decode overlaps the device: each tile's blend into the frame runs on a +worker thread while the next tile computes (`SD_TILE_ASYNC_MERGE=0` blends inline; this applies +to every tiled VAE decode), each temporal chunk is trimmed, cross-faded and copied into the final +frames on a worker thread while the next chunk decodes (`SD_H3_VAE_ASYNC_ASSEMBLY=0` restores +the concatenate-at-the-end path), the rotary tables are built once per tile shape, and a tile +that allocates nothing new reuses the device free-memory reading taken earlier in the decode +instead of querying the device for every capacity check (`SD_H3_VAE_REUSE_MEMQUERY=0`). The +decoded frames are bit-identical either way. + +The DiT blocks use fused ggml ops (CPU and CUDA) for the work around the matmuls and attention: +partial RoPE with the attention relayout and the K/V scale and F16 cast, per-segment adaLN +modulation and gated residuals written in place, and the MLP Linear scales folded into those ops +and the swiglu. The result is bit-identical to the unfused graph. `SD_H3_GRAPH_FAST=0` restores +the unfused graph; `SD_H3_FAST_QKV=0`, `SD_H3_FAST_MLP=0`, `SD_H3_FAST_SEGMENTS=0` and +`SD_H3_FAST_VIEWS=0` turn off one part each. With `--sage-attn` the same fused RoPE op writes the +F32 Q/K and F16 V layout the sage kernel reads (K and V carry the kv scale), so the sage path no +longer pays for the chunk / slice / rope / concat / scale / cast chain; the output is bit-identical +to the unfused sage graph. `SD_H3_FAST_SAGE_QKV=0` restores that chain. + +The audio VAE's anti-aliased activations run their up/down-sampling filters as direct F32 +depthwise convolutions (`CONV_2D_DW`) instead of an F16 im2col plus a matrix-vector product, +which makes the audio decode several times faster. The input is no longer rounded to F16, so +the waveform differs slightly from the previous graph (about 41 dB SNR on a 5 s clip). +`SD_H3_AUDIO_DIRECT_DW=0` restores the previous graph; backends without `CONV_2D_DW` keep it. + +The long-sequence flash attention kernel, the fused cuBLAS epilogues and the fused DiT ops come from +the ggml patches in `scripts/unsloth/ggml-patches`, which the Unsloth prebuilt binaries carry. A +source build gets them by applying the patches to the `ggml` submodule before configuring +(`for p in scripts/unsloth/ggml-patches/*.patch; do git -C ggml apply "../$p"; done`); without +them the build uses the stock ggml kernels and the unfused DiT graph. + ## First/last-frame conditioning Add `--init-img` for I2VA, or both `--init-img` and `--end-img` for FL2VA: diff --git a/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch b/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch new file mode 100644 index 0000000000..32b2c5378f --- /dev/null +++ b/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch @@ -0,0 +1,563 @@ +From ed3f840655078aef3c2bb4e6b9a1d26f8e868e64 Mon Sep 17 00:00:00 2001 +From: Daniel Han +Date: Thu, 1 Oct 2026 05:41:03 -0700 +Subject: [PATCH 1/3] ggml-cuda: pipelined long-sequence FlashAttention for + unmasked attention without GQA packing + +The mma FlashAttention kernels without GQA packing (ncols2 == 1, used whenever +there is no mask) load every K/V tile synchronously with an out-of-bounds check, +because the KV length need not be a multiple of the KV tile. For diffusion +transformers (no mask, 10^4 tokens) the kernel then spends most of its time +waiting on global memory. + +The long-sequence variant pipelines all full KV tiles with cp.async (the +existing two-stage path) and loads only the final, partial tile synchronously +with the bounds check. Per output element the arithmetic and its order are +unchanged, so results are bit-identical to the stock kernel. + +It is used for head sizes 64 and 128 with K == V head size, no mask, no sinks, +no logit softcap and KV length >= 256, on NVIDIA GPUs with cp.async (Ampere and +newer). GGML_CUDA_FA_LONGSEQ=0 restores the stock kernels. + +GGML_CUDA_FA_LONGSEQ_NCOLS=128 additionally selects a 128-column tile with 8 +warps for Q length >= 1024 (stream-k splits then differ, so results are not +bit-identical); off by default. + +test-backend-ops gains unmasked FLASH_ATTN_EXT cases with partial last KV +tiles, several Q tiles and GQA broadcast. +--- + src/ggml-cuda/fattn-mma-f16.cuh | 167 ++++++++++++++++-- + src/ggml-cuda/fattn.cu | 9 + + ...tn-mma-f16-instance-ncols1_128-ncols2_1.cu | 8 + + tests/test-backend-ops.cpp | 15 ++ + 4 files changed, 184 insertions(+), 15 deletions(-) + create mode 100644 src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_128-ncols2_1.cu + +diff --git a/src/ggml-cuda/fattn-mma-f16.cuh b/src/ggml-cuda/fattn-mma-f16.cuh +index 7f4cfd55..b4496734 100644 +--- a/src/ggml-cuda/fattn-mma-f16.cuh ++++ b/src/ggml-cuda/fattn-mma-f16.cuh +@@ -61,6 +61,10 @@ static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_co + GGML_CUDA_FATTN_MMA_CONFIG_CASE(128, 128, 32, 128, 2, 64, 64, 64, 64, 2, true); + GGML_CUDA_FATTN_MMA_CONFIG_CASE(128, 128, 64, 128, 2, 64, 64, 64, 64, 2, true); + ++ // ncols == 128 is only used by the long-sequence variant (ncols2 == 1, no mask), see ggml_cuda_fattn_mma_use_longseq: ++ GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64, 64, 128, 256, 1, 64, 32, 32, 32, 2, true); ++ GGML_CUDA_FATTN_MMA_CONFIG_CASE(128, 128, 128, 256, 1, 64, 64, 64, 64, 2, true); ++ + GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 8, 64, 4, 64, 96, 64, 64, 2, true); + GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 16, 64, 4, 32, 96, 64, 64, 2, true); + GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 32, 128, 2, 32, 96, 64, 64, 2, true); +@@ -358,6 +362,24 @@ static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages(const int DKQ, c + #endif // CP_ASYNC_AVAILABLE + } + ++// Long-sequence variant of the kernels without GQA packing (ncols2 == 1) and without a mask, used for diffusion models. ++// The stock ncols2 == 1 kernels load K/V synchronously with an out-of-bounds check on every KV tile because the ++// KV length need not be a multiple of nbatch_fa. The long-sequence variant pipelines all full KV tiles with cp.async ++// and only loads the final, partial tile synchronously with the OOB check. The arithmetic per output element is the ++// same as in the stock kernel, so the results are bit-identical. ++static __host__ int ggml_cuda_fattn_mma_get_nstages_longseq(const int DKQ, const int DV, const int ncols, const int cc) { ++ return cp_async_available(cc) ? ggml_cuda_fattn_mma_get_nstages_target(DKQ, DV, ncols, cc) : 0; ++} ++ ++static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages_longseq(const int DKQ, const int DV, const int ncols) { ++#ifdef CP_ASYNC_AVAILABLE ++ return ggml_cuda_fattn_mma_get_nstages_target(DKQ, DV, ncols); ++#else ++ GGML_UNUSED_VARS(DKQ, DV, ncols); ++ return 0; ++#endif // CP_ASYNC_AVAILABLE ++} ++ + // ------------------------------------------------------------------------------------------------------------------ + + template +@@ -529,7 +551,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( + + template ++ typename T_A_KQ, typename T_B_KQ, typename T_C_KQ, typename T_A_VKQ, typename T_B_VKQ, typename T_C_VKQ, int nstages_override = -1> + static __device__ __forceinline__ void flash_attn_ext_f16_iter( + const float2 * const __restrict__ Q_f2, + const half2 * const __restrict__ K_h2, +@@ -566,7 +588,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( + constexpr int nbatch_K2 = ggml_cuda_fattn_mma_get_nbatch_K2(DKQ, DV, ncols); + constexpr int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2(DKQ, DV, ncols); + constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols); +- constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2); ++ constexpr int nstages = nstages_override >= 0 ? nstages_override : ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2); + + constexpr int stride_tile_K = nbatch_K2 + 4; + +@@ -1113,7 +1135,8 @@ template struct mma_tile_sizes { + }; + #endif // defined(TURING_MMA_AVAILABLE) + +-template ++template + static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( + const float2 * const __restrict__ Q_f2, + const half2 * const __restrict__ K_h2, +@@ -1158,7 +1181,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( + constexpr int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2 (DKQ, DV, ncols); + constexpr int nbatch_combine = ggml_cuda_fattn_mma_get_nbatch_combine(DKQ, DV, ncols); + constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols); +- constexpr int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2); ++ constexpr int nstages = longseq ? ggml_cuda_fattn_mma_get_nstages_longseq(DKQ, DV, ncols) : ++ ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2); ++ static_assert(!longseq || (ncols2 == 1 && !V_is_K_view), "bad longseq variant"); + + if (cols_per_warp > ncols) { + NO_DEVICE_CODE; +@@ -1255,22 +1280,63 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( + + int kb0 = kb0_start; + ++ // Long-sequence variant: only the last KV tile can extend past ne11, all tiles before it are pipelined. ++ const int kb0_stop_pipe = longseq && kb0_stop*nbatch_fa > ne11 ? kb0_stop - 1 : kb0_stop; ++ + // Preload mask and K data for first iteration when using cp_async with multiple stages: + if constexpr (nstages > 1) { + static_assert(nbatch_K2 == DKQ/2, "batching not implemented for multi-stage pipeline"); + constexpr bool use_cp_async = true; + constexpr bool oob_check = false; + constexpr int k_VKQ_sup = nbatch_fa; +- if (ncols2 > 1 || mask_h) { +- flash_attn_ext_f16_load_mask +- (mask_h + kb0*nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01); ++ if (!longseq || kb0 < kb0_stop_pipe) { ++ if (ncols2 > 1 || mask_h) { ++ flash_attn_ext_f16_load_mask ++ (mask_h + kb0*nbatch_fa, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01); ++ } ++ flash_attn_ext_f16_load_tile ++ (K_h2 + int64_t(kb0)*nbatch_fa*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup); + } +- flash_attn_ext_f16_load_tile +- (K_h2 + int64_t(kb0)*nbatch_fa*stride_K, tile_K, nbatch_K2, stride_K, k_VKQ_sup); + } + + // kb0_start is always < kb0_stop so the last iter can be executed unconditionally. +- if constexpr (ncols2 == 1) { ++ if constexpr (longseq && nstages > 1) { ++ for (; kb0 < kb0_stop_pipe-1; ++kb0) { ++ constexpr bool last_iter = false; ++ constexpr bool oob_check = false; ++ constexpr int k_VKQ_sup = nbatch_fa; ++ flash_attn_ext_f16_iter ++ ++ (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, ++ ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, ++ KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); ++ } ++ if (kb0 < kb0_stop_pipe) { ++ constexpr bool last_iter = true; // Do not preload past the pipelined tiles. ++ constexpr bool oob_check = false; ++ constexpr int k_VKQ_sup = nbatch_fa; ++ flash_attn_ext_f16_iter ++ ++ (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, ++ ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, ++ KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); ++ ++kb0; ++ } ++ if (kb0 < kb0_stop) { ++ // Partial last tile: synchronous loads with an out-of-bounds check, as in the stock ncols2 == 1 kernel. ++ constexpr bool last_iter = true; ++ constexpr bool oob_check = true; ++ const int k_VKQ_sup = ne11 - kb0*nbatch_fa; ++ flash_attn_ext_f16_iter ++ ++ (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, ++ ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, ++ KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); ++ } ++ } else if constexpr (ncols2 == 1) { + constexpr bool oob_check = true; + for (; kb0 < kb0_stop-1; ++kb0) { + constexpr bool last_iter = false; +@@ -1700,7 +1766,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( + #endif // defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) + } + +-template ++template + __launch_bounds__(ggml_cuda_fattn_mma_get_nthreads(DKQ, DV, ncols1*ncols2), ggml_cuda_fattn_mma_get_occupancy(DKQ, DV, ncols1*ncols2)) + static __global__ void flash_attn_ext_f16( + const char * Q_ptr, +@@ -1829,12 +1895,12 @@ static __global__ void flash_attn_ext_f16( + constexpr bool is_fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer. + if (kb0_start == 0) { + constexpr bool needs_fixup = false; // CUDA block is working on an entire tile. +- flash_attn_ext_f16_process_tile ++ flash_attn_ext_f16_process_tile + (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, + ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); + } else { + constexpr bool needs_fixup = true; // CUDA block is missing the beginning of a tile. +- flash_attn_ext_f16_process_tile ++ flash_attn_ext_f16_process_tile + (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, + ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); + } +@@ -1875,7 +1941,7 @@ static __global__ void flash_attn_ext_f16( + + constexpr bool is_fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks. + constexpr bool needs_fixup = false; +- flash_attn_ext_f16_process_tile ++ flash_attn_ext_f16_process_tile + (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, + ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); + #else +@@ -1892,6 +1958,47 @@ static __global__ void flash_attn_ext_f16( + #endif // defined(FLASH_ATTN_AVAILABLE) && (defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)) + } + ++// Whether the long-sequence variant (see ggml_cuda_fattn_mma_get_nstages_longseq) can and should be used for dst. ++// GGML_CUDA_FA_LONGSEQ=0 disables it, =1 enables it on every NVIDIA GPU with cp.async, unset = per-architecture default. ++static bool ggml_cuda_fattn_mma_longseq_arch_default(const int cc) { ++ GGML_UNUSED(cc); ++ return true; ++} ++ ++// Q columns per CUDA block of the long-sequence variant for long Q: 64 (the stock tile) or 128 (8 warps sharing each K/V tile). ++// GGML_CUDA_FA_LONGSEQ_NCOLS=64|128 overrides the per-architecture default. ++static int ggml_cuda_fattn_mma_longseq_ncols(const int cc) { ++ static const int env = [] { ++ const char * e = getenv("GGML_CUDA_FA_LONGSEQ_NCOLS"); ++ return e != nullptr ? atoi(e) : -1; ++ }(); ++ if (env == 64 || env == 128) { ++ return env; ++ } ++ GGML_UNUSED(cc); ++ return 64; ++} ++ ++static bool ggml_cuda_fattn_mma_use_longseq(const int cc, const ggml_tensor * dst) { ++ static const int env = [] { ++ const char * e = getenv("GGML_CUDA_FA_LONGSEQ"); ++ return e != nullptr ? atoi(e) : -1; ++ }(); ++ if (env == 0 || !GGML_CUDA_CC_IS_NVIDIA(cc) || !cp_async_available(cc)) { ++ return false; ++ } ++ const ggml_tensor * K = dst->src[1]; ++ const ggml_tensor * V = dst->src[2]; ++ const ggml_tensor * mask = dst->src[3]; ++ const ggml_tensor * sinks = dst->src[4]; ++ float logit_softcap; ++ memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); ++ if (mask != nullptr || sinks != nullptr || logit_softcap != 0.0f || K->ne[0] != V->ne[0] || K->ne[1] < 256) { ++ return false; ++ } ++ return env == 1 || ggml_cuda_fattn_mma_longseq_arch_default(cc); ++} ++ + template + void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * KQV = dst; +@@ -1900,13 +2007,22 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml + + constexpr int ncols = ncols1 * ncols2; + ++#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) ++ constexpr bool longseq_instantiated = ncols2 == 1 && DKQ == DV && (DKQ == 64 || DKQ == 128); ++ const bool longseq = longseq_instantiated && ggml_cuda_fattn_mma_use_longseq(cc, dst); ++#else ++ constexpr bool longseq_instantiated = false; ++ const bool longseq = false; ++#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) ++ + const int nthreads = ggml_cuda_fattn_mma_get_nthreads (DKQ, DV, ncols, cc); + const int nbatch_fa = ggml_cuda_fattn_mma_get_nbatch_fa (DKQ, DV, ncols, cc); + const int nbatch_K2 = ggml_cuda_fattn_mma_get_nbatch_K2 (DKQ, DV, ncols, cc); + const int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2 (DKQ, DV, ncols, cc); + const int nbatch_combine = ggml_cuda_fattn_mma_get_nbatch_combine(DKQ, DV, ncols, cc); + const bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols, cc); +- const int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, cc); ++ const int nstages = longseq ? ggml_cuda_fattn_mma_get_nstages_longseq(DKQ, DV, ncols, cc) : ++ ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, cc); + + const int cols_per_warp = std::min(ncols, get_cols_per_warp(cc)); + const int warp_size_host = ggml_cuda_info().devices[ctx.device].warp_size; +@@ -1935,6 +2051,21 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml + using fattn_kernel_ptr_t = fattn_kernel_t; + #endif // defined(GGML_USE_HIP) + fattn_kernel_t fattn_kernel; ++ if constexpr (longseq_instantiated) { ++ if (longseq) { ++ GGML_ASSERT(logit_softcap == 0.0f); ++ fattn_kernel = flash_attn_ext_f16; ++ ++ static bool shared_memory_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; ++ if (!shared_memory_limit_raised[id]) { ++ CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); ++ shared_memory_limit_raised[id] = true; ++ } ++ launch_fattn ++ (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, warp_size_host); ++ return; ++ } ++ } + if (logit_softcap == 0.0f) { + constexpr bool use_logit_softcap = false; + fattn_kernel = flash_attn_ext_f16; +@@ -2003,6 +2134,12 @@ DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(112, 112, 64) + DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(128, 128, 64) + DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(256, 256, 64) + ++#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) ++// Long-sequence variant only: ++extern DECL_FATTN_MMA_F16_CASE( 64, 64, 128, 1); ++extern DECL_FATTN_MMA_F16_CASE(128, 128, 128, 1); ++#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) ++ + extern DECL_FATTN_MMA_F16_CASE(512, 512, 4, 2); + extern DECL_FATTN_MMA_F16_CASE(512, 512, 8, 2); + extern DECL_FATTN_MMA_F16_CASE(512, 512, 16, 2); +diff --git a/src/ggml-cuda/fattn.cu b/src/ggml-cuda/fattn.cu +index ab7a3b29..00e9554a 100644 +--- a/src/ggml-cuda/fattn.cu ++++ b/src/ggml-cuda/fattn.cu +@@ -30,6 +30,15 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con + return; + } + ++#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) ++ if constexpr (ncols2 == 1 && DKQ == DV && (DKQ == 64 || DKQ == 128)) { ++ if (Q->ne[1] >= 1024 && ggml_cuda_fattn_mma_longseq_ncols(cc) == 128 && ggml_cuda_fattn_mma_use_longseq(cc, dst)) { ++ ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); ++ return; ++ } ++ } ++#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) ++ + ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); + } + +diff --git a/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_128-ncols2_1.cu b/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_128-ncols2_1.cu +new file mode 100644 +index 00000000..b9bf54a1 +--- /dev/null ++++ b/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_128-ncols2_1.cu +@@ -0,0 +1,8 @@ ++// Long-sequence FlashAttention variant (CUDA only), see ggml_cuda_fattn_mma_use_longseq in fattn-mma-f16.cuh. ++ ++#include "../fattn-mma-f16.cuh" ++ ++#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) ++DECL_FATTN_MMA_F16_CASE(64, 64, 128, 1); ++DECL_FATTN_MMA_F16_CASE(128, 128, 128, 1); ++#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp +index fdf505cd..dd290caf 100644 +--- a/tests/test-backend-ops.cpp ++++ b/tests/test-backend-ops.cpp +@@ -9588,6 +9588,21 @@ static std::vector> make_test_cases_eval() { + } + } + ++ // diffusion-shaped cases without a mask (CUDA long-sequence FA variant): KV lengths that are not a multiple ++ // of the KV tile (partial last tile), several Q tiles per head, and with/without GQA broadcast. ++ for (int hs : { 64, 128, }) { ++ for (int kv : { 260, 1000, 4100, }) { ++ for (int nb : { 64, 200, 1000, }) { ++ for (int nr2 : { 1, 4, }) { ++ for (ggml_type type_KV : { GGML_TYPE_F16, GGML_TYPE_F32, }) { ++ if (type_KV == GGML_TYPE_F32 && (nr2 != 1 || nb != 200)) continue; ++ test_cases.emplace_back(new test_flash_attn_ext(hs, hs, 4, {nr2, 1}, kv, nb, false, false, 0, 0, GGML_PREC_F32, type_KV, type_KV)); ++ } ++ } ++ } ++ } ++ } ++ + // mixed quant and Q1_0 test cases + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_F16)); +-- +2.43.0 + + +From 7f7c132c57dbfd8fd44a13dd29031c6724589efc Mon Sep 17 00:00:00 2001 +From: Daniel Han +Date: Thu, 1 Oct 2026 06:27:44 -0700 +Subject: [PATCH 2/3] ggml-cuda: per-architecture tile width for the + long-sequence FlashAttention variant + +Use 128 Q columns per block (8 warps sharing each K/V tile) by default for head +size 128 on sm80 and sm89, where it was measured faster at MiniMax-H3 DiT shapes +(A100: 75 -> 70 ms, L4: 460 -> 365 ms per 19108-token, 56-head call); keep 64 +elsewhere (neutral on sm100/sm120, slower for head size 64 on A100). +GGML_CUDA_FA_LONGSEQ_NCOLS=64|128 overrides. +--- + src/ggml-cuda/fattn-mma-f16.cuh | 13 ++++++++----- + src/ggml-cuda/fattn.cu | 2 +- + 2 files changed, 9 insertions(+), 6 deletions(-) + +diff --git a/src/ggml-cuda/fattn-mma-f16.cuh b/src/ggml-cuda/fattn-mma-f16.cuh +index b4496734..57f284c6 100644 +--- a/src/ggml-cuda/fattn-mma-f16.cuh ++++ b/src/ggml-cuda/fattn-mma-f16.cuh +@@ -1965,9 +1965,11 @@ static bool ggml_cuda_fattn_mma_longseq_arch_default(const int cc) { + return true; + } + +-// Q columns per CUDA block of the long-sequence variant for long Q: 64 (the stock tile) or 128 (8 warps sharing each K/V tile). +-// GGML_CUDA_FA_LONGSEQ_NCOLS=64|128 overrides the per-architecture default. +-static int ggml_cuda_fattn_mma_longseq_ncols(const int cc) { ++// Q columns per CUDA block of the long-sequence variant for long Q: 64 (the stock tile) or 128 (8 warps sharing each K/V tile, ++// half the K/V traffic per FLOP). GGML_CUDA_FA_LONGSEQ_NCOLS=64|128 overrides the per-architecture default. ++// Measured at MiniMax-H3 DiT shapes (19108 tokens, 56 heads x 128): 128 is faster on A100 (sm80) and L4 (sm89), neutral on ++// RTX PRO 6000 Blackwell (sm120) and B200 (sm100), and slower for head size 64 on A100, so it is only the default there. ++static int ggml_cuda_fattn_mma_longseq_ncols(const int cc, const int DKQ) { + static const int env = [] { + const char * e = getenv("GGML_CUDA_FA_LONGSEQ_NCOLS"); + return e != nullptr ? atoi(e) : -1; +@@ -1975,8 +1977,9 @@ static int ggml_cuda_fattn_mma_longseq_ncols(const int cc) { + if (env == 64 || env == 128) { + return env; + } +- GGML_UNUSED(cc); +- return 64; ++ const int arch = ggml_cuda_highest_compiled_arch(cc); ++ return DKQ == 128 && (cc == GGML_CUDA_CC_AMPERE || cc == GGML_CUDA_CC_ADA_LOVELACE) && ++ (arch == GGML_CUDA_CC_AMPERE || arch == GGML_CUDA_CC_ADA_LOVELACE) ? 128 : 64; + } + + static bool ggml_cuda_fattn_mma_use_longseq(const int cc, const ggml_tensor * dst) { +diff --git a/src/ggml-cuda/fattn.cu b/src/ggml-cuda/fattn.cu +index 00e9554a..b8be30a0 100644 +--- a/src/ggml-cuda/fattn.cu ++++ b/src/ggml-cuda/fattn.cu +@@ -32,7 +32,7 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con + + #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + if constexpr (ncols2 == 1 && DKQ == DV && (DKQ == 64 || DKQ == 128)) { +- if (Q->ne[1] >= 1024 && ggml_cuda_fattn_mma_longseq_ncols(cc) == 128 && ggml_cuda_fattn_mma_use_longseq(cc, dst)) { ++ if (Q->ne[1] >= 1024 && ggml_cuda_fattn_mma_longseq_ncols(cc, DKQ) == 128 && ggml_cuda_fattn_mma_use_longseq(cc, dst)) { + ggml_cuda_flash_attn_ext_mma_f16_case(ctx, dst); + return; + } +-- +2.43.0 + + +From e6e78b0ba7477be2b0e6869e19bf86ac89b0381b Mon Sep 17 00:00:00 2001 +From: Daniel Han +Date: Thu, 1 Oct 2026 07:36:53 -0700 +Subject: [PATCH 3/3] ggml-cuda: long-sequence FlashAttention takes the stock + grid; 128-column tile opt-in only + +The stream-k grid is sized from occupancy, and the pipelined variant can have a +different occupancy than the stock kernel (two-stage shared memory), which +moves the stream-k seams and with them the order in which partial results are +combined (measured on L4: relative L2 6e-4 at the op, 36 dB PSNR end to end +over 4 diffusion steps). Size the grid from the stock kernel of the same tile +instead, so the pipelined variant is a bit-identical drop-in on every GPU. + +The 128-column tile cannot match the stock seams, so it is no longer a default +on any architecture; GGML_CUDA_FA_LONGSEQ_NCOLS=128 opts in. +--- + src/ggml-cuda/fattn-common.cuh | 8 +++++-- + src/ggml-cuda/fattn-mma-f16.cuh | 39 +++++++++++++++++++++++---------- + 2 files changed, 33 insertions(+), 14 deletions(-) + +diff --git a/src/ggml-cuda/fattn-common.cuh b/src/ggml-cuda/fattn-common.cuh +index e67cc7fd..8a8d4023 100644 +--- a/src/ggml-cuda/fattn-common.cuh ++++ b/src/ggml-cuda/fattn-common.cuh +@@ -972,7 +972,10 @@ static __global__ void flash_attn_combine_results( + template + void launch_fattn( + ggml_backend_cuda_context & ctx, ggml_tensor * dst, fattn_kernel_t fattn_kernel, const int nwarps, const size_t nbytes_shared, +- const int nbatch_fa, const bool need_f16_K, const bool need_f16_V, const bool stream_k, const int warp_size = WARP_SIZE ++ const int nbatch_fa, const bool need_f16_K, const bool need_f16_V, const bool stream_k, const int warp_size = WARP_SIZE, ++ // Optional: size the grid from the occupancy of this kernel instead of fattn_kernel. A drop-in kernel variant launched ++ // with the grid of the kernel it replaces splits the work (stream-k seams) identically and so gives identical results. ++ fattn_kernel_t occupancy_kernel = nullptr, const size_t occupancy_nbytes_shared = 0 + ) { + constexpr int ncols = ncols1 * ncols2; + +@@ -1110,7 +1113,8 @@ void launch_fattn( + + const dim3 block_dim(warp_size, nwarps, 1); + int max_blocks_per_sm = 1; // Max. number of active blocks limited by occupancy. +- CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&max_blocks_per_sm, fattn_kernel, block_dim.x * block_dim.y * block_dim.z, nbytes_shared)); ++ CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&max_blocks_per_sm, occupancy_kernel ? occupancy_kernel : fattn_kernel, ++ block_dim.x * block_dim.y * block_dim.z, occupancy_kernel ? occupancy_nbytes_shared : nbytes_shared)); + GGML_ASSERT(max_blocks_per_sm > 0); + int parallel_blocks = max_blocks_per_sm; + +diff --git a/src/ggml-cuda/fattn-mma-f16.cuh b/src/ggml-cuda/fattn-mma-f16.cuh +index 57f284c6..d5fd8eea 100644 +--- a/src/ggml-cuda/fattn-mma-f16.cuh ++++ b/src/ggml-cuda/fattn-mma-f16.cuh +@@ -1965,21 +1965,18 @@ static bool ggml_cuda_fattn_mma_longseq_arch_default(const int cc) { + return true; + } + +-// Q columns per CUDA block of the long-sequence variant for long Q: 64 (the stock tile) or 128 (8 warps sharing each K/V tile, +-// half the K/V traffic per FLOP). GGML_CUDA_FA_LONGSEQ_NCOLS=64|128 overrides the per-architecture default. +-// Measured at MiniMax-H3 DiT shapes (19108 tokens, 56 heads x 128): 128 is faster on A100 (sm80) and L4 (sm89), neutral on +-// RTX PRO 6000 Blackwell (sm120) and B200 (sm100), and slower for head size 64 on A100, so it is only the default there. ++// Q columns per CUDA block of the long-sequence variant for long Q: 64 (the stock tile, default) or 128 (8 warps sharing ++// each K/V tile, half the K/V traffic per FLOP). 128 is opt-in via GGML_CUDA_FA_LONGSEQ_NCOLS=128: it was measured faster at ++// MiniMax-H3 DiT shapes on A100 (sm80, 75 -> 70 ms) and L4 (sm89, 460 -> 365 ms) but splits the work differently from the ++// stock kernel, so its results are not bit-identical. + static int ggml_cuda_fattn_mma_longseq_ncols(const int cc, const int DKQ) { + static const int env = [] { + const char * e = getenv("GGML_CUDA_FA_LONGSEQ_NCOLS"); + return e != nullptr ? atoi(e) : -1; + }(); +- if (env == 64 || env == 128) { +- return env; +- } +- const int arch = ggml_cuda_highest_compiled_arch(cc); +- return DKQ == 128 && (cc == GGML_CUDA_CC_AMPERE || cc == GGML_CUDA_CC_ADA_LOVELACE) && +- (arch == GGML_CUDA_CC_AMPERE || arch == GGML_CUDA_CC_ADA_LOVELACE) ? 128 : 64; ++ GGML_UNUSED(cc); ++ GGML_UNUSED(DKQ); ++ return env == 128 ? 128 : 64; + } + + static bool ggml_cuda_fattn_mma_use_longseq(const int cc, const ggml_tensor * dst) { +@@ -2064,8 +2061,26 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml + CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(fattn_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_shared_total)); + shared_memory_limit_raised[id] = true; + } +- launch_fattn +- (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, warp_size_host); ++ if (ncols1 <= 64) { ++ // Same grid as the stock kernel of this tile size: identical stream-k seams, so results are bit-identical. ++ const int nstages_stock = ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2, cc); ++ const size_t nbytes_KV_stock = nstages_stock <= 1 ? nbytes_shared_KV_1stage : nbytes_shared_KV_2stage; ++ const size_t nbytes_total_stock = std::max(nbytes_shared_combine, Q_in_reg ? ++ std::max(nbytes_shared_Q, nbytes_KV_stock + nbytes_shared_mask) : ++ nbytes_shared_Q + nbytes_KV_stock + nbytes_shared_mask); ++ fattn_kernel_t stock_kernel = flash_attn_ext_f16; ++ static bool stock_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; ++ if (!stock_limit_raised[id]) { ++ CUDA_CHECK(cudaFuncSetAttribute(reinterpret_cast(stock_kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes_total_stock)); ++ stock_limit_raised[id] = true; ++ } ++ launch_fattn ++ (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, warp_size_host, ++ stock_kernel, nbytes_total_stock); ++ } else { ++ launch_fattn ++ (ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, warp_size_host); ++ } + return; + } + } +-- +2.43.0 + diff --git a/scripts/unsloth/ggml-patches/0003-ggml-cuda-h3-vae-epilogue-fusion.patch b/scripts/unsloth/ggml-patches/0003-ggml-cuda-h3-vae-epilogue-fusion.patch new file mode 100644 index 0000000000..1e14c53584 --- /dev/null +++ b/scripts/unsloth/ggml-patches/0003-ggml-cuda-h3-vae-epilogue-fusion.patch @@ -0,0 +1,348 @@ +From 76b49e564c23b99758642d5bc5013d09774a1e24 Mon Sep 17 00:00:00 2001 +From: Daniel Han +Date: Thu, 1 Oct 2026 05:46:18 -0700 +Subject: [PATCH] ggml-cuda: fuse bias epilogues into F16 cuBLAS matmuls, + narrow blocks for short-row RMS norm + +An F16 matmul on the cuBLAS path writes an F16 temporary, converts it to F32, and the +graph then adds a bias (and in transformer blocks multiplies by a per-channel scale and +adds the residual, or applies SWIGLU) in separate passes over the F32 result. When those +nodes directly follow the matmul and have no other users, they now run inside the F16 -> F32 +conversion: MUL_MAT+ADD, MUL_MAT+ADD+MUL+ADD and MUL_MAT+ADD+GLU(SWIGLU). Every element +goes through the same rounded operations in the same order (explicit _rn intrinsics stop +nvcc from contracting them into FMAs), so results are bit-identical. +GGML_CUDA_CUBLAS_EPILOGUE_FUSION=0 disables it. + +RMS norm over rows of at most 128 columns (per-head q/k norms) launched 256-thread blocks +in which most threads had no element. It now uses the next multiple of 32 threads: each +thread still owns the same element and each warp the same partial sum, and the wider block +only added zero partials, so the result is unchanged. GGML_CUDA_NORM_SMALL_ROWS=0 restores +the 256-thread launch. +--- + src/ggml-cuda/ggml-cuda.cu | 200 ++++++++++++++++++++++++++++++++++++- + src/ggml-cuda/norm.cu | 63 ++++++++++++ + 2 files changed, 262 insertions(+), 1 deletion(-) + +diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu +index 517d089..0e61996 100644 +--- a/src/ggml-cuda/ggml-cuda.cu ++++ b/src/ggml-cuda/ggml-cuda.cu +@@ -2017,8 +2017,73 @@ struct batched_mul_mat_traits { + static inline auto convert_nc(ggml_type src_type) { return ggml_get_to_fp16_nc_cuda(src_type); } + }; + ++// Elementwise ops that directly follow an F16 cuBLAS matmul, applied while converting its F16 ++// output to F32 instead of in separate passes over the F32 result. Each element sees the same ++// rounded operations in the same order as the unfused ADD / MUL / ADD or ADD / SWIGLU nodes ++// (explicit _rn intrinsics keep nvcc from contracting them into FMAs), so results are bit-identical. ++enum ggml_cuda_cublas_epilogue_mode { ++ GGML_CUDA_CUBLAS_EPILOGUE_BIAS, // dst = mm + bias ++ GGML_CUDA_CUBLAS_EPILOGUE_BIAS_SCALE_RESIDUAL, // dst = residual + (mm + bias) * scale ++ GGML_CUDA_CUBLAS_EPILOGUE_BIAS_SWIGLU, // dst = silu(mm[:n] + bias[:n]) * (mm[n:] + bias[n:]) ++}; ++ ++struct ggml_cuda_cublas_epilogue { ++ ggml_cuda_cublas_epilogue_mode mode; ++ const float * bias; ++ const float * scale; ++ const float * residual; ++ float * dst; ++ int64_t ne0_out; ++ int64_t n_out; ++}; ++ ++template ++static __global__ void k_cublas_f16_epilogue(const half * __restrict__ mm, const float * __restrict__ bias, ++ const float * __restrict__ scale, const float * residual, float * dst, ++ const int64_t ne0_out, const int64_t n_out) { ++ const int64_t i = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; ++ if (i >= n_out) { ++ return; ++ } ++ const int64_t row = i / ne0_out; ++ const int64_t col = i - row*ne0_out; ++ if constexpr (mode == GGML_CUDA_CUBLAS_EPILOGUE_BIAS_SWIGLU) { ++ const int64_t j = row*2*ne0_out + col; ++ const float x = __fadd_rn(__half2float(mm[j]), bias[col]); ++ const float g = __fadd_rn(__half2float(mm[j + ne0_out]), bias[col + ne0_out]); ++ dst[i] = __fmul_rn(ggml_cuda_op_silu_single(x), g); ++ } else { ++ float v = __fadd_rn(__half2float(mm[i]), bias[col]); ++ if constexpr (mode == GGML_CUDA_CUBLAS_EPILOGUE_BIAS_SCALE_RESIDUAL) { ++ v = __fadd_rn(residual[i], __fmul_rn(v, scale[col])); ++ } ++ dst[i] = v; ++ } ++} ++ ++static void ggml_cuda_cublas_f16_epilogue(const half * mm, const ggml_cuda_cublas_epilogue & epi, cudaStream_t stream) { ++ const int block = 256; ++ const int64_t grid = (epi.n_out + block - 1) / block; ++ switch (epi.mode) { ++ case GGML_CUDA_CUBLAS_EPILOGUE_BIAS: ++ k_cublas_f16_epilogue<<>>( ++ mm, epi.bias, epi.scale, epi.residual, epi.dst, epi.ne0_out, epi.n_out); ++ break; ++ case GGML_CUDA_CUBLAS_EPILOGUE_BIAS_SCALE_RESIDUAL: ++ k_cublas_f16_epilogue<<>>( ++ mm, epi.bias, epi.scale, epi.residual, epi.dst, epi.ne0_out, epi.n_out); ++ break; ++ case GGML_CUDA_CUBLAS_EPILOGUE_BIAS_SWIGLU: ++ k_cublas_f16_epilogue<<>>( ++ mm, epi.bias, epi.scale, epi.residual, epi.dst, epi.ne0_out, epi.n_out); ++ break; ++ } ++ CUDA_CHECK(cudaGetLastError()); ++} ++ + template +-static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { ++static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, ++ const ggml_cuda_cublas_epilogue * epilogue = nullptr) { + using traits = batched_mul_mat_traits; + using cuda_t = typename traits::cuda_type; + +@@ -2226,8 +2291,17 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const + + // Convert output back to F32 if needed + if (cu_data_type != CUDA_R_32F) { ++ if constexpr (compute_type == GGML_TYPE_F16) { ++ if (epilogue != nullptr) { ++ ggml_cuda_cublas_f16_epilogue((const half *) dst_temp.get(), *epilogue, main_stream); ++ return; ++ } ++ } ++ GGML_ASSERT(epilogue == nullptr); + const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(traits::ggml_type_val); + to_fp32_cuda(dst_temp.get(), dst_ddf, ne_dst, main_stream); ++ } else { ++ GGML_ASSERT(epilogue == nullptr); + } + } + +@@ -3838,6 +3912,123 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, + } + + // try and fuse nodes and return the number of nodes to skip ++static bool ggml_cuda_cublas_epilogue_enabled() { ++ static const bool enabled = [] { ++ const char * e = getenv("GGML_CUDA_CUBLAS_EPILOGUE_FUSION"); ++ return e == nullptr || atoi(e) != 0; ++ }(); ++ return enabled; ++} ++ ++// True when ggml_cuda_mul_mat sends this F16 matmul to cuBLAS with F16 compute and an F16 ++// temporary output (the path the epilogue fusion replaces the F16 -> F32 conversion of). ++static bool ggml_cuda_mul_mat_uses_cublas_f16_temp(const ggml_backend_cuda_context & ctx, const ggml_tensor * mm) { ++ const ggml_tensor * src0 = mm->src[0]; ++ const ggml_tensor * src1 = mm->src[1]; ++ if (src0->type != GGML_TYPE_F16 || src1->type != GGML_TYPE_F32 || mm->type != GGML_TYPE_F32 || ++ !ggml_is_contiguous(mm) || ggml_get_op_params_i32(mm, 1) == GGML_HINT_SRC0_IS_HADAMARD || ++ mm->op_params[0] == GGML_PREC_F32 || getenv("GGML_CUDA_CUBLAS_COMPUTE_TYPE") != nullptr) { ++ return false; ++ } ++ const int cc = ggml_cuda_info().devices[ctx.device].cc; ++ const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size; ++ if (!fast_fp16_hardware_available(cc) || cc == GGML_CUDA_CC_VOLTA || GGML_CUDA_CC_IS_RDNA4(cc) || GGML_CUDA_CC_IS_CDNA(cc)) { ++ return false; ++ } ++ const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE ++ && ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && src0->view_src; ++ if (bad_padding_clear) { ++ return true; ++ } ++ const int64_t ne11 = src1->ne[1]; ++ return !ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, ne11) && ++ !ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, ne11, /*mul_mat_id =*/ false); ++} ++ ++static const ggml_tensor * ggml_cuda_other_src(const ggml_tensor * node, const ggml_tensor * operand) { ++ if (node->src[0] == operand) { ++ return node->src[1]; ++ } ++ if (node->src[1] == operand) { ++ return node->src[0]; ++ } ++ return nullptr; ++} ++ ++// Like ggml_node_has_n_uses(cgraph, idx, 1), but also accepts an in-place node whose view ++// source is the (single-use) node before it, e.g. the bias ADD that ggml_add_inplace builds. ++static bool ggml_cuda_single_use_intermediate(const ggml_cgraph * cgraph, int idx, const ggml_tensor * prev) { ++ const ggml_tensor * node = cgraph->nodes[idx]; ++ if (ggml_node_get_use_count(cgraph, idx) != 1 || (node->flags & GGML_TENSOR_FLAG_OUTPUT)) { ++ return false; ++ } ++ return node->view_src == nullptr || node->view_src == prev; ++} ++ ++static bool ggml_cuda_is_row_vector(const ggml_tensor * t, int64_t ne0) { ++ return t != nullptr && t->type == GGML_TYPE_F32 && ggml_is_contiguous(t) && t->ne[0] == ne0 && ggml_nelements(t) == ne0; ++} ++ ++// MUL_MAT (F16 weights, cuBLAS) followed by a row bias ADD, optionally then MUL by a row scale and ++// ADD of a residual, or SWIGLU. Returns the number of extra nodes consumed, 0 if not fused. ++static int ggml_cuda_try_fuse_cublas_epilogue(ggml_backend_cuda_context & ctx, ggml_cgraph * cgraph, int i) { ++ if (!ggml_cuda_cublas_epilogue_enabled() || i + 1 >= cgraph->n_nodes) { ++ return 0; ++ } ++ ggml_tensor * mm = cgraph->nodes[i]; ++ ggml_tensor * add = cgraph->nodes[i + 1]; ++ if (mm->op != GGML_OP_MUL_MAT || add->op != GGML_OP_ADD || (add->flags & GGML_TENSOR_FLAG_COMPUTE) == 0 || ++ !ggml_node_has_n_uses(cgraph, i, 1) || !ggml_are_same_shape(mm, add) || add->type != GGML_TYPE_F32 || ++ !ggml_is_contiguous(add)) { ++ return 0; ++ } ++ const ggml_tensor * bias = ggml_cuda_other_src(add, mm); ++ if (!ggml_cuda_is_row_vector(bias, mm->ne[0]) || !ggml_cuda_mul_mat_uses_cublas_f16_temp(ctx, mm)) { ++ return 0; ++ } ++ ++ ggml_cuda_cublas_epilogue epi = {}; ++ epi.mode = GGML_CUDA_CUBLAS_EPILOGUE_BIAS; ++ epi.bias = (const float *) bias->data; ++ epi.dst = (float *) add->data; ++ epi.ne0_out = add->ne[0]; ++ epi.n_out = ggml_nelements(add); ++ int consumed = 1; ++ ++ if (i + 2 < cgraph->n_nodes && ggml_cuda_single_use_intermediate(cgraph, i + 1, mm)) { ++ ggml_tensor * mul = cgraph->nodes[i + 2]; ++ ggml_tensor * res = i + 3 < cgraph->n_nodes ? cgraph->nodes[i + 3] : nullptr; ++ const ggml_tensor * glu = cgraph->nodes[i + 2]; ++ if (res != nullptr && mul->op == GGML_OP_MUL && res->op == GGML_OP_ADD && (mul->flags & GGML_TENSOR_FLAG_COMPUTE) && ++ (res->flags & GGML_TENSOR_FLAG_COMPUTE) && ggml_node_has_n_uses(cgraph, i + 2, 1) && ++ ggml_are_same_shape(mul, add) && ggml_are_same_shape(res, add) && mul->type == GGML_TYPE_F32 && ++ res->type == GGML_TYPE_F32 && ggml_is_contiguous(res)) { ++ const ggml_tensor * scale = ggml_cuda_other_src(mul, add); ++ const ggml_tensor * residual = ggml_cuda_other_src(res, mul); ++ if (mul->src[0] == add && ggml_cuda_is_row_vector(scale, add->ne[0]) && residual != nullptr && ++ residual->type == GGML_TYPE_F32 && ggml_is_contiguous(residual) && ggml_are_same_shape(residual, res)) { ++ epi.mode = GGML_CUDA_CUBLAS_EPILOGUE_BIAS_SCALE_RESIDUAL; ++ epi.scale = (const float *) scale->data; ++ epi.residual = (const float *) residual->data; ++ epi.dst = (float *) res->data; ++ consumed = 3; ++ } ++ } else if (glu->op == GGML_OP_GLU && (glu->flags & GGML_TENSOR_FLAG_COMPUTE) && glu->src[0] == add && ++ glu->src[1] == nullptr && ggml_get_glu_op(glu) == GGML_GLU_OP_SWIGLU && ++ ggml_get_op_params_i32(glu, 1) == 0 && glu->type == GGML_TYPE_F32 && ggml_is_contiguous(glu) && ++ add->ne[0] % 2 == 0 && glu->ne[0] * 2 == add->ne[0] && ggml_nrows(glu) == ggml_nrows(add)) { ++ epi.mode = GGML_CUDA_CUBLAS_EPILOGUE_BIAS_SWIGLU; ++ epi.dst = (float *) glu->data; ++ epi.ne0_out = glu->ne[0]; ++ epi.n_out = ggml_nelements(glu); ++ consumed = 2; ++ } ++ } ++ ++ ggml_cuda_mul_mat_cublas_impl(ctx, mm->src[0], mm->src[1], mm, &epi); ++ return consumed; ++} ++ + static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { + + static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); +@@ -4523,6 +4714,13 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph + return fused_node_count - 1; + } + ++ if (node->op == GGML_OP_MUL_MAT) { ++ const int consumed = ggml_cuda_try_fuse_cublas_epilogue(*cuda_ctx, cgraph, i); ++ if (consumed > 0) { ++ return consumed; ++ } ++ } ++ + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD }, {})) { + ggml_cuda_op_rms_norm_fused_add(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]); + return 2; +diff --git a/src/ggml-cuda/norm.cu b/src/ggml-cuda/norm.cu +index c3758cd..9c9385f 100644 +--- a/src/ggml-cuda/norm.cu ++++ b/src/ggml-cuda/norm.cu +@@ -301,10 +301,68 @@ static void group_norm_f32_cuda( + } + } + ++// Rows of at most 128 columns (per-head q/k norms) leave most of a 256-thread block idle. A block ++// of the next multiple of 32 gives every thread the same element and every warp the same partial ++// sum; the 256-thread reduction only adds zero partials on top, so the result is bit-identical. ++// GGML_CUDA_NORM_SMALL_ROWS=0 restores the 256-thread launch. ++static bool rms_norm_small_rows_enabled() { ++#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) ++ return false; ++#else ++ static const bool enabled = [] { ++ const char * e = getenv("GGML_CUDA_NORM_SMALL_ROWS"); ++ return e == nullptr || atoi(e) != 0; ++ }(); ++ return enabled && WARP_SIZE == 32; ++#endif ++} ++ ++template ++static void rms_norm_f32_launch_small( ++ const float * x, float * dst, const int ncols, const dim3 blocks_num, ++ const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, ++ const float * mul, const int64_t mul_stride_row, const int64_t mul_stride_channel, const int64_t mul_stride_sample, ++ const uint3 mul_ncols_packed, const uint3 mul_nrows_packed, const uint3 mul_nchannels_packed, const uint3 mul_nsamples_packed, ++ cudaStream_t stream) { ++ const dim3 block_dims(block_size, 1, 1); ++ const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, block_size > WARP_SIZE ? 32 * sizeof(float) : 0, stream}; ++ ggml_cuda_kernel_launch(rms_norm_f32, launch_params, ++ x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, ++ mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, ++ nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); ++} ++ ++template ++static bool rms_norm_f32_try_small( ++ const float * x, float * dst, const int ncols, const dim3 blocks_num, ++ const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, ++ const float * mul, const int64_t mul_stride_row, const int64_t mul_stride_channel, const int64_t mul_stride_sample, ++ const uint3 mul_ncols_packed, const uint3 mul_nrows_packed, const uint3 mul_nchannels_packed, const uint3 mul_nsamples_packed, ++ cudaStream_t stream) { ++ if (ncols > 128 || !rms_norm_small_rows_enabled()) { ++ return false; ++ } ++ if (ncols <= 32) { ++ rms_norm_f32_launch_small<32, do_multiply>(x, dst, ncols, blocks_num, stride_row, stride_channel, stride_sample, eps, mul, ++ mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, stream); ++ } else if (ncols <= 64) { ++ rms_norm_f32_launch_small<64, do_multiply>(x, dst, ncols, blocks_num, stride_row, stride_channel, stride_sample, eps, mul, ++ mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, stream); ++ } else { ++ rms_norm_f32_launch_small<128, do_multiply>(x, dst, ncols, blocks_num, stride_row, stride_channel, stride_sample, eps, mul, ++ mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, stream); ++ } ++ return true; ++} ++ + static void rms_norm_f32_cuda( + const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples, + const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, cudaStream_t stream) { + const dim3 blocks_num(nrows, nchannels, nsamples); ++ if (rms_norm_f32_try_small(x, dst, ncols, blocks_num, stride_row, stride_channel, stride_sample, eps, ++ nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), stream)) { ++ return; ++ } + if (ncols < 1024) { + const dim3 block_dims(256, 1, 1); + const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; +@@ -360,6 +418,11 @@ static void rms_norm_mul_f32_cuda(const float * x, + const uint3 mul_nrows_packed = init_fastdiv_values(mul_nrows); + const uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels); + const uint3 mul_nsamples_packed = init_fastdiv_values(mul_nsamples); ++ if (rms_norm_f32_try_small(x, dst, ncols, blocks_num, stride_row, stride_channel, stride_sample, eps, ++ mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, ++ mul_nchannels_packed, mul_nsamples_packed, stream)) { ++ return; ++ } + if (ncols < 1024) { + const dim3 block_dims(256, 1, 1); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; +-- +2.43.0 + diff --git a/scripts/unsloth/ggml-patches/0004-ggml-h3-dit-fused-ops.patch b/scripts/unsloth/ggml-patches/0004-ggml-h3-dit-fused-ops.patch new file mode 100644 index 0000000000..689dfa5f67 --- /dev/null +++ b/scripts/unsloth/ggml-patches/0004-ggml-h3-dit-fused-ops.patch @@ -0,0 +1,1020 @@ +From 04e9fe6c85b680c2f8e29c22e084dab6012babbe Mon Sep 17 00:00:00 2001 +From: Daniel Han +Date: Thu, 1 Oct 2026 06:48:01 -0700 +Subject: [PATCH] ggml : fused value-preserving ops for adaLN DiT blocks + +Three ops that replace long chains of elementwise ops, copies and concats in +diffusion transformers such as MiniMax-H3 while producing the same bits as the +chains they replace (every product, sum, scale and conversion is rounded +separately, never contracted into an FMA): + +- ggml_rope_pe_permute: rotary embedding with a precomputed [cos, -sin; sin, + cos] table on the first n_rot channels (non-interleaved pairing), pass-through + for the rest, written in the head-major [d, tokens, heads] layout attention + reads, with an optional scale and F16 cast. Reads a strided view, so the + q/k/v chunks of a fused qkv projection need no copies. +- ggml_modulate_rows / ggml_gated_add_rows: (x + x*scale) + shift and + base + u*gate on a row range, written in place into an accumulator so a + sequence split into segments with different modulation vectors needs no + slice copies and no concat. Optional output / input scale. +- ggml_swiglu_scaled: silu(s_in*a) * (s_in*b) scaled by s_out, i.e. a Linear + output scale, swiglu and the next Linear input scale in one pass. + +CPU and CUDA implementations; other backends report the ops unsupported so +callers keep their unfused graphs. test-backend-ops covers strided inputs, +partial and full rotation, F32/F16 outputs, chained segments and the scales. +--- + include/ggml.h | 62 ++++++++++++ + src/ggml-backend-meta.cpp | 3 +- + src/ggml-cpu/ggml-cpu.c | 15 +++ + src/ggml-cpu/ops.cpp | 154 +++++++++++++++++++++++++++++ + src/ggml-cpu/ops.h | 3 + + src/ggml-cuda/ggml-cuda.cu | 17 ++++ + src/ggml-cuda/modulate-rows.cu | 168 ++++++++++++++++++++++++++++++++ + src/ggml-cuda/modulate-rows.cuh | 7 ++ + src/ggml-cuda/rope-pe.cu | 102 +++++++++++++++++++ + src/ggml-cuda/rope-pe.cuh | 5 + + src/ggml.c | 129 +++++++++++++++++++++++- + tests/test-backend-ops.cpp | 139 ++++++++++++++++++++++++++ + 12 files changed, 801 insertions(+), 3 deletions(-) + create mode 100644 src/ggml-cuda/modulate-rows.cu + create mode 100644 src/ggml-cuda/modulate-rows.cuh + create mode 100644 src/ggml-cuda/rope-pe.cu + create mode 100644 src/ggml-cuda/rope-pe.cuh + +diff --git a/include/ggml.h b/include/ggml.h +index 991cc4fc..153eed2a 100644 +--- a/include/ggml.h ++++ b/include/ggml.h +@@ -594,6 +594,9 @@ extern "C" { + + GGML_OP_QUANTIZE_I8_CONVROT, + GGML_OP_SAGE_ATTN, ++ GGML_OP_ROPE_PE_PERMUTE, ++ GGML_OP_MODULATE_ROWS, ++ GGML_OP_SWIGLU_SCALED, + + GGML_OP_COUNT, + }; +@@ -2461,6 +2464,65 @@ extern "C" { + float scale, + enum ggml_sage_attn_mode mode); + ++ // Rotary embedding with a precomputed cos/sin table, fused with the head-major relayout ++ // attention kernels consume, an optional scale and an optional cast to F16. ++ // a: F32 [head_dim, n_head, n_tokens, n_batch], rows contiguous (other strides free) ++ // pe: F32 contiguous [2, 2, n_rot/2, n_tokens], pe[c, r, j, t] = M[r][c] of the 2x2 rotation ++ // for pair j of token t (non-interleaved pairing j <-> j + n_rot/2); NULL when n_rot == 0 ++ // result: `type` (F32 or F16) contiguous [head_dim, n_tokens, n_head, n_batch] ++ // i < n_rot: y = a[j]*pe[0, r, j, t] + a[j + n_rot/2]*pe[1, r, j, t] (r = i / (n_rot/2), j = i % (n_rot/2)) ++ // i >= n_rot: y = a[i] ++ // y = scale*y when scale != 1, then y is converted to `type`. ++ // Each product, the sum, the scale and the conversion are rounded separately, so the result ++ // matches the equivalent chain of ggml_mul/ggml_add/ggml_concat/ggml_scale/ggml_cpy bit for bit. ++ GGML_API struct ggml_tensor * ggml_rope_pe_permute( ++ struct ggml_context * ctx, ++ struct ggml_tensor * a, ++ struct ggml_tensor * pe, ++ int n_rot, ++ float scale, ++ enum ggml_type type); ++ ++ // Row-range modulation for adaLN-style blocks whose sequence is split into segments that use ++ // different modulation vectors. Writes rows [row0, row1) along dim 1 (all of dims 2 and 3) of the ++ // result and leaves every other row as it is in `acc`: chain one call per segment, passing the ++ // previous result as `acc` (written in place); the first call passes acc = NULL and gets a new ++ // tensor shaped like x / base. x, base, u: F32 with contiguous rows; scale, shift, gate: F32 ++ // with ne[0] == x->ne[0] elements, broadcast over rows. ++ // ggml_modulate_rows: y = (x + x*scale) + shift; then y = out_scale*y when out_scale != 1 ++ // ggml_gated_add_rows: w = u_scale*u when u_scale != 1 (else u); y = base + w*gate ++ // Each product, sum and scale is rounded separately, so the result equals the slice + ++ // ggml_mul/ggml_add/ggml_scale + ggml_concat chain bit for bit. ++ GGML_API struct ggml_tensor * ggml_modulate_rows( ++ struct ggml_context * ctx, ++ struct ggml_tensor * acc, ++ struct ggml_tensor * x, ++ struct ggml_tensor * scale, ++ struct ggml_tensor * shift, ++ int64_t row0, ++ int64_t row1, ++ float out_scale); ++ ++ GGML_API struct ggml_tensor * ggml_gated_add_rows( ++ struct ggml_context * ctx, ++ struct ggml_tensor * acc, ++ struct ggml_tensor * base, ++ struct ggml_tensor * u, ++ struct ggml_tensor * gate, ++ int64_t row0, ++ int64_t row1, ++ float u_scale); ++ ++ // swiglu on two separately scaled halves, then scaled: with a' = in_scale*a, b' = in_scale*b, ++ // y = out_scale*(silu(a')*b') (each scale skipped when it is 1). a, b: F32, same shape, rows ++ // contiguous. Matches ggml_scale -> ggml_swiglu_split -> ggml_scale bit for bit. ++ GGML_API struct ggml_tensor * ggml_swiglu_scaled( ++ struct ggml_context * ctx, ++ struct ggml_tensor * a, ++ struct ggml_tensor * b, ++ float in_scale, ++ float out_scale); ++ + GGML_API void ggml_flash_attn_ext_set_prec( + struct ggml_tensor * a, + enum ggml_prec prec); +diff --git a/src/ggml-backend-meta.cpp b/src/ggml-backend-meta.cpp +index 539b4f70..5a34e748 100644 +--- a/src/ggml-backend-meta.cpp ++++ b/src/ggml-backend-meta.cpp +@@ -151,7 +151,8 @@ static ggml_backend_buffer_type_t ggml_backend_meta_device_get_host_buffer_type( + + static bool ggml_backend_meta_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { + GGML_ASSERT(ggml_backend_dev_is_meta(dev)); +- if (op->op == GGML_OP_SAGE_ATTN) { ++ if (op->op == GGML_OP_SAGE_ATTN || op->op == GGML_OP_ROPE_PE_PERMUTE || ++ op->op == GGML_OP_MODULATE_ROWS || op->op == GGML_OP_SWIGLU_SCALED) { + return false; + } + const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; +diff --git a/src/ggml-cpu/ggml-cpu.c b/src/ggml-cpu/ggml-cpu.c +index 350777c1..d83f15f9 100644 +--- a/src/ggml-cpu/ggml-cpu.c ++++ b/src/ggml-cpu/ggml-cpu.c +@@ -2190,6 +2190,18 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm + { + ggml_compute_forward_glu(params, tensor); + } break; ++ case GGML_OP_ROPE_PE_PERMUTE: ++ { ++ ggml_compute_forward_rope_pe_permute(params, tensor); ++ } break; ++ case GGML_OP_MODULATE_ROWS: ++ { ++ ggml_compute_forward_modulate_rows(params, tensor); ++ } break; ++ case GGML_OP_SWIGLU_SCALED: ++ { ++ ggml_compute_forward_swiglu_scaled(params, tensor); ++ } break; + case GGML_OP_GET_REL_POS: + { + ggml_compute_forward_get_rel_pos(params, tensor); +@@ -2522,6 +2534,9 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { + case GGML_OP_ROPE: + case GGML_OP_ROPE_BACK: + case GGML_OP_ADD_REL_POS: ++ case GGML_OP_ROPE_PE_PERMUTE: ++ case GGML_OP_MODULATE_ROWS: ++ case GGML_OP_SWIGLU_SCALED: + { + n_tasks = n_threads; + } break; +diff --git a/src/ggml-cpu/ops.cpp b/src/ggml-cpu/ops.cpp +index 727a4058..40fe1949 100644 +--- a/src/ggml-cpu/ops.cpp ++++ b/src/ggml-cpu/ops.cpp +@@ -12,6 +12,7 @@ + #include + #include + #include ++#include + + struct ggml_fp8_e4m3_cpu { + uint8_t value; +@@ -12136,3 +12137,156 @@ void ggml_compute_forward_quantize_i8_convrot(const ggml_compute_params * params + memset(dst_data + rows * dst->nb[1], 0, (size_t)(rows_padded - rows) * dst->nb[1]); + } + } ++ ++// ggml_compute_forward_rope_pe_permute ++ ++void ggml_compute_forward_rope_pe_permute( ++ const ggml_compute_params * params, ++ ggml_tensor * dst) { ++ const ggml_tensor * src0 = dst->src[0]; ++ const ggml_tensor * pe = dst->src[1]; ++ ++ GGML_ASSERT(src0->type == GGML_TYPE_F32); ++ GGML_ASSERT(dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_F16); ++ ++ const int n_rot = ggml_get_op_params_i32(dst, 0); ++ const float scale = ggml_get_op_params_f32(dst, 1); ++ const int half = n_rot / 2; ++ ++ const int64_t D = src0->ne[0]; ++ const int64_t H = src0->ne[1]; ++ const int64_t L = src0->ne[2]; ++ const int64_t N = src0->ne[3]; ++ ++ const int64_t nrows = H * L * N; ++ const int64_t dr = (nrows + params->nth - 1) / params->nth; ++ const int64_t ir0 = dr * params->ith; ++ const int64_t ir1 = MIN(ir0 + dr, nrows); ++ ++ // Same rounding points as the unfused graph: each product, the sum, the scale and the ++ // conversion are separate passes, so no contraction into fused multiply-adds can happen. ++ std::vector p0(n_rot), p1(n_rot), y(D); ++ ++ for (int64_t ir = ir0; ir < ir1; ++ir) { ++ const int64_t h = ir % H; ++ const int64_t l = (ir / H) % L; ++ const int64_t n = ir / (H * L); ++ ++ const float * x = (const float *) ((const char *) src0->data + h * src0->nb[1] + l * src0->nb[2] + n * src0->nb[3]); ++ ++ if (n_rot > 0) { ++ const float * t = (const float *) pe->data + l * (int64_t) half * 4; ++ for (int i = 0; i < n_rot; ++i) { ++ const int r = i / half; ++ const int j = i - r * half; ++ p0[i] = x[j] * t[4 * j + 2 * r + 0]; ++ } ++ for (int i = 0; i < n_rot; ++i) { ++ const int r = i / half; ++ const int j = i - r * half; ++ p1[i] = x[j + half] * t[4 * j + 2 * r + 1]; ++ } ++ ggml_vec_add_f32(n_rot, y.data(), p0.data(), p1.data()); ++ } ++ if (D > n_rot) { ++ memcpy(y.data() + n_rot, x + n_rot, (D - n_rot) * sizeof(float)); ++ } ++ if (scale != 1.0f) { ++ ggml_vec_scale_f32((int) D, y.data(), scale); ++ } ++ ++ char * out = (char *) dst->data + l * dst->nb[1] + h * dst->nb[2] + n * dst->nb[3]; ++ if (dst->type == GGML_TYPE_F32) { ++ memcpy(out, y.data(), D * sizeof(float)); ++ } else { ++ ggml_cpu_fp32_to_fp16(y.data(), (ggml_fp16_t *) out, D); ++ } ++ } ++} ++ ++// ggml_compute_forward_modulate_rows ++ ++void ggml_compute_forward_modulate_rows( ++ const ggml_compute_params * params, ++ ggml_tensor * dst) { ++ const ggml_tensor * a = dst->src[0]; ++ const ggml_tensor * b = dst->src[1]; ++ const ggml_tensor * c = dst->src[2]; ++ ++ const int mode = ggml_get_op_params_i32(dst, 0); ++ const int64_t row0 = ggml_get_op_params_i32(dst, 1); ++ const int64_t row1 = ggml_get_op_params_i32(dst, 2); ++ const float s = ggml_get_op_params_f32(dst, 3); ++ ++ const int64_t nc = a->ne[0]; ++ const int64_t nr1 = row1 - row0; ++ const int64_t nrows = nr1 * a->ne[2] * a->ne[3]; ++ const int64_t dr = (nrows + params->nth - 1) / params->nth; ++ const int64_t ir0 = dr * params->ith; ++ const int64_t ir1 = MIN(ir0 + dr, nrows); ++ ++ // separate passes keep every product, sum and scale individually rounded, as in the unfused graph ++ std::vector t(nc), w(nc); ++ for (int64_t ir = ir0; ir < ir1; ++ir) { ++ const int64_t i1 = row0 + ir % nr1; ++ const int64_t i2 = (ir / nr1) % a->ne[2]; ++ const int64_t i3 = ir / (nr1 * a->ne[2]); ++ ++ const float * x = (const float *) ((const char *) a->data + i1 * a->nb[1] + i2 * a->nb[2] + i3 * a->nb[3]); ++ float * y = (float *) ((char *) dst->data + i1 * dst->nb[1] + i2 * dst->nb[2] + i3 * dst->nb[3]); ++ ++ if (mode == 0) { ++ ggml_vec_mul_f32((int) nc, t.data(), x, (const float *) b->data); ++ ggml_vec_add_f32((int) nc, w.data(), x, t.data()); ++ ggml_vec_add_f32((int) nc, y, w.data(), (const float *) c->data); ++ if (s != 1.0f) { ++ ggml_vec_scale_f32((int) nc, y, s); ++ } ++ } else { ++ const float * u = (const float *) ((const char *) b->data + i1 * b->nb[1] + i2 * b->nb[2] + i3 * b->nb[3]); ++ memcpy(w.data(), u, nc * sizeof(float)); ++ if (s != 1.0f) { ++ ggml_vec_scale_f32((int) nc, w.data(), s); ++ } ++ ggml_vec_mul_f32((int) nc, t.data(), w.data(), (const float *) c->data); ++ ggml_vec_add_f32((int) nc, y, x, t.data()); ++ } ++ } ++} ++ ++// ggml_compute_forward_swiglu_scaled ++ ++void ggml_compute_forward_swiglu_scaled( ++ const ggml_compute_params * params, ++ ggml_tensor * dst) { ++ const ggml_tensor * a = dst->src[0]; ++ const ggml_tensor * b = dst->src[1]; ++ ++ const float in_scale = ggml_get_op_params_f32(dst, 0); ++ const float out_scale = ggml_get_op_params_f32(dst, 1); ++ ++ const int64_t nc = a->ne[0]; ++ const int64_t nrows = ggml_nrows(a); ++ const int64_t dr = (nrows + params->nth - 1) / params->nth; ++ const int64_t ir0 = dr * params->ith; ++ const int64_t ir1 = MIN(ir0 + dr, nrows); ++ ++ std::vector xa(nc), xb(nc); ++ for (int64_t ir = ir0; ir < ir1; ++ir) { ++ const int64_t i1 = ir % a->ne[1]; ++ const int64_t i2 = (ir / a->ne[1]) % a->ne[2]; ++ const int64_t i3 = ir / (a->ne[1] * a->ne[2]); ++ ++ memcpy(xa.data(), (const char *) a->data + i1 * a->nb[1] + i2 * a->nb[2] + i3 * a->nb[3], nc * sizeof(float)); ++ memcpy(xb.data(), (const char *) b->data + i1 * b->nb[1] + i2 * b->nb[2] + i3 * b->nb[3], nc * sizeof(float)); ++ if (in_scale != 1.0f) { ++ ggml_vec_scale_f32((int) nc, xa.data(), in_scale); ++ ggml_vec_scale_f32((int) nc, xb.data(), in_scale); ++ } ++ float * y = (float *) ((char *) dst->data + i1 * dst->nb[1] + i2 * dst->nb[2] + i3 * dst->nb[3]); ++ ggml_vec_swiglu_f32((int) nc, y, xa.data(), xb.data()); ++ if (out_scale != 1.0f) { ++ ggml_vec_scale_f32((int) nc, y, out_scale); ++ } ++ } ++} +diff --git a/src/ggml-cpu/ops.h b/src/ggml-cpu/ops.h +index 3b579167..e09cc11d 100644 +--- a/src/ggml-cpu/ops.h ++++ b/src/ggml-cpu/ops.h +@@ -98,6 +98,9 @@ void ggml_compute_forward_win_part(const struct ggml_compute_params * params, st + void ggml_compute_forward_win_unpart(const struct ggml_compute_params * params, struct ggml_tensor * dst); + void ggml_compute_forward_unary(const struct ggml_compute_params * params, struct ggml_tensor * dst); + void ggml_compute_forward_glu(const struct ggml_compute_params * params, struct ggml_tensor * dst); ++void ggml_compute_forward_rope_pe_permute(const struct ggml_compute_params * params, struct ggml_tensor * dst); ++void ggml_compute_forward_modulate_rows(const struct ggml_compute_params * params, struct ggml_tensor * dst); ++void ggml_compute_forward_swiglu_scaled(const struct ggml_compute_params * params, struct ggml_tensor * dst); + void ggml_compute_forward_get_rel_pos(const struct ggml_compute_params * params, struct ggml_tensor * dst); + void ggml_compute_forward_add_rel_pos(const struct ggml_compute_params * params, struct ggml_tensor * dst); + void ggml_compute_forward_rwkv_wkv6(const struct ggml_compute_params * params, struct ggml_tensor * dst); +diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu +index 517d0894..2c0984a0 100644 +--- a/src/ggml-cuda/ggml-cuda.cu ++++ b/src/ggml-cuda/ggml-cuda.cu +@@ -44,6 +44,8 @@ + #include "ggml-cuda/rope.cuh" + #include "ggml-cuda/roll.cuh" + #include "ggml-cuda/scale.cuh" ++#include "ggml-cuda/rope-pe.cuh" ++#include "ggml-cuda/modulate-rows.cuh" + #include "ggml-cuda/snake.cuh" + #include "ggml-cuda/softcap.cuh" + #include "ggml-cuda/softmax.cuh" +@@ -3002,6 +3004,15 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg + case GGML_OP_SAGE_ATTN: + ggml_cuda_sage_attn(ctx, dst); + break; ++ case GGML_OP_ROPE_PE_PERMUTE: ++ ggml_cuda_op_rope_pe_permute(ctx, dst); ++ break; ++ case GGML_OP_MODULATE_ROWS: ++ ggml_cuda_op_modulate_rows(ctx, dst); ++ break; ++ case GGML_OP_SWIGLU_SCALED: ++ ggml_cuda_op_swiglu_scaled(ctx, dst); ++ break; + case GGML_OP_CROSS_ENTROPY_LOSS: + ggml_cuda_cross_entropy_loss(ctx, dst); + break; +@@ -5886,6 +5897,12 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g + return ggml_cuda_flash_attn_ext_supported(dev_ctx->device, op); + case GGML_OP_SAGE_ATTN: + return ggml_cuda_sage_attn_supported(dev_ctx->device, op); ++ case GGML_OP_ROPE_PE_PERMUTE: ++ return ggml_cuda_rope_pe_permute_supported(op); ++ case GGML_OP_MODULATE_ROWS: ++ return ggml_cuda_modulate_rows_supported(op); ++ case GGML_OP_SWIGLU_SCALED: ++ return ggml_cuda_swiglu_scaled_supported(op); + case GGML_OP_CROSS_ENTROPY_LOSS: + case GGML_OP_CROSS_ENTROPY_LOSS_BACK: + case GGML_OP_OPT_STEP_ADAMW: +diff --git a/src/ggml-cuda/modulate-rows.cu b/src/ggml-cuda/modulate-rows.cu +new file mode 100644 +index 00000000..1b639d28 +--- /dev/null ++++ b/src/ggml-cuda/modulate-rows.cu +@@ -0,0 +1,168 @@ ++#include "modulate-rows.cuh" ++#include "unary.cuh" ++ ++// Rounding matches the unfused graph: __fmul_rn/__fadd_rn are the separately rounded MUL and ADD ++// kernels (never contracted into an FMA), `s*y + bias` is the expression the SCALE kernel ++// evaluates, and the swiglu product is the GLU kernel's silu(a)*b. ++ ++template ++static __global__ void modulate_rows_kernel( ++ const char * __restrict__ a, const char * __restrict__ b, const char * __restrict__ c, char * __restrict__ dst, ++ const int nc, const int64_t nr1, const int64_t ne2, const int64_t row0, ++ const int64_t nba1, const int64_t nba2, const int64_t nba3, ++ const int64_t nbb1, const int64_t nbb2, const int64_t nbb3, ++ const int64_t nbd1, const int64_t nbd2, const int64_t nbd3, ++ const float s, const float bias) { ++ const int64_t ir = blockIdx.x; ++ const int64_t i1 = row0 + ir % nr1; ++ const int64_t i2 = (ir / nr1) % ne2; ++ const int64_t i3 = ir / (nr1 * ne2); ++ ++ const float * xa = (const float *) (a + i1 * nba1 + i2 * nba2 + i3 * nba3); ++ float * y = (float *) (dst + i1 * nbd1 + i2 * nbd2 + i3 * nbd3); ++ ++ if constexpr (mode == 0) { ++ const float * scale = (const float *) b; ++ const float * shift = (const float *) c; ++ for (int i = threadIdx.x; i < nc; i += blockDim.x) { ++ const float x = xa[i]; ++ float v = __fadd_rn(__fadd_rn(x, __fmul_rn(x, scale[i])), shift[i]); ++ if constexpr (has_scale) { ++ v = s * v + bias; ++ } ++ y[i] = v; ++ } ++ } else { ++ const float * u = (const float *) (b + i1 * nbb1 + i2 * nbb2 + i3 * nbb3); ++ const float * gate = (const float *) c; ++ for (int i = threadIdx.x; i < nc; i += blockDim.x) { ++ float w = u[i]; ++ if constexpr (has_scale) { ++ w = s * w + bias; ++ } ++ y[i] = __fadd_rn(xa[i], __fmul_rn(w, gate[i])); ++ } ++ } ++} ++ ++void ggml_cuda_op_modulate_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ++ const ggml_tensor * a = dst->src[0]; ++ const ggml_tensor * b = dst->src[1]; ++ const ggml_tensor * c = dst->src[2]; ++ ++ const int mode = ggml_get_op_params_i32(dst, 0); ++ const int64_t row0 = ggml_get_op_params_i32(dst, 1); ++ const int64_t row1 = ggml_get_op_params_i32(dst, 2); ++ const float s = ggml_get_op_params_f32(dst, 3); ++ ++ const int64_t nr1 = row1 - row0; ++ const int64_t nrows = nr1 * a->ne[2] * a->ne[3]; ++ const int nc = (int) a->ne[0]; ++ if (nrows == 0) { ++ return; ++ } ++ ++ cudaStream_t stream = ctx.stream(); ++ const dim3 grid((unsigned int) nrows, 1, 1); ++ const dim3 block(256, 1, 1); ++ ++#define MODULATE_ROWS_LAUNCH(MODE, HAS_SCALE) \ ++ modulate_rows_kernel<<>>( \ ++ (const char *) a->data, (const char *) b->data, (const char *) c->data, (char *) dst->data, nc, nr1, \ ++ a->ne[2], row0, a->nb[1], a->nb[2], a->nb[3], b->nb[1], b->nb[2], b->nb[3], dst->nb[1], dst->nb[2], \ ++ dst->nb[3], s, 0.0f) ++ ++ if (mode == 0) { ++ if (s != 1.0f) { ++ MODULATE_ROWS_LAUNCH(0, true); ++ } else { ++ MODULATE_ROWS_LAUNCH(0, false); ++ } ++ } else { ++ if (s != 1.0f) { ++ MODULATE_ROWS_LAUNCH(1, true); ++ } else { ++ MODULATE_ROWS_LAUNCH(1, false); ++ } ++ } ++#undef MODULATE_ROWS_LAUNCH ++} ++ ++bool ggml_cuda_modulate_rows_supported(const ggml_tensor * op) { ++ const int64_t nrows = (int64_t) (ggml_get_op_params_i32(op, 2) - ggml_get_op_params_i32(op, 1)) * op->src[0]->ne[2] * op->src[0]->ne[3]; ++ return op->type == GGML_TYPE_F32 && ggml_is_contiguous_rows(op) && op->src[0]->ne[0] <= INT_MAX && nrows <= INT_MAX && ++ ggml_is_contiguous(op->src[ggml_get_op_params_i32(op, 0) == 0 ? 1 : 2]) && ++ (ggml_get_op_params_i32(op, 0) != 0 || ggml_is_contiguous(op->src[2])); ++} ++ ++template ++static __global__ void swiglu_scaled_kernel( ++ const char * __restrict__ a, const char * __restrict__ b, float * __restrict__ dst, ++ const int nc, const int64_t ne1, const int64_t ne2, ++ const int64_t nba1, const int64_t nba2, const int64_t nba3, ++ const int64_t nbb1, const int64_t nbb2, const int64_t nbb3, ++ const float in_scale, const float out_scale, const float bias) { ++ const int64_t ir = blockIdx.x; ++ const int64_t i1 = ir % ne1; ++ const int64_t i2 = (ir / ne1) % ne2; ++ const int64_t i3 = ir / (ne1 * ne2); ++ ++ const float * xa = (const float *) (a + i1 * nba1 + i2 * nba2 + i3 * nba3); ++ const float * xb = (const float *) (b + i1 * nbb1 + i2 * nbb2 + i3 * nbb3); ++ float * y = dst + ir * nc; ++ ++ for (int i = threadIdx.x; i < nc; i += blockDim.x) { ++ float va = xa[i]; ++ float vb = xb[i]; ++ if constexpr (has_in) { ++ va = in_scale * va + bias; ++ vb = in_scale * vb + bias; ++ } ++ float v = ggml_cuda_op_silu_single(va) * vb; ++ if constexpr (has_out) { ++ v = out_scale * v + bias; ++ } ++ y[i] = v; ++ } ++} ++ ++void ggml_cuda_op_swiglu_scaled(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ++ const ggml_tensor * a = dst->src[0]; ++ const ggml_tensor * b = dst->src[1]; ++ ++ const float in_scale = ggml_get_op_params_f32(dst, 0); ++ const float out_scale = ggml_get_op_params_f32(dst, 1); ++ ++ GGML_ASSERT(ggml_is_contiguous(dst)); ++ const int64_t nrows = ggml_nrows(a); ++ const int nc = (int) a->ne[0]; ++ if (nrows == 0) { ++ return; ++ } ++ ++ cudaStream_t stream = ctx.stream(); ++ const dim3 grid((unsigned int) nrows, 1, 1); ++ const dim3 block(256, 1, 1); ++ ++#define SWIGLU_SCALED_LAUNCH(HAS_IN, HAS_OUT) \ ++ swiglu_scaled_kernel<<>>( \ ++ (const char *) a->data, (const char *) b->data, (float *) dst->data, nc, a->ne[1], a->ne[2], a->nb[1], \ ++ a->nb[2], a->nb[3], b->nb[1], b->nb[2], b->nb[3], in_scale, out_scale, 0.0f) ++ ++ const bool has_in = in_scale != 1.0f; ++ const bool has_out = out_scale != 1.0f; ++ if (has_in && has_out) { ++ SWIGLU_SCALED_LAUNCH(true, true); ++ } else if (has_in) { ++ SWIGLU_SCALED_LAUNCH(true, false); ++ } else if (has_out) { ++ SWIGLU_SCALED_LAUNCH(false, true); ++ } else { ++ SWIGLU_SCALED_LAUNCH(false, false); ++ } ++#undef SWIGLU_SCALED_LAUNCH ++} ++ ++bool ggml_cuda_swiglu_scaled_supported(const ggml_tensor * op) { ++ return op->type == GGML_TYPE_F32 && op->src[0]->ne[0] <= INT_MAX && ggml_nrows(op->src[0]) <= INT_MAX; ++} +diff --git a/src/ggml-cuda/modulate-rows.cuh b/src/ggml-cuda/modulate-rows.cuh +new file mode 100644 +index 00000000..216d0767 +--- /dev/null ++++ b/src/ggml-cuda/modulate-rows.cuh +@@ -0,0 +1,7 @@ ++#include "common.cuh" ++ ++void ggml_cuda_op_modulate_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst); ++bool ggml_cuda_modulate_rows_supported(const ggml_tensor * op); ++ ++void ggml_cuda_op_swiglu_scaled(ggml_backend_cuda_context & ctx, ggml_tensor * dst); ++bool ggml_cuda_swiglu_scaled_supported(const ggml_tensor * op); +diff --git a/src/ggml-cuda/rope-pe.cu b/src/ggml-cuda/rope-pe.cu +new file mode 100644 +index 00000000..3afe4423 +--- /dev/null ++++ b/src/ggml-cuda/rope-pe.cu +@@ -0,0 +1,102 @@ ++#include "rope-pe.cuh" ++#include "convert.cuh" ++ ++// Fused rotary embedding (precomputed cos/sin table) + head-major relayout + scale + cast. ++// One block per (token, batch); the block walks every head so the table entries are read once. ++// Rounding matches the unfused graph exactly: __fmul_rn/__fadd_rn keep the two products and the ++// sum as separately rounded operations (no FMA contraction), the scale is the same `scale*y + bias` ++// expression the SCALE kernel evaluates, and the conversion uses the same ggml_cuda_cast as CPY. ++template ++static __global__ void rope_pe_permute_kernel( ++ const char * __restrict__ x, const float * __restrict__ pe, dst_t * __restrict__ dst, ++ const int D, const int H, const int64_t L, const int n_rot, ++ const int64_t nb01, const int64_t nb02, const int64_t nb03, ++ const float scale, const float bias) { ++ const int64_t l = blockIdx.x; ++ const int64_t n = blockIdx.y; ++ const int half = n_rot / 2; ++ ++ for (int i = threadIdx.x; i < D; i += blockDim.x) { ++ const bool rot = i < n_rot; ++ int j = i; ++ float c0 = 0.0f; ++ float c1 = 0.0f; ++ if (rot) { ++ const int r = i >= half; ++ j = i - r * half; ++ const float * t = pe + (l * half + j) * 4 + 2 * r; ++ c0 = t[0]; ++ c1 = t[1]; ++ } ++ for (int h = 0; h < H; ++h) { ++ const float * xr = (const float *) (x + h * nb01 + l * nb02 + n * nb03); ++ float y; ++ if (rot) { ++ y = __fadd_rn(__fmul_rn(xr[j], c0), __fmul_rn(xr[j + half], c1)); ++ } else { ++ y = xr[i]; ++ } ++ if constexpr (has_scale) { ++ y = scale * y + bias; ++ } ++ dst[((n * H + h) * L + l) * D + i] = ggml_cuda_cast(y); ++ } ++ } ++} ++ ++template ++static void rope_pe_permute_cuda(const ggml_tensor * src0, const float * pe, dst_t * dst, const int n_rot, ++ const float scale, cudaStream_t stream) { ++ const int D = (int) src0->ne[0]; ++ const int H = (int) src0->ne[1]; ++ const int64_t L = src0->ne[2]; ++ const int64_t N = src0->ne[3]; ++ ++ const dim3 grid((unsigned int) L, (unsigned int) N, 1); ++ const dim3 block(D >= 128 ? 128 : ((D + 31) / 32) * 32, 1, 1); ++ const char * x = (const char *) src0->data; ++ if (scale != 1.0f) { ++ rope_pe_permute_kernel<<>>( ++ x, pe, dst, D, H, L, n_rot, src0->nb[1], src0->nb[2], src0->nb[3], scale, 0.0f); ++ } else { ++ rope_pe_permute_kernel<<>>( ++ x, pe, dst, D, H, L, n_rot, src0->nb[1], src0->nb[2], src0->nb[3], scale, 0.0f); ++ } ++} ++ ++void ggml_cuda_op_rope_pe_permute(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ++ const ggml_tensor * src0 = dst->src[0]; ++ const ggml_tensor * pe = dst->src[1]; ++ ++ const int n_rot = ggml_get_op_params_i32(dst, 0); ++ const float scale = ggml_get_op_params_f32(dst, 1); ++ ++ GGML_ASSERT(src0->type == GGML_TYPE_F32 && src0->nb[0] == sizeof(float)); ++ GGML_ASSERT(ggml_is_contiguous(dst)); ++ GGML_ASSERT(n_rot == 0 || (pe != nullptr && ggml_is_contiguous(pe))); ++ ++ cudaStream_t stream = ctx.stream(); ++ const float * pe_d = pe ? (const float *) pe->data : nullptr; ++ ++ if (dst->type == GGML_TYPE_F16) { ++ rope_pe_permute_cuda(src0, pe_d, (half *) dst->data, n_rot, scale, stream); ++ } else { ++ GGML_ASSERT(dst->type == GGML_TYPE_F32); ++ rope_pe_permute_cuda(src0, pe_d, (float *) dst->data, n_rot, scale, stream); ++ } ++} ++ ++bool ggml_cuda_rope_pe_permute_supported(const ggml_tensor * op) { ++ const ggml_tensor * src0 = op->src[0]; ++ const ggml_tensor * pe = op->src[1]; ++ if (src0->type != GGML_TYPE_F32 || src0->nb[0] != sizeof(float)) { ++ return false; ++ } ++ if (op->type != GGML_TYPE_F32 && op->type != GGML_TYPE_F16) { ++ return false; ++ } ++ if (src0->ne[2] > INT_MAX || src0->ne[3] > 65535 || src0->ne[1] > INT_MAX || src0->ne[0] > INT_MAX) { ++ return false; ++ } ++ return pe == nullptr || (pe->type == GGML_TYPE_F32 && ggml_is_contiguous(pe)); ++} +diff --git a/src/ggml-cuda/rope-pe.cuh b/src/ggml-cuda/rope-pe.cuh +new file mode 100644 +index 00000000..b68b6af4 +--- /dev/null ++++ b/src/ggml-cuda/rope-pe.cuh +@@ -0,0 +1,5 @@ ++#include "common.cuh" ++ ++void ggml_cuda_op_rope_pe_permute(ggml_backend_cuda_context & ctx, ggml_tensor * dst); ++ ++bool ggml_cuda_rope_pe_permute_supported(const ggml_tensor * op); +diff --git a/src/ggml.c b/src/ggml.c +index dc4a29f9..fec933bf 100644 +--- a/src/ggml.c ++++ b/src/ggml.c +@@ -1115,9 +1115,12 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { + + "QUANTIZE_I8_CONVROT", + "SAGE_ATTN", ++ "ROPE_PE_PERMUTE", ++ "MODULATE_ROWS", ++ "SWIGLU_SCALED", + }; + +-static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103"); ++static_assert(GGML_OP_COUNT == 106, "GGML_OP_COUNT != 106"); + + static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { + "none", +@@ -1233,9 +1236,12 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { + + "quantize_i8_convrot(x)", + "sage_attn(x,y,z)", ++ "rope_pe_permute(x,pe)", ++ "modulate_rows(x,a,b)", ++ "swiglu_scaled(x,y)", + }; + +-static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103"); ++static_assert(GGML_OP_COUNT == 106, "GGML_OP_COUNT != 106"); + + static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); + +@@ -5557,6 +5563,125 @@ struct ggml_tensor * ggml_sage_attn( + return result; + } + ++struct ggml_tensor * ggml_rope_pe_permute( ++ struct ggml_context * ctx, ++ struct ggml_tensor * a, ++ struct ggml_tensor * pe, ++ int n_rot, ++ float scale, ++ enum ggml_type type) { ++ GGML_ASSERT(a->type == GGML_TYPE_F32); ++ GGML_ASSERT(a->nb[0] == ggml_type_size(a->type)); ++ GGML_ASSERT(type == GGML_TYPE_F32 || type == GGML_TYPE_F16); ++ GGML_ASSERT(n_rot >= 0 && n_rot % 2 == 0 && n_rot <= a->ne[0]); ++ if (n_rot > 0) { ++ GGML_ASSERT(pe != NULL && pe->type == GGML_TYPE_F32 && ggml_is_contiguous(pe)); ++ GGML_ASSERT(pe->ne[0] == 2 && pe->ne[1] == 2 && pe->ne[2] == n_rot / 2 && pe->ne[3] == a->ne[2]); ++ } else { ++ GGML_ASSERT(pe == NULL); ++ } ++ GGML_ASSERT(isfinite(scale)); ++ ++ struct ggml_tensor * result = ggml_new_tensor_4d(ctx, type, a->ne[0], a->ne[2], a->ne[1], a->ne[3]); ++ ++ ggml_set_op_params_i32(result, 0, n_rot); ++ ggml_set_op_params_f32(result, 1, scale); ++ ++ result->op = GGML_OP_ROPE_PE_PERMUTE; ++ result->src[0] = a; ++ result->src[1] = pe; ++ ++ return result; ++} ++ ++static struct ggml_tensor * ggml_modulate_rows_impl( ++ struct ggml_context * ctx, ++ struct ggml_tensor * acc, ++ struct ggml_tensor * a, ++ struct ggml_tensor * b, ++ struct ggml_tensor * c, ++ int64_t row0, ++ int64_t row1, ++ float s, ++ int32_t mode) { ++ GGML_ASSERT(a->type == GGML_TYPE_F32 && a->nb[0] == sizeof(float)); ++ GGML_ASSERT(row0 >= 0 && row0 < row1 && row1 <= a->ne[1] && row1 <= INT32_MAX); ++ GGML_ASSERT(isfinite(s)); ++ // mode 0: a = x, b = scale, c = shift; mode 1: a = base, b = u, c = gate ++ struct ggml_tensor * vec0 = mode == 0 ? b : c; ++ GGML_ASSERT(vec0->type == GGML_TYPE_F32 && vec0->nb[0] == sizeof(float) && ggml_nelements(vec0) == a->ne[0]); ++ if (mode == 0) { ++ GGML_ASSERT(c->type == GGML_TYPE_F32 && c->nb[0] == sizeof(float) && ggml_nelements(c) == a->ne[0]); ++ } else { ++ GGML_ASSERT(b->type == GGML_TYPE_F32 && b->nb[0] == sizeof(float) && ggml_are_same_shape(a, b)); ++ } ++ if (acc != NULL) { ++ GGML_ASSERT(acc->type == GGML_TYPE_F32 && ggml_is_contiguous(acc) && ggml_are_same_shape(acc, a)); ++ } ++ ++ struct ggml_tensor * result = acc != NULL ? ggml_view_tensor(ctx, acc) : ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, a->ne); ++ ++ ggml_set_op_params_i32(result, 0, mode); ++ ggml_set_op_params_i32(result, 1, (int32_t) row0); ++ ggml_set_op_params_i32(result, 2, (int32_t) row1); ++ ggml_set_op_params_f32(result, 3, s); ++ ++ result->op = GGML_OP_MODULATE_ROWS; ++ result->src[0] = a; ++ result->src[1] = b; ++ result->src[2] = c; ++ result->src[3] = acc; ++ ++ return result; ++} ++ ++struct ggml_tensor * ggml_modulate_rows( ++ struct ggml_context * ctx, ++ struct ggml_tensor * acc, ++ struct ggml_tensor * x, ++ struct ggml_tensor * scale, ++ struct ggml_tensor * shift, ++ int64_t row0, ++ int64_t row1, ++ float out_scale) { ++ return ggml_modulate_rows_impl(ctx, acc, x, scale, shift, row0, row1, out_scale, 0); ++} ++ ++struct ggml_tensor * ggml_gated_add_rows( ++ struct ggml_context * ctx, ++ struct ggml_tensor * acc, ++ struct ggml_tensor * base, ++ struct ggml_tensor * u, ++ struct ggml_tensor * gate, ++ int64_t row0, ++ int64_t row1, ++ float u_scale) { ++ return ggml_modulate_rows_impl(ctx, acc, base, u, gate, row0, row1, u_scale, 1); ++} ++ ++struct ggml_tensor * ggml_swiglu_scaled( ++ struct ggml_context * ctx, ++ struct ggml_tensor * a, ++ struct ggml_tensor * b, ++ float in_scale, ++ float out_scale) { ++ GGML_ASSERT(a->type == GGML_TYPE_F32 && b->type == GGML_TYPE_F32); ++ GGML_ASSERT(a->nb[0] == sizeof(float) && b->nb[0] == sizeof(float)); ++ GGML_ASSERT(ggml_are_same_shape(a, b)); ++ GGML_ASSERT(isfinite(in_scale) && isfinite(out_scale)); ++ ++ struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, a->ne); ++ ++ ggml_set_op_params_f32(result, 0, in_scale); ++ ggml_set_op_params_f32(result, 1, out_scale); ++ ++ result->op = GGML_OP_SWIGLU_SCALED; ++ result->src[0] = a; ++ result->src[1] = b; ++ ++ return result; ++} ++ + void ggml_flash_attn_ext_set_prec( + struct ggml_tensor * a, + enum ggml_prec prec) { +diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp +index dd290caf..1e46361c 100644 +--- a/tests/test-backend-ops.cpp ++++ b/tests/test-backend-ops.cpp +@@ -2125,6 +2125,118 @@ struct test_glu : public test_case { + } + }; + ++// GGML_OP_ROPE_PE_PERMUTE ++struct test_rope_pe_permute : public test_case { ++ const ggml_type type; ++ const std::array ne_a; // [head_dim, n_head, n_tokens, n_batch] ++ const int n_rot; ++ const float scale; ++ const bool view; // a is a strided view (one third of a fused qkv projection) ++ ++ std::string vars() override { ++ return VARS_TO_STR5(type, ne_a, n_rot, scale, view); ++ } ++ ++ test_rope_pe_permute(ggml_type type = GGML_TYPE_F32, ++ std::array ne_a = {128, 4, 9, 1}, ++ int n_rot = 96, float scale = 1.0f, bool view = false) ++ : type(type), ne_a(ne_a), n_rot(n_rot), scale(scale), view(view) {} ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * a; ++ if (view) { ++ ggml_tensor * qkv = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 3 * ne_a[0] * ne_a[1], ne_a[2], ne_a[3], 1); ++ ggml_set_name(qkv, "qkv"); ++ a = ggml_view_4d(ctx, qkv, ne_a[0], ne_a[1], ne_a[2], ne_a[3], ++ ne_a[0] * sizeof(float), qkv->nb[1], qkv->nb[2], ne_a[0] * ne_a[1] * sizeof(float)); ++ ggml_set_name(a, "view_of_qkv"); ++ } else { ++ a = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne_a.data()); ++ ggml_set_name(a, "a"); ++ } ++ ggml_tensor * pe = nullptr; ++ if (n_rot > 0) { ++ pe = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 2, 2, n_rot / 2, ne_a[2]); ++ ggml_set_name(pe, "pe"); ++ } ++ ggml_tensor * out = ggml_rope_pe_permute(ctx, a, pe, n_rot, scale, type); ++ ggml_set_name(out, "out"); ++ return out; ++ } ++}; ++ ++// GGML_OP_MODULATE_ROWS (a chain of segments written in place) ++struct test_modulate_rows : public test_case { ++ const int mode; // 0: modulate, 1: gated add ++ const std::array ne; ++ const float s; ++ const bool view; ++ ++ std::string vars() override { ++ return VARS_TO_STR4(mode, ne, s, view); ++ } ++ ++ test_modulate_rows(int mode = 0, std::array ne = {64, 12, 2, 1}, float s = 1.0f, bool view = false) ++ : mode(mode), ne(ne), s(s), view(view) {} ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ auto make = [&](const char * name) { ++ ggml_tensor * t; ++ if (view) { ++ ggml_tensor * big = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 2 * ne[0], ne[1], ne[2], ne[3]); ++ t = ggml_view_4d(ctx, big, ne[0], ne[1], ne[2], ne[3], big->nb[1], big->nb[2], big->nb[3], ne[0] * sizeof(float)); ++ } else { ++ t = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); ++ } ++ ggml_set_name(t, name); ++ return t; ++ }; ++ ggml_tensor * a = make("a"); ++ ggml_tensor * u = mode == 1 ? make("u") : nullptr; ++ // modulation table: 3 rows of 2 vectors each, taken as views like an adaLN projection ++ ggml_tensor * table = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2 * ne[0], 3); ++ ggml_set_name(table, "table"); ++ const int64_t bounds[4] = {0, ne[1] / 4, ne[1] - 2, ne[1]}; ++ ggml_tensor * out = nullptr; ++ for (int seg = 0; seg < 3; ++seg) { ++ const int row = (seg + 1) % 3; ++ ggml_tensor * v0 = ggml_view_1d(ctx, table, ne[0], row * table->nb[1]); ++ ggml_tensor * v1 = ggml_view_1d(ctx, table, ne[0], row * table->nb[1] + ne[0] * sizeof(float)); ++ if (mode == 0) { ++ out = ggml_modulate_rows(ctx, out, a, v0, v1, bounds[seg], bounds[seg + 1], s); ++ } else { ++ out = ggml_gated_add_rows(ctx, out, a, u, v0, bounds[seg], bounds[seg + 1], s); ++ } ++ } ++ ggml_set_name(out, "out"); ++ return out; ++ } ++}; ++ ++// GGML_OP_SWIGLU_SCALED ++struct test_swiglu_scaled : public test_case { ++ const std::array ne; ++ const float in_scale; ++ const float out_scale; ++ ++ std::string vars() override { ++ return VARS_TO_STR3(ne, in_scale, out_scale); ++ } ++ ++ test_swiglu_scaled(std::array ne = {64, 5, 3, 1}, float in_scale = 1.0f, float out_scale = 1.0f) ++ : ne(ne), in_scale(in_scale), out_scale(out_scale) {} ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * h = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 2 * ne[0], ne[1], ne[2], ne[3]); ++ ggml_set_name(h, "h"); ++ ggml_tensor * a = ggml_view_4d(ctx, h, ne[0], ne[1], ne[2], ne[3], h->nb[1], h->nb[2], h->nb[3], 0); ++ ggml_tensor * b = ggml_view_4d(ctx, h, ne[0], ne[1], ne[2], ne[3], h->nb[1], h->nb[2], h->nb[3], ne[0] * sizeof(float)); ++ ggml_tensor * out = ggml_swiglu_scaled(ctx, a, b, in_scale, out_scale); ++ ggml_set_name(out, "out"); ++ return out; ++ } ++}; ++ + struct test_glu_split : public test_case { + const ggml_glu_op op; + const ggml_type type; +@@ -8101,6 +8213,33 @@ static std::vector> make_test_cases_eval() { + } + } + ++ for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_F16}) { ++ for (bool view : {false, true}) { ++ for (float scale : {1.0f, 1.0f / 128.0f}) { ++ test_cases.emplace_back(new test_rope_pe_permute(type, {128, 4, 9, 1}, 96, scale, view)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {128, 3, 5, 2}, 128, scale, view)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {128, 4, 9, 1}, 0, scale, view)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {64, 3, 7, 2}, 32, scale, view)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {200, 2, 3, 1}, 100, scale, view)); ++ } ++ } ++ } ++ ++ for (int mode : {0, 1}) { ++ for (bool view : {false, true}) { ++ for (float sc : {1.0f, mode == 0 ? 1.0f / 128.0f : 128.0f}) { ++ test_cases.emplace_back(new test_modulate_rows(mode, {64, 12, 2, 1}, sc, view)); ++ test_cases.emplace_back(new test_modulate_rows(mode, {300, 9, 1, 2}, sc, view)); ++ } ++ } ++ } ++ for (float in_scale : {1.0f, 128.0f}) { ++ for (float out_scale : {1.0f, 1.0f / 128.0f}) { ++ test_cases.emplace_back(new test_swiglu_scaled({64, 5, 3, 1}, in_scale, out_scale)); ++ test_cases.emplace_back(new test_swiglu_scaled({300, 7, 1, 2}, in_scale, out_scale)); ++ } ++ } ++ + for (int v : {0, 1}) { + for (float alpha : {.5f, 1.702f}) { + for (float limit : {2.0f, 7.0f}) { +-- +2.43.0 + diff --git a/scripts/unsloth/ggml-patches/0005-ggml-rope-pe-permute-rms.patch b/scripts/unsloth/ggml-patches/0005-ggml-rope-pe-permute-rms.patch new file mode 100644 index 0000000000..65b336de54 --- /dev/null +++ b/scripts/unsloth/ggml-patches/0005-ggml-rope-pe-permute-rms.patch @@ -0,0 +1,304 @@ +From 092f5d51d0844f7cf282ca115a81fccdde9fafad Mon Sep 17 00:00:00 2001 +From: Daniel Han +Date: Thu, 1 Oct 2026 23:16:25 -0700 +Subject: [PATCH] ggml : rope_pe_permute with the per-head RMS norm folded in + +ggml_rope_pe_permute_rms(a, pe, n_rot, scale, type, eps) is ggml_rope_pe_permute of +ggml_rms_norm(a, eps) without a weight. CPU uses rms_norm's arithmetic (double sum, +1/sqrtf, vec scale); CUDA normalises each head row with the block size and block +reduction rms_norm_f32_cuda's small-row path uses for that row length, so both match +the two-op chain bit for bit. CUDA supports it for head_dim <= 128 when that small-row +path is active (not HIP/MUSA, GGML_CUDA_NORM_SMALL_ROWS not 0); elsewhere the op is +reported unsupported and callers keep the separate norm. test-backend-ops cases added. +--- + include/ggml.h | 12 ++++ + src/ggml-cpu/ops.cpp | 17 +++++- + src/ggml-cuda/rope-pe.cu | 120 +++++++++++++++++++++++++++++++++++++ + src/ggml.c | 15 +++++ + tests/test-backend-ops.cpp | 15 +++-- + 5 files changed, 174 insertions(+), 5 deletions(-) + +diff --git a/include/ggml.h b/include/ggml.h +index 153eed2a..8134626f 100644 +--- a/include/ggml.h ++++ b/include/ggml.h +@@ -2483,6 +2483,18 @@ extern "C" { + float scale, + enum ggml_type type); + ++ // ggml_rope_pe_permute of ggml_rms_norm(a, eps) (no weight): every head row of `a` is ++ // RMS-normalised first, with the same arithmetic as GGML_OP_RMS_NORM on the same backend, so ++ // the result matches the two-op chain bit for bit where the backend supports it. ++ GGML_API struct ggml_tensor * ggml_rope_pe_permute_rms( ++ struct ggml_context * ctx, ++ struct ggml_tensor * a, ++ struct ggml_tensor * pe, ++ int n_rot, ++ float scale, ++ enum ggml_type type, ++ float eps); ++ + // Row-range modulation for adaLN-style blocks whose sequence is split into segments that use + // different modulation vectors. Writes rows [row0, row1) along dim 1 (all of dims 2 and 3) of the + // result and leaves every other row as it is in `acc`: chain one call per segment, passing the +diff --git a/src/ggml-cpu/ops.cpp b/src/ggml-cpu/ops.cpp +index 40fe1949..1d67b482 100644 +--- a/src/ggml-cpu/ops.cpp ++++ b/src/ggml-cpu/ops.cpp +@@ -12151,6 +12151,8 @@ void ggml_compute_forward_rope_pe_permute( + + const int n_rot = ggml_get_op_params_i32(dst, 0); + const float scale = ggml_get_op_params_f32(dst, 1); ++ const bool norm = ggml_get_op_params_i32(dst, 2) != 0; ++ const float eps = ggml_get_op_params_f32(dst, 3); + const int half = n_rot / 2; + + const int64_t D = src0->ne[0]; +@@ -12165,7 +12167,7 @@ void ggml_compute_forward_rope_pe_permute( + + // Same rounding points as the unfused graph: each product, the sum, the scale and the + // conversion are separate passes, so no contraction into fused multiply-adds can happen. +- std::vector p0(n_rot), p1(n_rot), y(D); ++ std::vector p0(n_rot), p1(n_rot), y(D), xn(norm ? D : 0); + + for (int64_t ir = ir0; ir < ir1; ++ir) { + const int64_t h = ir % H; +@@ -12174,6 +12176,19 @@ void ggml_compute_forward_rope_pe_permute( + + const float * x = (const float *) ((const char *) src0->data + h * src0->nb[1] + l * src0->nb[2] + n * src0->nb[3]); + ++ if (norm) { ++ // ggml_compute_forward_rms_norm_f32 without a weight ++ ggml_float sum = 0.0; ++ for (int64_t i = 0; i < D; i++) { ++ sum += (ggml_float)(x[i] * x[i]); ++ } ++ const float mean = sum/D; ++ const float s = 1.0f/sqrtf(mean + eps); ++ memcpy(xn.data(), x, D * sizeof(float)); ++ ggml_vec_scale_f32((int) D, xn.data(), s); ++ x = xn.data(); ++ } ++ + if (n_rot > 0) { + const float * t = (const float *) pe->data + l * (int64_t) half * 4; + for (int i = 0; i < n_rot; ++i) { +diff --git a/src/ggml-cuda/rope-pe.cu b/src/ggml-cuda/rope-pe.cu +index 3afe4423..b6dd7a69 100644 +--- a/src/ggml-cuda/rope-pe.cu ++++ b/src/ggml-cuda/rope-pe.cu +@@ -44,6 +44,112 @@ static __global__ void rope_pe_permute_kernel( + } + } + ++// Same op with every head row RMS-normalised first (no weight). One block of block_size threads ++// per (token, batch), the size and reduction rms_norm_f32_cuda's small-row path uses for this row ++// length, so each normalised value is bit-identical to GGML_OP_RMS_NORM's; the normalised row is ++// staged in shared memory for the RoPE partner reads. ++template ++static __global__ void rope_pe_permute_rms_kernel( ++ const char * __restrict__ x, const float * __restrict__ pe, dst_t * __restrict__ dst, ++ const int D, const int H, const int64_t L, const int n_rot, ++ const int64_t nb01, const int64_t nb02, const int64_t nb03, ++ const float scale, const float bias, const float eps) { ++ const int64_t l = blockIdx.x; ++ const int64_t n = blockIdx.y; ++ const int half = n_rot / 2; ++ const int i = threadIdx.x; ++ ++ __shared__ float s_sum[32]; ++ __shared__ float xs[block_size]; ++ ++ const bool rot = i < n_rot; ++ int j = i; ++ float c0 = 0.0f; ++ float c1 = 0.0f; ++ if (rot) { ++ const int r = i >= half; ++ j = i - r * half; ++ const float * t = pe + (l * half + j) * 4 + 2 * r; ++ c0 = t[0]; ++ c1 = t[1]; ++ } ++ for (int h = 0; h < H; ++h) { ++ const float * xr = (const float *) (x + h * nb01 + l * nb02 + n * nb03); ++ float xi = 0.0f; ++ float tmp = 0.0f; ++ if (i < D) { ++ xi = xr[i]; ++ tmp += xi * xi; ++ } ++ tmp = block_reduce(tmp, s_sum); ++ const float mean = tmp / D; ++ const float sc = rsqrtf(mean + eps); ++ if (i < D) { ++ xs[i] = sc * xi; ++ } ++ __syncthreads(); ++ if (i < D) { ++ float y; ++ if (rot) { ++ y = __fadd_rn(__fmul_rn(xs[j], c0), __fmul_rn(xs[j + half], c1)); ++ } else { ++ y = xs[i]; ++ } ++ if constexpr (has_scale) { ++ y = scale * y + bias; ++ } ++ dst[((n * H + h) * L + l) * D + i] = ggml_cuda_cast(y); ++ } ++ __syncthreads(); ++ } ++} ++ ++// rms_norm_f32_cuda takes its small-row path (and so its reduction order) only here. ++static bool rope_pe_rms_small_rows_enabled() { ++#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) ++ return false; ++#else ++ static const bool enabled = [] { ++ const char * e = getenv("GGML_CUDA_NORM_SMALL_ROWS"); ++ return e == nullptr || atoi(e) != 0; ++ }(); ++ return enabled && WARP_SIZE == 32; ++#endif ++} ++ ++template ++static void rope_pe_permute_rms_launch(const ggml_tensor * src0, const float * pe, dst_t * dst, const int n_rot, ++ const float scale, const float eps, cudaStream_t stream) { ++ const int D = (int) src0->ne[0]; ++ const int H = (int) src0->ne[1]; ++ const int64_t L = src0->ne[2]; ++ const int64_t N = src0->ne[3]; ++ const dim3 grid((unsigned int) L, (unsigned int) N, 1); ++ const dim3 block(block_size, 1, 1); ++ const char * x = (const char *) src0->data; ++ if (scale != 1.0f) { ++ rope_pe_permute_rms_kernel<<>>( ++ x, pe, dst, D, H, L, n_rot, src0->nb[1], src0->nb[2], src0->nb[3], scale, 0.0f, eps); ++ } else { ++ rope_pe_permute_rms_kernel<<>>( ++ x, pe, dst, D, H, L, n_rot, src0->nb[1], src0->nb[2], src0->nb[3], scale, 0.0f, eps); ++ } ++} ++ ++template ++static void rope_pe_permute_rms_cuda(const ggml_tensor * src0, const float * pe, dst_t * dst, const int n_rot, ++ const float scale, const float eps, cudaStream_t stream) { ++ const int D = (int) src0->ne[0]; ++ GGML_ASSERT(D <= 128); ++ if (D <= 32) { ++ rope_pe_permute_rms_launch(src0, pe, dst, n_rot, scale, eps, stream); ++ } else if (D <= 64) { ++ rope_pe_permute_rms_launch(src0, pe, dst, n_rot, scale, eps, stream); ++ } else { ++ rope_pe_permute_rms_launch(src0, pe, dst, n_rot, scale, eps, stream); ++ } ++} ++ + template + static void rope_pe_permute_cuda(const ggml_tensor * src0, const float * pe, dst_t * dst, const int n_rot, + const float scale, cudaStream_t stream) { +@@ -78,6 +184,17 @@ void ggml_cuda_op_rope_pe_permute(ggml_backend_cuda_context & ctx, ggml_tensor * + cudaStream_t stream = ctx.stream(); + const float * pe_d = pe ? (const float *) pe->data : nullptr; + ++ if (ggml_get_op_params_i32(dst, 2) != 0) { ++ const float eps = ggml_get_op_params_f32(dst, 3); ++ if (dst->type == GGML_TYPE_F16) { ++ rope_pe_permute_rms_cuda(src0, pe_d, (half *) dst->data, n_rot, scale, eps, stream); ++ } else { ++ GGML_ASSERT(dst->type == GGML_TYPE_F32); ++ rope_pe_permute_rms_cuda(src0, pe_d, (float *) dst->data, n_rot, scale, eps, stream); ++ } ++ return; ++ } ++ + if (dst->type == GGML_TYPE_F16) { + rope_pe_permute_cuda(src0, pe_d, (half *) dst->data, n_rot, scale, stream); + } else { +@@ -98,5 +215,8 @@ bool ggml_cuda_rope_pe_permute_supported(const ggml_tensor * op) { + if (src0->ne[2] > INT_MAX || src0->ne[3] > 65535 || src0->ne[1] > INT_MAX || src0->ne[0] > INT_MAX) { + return false; + } ++ if (ggml_get_op_params_i32(op, 2) != 0 && (src0->ne[0] > 128 || !rope_pe_rms_small_rows_enabled())) { ++ return false; ++ } + return pe == nullptr || (pe->type == GGML_TYPE_F32 && ggml_is_contiguous(pe)); + } +diff --git a/src/ggml.c b/src/ggml.c +index fec933bf..62a157c6 100644 +--- a/src/ggml.c ++++ b/src/ggml.c +@@ -5594,6 +5594,21 @@ struct ggml_tensor * ggml_rope_pe_permute( + return result; + } + ++struct ggml_tensor * ggml_rope_pe_permute_rms( ++ struct ggml_context * ctx, ++ struct ggml_tensor * a, ++ struct ggml_tensor * pe, ++ int n_rot, ++ float scale, ++ enum ggml_type type, ++ float eps) { ++ GGML_ASSERT(eps >= 0.0f); ++ struct ggml_tensor * result = ggml_rope_pe_permute(ctx, a, pe, n_rot, scale, type); ++ ggml_set_op_params_i32(result, 2, 1); ++ ggml_set_op_params_f32(result, 3, eps); ++ return result; ++} ++ + static struct ggml_tensor * ggml_modulate_rows_impl( + struct ggml_context * ctx, + struct ggml_tensor * acc, +diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp +index 1e46361c..b58c815c 100644 +--- a/tests/test-backend-ops.cpp ++++ b/tests/test-backend-ops.cpp +@@ -2132,15 +2132,16 @@ struct test_rope_pe_permute : public test_case { + const int n_rot; + const float scale; + const bool view; // a is a strided view (one third of a fused qkv projection) ++ const float eps; // >= 0: ggml_rope_pe_permute_rms (rows RMS-normalised first) + + std::string vars() override { +- return VARS_TO_STR5(type, ne_a, n_rot, scale, view); ++ return VARS_TO_STR6(type, ne_a, n_rot, scale, view, eps); + } + + test_rope_pe_permute(ggml_type type = GGML_TYPE_F32, + std::array ne_a = {128, 4, 9, 1}, +- int n_rot = 96, float scale = 1.0f, bool view = false) +- : type(type), ne_a(ne_a), n_rot(n_rot), scale(scale), view(view) {} ++ int n_rot = 96, float scale = 1.0f, bool view = false, float eps = -1.0f) ++ : type(type), ne_a(ne_a), n_rot(n_rot), scale(scale), view(view), eps(eps) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * a; +@@ -2159,7 +2160,8 @@ struct test_rope_pe_permute : public test_case { + pe = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 2, 2, n_rot / 2, ne_a[2]); + ggml_set_name(pe, "pe"); + } +- ggml_tensor * out = ggml_rope_pe_permute(ctx, a, pe, n_rot, scale, type); ++ ggml_tensor * out = eps >= 0.0f ? ggml_rope_pe_permute_rms(ctx, a, pe, n_rot, scale, type, eps) ++ : ggml_rope_pe_permute(ctx, a, pe, n_rot, scale, type); + ggml_set_name(out, "out"); + return out; + } +@@ -8221,6 +8223,11 @@ static std::vector> make_test_cases_eval() { + test_cases.emplace_back(new test_rope_pe_permute(type, {128, 4, 9, 1}, 0, scale, view)); + test_cases.emplace_back(new test_rope_pe_permute(type, {64, 3, 7, 2}, 32, scale, view)); + test_cases.emplace_back(new test_rope_pe_permute(type, {200, 2, 3, 1}, 100, scale, view)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {64, 32, 37, 1}, 48, scale, view, 1e-5f)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {64, 3, 7, 2}, 0, scale, view, 1e-5f)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {128, 4, 9, 1}, 96, scale, view, 1e-6f)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {32, 2, 5, 1}, 16, scale, view, 1e-5f)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {48, 2, 5, 1}, 32, scale, view, 1e-5f)); + } + } + } +-- +2.43.0 + diff --git a/scripts/unsloth/ggml-patches/0006-ggml-cuda-cudnn-attention.patch b/scripts/unsloth/ggml-patches/0006-ggml-cuda-cudnn-attention.patch new file mode 100644 index 0000000000..2b23434e88 --- /dev/null +++ b/scripts/unsloth/ggml-patches/0006-ggml-cuda-cudnn-attention.patch @@ -0,0 +1,994 @@ +From 784b441f8b098863343e69e0345ef62d2ff580af Mon Sep 17 00:00:00 2001 +From: Daniel Han +Date: Sat, 3 Oct 2026 18:27:29 -0700 +Subject: [PATCH] ggml-cuda: optional cuDNN attention for unmasked + FLASH_ATTN_EXT and SAGE_ATTN + +GGML_CUDA_CUDNN=ON (default OFF, Linux) builds a cuDNN fused-attention path +through the header-only cudnn-frontend v1.26.0 (MIT; fetched by CMake or taken +from GGML_CUDA_CUDNN_FRONTEND_DIR) with NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING: +only cudnn.h is needed at build time and libcudnn.so.9 is dlopen'ed on first use +(GGML_CUDA_CUDNN_LIB overrides the name), so the binary has no cuDNN link +dependency and runs unchanged where cuDNN is missing. + +On NVIDIA sm80+ it can take unmasked, sink/ALiBi/softcap-free attention with equal +Q/K/V head sizes (multiple of 8, up to 256) and at least 2^20 attention scores, +GQA and batch included, for FLASH_ATTN_EXT and for SAGE_ATTN in AUTO mode (an +explicitly requested Sage variant keeps its kernel). Q (F32) is narrowed to F16 +into dst's own bytes, F16 K/V are read in place with their strides (F32 K/V are +narrowed), cuDNN writes F16 O to scratch after dst and one kernel widens it into +dst. The scratch is part of the op's graph allocation (max with the existing +kernels' needs). Plans are built once per shape and cached; a shape cuDNN +declines, a failed build or execute, stream capture, Turing and older, or a +missing library keep the existing kernels. + +By default FLASH_ATTN_EXT takes it on sm89 and sm100 and SAGE_ATTN on sm100, +where it measured faster; the existing kernels were faster on sm86 (both ops) +and sm89 (Sage), and other architectures were not measured. +GGML_CUDA_CUDNN_ATTN=0|1 and GGML_CUDA_CUDNN_SAGE=0|1 turn it off or on for every +sm80+ GPU, GGML_CUDA_CUDNN_ATTN_BF16=1 runs BF16, GGML_CUDA_CUDNN_ATTN_LOG=1 +logs each call. + +test-backend-ops: unmasked cases at the MiniMax-H3 head layout (odd and aligned +lengths, GQA with batch, F32 K/V, permuted Q) and unmasked perf cases. +--- + CMakeLists.txt | 3 + + src/ggml-cuda/CMakeLists.txt | 49 ++++ + src/ggml-cuda/cudnn-attn-graph.cpp | 258 +++++++++++++++++++ + src/ggml-cuda/cudnn-attn-graph.h | 31 +++ + src/ggml-cuda/cudnn-attn.cu | 384 +++++++++++++++++++++++++++++ + src/ggml-cuda/cudnn-attn.cuh | 26 ++ + src/ggml-cuda/fattn.cu | 7 +- + src/ggml-cuda/ggml-cuda.cu | 2 +- + src/ggml-cuda/sage-attn.cu | 12 +- + src/ggml-cuda/sage-attn.cuh | 2 +- + tests/test-backend-ops.cpp | 25 ++ + 11 files changed, 793 insertions(+), 6 deletions(-) + create mode 100644 src/ggml-cuda/cudnn-attn-graph.cpp + create mode 100644 src/ggml-cuda/cudnn-attn-graph.h + create mode 100644 src/ggml-cuda/cudnn-attn.cu + create mode 100644 src/ggml-cuda/cudnn-attn.cuh + +diff --git a/CMakeLists.txt b/CMakeLists.txt +index b6843578..1c32cf3b 100644 +--- a/CMakeLists.txt ++++ b/CMakeLists.txt +@@ -208,6 +208,9 @@ option(GGML_CUDA_FA "ggml: compile ggml FlashAttention C + option(GGML_CUDA_FA_ALL_QUANTS "ggml: compile all quants for FlashAttention" OFF) + option(GGML_CUDA_GRAPHS "ggml: use CUDA graphs (llama.cpp only)" ${GGML_CUDA_GRAPHS_DEFAULT}) + option(GGML_CUDA_NCCL "ggml: use NVIDIA Collective Comm. Library" ON) ++option(GGML_CUDA_CUDNN "ggml: cuDNN attention (cudnn-frontend, libcudnn opened at run time)" OFF) ++set (GGML_CUDA_CUDNN_FRONTEND_DIR "" CACHE PATH "ggml: cudnn-frontend checkout for GGML_CUDA_CUDNN (fetched when empty)") ++set (GGML_CUDA_CUDNN_INCLUDE_DIR "" CACHE PATH "ggml: directory holding cudnn.h for GGML_CUDA_CUDNN") + set (GGML_CUDA_COMPRESSION_MODE "size" CACHE STRING + "ggml: cuda link binary compression mode; requires cuda 12.8+") + set_property(CACHE GGML_CUDA_COMPRESSION_MODE PROPERTY STRINGS "none;speed;balance;size") +diff --git a/src/ggml-cuda/CMakeLists.txt b/src/ggml-cuda/CMakeLists.txt +index 05b81487..0ddeb81e 100644 +--- a/src/ggml-cuda/CMakeLists.txt ++++ b/src/ggml-cuda/CMakeLists.txt +@@ -124,11 +124,60 @@ if (CUDAToolkit_FOUND) + template-instances/fattn-vec-instance-bf16-bf16.cu) + endif() + ++ if (GGML_CUDA_CUDNN AND WIN32) ++ message(WARNING "GGML_CUDA_CUDNN: the run-time cuDNN loader is Linux only, building without it") ++ set(GGML_CUDA_CUDNN OFF) ++ endif() ++ if (GGML_CUDA_CUDNN) ++ # Headers only at build time: cudnn.h for the types, cudnn-frontend for the graph API. ++ # libcudnn itself is opened with dlopen when the first attention op runs. ++ if (GGML_CUDA_CUDNN_INCLUDE_DIR) ++ # an explicit directory wins over any earlier cached search result ++ set(GGML_CUDNN_HEADER_DIR ${GGML_CUDA_CUDNN_INCLUDE_DIR}) ++ else() ++ find_path(GGML_CUDNN_HEADER_DIR cudnn.h ++ HINTS $ENV{CUDNN_INCLUDE_DIR} $ENV{CUDNN_PATH} ${CUDAToolkit_INCLUDE_DIRS} ++ PATH_SUFFIXES include) ++ endif() ++ if (NOT GGML_CUDNN_HEADER_DIR OR NOT EXISTS ${GGML_CUDNN_HEADER_DIR}/cudnn.h) ++ message(FATAL_ERROR "GGML_CUDA_CUDNN: cudnn.h not found, set GGML_CUDA_CUDNN_INCLUDE_DIR") ++ endif() ++ if (GGML_CUDA_CUDNN_FRONTEND_DIR) ++ set(GGML_CUDNN_FRONTEND_INCLUDE ${GGML_CUDA_CUDNN_FRONTEND_DIR}/include) ++ else() ++ include(FetchContent) ++ # v1.26.0, the last MIT-licensed release (later ones are Apache-2.0). Only the headers are ++ # used; SOURCE_SUBDIR points at a directory that does not exist so its CMake project stays out. ++ FetchContent_Declare( ++ ggml_cudnn_frontend ++ URL https://github.com/NVIDIA/cudnn-frontend/archive/refs/tags/v1.26.0.tar.gz ++ URL_HASH SHA256=f88361fb714589c970bac2ab56d6a7843722be08ca5ba226484de3fa46fd22cf ++ DOWNLOAD_EXTRACT_TIMESTAMP TRUE ++ SOURCE_SUBDIR ggml-headers-only ++ ) ++ FetchContent_MakeAvailable(ggml_cudnn_frontend) ++ set(GGML_CUDNN_FRONTEND_INCLUDE ${ggml_cudnn_frontend_SOURCE_DIR}/include) ++ endif() ++ if (NOT EXISTS ${GGML_CUDNN_FRONTEND_INCLUDE}/cudnn_frontend.h) ++ message(FATAL_ERROR "GGML_CUDA_CUDNN: cudnn_frontend.h not found under ${GGML_CUDNN_FRONTEND_INCLUDE}") ++ endif() ++ message(STATUS "cuDNN attention: cudnn.h in ${GGML_CUDNN_HEADER_DIR}, cudnn-frontend in ${GGML_CUDNN_FRONTEND_INCLUDE}") ++ list(APPEND GGML_SOURCES_CUDA cudnn-attn-graph.cpp) ++ endif() ++ + ggml_add_backend_library(ggml-cuda + ${GGML_HEADERS_CUDA} + ${GGML_SOURCES_CUDA} + ) + ++ if (GGML_CUDA_CUDNN) ++ target_compile_definitions(ggml-cuda PRIVATE GGML_CUDA_USE_CUDNN NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING) ++ # -isystem: third-party headers, kept out of ggml's warning flags. ++ set_source_files_properties(cudnn-attn-graph.cpp PROPERTIES ++ COMPILE_OPTIONS "-isystem;${GGML_CUDNN_FRONTEND_INCLUDE};-isystem;${GGML_CUDNN_HEADER_DIR}") ++ target_link_libraries(ggml-cuda PRIVATE ${CMAKE_DL_LIBS}) ++ endif() ++ + if (CUDAToolkit_VERSION VERSION_GREATER_EQUAL "11.8") + target_compile_definitions(ggml-cuda PRIVATE GGML_CUDA_USE_CUBLASLT_FP8) + if (NOT GGML_STATIC OR WIN32) +diff --git a/src/ggml-cuda/cudnn-attn-graph.cpp b/src/ggml-cuda/cudnn-attn-graph.cpp +new file mode 100644 +index 00000000..b0f1dc11 +--- /dev/null ++++ b/src/ggml-cuda/cudnn-attn-graph.cpp +@@ -0,0 +1,258 @@ ++// cuDNN scaled dot product attention through cudnn-frontend (header-only), with libcudnn opened at ++// run time (NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING): a build with GGML_CUDA_CUDNN still starts and ++// runs on hosts without cuDNN, it just keeps the ggml kernels. ++ ++#include "cudnn-attn-graph.h" ++ ++#include "ggml-impl.h" ++ ++#include ++ ++#include ++ ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++ ++namespace cudnn_frontend { ++// cudnn-frontend resolves every backend call through this handle when dynamic loading is on. ++void * cudnn_dlhandle = nullptr; ++} // namespace cudnn_frontend ++ ++namespace fe = cudnn_frontend; ++ ++namespace { ++ ++enum : int64_t { UID_Q = 1, UID_K = 2, UID_V = 3, UID_O = 4 }; ++ ++using plan_key = std::array; ++ ++struct plan_entry { ++ std::shared_ptr graph; // null: cuDNN has no plan for this key ++ int64_t workspace = 0; ++}; ++ ++std::mutex g_mutex; ++int g_load_state = 0; // 0 not tried, 1 loaded, -1 unavailable ++std::map g_handles; ++std::map g_plans; ++ ++plan_key make_key(const ggml_cuda_cudnn_sdpa_params & p) { ++ plan_key k{}; ++ int i = 0; ++ k[i++] = p.device; ++ k[i++] = p.b; ++ k[i++] = p.hq; ++ k[i++] = p.hk; ++ k[i++] = p.sq; ++ k[i++] = p.skv; ++ k[i++] = p.d; ++ for (int j = 0; j < 4; ++j) { ++ k[i++] = p.q_stride[j]; ++ } ++ for (int j = 0; j < 4; ++j) { ++ k[i++] = p.k_stride[j]; ++ } ++ for (int j = 0; j < 4; ++j) { ++ k[i++] = p.v_stride[j]; ++ } ++ for (int j = 0; j < 4; ++j) { ++ k[i++] = p.o_stride[j]; ++ } ++ uint32_t scale_bits = 0; ++ std::memcpy(&scale_bits, &p.scale, sizeof(scale_bits)); ++ k[i++] = scale_bits; ++ k[i++] = p.bf16 ? 1 : 0; ++ GGML_ASSERT(i == (int) k.size()); ++ return k; ++} ++ ++bool load_locked() { ++ if (g_load_state != 0) { ++ return g_load_state == 1; ++ } ++ g_load_state = -1; ++ ++ // cudnn-frontend opens libcudart itself for a few runtime calls; point it at the runtime major ++ // this backend was built against so it never picks a second, mismatched copy. ++ char cudart_name[32]; ++ std::snprintf(cudart_name, sizeof(cudart_name), "libcudart.so.%d", CUDART_VERSION / 1000); ++ setenv("CUDNN_FRONTEND_CUDART_LIB_NAME", cudart_name, 0); ++ ++ const char * user = std::getenv("GGML_CUDA_CUDNN_LIB"); ++ const char * name = user != nullptr && user[0] != '\0' ? user : "libcudnn.so.9"; ++ void * h = dlopen(name, RTLD_NOW | RTLD_LOCAL); ++ if (h == nullptr) { ++ GGML_LOG_DEBUG("%s: cuDNN not loaded (%s), attention keeps the ggml kernels\n", __func__, dlerror()); ++ return false; ++ } ++ using get_version_t = size_t (*)(); ++ auto get_version = reinterpret_cast(dlsym(h, "cudnnGetVersion")); ++ const size_t version = get_version != nullptr ? get_version() : 0; ++ if (version < 90000) { ++ GGML_LOG_WARN("%s: %s reports cuDNN %zu, need 9.0 or newer; attention keeps the ggml kernels\n", __func__, ++ name, version); ++ dlclose(h); ++ return false; ++ } ++ fe::cudnn_dlhandle = h; ++ g_load_state = 1; ++ GGML_LOG_INFO("%s: cuDNN %zu loaded from %s for attention\n", __func__, version, name); ++ return true; ++} ++ ++cudnnHandle_t handle_locked(int device) { ++ auto it = g_handles.find(device); ++ if (it != g_handles.end()) { ++ return it->second; ++ } ++ cudnnHandle_t h = nullptr; ++ if (fe::detail::create_handle(&h) != CUDNN_STATUS_SUCCESS) { ++ h = nullptr; ++ } ++ g_handles[device] = h; ++ return h; ++} ++ ++std::shared_ptr build_graph(const ggml_cuda_cudnn_sdpa_params & p, cudnnHandle_t handle, ++ std::string & why) { ++ auto g = std::make_shared(); ++ g->set_io_data_type(p.bf16 ? fe::DataType_t::BFLOAT16 : fe::DataType_t::HALF) ++ .set_intermediate_data_type(fe::DataType_t::FLOAT) ++ .set_compute_data_type(fe::DataType_t::FLOAT); ++ ++ auto tensor = [&](const char * name, int64_t uid, int64_t h, int64_t s, const int64_t * stride) { ++ return g->tensor(fe::graph::Tensor_attributes() ++ .set_name(name) ++ .set_uid(uid) ++ .set_dim({ p.b, h, s, p.d }) ++ .set_stride({ stride[0], stride[1], stride[2], stride[3] })); ++ }; ++ auto Q = tensor("Q", UID_Q, p.hq, p.sq, p.q_stride); ++ auto K = tensor("K", UID_K, p.hk, p.skv, p.k_stride); ++ auto V = tensor("V", UID_V, p.hk, p.skv, p.v_stride); ++ ++ auto options = fe::graph::SDPA_attributes().set_name("ggml_sdpa").set_generate_stats(false).set_attn_scale(p.scale); ++ auto [O, stats] = g->sdpa(Q, K, V, options); ++ (void) stats; ++ O->set_output(true) ++ .set_uid(UID_O) ++ .set_dim({ p.b, p.hq, p.sq, p.d }) ++ .set_stride({ p.o_stride[0], p.o_stride[1], p.o_stride[2], p.o_stride[3] }); ++ ++ fe::error_t st = g->validate(); ++ if (st.is_good()) { ++ st = g->build_operation_graph(handle); ++ } ++ if (st.is_good()) { ++ st = g->create_execution_plans({ fe::HeurMode_t::A, fe::HeurMode_t::FALLBACK }); ++ } ++ if (st.is_good()) { ++ st = g->check_support(handle); ++ } ++ if (st.is_good()) { ++ st = g->build_plans(handle); ++ } ++ if (!st.is_good()) { ++ why = st.get_message(); ++ return nullptr; ++ } ++ return g; ++} ++ ++} // namespace ++ ++bool ggml_cuda_cudnn_available() { ++ std::lock_guard lock(g_mutex); ++ try { ++ return load_locked(); ++ } catch (const std::exception & e) { ++ GGML_LOG_WARN("%s: %s\n", __func__, e.what()); ++ g_load_state = -1; ++ return false; ++ } ++} ++ ++bool ggml_cuda_cudnn_sdpa_prepare(const ggml_cuda_cudnn_sdpa_params & p, size_t * workspace_size) { ++ std::lock_guard lock(g_mutex); ++ if (g_load_state != 1) { ++ return false; ++ } ++ const plan_key key = make_key(p); ++ auto it = g_plans.find(key); ++ if (it == g_plans.end()) { ++ plan_entry entry; ++ std::string why; ++ const auto t0 = std::chrono::steady_clock::now(); ++ try { ++ cudnnHandle_t handle = handle_locked(p.device); ++ if (handle == nullptr) { ++ why = "cudnnCreate failed"; ++ } else { ++ entry.graph = build_graph(p, handle, why); ++ if (entry.graph && !entry.graph->get_workspace_size(entry.workspace).is_good()) { ++ entry.graph = nullptr; ++ why = "workspace query failed"; ++ } ++ } ++ } catch (const std::exception & e) { ++ entry.graph = nullptr; ++ why = e.what(); ++ } ++ const double ms = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); ++ if (entry.graph) { ++ GGML_LOG_INFO("%s: cuDNN SDPA plan b=%lld hq=%lld hk=%lld sq=%lld skv=%lld d=%lld %s, workspace %lld B, " ++ "built in %.0f ms\n", ++ __func__, (long long) p.b, (long long) p.hq, (long long) p.hk, (long long) p.sq, ++ (long long) p.skv, (long long) p.d, p.bf16 ? "bf16" : "f16", (long long) entry.workspace, ++ ms); ++ } else { ++ GGML_LOG_INFO("%s: no cuDNN SDPA plan for b=%lld hq=%lld hk=%lld sq=%lld skv=%lld d=%lld, keeping the ggml " ++ "kernel (%s)\n", ++ __func__, (long long) p.b, (long long) p.hq, (long long) p.hk, (long long) p.sq, ++ (long long) p.skv, (long long) p.d, why.c_str()); ++ } ++ it = g_plans.emplace(key, std::move(entry)).first; ++ } ++ if (!it->second.graph) { ++ return false; ++ } ++ *workspace_size = (size_t) it->second.workspace; ++ return true; ++} ++ ++bool ggml_cuda_cudnn_sdpa_execute(const ggml_cuda_cudnn_sdpa_params & p, cudaStream_t stream, void * q, void * k, ++ void * v, void * o, void * workspace) { ++ std::lock_guard lock(g_mutex); ++ auto it = g_plans.find(make_key(p)); ++ if (it == g_plans.end() || !it->second.graph) { ++ return false; ++ } ++ try { ++ cudnnHandle_t handle = handle_locked(p.device); ++ if (handle == nullptr || fe::detail::set_stream(handle, stream) != CUDNN_STATUS_SUCCESS) { ++ return false; ++ } ++ std::unordered_map pack = { ++ { UID_Q, q }, { UID_K, k }, { UID_V, v }, { UID_O, o } ++ }; ++ fe::error_t st = it->second.graph->execute(handle, pack, workspace); ++ if (!st.is_good()) { ++ GGML_LOG_WARN("%s: cuDNN SDPA execute failed, keeping the ggml kernel for this shape: %s\n", __func__, ++ st.get_message().c_str()); ++ it->second.graph = nullptr; ++ return false; ++ } ++ } catch (const std::exception & e) { ++ GGML_LOG_WARN("%s: cuDNN SDPA execute failed, keeping the ggml kernel for this shape: %s\n", __func__, e.what()); ++ it->second.graph = nullptr; ++ return false; ++ } ++ return true; ++} +diff --git a/src/ggml-cuda/cudnn-attn-graph.h b/src/ggml-cuda/cudnn-attn-graph.h +new file mode 100644 +index 00000000..9749f2ab +--- /dev/null ++++ b/src/ggml-cuda/cudnn-attn-graph.h +@@ -0,0 +1,31 @@ ++#pragma once ++ ++// Host-side half of the cuDNN attention path (cudnn-attn.cu is the device side). Only compiled ++// with GGML_CUDA_CUDNN. libcudnn is opened at run time, so the binary never links against it. ++ ++#include ++#include ++ ++#include ++ ++struct ggml_cuda_cudnn_sdpa_params { ++ int device; ++ int64_t b, hq, hk, sq, skv, d; // batch, query heads, key/value heads, query tokens, key tokens, head size ++ int64_t q_stride[4]; // element strides in (b, h, s, d) order ++ int64_t k_stride[4]; ++ int64_t v_stride[4]; ++ int64_t o_stride[4]; ++ float scale; ++ bool bf16; // io type: BF16 when set, F16 otherwise ++}; ++ ++// Opens libcudnn once per process. False when it cannot be loaded or is older than cuDNN 9. ++bool ggml_cuda_cudnn_available(); ++ ++// Builds (or finds) the execution plan for these parameters. False when cuDNN has no plan for ++// them; the result is cached either way, so a rejected shape costs one attempt per process. ++bool ggml_cuda_cudnn_sdpa_prepare(const ggml_cuda_cudnn_sdpa_params & p, size_t * workspace_size); ++ ++// Runs a plan built by ggml_cuda_cudnn_sdpa_prepare on the given stream. ++bool ggml_cuda_cudnn_sdpa_execute(const ggml_cuda_cudnn_sdpa_params & p, cudaStream_t stream, ++ void * q, void * k, void * v, void * o, void * workspace); +diff --git a/src/ggml-cuda/cudnn-attn.cu b/src/ggml-cuda/cudnn-attn.cu +new file mode 100644 +index 00000000..a5b9de14 +--- /dev/null ++++ b/src/ggml-cuda/cudnn-attn.cu +@@ -0,0 +1,384 @@ ++#include "cudnn-attn.cuh" ++ ++#if defined(GGML_CUDA_USE_CUDNN) && !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) ++ ++#include "convert.cuh" ++#include "cudnn-attn-graph.h" ++ ++#include ++#include ++#include ++ ++static bool cudnn_attn_env_set(const char * name, bool default_value) { ++ const char * v = std::getenv(name); ++ if (v == nullptr || v[0] == '\0') { ++ return default_value; ++ } ++ return !(std::strcmp(v, "0") == 0 || std::strcmp(v, "off") == 0 || std::strcmp(v, "false") == 0); ++} ++ ++// Where a staged F16 / BF16 tensor lives: inside dst's bytes or in the scratch after dst. ++struct cudnn_attn_slot { ++ bool in_dst; ++ size_t off; ++}; ++ ++// Everything the cuDNN path needs to know about one op, from shapes and types alone. ++struct cudnn_attn_plan { ++ ggml_cuda_cudnn_sdpa_params p; ++ bool convert_k; ++ bool convert_v; ++ cudnn_attn_slot q, k, v, o; ++ size_t scratch; // scratch bytes after dst ++}; ++ ++static bool cudnn_attn_strides_ok(const ggml_tensor * t) { ++ // cuDNN wants a unit stride on the head dimension and 16-byte aligned rows. ++ return t->nb[0] == ggml_type_size(t->type) && t->nb[1] % 16 == 0 && t->nb[2] % 16 == 0 && t->nb[3] % 16 == 0; ++} ++ ++// -1 when unset, otherwise 0 or 1. ++static int cudnn_attn_env_mode(const char * name) { ++ const char * v = std::getenv(name); ++ if (v == nullptr || v[0] == '\0') { ++ return -1; ++ } ++ return cudnn_attn_env_set(name, true) ? 1 : 0; ++} ++ ++static bool cudnn_attn_make_plan(int device, const ggml_tensor * dst, cudnn_attn_plan & plan) { ++ static const int attn_mode = cudnn_attn_env_mode("GGML_CUDA_CUDNN_ATTN"); ++ static const int sage_mode = cudnn_attn_env_mode("GGML_CUDA_CUDNN_SAGE"); ++ static const bool want_bf16 = cudnn_attn_env_set("GGML_CUDA_CUDNN_ATTN_BF16", false); ++ if (attn_mode == 0 || (dst->op == GGML_OP_SAGE_ATTN && sage_mode == 0)) { ++ return false; ++ } ++ if (dst->op != GGML_OP_FLASH_ATTN_EXT && dst->op != GGML_OP_SAGE_ATTN) { ++ return false; ++ } ++ // An explicitly requested SageAttention variant keeps its kernel; AUTO takes the cuDNN path. ++ if (dst->op == GGML_OP_SAGE_ATTN && ggml_get_op_params_i32(dst, 1) != GGML_SAGE_ATTN_AUTO) { ++ return false; ++ } ++ const int cc = ggml_cuda_info().devices[device].cc; ++ if (!GGML_CUDA_CC_IS_NVIDIA(cc) || cc < GGML_CUDA_CC_AMPERE) { ++ return false; ++ } ++ // Unset switches keep the existing kernels except where cuDNN measured faster at H3 shapes: against the flash ++ // attention kernel on sm89 and sm100 (slower on sm86), against the Sage kernel on sm100 (slower on sm86 and ++ // sm89). Other architectures were not measured. =1 opts in on any sm80+ GPU. ++ const bool sm100 = cc >= 1000 && cc < GGML_CUDA_CC_BLACKWELL; ++ const bool sm89 = cc >= GGML_CUDA_CC_ADA_LOVELACE && cc < GGML_CUDA_CC_HOPPER; ++ if (dst->op == GGML_OP_SAGE_ATTN ? sage_mode < 0 && !sm100 : attn_mode < 0 && !sm89 && !sm100) { ++ return false; ++ } ++ ++ const ggml_tensor * Q = dst->src[0]; ++ const ggml_tensor * K = dst->src[1]; ++ const ggml_tensor * V = dst->src[2]; ++ if (Q == nullptr || K == nullptr || V == nullptr) { ++ return false; ++ } ++ if (dst->op == GGML_OP_FLASH_ATTN_EXT) { ++ float max_bias = 0.0f, softcap = 0.0f; ++ std::memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); ++ std::memcpy(&softcap, (const float *) dst->op_params + 2, sizeof(float)); ++ if (dst->src[3] != nullptr || dst->src[4] != nullptr || max_bias != 0.0f || softcap != 0.0f) { ++ return false; ++ } ++ } ++ ++ const ggml_type io = want_bf16 ? GGML_TYPE_BF16 : GGML_TYPE_F16; ++ auto kv_type_ok = [](ggml_type t) { ++ return t == GGML_TYPE_F16 || t == GGML_TYPE_F32 || t == GGML_TYPE_BF16; ++ }; ++ if (Q->type != GGML_TYPE_F32 || !kv_type_ok(K->type) || !kv_type_ok(V->type) || dst->type != GGML_TYPE_F32) { ++ return false; ++ } ++ // BF16 -> F16 can overflow, so BF16 K/V only run in BF16 mode. F16 K/V in BF16 mode are narrowed ++ // (3 mantissa bits lost); that mode exists to measure BF16 against the F16 default. ++ if ((K->type == GGML_TYPE_BF16 || V->type == GGML_TYPE_BF16) && io != GGML_TYPE_BF16) { ++ return false; ++ } ++ ++ const int64_t d = Q->ne[0]; ++ if (d % 8 != 0 || d > 256 || K->ne[0] != d || V->ne[0] != d) { ++ return false; ++ } ++ const int64_t sq = Q->ne[1], hq = Q->ne[2], b = Q->ne[3]; ++ const int64_t skv = K->ne[1], hk = K->ne[2]; ++ if (V->ne[1] != skv || V->ne[2] != hk || K->ne[3] != b || V->ne[3] != b || hk == 0 || hq % hk != 0) { ++ return false; ++ } ++ if (!ggml_is_contiguous(dst) || dst->ne[0] != d || dst->ne[1] != hq || dst->ne[2] != sq || dst->ne[3] != b) { ++ return false; ++ } ++ if (Q->nb[0] != ggml_type_size(Q->type) || !cudnn_attn_strides_ok(K) || !cudnn_attn_strides_ok(V)) { ++ return false; ++ } ++ if (hq * b > INT32_MAX || sq > INT32_MAX || skv > INT32_MAX) { ++ return false; ++ } ++ // Below about 2^20 attention scores the two extra conversion launches cost more than cuDNN saves ++ // (B200: 31 tokens x 56 heads 10.6 -> 14-16 us, one query row x 32 heads x 4096 keys 12.9 -> 20 us, ++ // break-even near 256 x 256 x 16, 2.3x faster at 512 x 512 x 16), so small shapes keep the ggml kernels. ++ if ((double) sq * (double) skv * (double) (hq * b) < (double) (1 << 20)) { ++ return false; ++ } ++ if (!ggml_cuda_cudnn_available()) { ++ return false; ++ } ++ ++ plan = {}; ++ plan.p.device = device; ++ plan.p.b = b; ++ plan.p.hq = hq; ++ plan.p.hk = hk; ++ plan.p.sq = sq; ++ plan.p.skv = skv; ++ plan.p.d = d; ++ plan.p.bf16 = io == GGML_TYPE_BF16; ++ std::memcpy(&plan.p.scale, (const float *) dst->op_params + 0, sizeof(float)); ++ ++ // Converted tensors are written contiguous: [d, s, h, b], i.e. (b, h, s, d) strides below. ++ auto set_contiguous = [d](int64_t * stride, int64_t s, int64_t h) { ++ stride[3] = 1; ++ stride[2] = d; ++ stride[1] = d * s; ++ stride[0] = d * s * h; ++ }; ++ auto set_from = [](int64_t * stride, const ggml_tensor * t) { ++ const size_t ts = ggml_type_size(t->type); ++ stride[3] = 1; ++ stride[2] = t->nb[1] / ts; ++ stride[1] = t->nb[2] / ts; ++ stride[0] = t->nb[3] / ts; ++ }; ++ ++ // Q (and K/V when they need converting) are staged in dst's own bytes: dst is F32 and at least ++ // twice the size of F16 Q, and it is only written after cuDNN has read them. O and whatever does ++ // not fit go to scratch after dst. ++ const size_t es = 2; // F16 / BF16 ++ const size_t dst_size = ggml_nbytes(dst); ++ size_t dst_off = 0; ++ size_t off = 0; ++ auto take = [&](size_t bytes, cudnn_attn_slot & slot, bool allow_dst) { ++ if (allow_dst && dst_off + bytes <= dst_size) { ++ slot = { true, dst_off }; ++ dst_off += GGML_PAD(bytes, 256); ++ } else { ++ slot = { false, off }; ++ off += GGML_PAD(bytes, 256); ++ } ++ }; ++ ++ take(ggml_nelements(Q) * es, plan.q, true); ++ set_contiguous(plan.p.q_stride, sq, hq); ++ ++ plan.convert_k = K->type != io; ++ if (plan.convert_k) { ++ take(ggml_nelements(K) * es, plan.k, true); ++ set_contiguous(plan.p.k_stride, skv, hk); ++ } else { ++ set_from(plan.p.k_stride, K); ++ } ++ plan.convert_v = V->type != io; ++ if (plan.convert_v) { ++ take(ggml_nelements(V) * es, plan.v, true); ++ set_contiguous(plan.p.v_stride, skv, hk); ++ } else { ++ set_from(plan.p.v_stride, V); ++ } ++ ++ // O is written in dst's layout, [d, hq, sq, b], then widened to F32 into dst. ++ take(ggml_nelements(dst) * es, plan.o, false); ++ plan.p.o_stride[3] = 1; ++ plan.p.o_stride[1] = d; ++ plan.p.o_stride[2] = d * hq; ++ plan.p.o_stride[0] = d * hq * sq; ++ plan.scratch = off; ++ return true; ++} ++ ++size_t ggml_cuda_cudnn_attn_alloc_size(int device, const ggml_tensor * dst) { ++ cudnn_attn_plan plan; ++ if (!cudnn_attn_make_plan(device, dst, plan)) { ++ return 0; ++ } ++ // dst (which also stages Q), up to 256 bytes to align the scratch start, then the scratch. ++ return ggml_nbytes(dst) + 256 + plan.scratch; ++} ++ ++// F32 rows of 8-element groups (strided rows, 16-byte aligned) -> contiguous F16 / BF16, 8 per thread. ++template ++static __global__ void cudnn_attn_narrow_f32(const char * __restrict__ x, dst_t * __restrict__ y, const int64_t n8, ++ const int64_t g0, const int64_t ne1, const int64_t ne2, const int64_t nb1, ++ const int64_t nb2, const int64_t nb3) { ++ for (int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; i < n8; i += (int64_t) gridDim.x * blockDim.x) { ++ const int64_t g = i % g0; ++ const int64_t row = i / g0; ++ const int64_t i1 = row % ne1; ++ const int64_t i2 = (row / ne1) % ne2; ++ const int64_t i3 = row / (ne1 * ne2); ++ const float4 * src = (const float4 *) (x + i3 * nb3 + i2 * nb2 + i1 * nb1) + 2 * g; ++ const float4 a = src[0]; ++ const float4 b = src[1]; ++ union { ++ dst_t h[8]; ++ uint4 u; ++ } out; ++ out.h[0] = ggml_cuda_cast(a.x); ++ out.h[1] = ggml_cuda_cast(a.y); ++ out.h[2] = ggml_cuda_cast(a.z); ++ out.h[3] = ggml_cuda_cast(a.w); ++ out.h[4] = ggml_cuda_cast(b.x); ++ out.h[5] = ggml_cuda_cast(b.y); ++ out.h[6] = ggml_cuda_cast(b.z); ++ out.h[7] = ggml_cuda_cast(b.w); ++ ((uint4 *) y)[i] = out.u; ++ } ++} ++ ++// Contiguous F16 / BF16 -> contiguous F32, 8 per thread. ++template ++static __global__ void cudnn_attn_widen(const src_t * __restrict__ x, float * __restrict__ y, const int64_t n8) { ++ for (int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; i < n8; i += (int64_t) gridDim.x * blockDim.x) { ++ union { ++ src_t h[8]; ++ uint4 u; ++ } in; ++ in.u = ((const uint4 *) x)[i]; ++ float4 a, b; ++ a.x = ggml_cuda_cast(in.h[0]); ++ a.y = ggml_cuda_cast(in.h[1]); ++ a.z = ggml_cuda_cast(in.h[2]); ++ a.w = ggml_cuda_cast(in.h[3]); ++ b.x = ggml_cuda_cast(in.h[4]); ++ b.y = ggml_cuda_cast(in.h[5]); ++ b.z = ggml_cuda_cast(in.h[6]); ++ b.w = ggml_cuda_cast(in.h[7]); ++ ((float4 *) y)[2 * i] = a; ++ ((float4 *) y)[2 * i + 1] = b; ++ } ++} ++ ++static int cudnn_attn_grid(int64_t n) { ++ return (int) std::min((n + 255) / 256, 65535LL * 8); ++} ++ ++static void cudnn_attn_convert(const ggml_tensor * src, void * dst, bool bf16, cudaStream_t stream) { ++ const size_t ts = ggml_type_size(src->type); ++ const bool vec_ok = src->type == GGML_TYPE_F32 && src->ne[0] % 8 == 0 && (uintptr_t) src->data % 16 == 0 && ++ src->nb[1] % 16 == 0 && src->nb[2] % 16 == 0 && src->nb[3] % 16 == 0; ++ if (vec_ok) { ++ const int64_t n8 = ggml_nelements(src) / 8; ++ if (bf16) { ++ cudnn_attn_narrow_f32<<>>( ++ (const char *) src->data, (nv_bfloat16 *) dst, n8, src->ne[0] / 8, src->ne[1], src->ne[2], src->nb[1], ++ src->nb[2], src->nb[3]); ++ } else { ++ cudnn_attn_narrow_f32<<>>( ++ (const char *) src->data, (half *) dst, n8, src->ne[0] / 8, src->ne[1], src->ne[2], src->nb[1], ++ src->nb[2], src->nb[3]); ++ } ++ return; ++ } ++ if (bf16) { ++ to_bf16_nc_cuda_t f = ggml_get_to_bf16_nc_cuda(src->type); ++ GGML_ASSERT(f != nullptr); ++ f(src->data, (nv_bfloat16 *) dst, src->ne[0], src->ne[1], src->ne[2], src->ne[3], src->nb[1] / ts, ++ src->nb[2] / ts, src->nb[3] / ts, stream); ++ } else { ++ to_fp16_nc_cuda_t f = ggml_get_to_fp16_nc_cuda(src->type); ++ GGML_ASSERT(f != nullptr); ++ f(src->data, (half *) dst, src->ne[0], src->ne[1], src->ne[2], src->ne[3], src->nb[1] / ts, src->nb[2] / ts, ++ src->nb[3] / ts, stream); ++ } ++} ++ ++bool ggml_cuda_cudnn_attn(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ++ static const bool log_calls = cudnn_attn_env_set("GGML_CUDA_CUDNN_ATTN_LOG", false); ++ static std::atomic n_calls{ 0 }; ++ ++ cudnn_attn_plan plan; ++ if (!cudnn_attn_make_plan(ctx.device, dst, plan)) { ++ return false; ++ } ++ cudaStream_t stream = ctx.stream(); ++ ++ // Plans are built outside stream capture only; a captured graph keeps the ggml kernels. ++ cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; ++ CUDA_CHECK(cudaStreamIsCapturing(stream, &capture)); ++ if (capture != cudaStreamCaptureStatusNone) { ++ return false; ++ } ++ ++ const ggml_tensor * Q = dst->src[0]; ++ const ggml_tensor * K = dst->src[1]; ++ const ggml_tensor * V = dst->src[2]; ++ if ((uintptr_t) dst->data % 16 != 0 || (!plan.convert_k && (uintptr_t) K->data % 16 != 0) || ++ (!plan.convert_v && (uintptr_t) V->data % 16 != 0)) { ++ return false; ++ } ++ ++ size_t workspace_size = 0; ++ if (!ggml_cuda_cudnn_sdpa_prepare(plan.p, &workspace_size)) { ++ if (log_calls) { ++ GGML_LOG_INFO("%s: %s %s -> ggml kernel (no cuDNN plan)\n", __func__, ggml_op_desc(dst), dst->name); ++ } ++ return false; ++ } ++ ++ char * base = (char *) dst->data; ++ char * scratch = (char *) GGML_PAD((uintptr_t) dst->data + ggml_nbytes(dst), 256); ++ auto where = [&](const cudnn_attn_slot & slot) -> void * { ++ return (slot.in_dst ? base : scratch) + slot.off; ++ }; ++ void * q = where(plan.q); ++ void * k = plan.convert_k ? where(plan.k) : K->data; ++ void * v = plan.convert_v ? where(plan.v) : V->data; ++ void * o = where(plan.o); ++ ++ cudnn_attn_convert(Q, q, plan.p.bf16, stream); ++ if (plan.convert_k) { ++ cudnn_attn_convert(K, k, plan.p.bf16, stream); ++ } ++ if (plan.convert_v) { ++ cudnn_attn_convert(V, v, plan.p.bf16, stream); ++ } ++ ++ ggml_cuda_pool_alloc workspace(ctx.pool()); ++ if (workspace_size > 0) { ++ workspace.alloc(workspace_size); ++ } ++ if (!ggml_cuda_cudnn_sdpa_execute(plan.p, stream, q, k, v, o, workspace_size > 0 ? workspace.get() : nullptr)) { ++ return false; ++ } ++ ++ // dst holds d * hq * sq * b values with d % 8 == 0, and both buffers are 256-byte aligned. ++ const int64_t n8 = ggml_nelements(dst) / 8; ++ if (plan.p.bf16) { ++ cudnn_attn_widen<<>>((const nv_bfloat16 *) o, (float *) dst->data, n8); ++ } else { ++ cudnn_attn_widen<<>>((const half *) o, (float *) dst->data, n8); ++ } ++ CUDA_CHECK(cudaGetLastError()); ++ ++ const long long n = ++n_calls; ++ if (log_calls) { ++ GGML_LOG_INFO("%s: %s %s -> cuDNN (call %lld)\n", __func__, ggml_op_desc(dst), dst->name, n); ++ } ++ return true; ++} ++ ++#else ++ ++size_t ggml_cuda_cudnn_attn_alloc_size(int, const ggml_tensor *) { ++ return 0; ++} ++ ++bool ggml_cuda_cudnn_attn(ggml_backend_cuda_context &, ggml_tensor *) { ++ return false; ++} ++ ++#endif +diff --git a/src/ggml-cuda/cudnn-attn.cuh b/src/ggml-cuda/cudnn-attn.cuh +new file mode 100644 +index 00000000..067bea63 +--- /dev/null ++++ b/src/ggml-cuda/cudnn-attn.cuh +@@ -0,0 +1,26 @@ ++#pragma once ++ ++#include "common.cuh" ++ ++// cuDNN scaled dot product attention for GGML_OP_FLASH_ATTN_EXT and GGML_OP_SAGE_ATTN on NVIDIA ++// sm80+ when the backend is built with GGML_CUDA_CUDNN and libcudnn 9 loads at run time. ++// Only unmasked, sink-free, ALiBi-free, softcap-free attention with equal Q/K/V head sizes and at ++// least 2^20 attention scores is taken; everything else, and every shape cuDNN rejects, keeps the ++// existing kernels. ++// ++// By default FLASH_ATTN_EXT takes it on sm89 and sm100 and SAGE_ATTN on sm100, where it measured faster. ++// Environment: ++// GGML_CUDA_CUDNN_ATTN=0|1 never use cuDNN for attention / use it on every sm80+ GPU ++// GGML_CUDA_CUDNN_SAGE=0|1 keep the SageAttention kernels for GGML_OP_SAGE_ATTN / use cuDNN on every ++// sm80+ GPU (AUTO mode; an explicitly requested Sage variant keeps its kernel) ++// GGML_CUDA_CUDNN_ATTN_BF16=1 run cuDNN in BF16 instead of F16 (F16 K/V are narrowed) ++// GGML_CUDA_CUDNN_LIB=path library to open instead of libcudnn.so.9 ++// GGML_CUDA_CUDNN_ATTN_LOG=1 log every call that takes (or falls back from) the cuDNN path ++ ++// Bytes the op needs from the graph allocation when the cuDNN path can run, 0 otherwise. The ++// region starts at dst->data and includes dst itself. ++size_t ggml_cuda_cudnn_attn_alloc_size(int device, const ggml_tensor * dst); ++ ++// Runs dst through cuDNN. Returns false without touching dst when the cuDNN path does not apply, ++// so the caller runs its own kernel. ++bool ggml_cuda_cudnn_attn(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +diff --git a/src/ggml-cuda/fattn.cu b/src/ggml-cuda/fattn.cu +index b8be30a0..ddc93e30 100644 +--- a/src/ggml-cuda/fattn.cu ++++ b/src/ggml-cuda/fattn.cu +@@ -4,6 +4,7 @@ + #include "fattn-tile.cuh" + #include "fattn-vec.cuh" + #include "fattn.cuh" ++#include "cudnn-attn.cuh" + + template + static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { +@@ -573,11 +574,15 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d + const ggml_cuda_flash_attn_ext_f16_extra_data f16_extra = + ggml_cuda_flash_attn_ext_get_f16_extra_data(dst, need_f16_K, need_f16_V); + +- return f16_extra.end - (uintptr_t) dst->data; ++ // The cuDNN path, when it can run, uses the same region after dst for its own scratch. ++ return std::max(f16_extra.end - (uintptr_t) dst->data, ggml_cuda_cudnn_attn_alloc_size(device, dst)); + } + + void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + ggml_cuda_set_device(ctx.device); ++ if (ggml_cuda_cudnn_attn(ctx, dst)) { ++ return; ++ } + switch (ggml_cuda_get_best_fattn_kernel(ggml_cuda_get_device(), dst)) { + case BEST_FATTN_KERNEL_NONE: + GGML_ABORT("fatal error"); +diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu +index 86a6cc5b..749aead9 100644 +--- a/src/ggml-cuda/ggml-cuda.cu ++++ b/src/ggml-cuda/ggml-cuda.cu +@@ -919,7 +919,7 @@ static size_t ggml_backend_cuda_buffer_type_get_alloc_size(ggml_backend_buffer_t + ? ggml_cuda_flash_attn_ext_get_alloc_size(buft_ctx->device, tensor) + : ggml_nbytes(tensor); + if (tensor->op == GGML_OP_SAGE_ATTN) { +- size = ggml_cuda_sage_attn_get_alloc_size(tensor); ++ size = ggml_cuda_sage_attn_get_alloc_size(buft_ctx->device, tensor); + } + int64_t ne0 = tensor->ne[0]; + +diff --git a/src/ggml-cuda/sage-attn.cu b/src/ggml-cuda/sage-attn.cu +index 2c103eb0..ceb29ed2 100644 +--- a/src/ggml-cuda/sage-attn.cu ++++ b/src/ggml-cuda/sage-attn.cu +@@ -1,4 +1,5 @@ + #include "sage-attn.cuh" ++#include "cudnn-attn.cuh" + + #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDACC_VER_MAJOR__ >= 12 + +@@ -332,13 +333,18 @@ bool ggml_cuda_sage_attn_supported(int device, const ggml_tensor* op) { + op->ne[2] == q->ne[1] && op->ne[3] == q->ne[3]; + } + +-size_t ggml_cuda_sage_attn_get_alloc_size(const ggml_tensor* op) { ++size_t ggml_cuda_sage_attn_get_alloc_size(int device, const ggml_tensor* op) { + // Scratch belongs to the graph allocation, so graph capture and VRAM planning see its full lifetime. +- return SageWorkspace(op).size; ++ // The cuDNN path, when it can run, reuses the same region for its own scratch. ++ return std::max(SageWorkspace(op).size, ggml_cuda_cudnn_attn_alloc_size(device, op)); + } + + void ggml_cuda_sage_attn(ggml_backend_cuda_context& ctx, ggml_tensor* dst) { + GGML_ASSERT(ggml_cuda_sage_attn_supported(ctx.device, dst)); ++ ggml_cuda_set_device(ctx.device); ++ if (ggml_cuda_cudnn_attn(ctx, dst)) { ++ return; ++ } + if (dst->src[0]->ne[0] == 64) { + sage_attn<64>(ctx, dst); + } else { +@@ -352,7 +358,7 @@ bool ggml_cuda_sage_attn_supported(int, const ggml_tensor*) { + return false; + } + +-size_t ggml_cuda_sage_attn_get_alloc_size(const ggml_tensor* op) { ++size_t ggml_cuda_sage_attn_get_alloc_size(int, const ggml_tensor* op) { + return ggml_nbytes(op); + } + +diff --git a/src/ggml-cuda/sage-attn.cuh b/src/ggml-cuda/sage-attn.cuh +index 126c7a74..e2648857 100644 +--- a/src/ggml-cuda/sage-attn.cuh ++++ b/src/ggml-cuda/sage-attn.cuh +@@ -4,7 +4,7 @@ + #include "common.cuh" + + bool ggml_cuda_sage_attn_supported(int device, const ggml_tensor* op); +-size_t ggml_cuda_sage_attn_get_alloc_size(const ggml_tensor* op); ++size_t ggml_cuda_sage_attn_get_alloc_size(int device, const ggml_tensor* op); + void ggml_cuda_sage_attn(ggml_backend_cuda_context& ctx, ggml_tensor* dst); + + #endif +diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp +index b58c815c..55db14a0 100644 +--- a/tests/test-backend-ops.cpp ++++ b/tests/test-backend-ops.cpp +@@ -9773,6 +9773,21 @@ static std::vector> make_test_cases_eval() { + GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + } + ++ // unmasked F16 K/V with the MiniMax-H3 DiT head layout (42 heads x 128), the shape class the CUDA ++ // cuDNN attention path takes: odd and tile-aligned lengths, GQA with batch, F32 K/V, head size 64. ++ for (int64_t n : { 1023, 4096 }) { ++ test_cases.emplace_back(new test_flash_attn_ext(128, 128, 42, {1, 1}, n, n, false, false, 0, 0, ++ GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); ++ } ++ test_cases.emplace_back(new test_flash_attn_ext(128, 128, 42, {1, 1}, 4096, 333, false, false, 0, 0, ++ GGML_PREC_DEFAULT, GGML_TYPE_F16, GGML_TYPE_F16)); ++ test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {4, 2}, 777, 333, false, false, 0, 0, ++ GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); ++ test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {1, 1}, 512, 512, false, false, 0, 0, ++ GGML_PREC_F32, GGML_TYPE_F32, GGML_TYPE_F32)); ++ test_cases.emplace_back(new test_flash_attn_ext(128, 128, 42, {1, 1}, 1023, 1023, false, false, 0, 0, ++ GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3})); ++ + test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3})); + test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1})); + test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, { 10, 5, 4, 3})); +@@ -10124,6 +10139,16 @@ static std::vector> make_test_cases_perf() { + // Qwen3-VL-8B https://github.com/ggml-org/llama.cpp/issues/17012 + test_cases.emplace_back(new test_flash_attn_ext(72, 72, 16, {1, 1}, 5776, 5776, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + ++ // unmasked diffusion attention: MiniMax-H3 DiT (19315 tokens x 56 heads at 960x544x124; 22440 x 42), ++ // its token refiner (31 tokens) and video VAE decoder (1797 x 32 x 64), FLUX-class 1024px (4096 x 24), ++ // and a single-query row ++ for (auto [hs, nh, n] : std::vector>{ { 128, 56, 19315 }, { 128, 42, 22440 }, { 128, 56, 31 }, ++ { 64, 32, 1797 }, { 128, 24, 4096 }, { 128, 16, 256 }, ++ { 128, 16, 512 }, { 128, 16, 1024 } }) { ++ test_cases.emplace_back(new test_flash_attn_ext(hs, hs, nh, {1, 1}, n, n, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); ++ } ++ test_cases.emplace_back(new test_flash_attn_ext(128, 128, 8, {4, 1}, 4096, 1, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); ++ + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0)); +-- +2.43.0 + diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 74324169ee..58130582d4 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -207,15 +207,10 @@ ggml_tensor* ggml_ext_gelu_quick(ggml_context* ctx, return x; } -ggml_tensor* ggml_ext_linear(ggml_context* ctx, - ggml_tensor* x, - ggml_tensor* w, - ggml_tensor* b, - bool force_prec_f32, - float scale) { - if (scale != 1.f) { - x = ggml_ext_scale(ctx, x, scale); - } +ggml_tensor* ggml_ext_linear_matmul(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + bool force_prec_f32) { if (x->ne[2] * x->ne[3] > 1024) { // workaround: avoid ggml cuda error int64_t ne2 = x->ne[2]; @@ -232,6 +227,19 @@ ggml_tensor* ggml_ext_linear(ggml_context* ctx, ggml_mul_mat_set_prec(x, GGML_PREC_F32); } } + return x; +} + +ggml_tensor* ggml_ext_linear(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* b, + bool force_prec_f32, + float scale) { + if (scale != 1.f) { + x = ggml_ext_scale(ctx, x, scale); + } + x = ggml_ext_linear_matmul(ctx, x, w, force_prec_f32); if (scale != 1.f) { x = ggml_ext_scale(ctx, x, 1.f / scale); } @@ -796,6 +804,47 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, return kqv; } +ggml_tensor* ggml_ext_attention_prepared(ggml_context* ctx, + ggml_backend_t backend, + ggml_tensor* q, + ggml_tensor* k, + ggml_tensor* v, + int64_t n_head, + int64_t N, + float kv_scale) { + GGML_ASSERT(q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16); + const int64_t L_q = q->ne[1]; + const int64_t d_head = v->ne[0]; + if (backend == nullptr || d_head < 64) { + return nullptr; + } + + // Same expressions as ggml_ext_attention_ext so the softmax scale rounds identically. + float scale = (1.0f / sqrt((float)d_head)); + + auto kqv = ggml_flash_attn_ext(ctx, q, k, v, nullptr, scale / kv_scale, 0, 0); + if (!ggml_backend_supports_op(backend, kqv)) { + return nullptr; + } + ggml_flash_attn_ext_set_prec(kqv, GGML_PREC_F32); + if (kv_scale != 1.0f) { + kqv = ggml_ext_scale(ctx, kqv, 1.0f / kv_scale); + } + kqv = ggml_view_4d(ctx, + kqv, + d_head, + n_head, + L_q, + N, + kqv->nb[1], + kqv->nb[2], + kqv->nb[1] * n_head, + 0); + kqv = ggml_ext_cont(ctx, kqv); + kqv = ggml_reshape_3d(ctx, kqv, d_head * n_head, L_q, N); // [N, L_q, C] + return kqv; +} + ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx, ggml_tensor* x, ggml_tensor* w, diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h index 33d8c3bf21..a5076262ff 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -65,6 +65,12 @@ ggml_tensor* ggml_ext_gelu_quick(ggml_context* ctx, ggml_tensor* x, bool inplace = false); +// The matrix multiply of ggml_ext_linear without its input/output scaling and bias. +ggml_tensor* ggml_ext_linear_matmul(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + bool force_prec_f32); + ggml_tensor* ggml_ext_linear(ggml_context* ctx, ggml_tensor* x, ggml_tensor* w, @@ -223,6 +229,21 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, float kv_scale = 1.0f, bool sage_attn = false); +// Flash attention on inputs that are already in the head-major layout the kernel reads: +// q: F32 [d_head, L_q, n_head * N], k/v: F16 [d_head, L_k, n_kv_head * N] with the kv_scale +// already applied. Builds the same flash_attn_ext + output scaling as ggml_ext_attention_ext +// (skip_reshape, unmasked), so the result is identical; returns nullptr when the backend lacks +// flash attention for these inputs (the caller then uses ggml_ext_attention_ext). +// return: [N, L_q, n_head * d_head] +ggml_tensor* ggml_ext_attention_prepared(ggml_context* ctx, + ggml_backend_t backend, + ggml_tensor* q, + ggml_tensor* k, + ggml_tensor* v, + int64_t n_head, + int64_t N, + float kv_scale); + ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx, ggml_tensor* x, ggml_tensor* w, diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp index abfa8639bd..3ac8e0b2dc 100644 --- a/src/core/ggml_runner.cpp +++ b/src/core/ggml_runner.cpp @@ -725,7 +725,7 @@ std::vector GGMLRunner::memory_requests( } } requests.push_back({size.backend, reinterpret_cast(this), pending, - retained, limit}); + retained, limit, reuse_device_query_}); } return requests; } diff --git a/src/core/ggml_runner.h b/src/core/ggml_runner.h index 29a6808bd8..c5e1a65164 100644 --- a/src/core/ggml_runner.h +++ b/src/core/ggml_runner.h @@ -156,6 +156,7 @@ struct GGMLRunner { bool graph_active_ = false; size_t max_graph_vram_bytes = 0; + bool reuse_device_query_ = false; bool graph_cut_layer_split_enabled = false; std::vector graph_cut_layer_split_backend_vram_limits_; @@ -369,6 +370,11 @@ struct GGMLRunner { weight_adapter = adapter; } + // See DeviceMemoryRequest::reuse_device_query; runner_end() drops the saved reading. + void set_reuse_device_query(bool enabled) { + reuse_device_query_ = enabled; + } + void set_max_graph_vram_bytes(size_t max_vram_bytes) { max_graph_vram_bytes = max_vram_bytes; } diff --git a/src/device_residency_manager.h b/src/device_residency_manager.h index 0a2a8895f3..0be0d2ee42 100644 --- a/src/device_residency_manager.h +++ b/src/device_residency_manager.h @@ -27,6 +27,8 @@ struct DeviceMemoryRequest { size_t pending_allocation_bytes = 0; size_t runtime_resident_bytes = 0; size_t max_backend_bytes = 0; + // Repeated identical compute: a check that allocates nothing may reuse this owner's last free-memory reading. + bool reuse_device_query = false; // Runtime buffers only; the manager accounts for weights separately. size_t runtime_peak_bytes() const { diff --git a/src/model/common/ggml_block.hpp b/src/model/common/ggml_block.hpp index 991b88f28b..b4c94b2d94 100644 --- a/src/model/common/ggml_block.hpp +++ b/src/model/common/ggml_block.hpp @@ -204,6 +204,30 @@ class Linear : public UnaryBlock { force_prec_f32 = force_prec_f32_; } + // For callers that fold this layer's input/output scaling into neighbouring fused ops: the + // weight when forward() is exactly ggml_ext_linear(x, weight, nullptr, prec_f32(), + // effective_scale()) (no bias, weight scale, adapter, INT8 or FP8 path), nullptr otherwise. + ggml_tensor* foldable_weight(GGMLRunnerContext* ctx) { + ggml_tensor* w = params["weight"]; + if (bias || has_weight_scale || ctx->weight_adapter || w->type == GGML_TYPE_I8) { + return nullptr; + } +#ifndef SD_USE_UPSTREAM_GGML + if (w->type == GGML_TYPE_F8_E4M3 || w->type == GGML_TYPE_F8_E5M2) { + return nullptr; + } +#endif + return w; + } + + float effective_scale(GGMLRunnerContext* ctx) const { + return ctx->linear_scale > 0.f ? ctx->linear_scale : scale; + } + + bool prec_f32() const { + return force_prec_f32; + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { ggml_tensor* w = params["weight"]; const float scale = ctx->linear_scale > 0.f ? ctx->linear_scale : this->scale; diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index 80c29ab528..b065f02bd1 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -120,6 +121,54 @@ namespace MiniMaxH3 { } }; + // Value-preserving graph rewrites for the DiT blocks. Every lever produces the same bits as + // the original graph (same arithmetic, same rounding points); each can be switched off: + // SD_H3_GRAPH_FAST=0 restores the original graph for all of them + // SD_H3_FAST_QKV=0 q/k/v: unfused chunk copies, partial RoPE via slice/rope/concat, K/V scale+cast + // SD_H3_FAST_MLP=0 MLP gate: copied chunks + silu + mul instead of one swiglu on views + // SD_H3_FAST_VIEWS=0 modulation/residual segments copied before use instead of read in place + // SD_H3_FAST_SEGMENTS=0 per-segment adaLN modulation / gated residual as mul+add+concat chains + // instead of one in-place op per segment, MLP Linear scales kept as + // separate scale ops instead of folded into those ops and the swiglu + static bool env_flag(const char* name, bool default_value) { + const char* value = std::getenv(name); + if (value == nullptr || value[0] == '\0') { + return default_value; + } + return std::atoi(value) != 0; + } + + static bool graph_fast() { + static const bool enabled = env_flag("SD_H3_GRAPH_FAST", true); + return enabled; + } + + static bool fast_qkv() { + static const bool enabled = graph_fast() && env_flag("SD_H3_FAST_QKV", true); + return enabled; + } + + static bool fast_mlp() { + static const bool enabled = graph_fast() && env_flag("SD_H3_FAST_MLP", true); + return enabled; + } + + // SD_H3_FAST_SAGE_QKV=0: --sage-attn keeps the unfused q/k/v chain. + static bool fast_sage_qkv() { + static const bool enabled = fast_qkv() && env_flag("SD_H3_FAST_SAGE_QKV", true); + return enabled; + } + + static bool fast_segments() { + static const bool enabled = graph_fast() && env_flag("SD_H3_FAST_SEGMENTS", true); + return enabled; + } + + static bool fast_views() { + static const bool enabled = graph_fast() && env_flag("SD_H3_FAST_VIEWS", true); + return enabled; + } + static float time_shift_sigma(float sigma, float from_shift, float to_shift) { float base = sigma / (from_shift + sigma * (1.f - from_shift)); return to_shift * base / (1.f + (to_shift - 1.f) * base); @@ -145,10 +194,57 @@ namespace MiniMaxH3 { blocks["fc2"] = std::make_shared(ffn_hidden_size, hidden_size, false, false, true, 1.f / 128.f); } + // Scales for forward_folded(): the caller applies `input_scale` to the input (fused into the + // preceding op) and `output_scale` to the result (fused into the following op). Returns false + // when the Linears cannot be folded (bias, adapter, weight scales, INT8/FP8 weights). + bool foldable(GGMLRunnerContext* ctx, float* input_scale, float* output_scale) { + auto fc1 = std::dynamic_pointer_cast(blocks["fc1"]); + auto fc2 = std::dynamic_pointer_cast(blocks["fc2"]); + if (fc1->foldable_weight(ctx) == nullptr || fc2->foldable_weight(ctx) == nullptr) { + return false; + } + const float s1 = fc1->effective_scale(ctx); + const float s2 = fc2->effective_scale(ctx); + // the same factors ggml_ext_linear applies around each matmul + *input_scale = s1; + *output_scale = s2 != 1.f ? 1.f / s2 : 1.f; + return true; + } + + // fc1 matmul -> swiglu with fc1's output scale and fc2's input scale fused in -> fc2 matmul. + // x must already carry fc1's input scale; the result still needs foldable()'s output_scale. + ggml_tensor* forward_folded(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto fc1 = std::dynamic_pointer_cast(blocks["fc1"]); + auto fc2 = std::dynamic_pointer_cast(blocks["fc2"]); + const float s1 = fc1->effective_scale(ctx); + const float s2 = fc2->effective_scale(ctx); + auto h = ggml_ext_linear_matmul(ctx->ggml_ctx, x, fc1->foldable_weight(ctx), fc1->prec_f32()); + auto halves = ggml_ext_chunk(ctx->ggml_ctx, h, 2, 0, false); +#ifndef SD_GGML_H3_FUSED_OPS + GGML_UNUSED(s1); + GGML_UNUSED(s2); + GGML_UNUSED(halves); + GGML_ABORT("forward_folded needs the fused swiglu op"); +#else + auto gated = ggml_swiglu_scaled(ctx->ggml_ctx, halves[0], halves[1], s1 != 1.f ? 1.f / s1 : 1.f, s2); + return ggml_ext_linear_matmul(ctx->ggml_ctx, gated, fc2->foldable_weight(ctx), fc2->prec_f32()); +#endif + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { auto fc1 = std::dynamic_pointer_cast(blocks["fc1"]); auto fc2 = std::dynamic_pointer_cast(blocks["fc2"]); - auto uv = ggml_ext_chunk(ctx->ggml_ctx, fc1->forward(ctx, x), 2, 0); + auto h = fc1->forward(ctx, x); + if (fast_mlp() && ctx->backend != nullptr && h->type == GGML_TYPE_F32 && ggml_is_contiguous(h)) { + // silu(u) * v straight from the two halves of fc1's output: same per-element + // silu and product as the copied-chunk graph, without the two chunk copies. + auto halves = ggml_ext_chunk(ctx->ggml_ctx, h, 2, 0, false); + auto gated = ggml_swiglu_split(ctx->ggml_ctx, halves[0], halves[1]); + if (ggml_backend_supports_op(ctx->backend, gated)) { + return fc2->forward(ctx, gated); + } + } + auto uv = ggml_ext_chunk(ctx->ggml_ctx, h, 2, 0); return fc2->forward(ctx, ggml_mul(ctx->ggml_ctx, ggml_silu(ctx->ggml_ctx, uv[0]), uv[1])); @@ -192,6 +288,80 @@ namespace MiniMaxH3 { blocks["out_proj"] = std::make_shared(inner, hidden_size, false); } + // q/k/v read in place from the fused projection, q/k RMS-normed, then one op per tensor + // applies the partial RoPE, writes the head-major layout flash attention reads, and for + // K/V applies the kv scale and the F16 cast. Returns nullptr when unsupported. + ggml_tensor* forward_head_major(GGMLRunnerContext* ctx, + ggml_tensor* qkv_out, + int64_t sequence, + int64_t batch, + ggml_tensor* pe) { +#ifndef SD_GGML_H3_FUSED_OPS + return nullptr; +#else + const bool sage = ctx->sage_attn_enabled; + if (ctx->backend == nullptr || (sage ? !fast_sage_qkv() : !ctx->flash_attn_enabled)) { + return nullptr; + } + // the sage kernel takes 64 or 128 channel heads without padding + if (sage && head_dim != 64 && head_dim != 128) { + return nullptr; + } + const int64_t inner = heads * head_dim; + if (qkv_out->type != GGML_TYPE_F32 || !ggml_is_contiguous(qkv_out) || qkv_out->ne[0] != inner * 3 || + qkv_out->ne[1] != sequence) { + return nullptr; + } + const int n_rot = pe != nullptr ? static_cast(pe->ne[2] * 2) : 0; + if (n_rot > head_dim || (pe != nullptr && (pe->ne[0] != 2 || pe->ne[1] != 2 || pe->ne[3] != sequence || + pe->type != GGML_TYPE_F32 || !ggml_is_contiguous(pe)))) { + return nullptr; + } + const float kv_scale = ctx->attn_scale > 0.f ? ctx->attn_scale : 1.f / 128.f; + + auto q_norm = std::dynamic_pointer_cast(blocks["q_norm"]); + auto k_norm = std::dynamic_pointer_cast(blocks["k_norm"]); + auto part = [&](int index) { + return ggml_view_4d(ctx->ggml_ctx, + qkv_out, + head_dim, + heads, + sequence, + batch, + head_dim * ggml_element_size(qkv_out), + qkv_out->nb[1], + qkv_out->nb[1] * sequence, + index * inner * ggml_element_size(qkv_out)); + }; + auto q = q_norm->forward(ctx, part(0)); + auto k = k_norm->forward(ctx, part(1)); + // sage reads F32 Q/K and F16 V in this layout; K and V are scaled in F32, as on the unfused path + q = ggml_rope_pe_permute(ctx->ggml_ctx, q, pe, n_rot, 1.f, GGML_TYPE_F32); + k = ggml_rope_pe_permute(ctx->ggml_ctx, k, pe, n_rot, kv_scale, sage ? GGML_TYPE_F32 : GGML_TYPE_F16); + auto v = ggml_rope_pe_permute(ctx->ggml_ctx, part(2), nullptr, 0, kv_scale, GGML_TYPE_F16); + if (!ggml_backend_supports_op(ctx->backend, q) || !ggml_backend_supports_op(ctx->backend, k) || + !ggml_backend_supports_op(ctx->backend, v)) { + return nullptr; + } + if (sage) { + // same softmax scale expression as ggml_ext_attention_ext's sage branch + const float scale = 1.0f / sqrt((float)head_dim); + auto out = ggml_sage_attn(ctx->ggml_ctx, q, k, v, scale / kv_scale, GGML_SAGE_ATTN_AUTO); + if (!ggml_backend_supports_op(ctx->backend, out)) { + return nullptr; + } + if (kv_scale != 1.0f) { + out = ggml_ext_scale(ctx->ggml_ctx, out, 1.0f / kv_scale); + } + return ggml_reshape_3d(ctx->ggml_ctx, out, head_dim * heads, sequence, batch); // [N, L, C] + } + q = ggml_reshape_3d(ctx->ggml_ctx, q, head_dim, sequence, heads * batch); + k = ggml_reshape_3d(ctx->ggml_ctx, k, head_dim, sequence, heads * batch); + v = ggml_reshape_3d(ctx->ggml_ctx, v, head_dim, sequence, heads * batch); + return ggml_ext_attention_prepared(ctx->ggml_ctx, ctx->backend, q, k, v, heads, batch, kv_scale); +#endif + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe = nullptr) { @@ -202,7 +372,14 @@ namespace MiniMaxH3 { int64_t sequence = x->ne[1]; int64_t batch = x->ne[2] * x->ne[3]; - auto qkv = ggml_ext_chunk(ctx->ggml_ctx, qkv_proj->forward(ctx, x), 3, 0); + auto qkv_out = qkv_proj->forward(ctx, x); + if (fast_qkv()) { + auto out = forward_head_major(ctx, qkv_out, sequence, batch, pe); + if (out != nullptr) { + return out_proj->forward(ctx, out); + } + } + auto qkv = ggml_ext_chunk(ctx->ggml_ctx, qkv_out, 3, 0); auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv[0], head_dim, heads, sequence, batch); auto k = ggml_reshape_4d(ctx->ggml_ctx, qkv[1], head_dim, heads, sequence, batch); auto v = ggml_reshape_4d(ctx->ggml_ctx, qkv[2], head_dim, heads, sequence, batch); @@ -331,8 +508,8 @@ namespace MiniMaxH3 { projection, hidden_size * expand, timestep_rows * modalities); - auto selected = ggml_ext_slice(ctx, reshaped, 1, row, row + 1); - return ggml_ext_chunk(ctx, selected, expand, 0); + auto selected = ggml_ext_slice(ctx, reshaped, 1, row, row + 1, !fast_views()); + return ggml_ext_chunk(ctx, selected, expand, 0, !fast_views()); } static ggml_tensor* modulate_segments(ggml_context* ctx, @@ -352,7 +529,7 @@ namespace MiniMaxH3 { expand, modalities, segment.modulation_row); - auto part = ggml_ext_slice(ctx, x, 1, segment.start, segment.end); + auto part = ggml_ext_slice(ctx, x, 1, segment.start, segment.end, !fast_views()); part = ggml_add(ctx, ggml_add(ctx, part, ggml_mul(ctx, part, mods[scale_index])), mods[shift_index]); @@ -371,8 +548,8 @@ namespace MiniMaxH3 { ggml_tensor* out = nullptr; for (const auto& segment : segments) { auto mods = modulation_row(ctx, projection, hidden_size, 6, 3, segment.modulation_row); - auto base = ggml_ext_slice(ctx, x, 1, segment.start, segment.end); - auto add = ggml_ext_slice(ctx, update, 1, segment.start, segment.end); + auto base = ggml_ext_slice(ctx, x, 1, segment.start, segment.end, !fast_views()); + auto add = ggml_ext_slice(ctx, update, 1, segment.start, segment.end, !fast_views()); auto part = ggml_add(ctx, base, ggml_mul(ctx, add, mods[gate_index])); out = out == nullptr ? part : ggml_concat(ctx, out, part, 1); } @@ -400,11 +577,95 @@ namespace MiniMaxH3 { config.uses_adaln_curves()); } + // The block with every per-segment modulation / gated residual as one in-place op per segment + // (no slice copies, no concat) and the MLP Linear scales folded into those ops and the + // swiglu. Same arithmetic and rounding as forward(); nullptr when unsupported. + ggml_tensor* forward_fused(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* t_emb, + const std::vector& segments, + ggml_tensor* pe) { +#ifndef SD_GGML_H3_FUSED_OPS + return nullptr; +#else + if (ctx->backend == nullptr || segments.empty() || x->type != GGML_TYPE_F32) { + return nullptr; + } + int64_t expected = 0; + for (const auto& segment : segments) { + if (segment.start != expected || segment.end <= segment.start) { + return nullptr; + } + expected = segment.end; + } + if (expected != x->ne[1]) { + return nullptr; + } + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto attn = std::dynamic_pointer_cast(blocks["attn"]); + auto mlp = std::dynamic_pointer_cast(blocks["mlp"]); + auto adaln = std::dynamic_pointer_cast(blocks["adaln_proj"]); + auto mods = adaln->forward(ctx, t_emb); + + auto vec = [&](const TokenModulationSpan& segment, int index) { + return modulation_row(ctx->ggml_ctx, mods, config.hidden_size, 6, 3, segment.modulation_row)[index]; + }; + { + auto probe_vec = vec(segments[0], 0); + auto probe = ggml_modulate_rows(ctx->ggml_ctx, nullptr, x, probe_vec, probe_vec, 0, 1, 1.f); + auto probe_g = ggml_gated_add_rows(ctx->ggml_ctx, nullptr, x, x, probe_vec, 0, 1, 1.f); + if (!ggml_backend_supports_op(ctx->backend, probe) || !ggml_backend_supports_op(ctx->backend, probe_g)) { + return nullptr; + } + } + auto modulate = [&](ggml_tensor* value, int shift_index, int scale_index, float out_scale) { + ggml_tensor* out = nullptr; + for (const auto& segment : segments) { + out = ggml_modulate_rows(ctx->ggml_ctx, out, value, vec(segment, scale_index), vec(segment, shift_index), + segment.start, segment.end, out_scale); + } + return out; + }; + auto gated = [&](ggml_tensor* base, ggml_tensor* update, int gate_index, float update_scale) { + ggml_tensor* out = nullptr; + for (const auto& segment : segments) { + out = ggml_gated_add_rows(ctx->ggml_ctx, out, base, update, vec(segment, gate_index), + segment.start, segment.end, update_scale); + } + return out; + }; + + auto h = modulate(norm1->forward(ctx, x), 0, 1, 1.f); + x = gated(x, attn->forward(ctx, h, pe), 2, 1.f); + + float mlp_in_scale = 1.f; + float mlp_out_scale = 1.f; + bool fold = fast_mlp() && mlp->foldable(ctx, &mlp_in_scale, &mlp_out_scale); + if (fold) { + auto probe = ggml_swiglu_scaled(ctx->ggml_ctx, x, x, 1.f, 1.f); + fold = ggml_backend_supports_op(ctx->backend, probe); + } + if (fold) { + h = modulate(norm2->forward(ctx, x), 3, 4, mlp_in_scale); + return gated(x, mlp->forward_folded(ctx, h), 5, mlp_out_scale); + } + h = modulate(norm2->forward(ctx, x), 3, 4, 1.f); + return gated(x, mlp->forward(ctx, h), 5, 1.f); +#endif + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* t_emb, const std::vector& segments, ggml_tensor* pe) { + if (fast_segments()) { + auto out = forward_fused(ctx, x, t_emb, segments, pe); + if (out != nullptr) { + return out; + } + } auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); auto attn = std::dynamic_pointer_cast(blocks["attn"]); diff --git a/src/model/vae/ltx_audio_vae.hpp b/src/model/vae/ltx_audio_vae.hpp index d2314b56a7..2ff8528ecc 100644 --- a/src/model/vae/ltx_audio_vae.hpp +++ b/src/model/vae/ltx_audio_vae.hpp @@ -396,6 +396,71 @@ namespace LTXV { return ggml_reshape_4d(ctx, out, out_time, channels, 1, 1); } + // F32 depthwise conv without im2col or F16 input rounding; nullptr when the backend lacks CONV_2D_DW. + static ggml_tensor* depthwise_conv1d_direct(GGMLRunnerContext* runner_ctx, + ggml_tensor* x, + ggml_tensor* filter, + int stride, + int padding, + bool reverse) { + auto ctx = runner_ctx->ggml_ctx; + if (runner_ctx->backend == nullptr || x->type != GGML_TYPE_F32 || x->ne[2] != 1 || x->ne[3] != 1) { + return nullptr; + } + const int64_t time = x->ne[0]; + const int64_t channels = x->ne[1]; + ggml_tensor* f = ggml_reshape_4d(ctx, filter, filter->ne[0], 1, 1, 1); + if (reverse) { + f = reverse_1d_filter(ctx, f); + } + if (f->type != GGML_TYPE_F32) { + f = ggml_cast(ctx, f, GGML_TYPE_F32); + } + auto kernel = repeat_with_vulkan_f32_workaround(runner_ctx->backend, ctx, f, f->ne[0], 1, 1, channels); // [K, 1, 1, C] + auto input = ggml_reshape_4d(ctx, ggml_ext_cont(ctx, x), time, 1, channels, 1); // [W, H=1, C, N=1] + auto out = ggml_conv_2d_dw_direct(ctx, kernel, input, stride, 1, padding, 0, 1, 1); + if (!ggml_backend_supports_op(runner_ctx->backend, out)) { + return nullptr; + } + return ggml_reshape_4d(ctx, out, out->ne[0], channels, 1, 1); + } + + // depthwise_conv_transpose1d via the direct conv: same zero-stuffing, reversed filter, padding and gain. + static ggml_tensor* depthwise_conv_transpose1d_direct(GGMLRunnerContext* runner_ctx, + ggml_tensor* x, + ggml_tensor* filter, + int stride) { + auto ctx = runner_ctx->ggml_ctx; + if (x->ne[2] != 1 || x->ne[3] != 1 || filter->ne[1] != 1 || filter->ne[2] != 1 || filter->ne[3] != 1) { + return nullptr; + } + const int64_t time = x->ne[0]; + const int64_t channels = x->ne[1]; + const int64_t kernel_size = filter->ne[0]; + const int64_t out_time = (time - 1) * stride + kernel_size; + + auto x_flat = ggml_reshape_3d(ctx, x, 1, time, channels); + if (stride > 1) { + auto zero_unit = ggml_ext_scale(ctx, x_flat, 0.0f); + auto zero_tail = zero_unit; + for (int i = 1; i < stride - 1; ++i) { + zero_tail = ggml_concat(ctx, zero_tail, zero_unit, 0); + } + x_flat = ggml_concat(ctx, x_flat, zero_tail, 0); + } + x_flat = ggml_reshape_4d(ctx, x_flat, time * stride, channels, 1, 1); + auto out = depthwise_conv1d_direct(runner_ctx, x_flat, filter, 1, static_cast(kernel_size - 1), true); + if (out == nullptr) { + return nullptr; + } + if (out->ne[0] > out_time) { + out = ggml_ext_slice(ctx, out, 0, 0, out_time); + } + GGML_ASSERT(out->ne[0] == out_time); + out = ggml_ext_scale(ctx, out, static_cast(stride)); + return ggml_reshape_4d(ctx, out, out_time, channels, 1, 1); + } + static ggml_tensor* upsample_waveform_hann(GGMLRunnerContext* runner_ctx, ggml_tensor* waveform, ggml_tensor* filter, @@ -667,8 +732,10 @@ namespace LTXV { int up_kernel_size = 12; int down_kernel_size = 12; - explicit Activation1D(int64_t channels) - : channels(channels) { + bool direct_dw = false; + + explicit Activation1D(int64_t channels, bool direct_dw = false) + : channels(channels), direct_dw(direct_dw) { blocks["act"] = std::make_shared(channels); } @@ -690,16 +757,17 @@ namespace LTXV { int up_pad_right = up_pad * up_ratio + (up_kernel_size - up_ratio + 1) / 2; x = replicate_pad_1d(ctx, x, up_pad, up_pad); - x = depthwise_conv_transpose1d(ctx, x, up_filter, up_ratio); - x = ggml_ext_slice(ctx->ggml_ctx, x, 0, up_pad_left, x->ne[0] - up_pad_right); + ggml_tensor* up = direct_dw ? depthwise_conv_transpose1d_direct(ctx, x, up_filter, up_ratio) : nullptr; + x = up != nullptr ? up : depthwise_conv_transpose1d(ctx, x, up_filter, up_ratio); + x = ggml_ext_slice(ctx->ggml_ctx, x, 0, up_pad_left, x->ne[0] - up_pad_right); x = act->forward(ctx, x); int down_pad_left = down_kernel_size / 2 - (down_kernel_size % 2 == 0 ? 1 : 0); int down_pad_right = down_kernel_size / 2; x = replicate_pad_1d(ctx, x, down_pad_left, down_pad_right); - x = depthwise_conv1d(ctx, x, down_filter, down_ratio, 0); - return x; + ggml_tensor* down = direct_dw ? depthwise_conv1d_direct(ctx, x, down_filter, down_ratio, 0, false) : nullptr; + return down != nullptr ? down : depthwise_conv1d(ctx, x, down_filter, down_ratio, 0); } }; diff --git a/src/model/vae/minimax_h3_audio_vae.hpp b/src/model/vae/minimax_h3_audio_vae.hpp index a4fc144736..300826a030 100644 --- a/src/model/vae/minimax_h3_audio_vae.hpp +++ b/src/model/vae/minimax_h3_audio_vae.hpp @@ -2,6 +2,7 @@ #define __SD_MODEL_VAE_MINIMAX_H3_AUDIO_VAE_HPP__ #include +#include #include #include @@ -10,6 +11,15 @@ namespace MiniMaxH3 { + // SD_H3_AUDIO_DIRECT_DW=0: anti-aliased activations go back to the im2col + mat-vec depthwise convs. + static bool audio_direct_dw() { + static const bool enabled = [] { + const char* v = std::getenv("SD_H3_AUDIO_DIRECT_DW"); + return v == nullptr || v[0] == '\0' || std::atoi(v) != 0; + }(); + return enabled; + } + struct AudioSnake1D : public UnaryBlock { int64_t channels; @@ -229,9 +239,9 @@ namespace MiniMaxH3 { : channels(channels) { for (int i = 0; i < 3; ++i) { blocks["activations." + std::to_string(i * 2)] = - std::make_shared(channels); + std::make_shared(channels, audio_direct_dw()); blocks["activations." + std::to_string(i * 2 + 1)] = - std::make_shared(channels); + std::make_shared(channels, audio_direct_dw()); blocks["convs1." + std::to_string(i)] = std::make_shared(channels, channels, @@ -298,7 +308,7 @@ namespace MiniMaxH3 { } channels = next_channels; } - blocks["activation_post"] = std::make_shared(channels); + blocks["activation_post"] = std::make_shared(channels, audio_direct_dw()); blocks["conv_post"] = std::make_shared(channels, 1, 7, diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index 777566438c..68cb9a9a27 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -1,11 +1,15 @@ #ifndef __SD_MODEL_VAE_MINIMAX_H3_VAE_HPP__ #define __SD_MODEL_VAE_MINIMAX_H3_VAE_HPP__ +#include #include #include #include +#include +#include #include #include +#include #include #include #include @@ -13,6 +17,7 @@ #include "model/common/rope.hpp" #include "model/diffusion/dit.hpp" #include "model/vae/vae.hpp" +#include "runtime/tiling.h" namespace MiniMaxH3VAE { @@ -244,6 +249,37 @@ namespace MiniMaxH3VAE { return ggml_concat(ctx, rotated, tail, 0); } + // Same arithmetic as apply_partial_rope: rotated channel i of each half is + // x_lo[i] * pe[0][j][i] + x_hi[i] * pe[1][j][i] as separate mul/mul/add ops, and the tail is + // copied unchanged. The halves are read as [half, L, n_head, N] views of the projection + // layout, so the per-half permutes, repeats and the tail permute + concat of the generic + // path disappear. rope_a / rope_b are the [head_dim, L] coefficient tables built by + // MiniMaxH3VideoVAERunner::build_rope_tables. Returns [head_dim, L, n_head * N]. + static ggml_tensor* apply_partial_rope_tables(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* rope_a, + ggml_tensor* rope_b, + int64_t rot_dim) { + const int64_t half = rot_dim / 2; + auto heads_view = [&](int64_t start, int64_t width) { + auto view = ggml_view_4d(ctx, x, width, x->ne[1], x->ne[2], x->ne[3], x->nb[1], x->nb[2], x->nb[3], start * x->nb[0]); + return ggml_permute(ctx, view, 0, 2, 1, 3); // [width, L, n_head, N] + }; + auto table = [&](ggml_tensor* t, int64_t start) { + return ggml_view_2d(ctx, t, half, t->ne[2], t->nb[2], start * t->nb[0]); // [half, L] + }; + auto lo = heads_view(0, half); + auto hi = heads_view(half, half); + auto out = ggml_concat(ctx, + ggml_add(ctx, ggml_mul(ctx, lo, table(rope_a, 0)), ggml_mul(ctx, hi, table(rope_b, 0))), + ggml_add(ctx, ggml_mul(ctx, lo, table(rope_b, half)), ggml_mul(ctx, hi, table(rope_a, half))), + 0); + if (rot_dim < x->ne[0]) { + out = ggml_concat(ctx, out, ggml_cont(ctx, heads_view(rot_dim, x->ne[0] - rot_dim)), 0); + } + return ggml_reshape_3d(ctx, out, out->ne[0], out->ne[1], out->ne[2] * out->ne[3]); + } + struct DecoderAttention : public GGMLBlock { static constexpr int num_head = 32; static constexpr int head_dim = 64; @@ -254,52 +290,156 @@ namespace MiniMaxH3VAE { blocks["to_out"] = std::make_shared(dim, dim, true); } + // Single-tile flash-attention decode: one ggml_rope_pe_permute per tensor replaces the table RoPE + + // chunk / permute / cast chain, bit-identically. nullptr when it does not apply (SD_H3_VAE_FUSED_QKV=0). + ggml_tensor* forward_head_major(GGMLRunnerContext* ctx, + ggml_tensor* qkv, + ggml_tensor* pe, + ggml_tensor* rope_a, + ggml_tensor* rope_b) { +#ifndef SD_GGML_H3_FUSED_OPS + return nullptr; +#else + static const bool enabled = [] { + const char* v = getenv("SD_H3_VAE_FUSED_QKV"); + return v == nullptr || v[0] == '\0' || atoi(v) != 0; + }(); + // rope_a / rope_b are only passed with SD_H3_VAE_GRAPH_OPT on; that path is what this replaces + if (!enabled || rope_a == nullptr || rope_b == nullptr || pe == nullptr || ctx->backend == nullptr || + !ctx->flash_attn_enabled || ctx->sage_attn_enabled || qkv->type != GGML_TYPE_F32 || + qkv->ne[3] != 1 || qkv->nb[0] != sizeof(float)) { + return nullptr; + } + const int64_t sequence = qkv->ne[2]; + const int n_rot = static_cast(pe->ne[2] * 2); + if (pe->type != GGML_TYPE_F32 || !ggml_is_contiguous(pe) || pe->ne[0] != 2 || pe->ne[1] != 2 || + pe->ne[3] != sequence || n_rot > head_dim) { + return nullptr; + } + const float kv_scale = ctx->attn_scale > 0.f ? ctx->attn_scale : 1.f; + auto part = [&](int index) { + return ggml_view_4d(ctx->ggml_ctx, qkv, head_dim, num_head, sequence, 1, + qkv->nb[1], qkv->nb[2], qkv->nb[3], index * head_dim * sizeof(float)); + }; + ggml_tensor* q = nullptr; + ggml_tensor* k = nullptr; +#ifdef SD_GGML_H3_FUSED_QK_NORM + // same reduction as GGML_OP_RMS_NORM where supported (SD_H3_VAE_FUSED_QK_NORM=0 keeps it separate) + static const bool fuse_norm = [] { + const char* v = getenv("SD_H3_VAE_FUSED_QK_NORM"); + return v == nullptr || v[0] == '\0' || atoi(v) != 0; + }(); + if (fuse_norm) { + q = ggml_rope_pe_permute_rms(ctx->ggml_ctx, part(0), pe, n_rot, 1.f, GGML_TYPE_F32, 1e-5f); + k = ggml_rope_pe_permute_rms(ctx->ggml_ctx, part(1), pe, n_rot, kv_scale, GGML_TYPE_F16, 1e-5f); + if (!ggml_backend_supports_op(ctx->backend, q) || !ggml_backend_supports_op(ctx->backend, k)) { + q = k = nullptr; + } + } +#endif + if (q == nullptr) { + q = ggml_rope_pe_permute(ctx->ggml_ctx, ggml_rms_norm(ctx->ggml_ctx, part(0), 1e-5f), pe, n_rot, 1.f, GGML_TYPE_F32); + k = ggml_rope_pe_permute(ctx->ggml_ctx, ggml_rms_norm(ctx->ggml_ctx, part(1), 1e-5f), pe, n_rot, kv_scale, GGML_TYPE_F16); + } + auto v = ggml_rope_pe_permute(ctx->ggml_ctx, part(2), nullptr, 0, kv_scale, GGML_TYPE_F16); + if (!ggml_backend_supports_op(ctx->backend, q) || !ggml_backend_supports_op(ctx->backend, k) || + !ggml_backend_supports_op(ctx->backend, v)) { + return nullptr; + } + q = ggml_reshape_3d(ctx->ggml_ctx, q, head_dim, sequence, num_head); + k = ggml_reshape_3d(ctx->ggml_ctx, k, head_dim, sequence, num_head); + v = ggml_reshape_3d(ctx->ggml_ctx, v, head_dim, sequence, num_head); + return ggml_ext_attention_prepared(ctx->ggml_ctx, ctx->backend, q, k, v, num_head, 1, kv_scale); +#endif + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, - ggml_tensor* pe) { + ggml_tensor* pe, + ggml_tensor* rope_a = nullptr, + ggml_tensor* rope_b = nullptr, + int64_t flat_tiles = 1) { auto to_qkv = std::dynamic_pointer_cast(blocks["to_qkv"]); auto to_out = std::dynamic_pointer_cast(blocks["to_out"]); auto qkv_projection = to_qkv->forward(ctx, x); - int64_t sequence = x->ne[1]; - int64_t batch_size = x->ne[2] * x->ne[3]; - qkv_projection = ggml_reshape_4d(ctx->ggml_ctx, - qkv_projection, - 3 * head_dim, - num_head, - sequence, - batch_size); - auto qkv = ggml_ext_chunk(ctx->ggml_ctx, qkv_projection, 3, 0); - auto q = ggml_reshape_4d(ctx->ggml_ctx, - qkv[0], - head_dim, - num_head, - sequence, - batch_size); - auto k = ggml_reshape_4d(ctx->ggml_ctx, - qkv[1], - head_dim, - num_head, - sequence, - batch_size); - auto v = ggml_reshape_4d(ctx->ggml_ctx, - qkv[2], - head_dim, - num_head, - sequence, - batch_size); - q = ggml_rms_norm(ctx->ggml_ctx, q, 1e-5f); - k = ggml_rms_norm(ctx->ggml_ctx, k, 1e-5f); - q = apply_partial_rope(ctx->ggml_ctx, q, pe); - k = apply_partial_rope(ctx->ggml_ctx, k, pe); - auto out = ggml_ext_attention_ext(ctx, - q, - k, - v, - num_head, - nullptr, - true, - ctx->flash_attn_enabled); - return to_out->forward(ctx, out); + int64_t sequence = x->ne[1] / flat_tiles; + int64_t batch_size = x->ne[2] * x->ne[3] * flat_tiles; + auto project_out = [&](ggml_tensor* attn) { + if (flat_tiles > 1) { + attn = ggml_reshape_2d(ctx->ggml_ctx, attn, attn->ne[0], attn->ne[1] * attn->ne[2] * attn->ne[3]); + } + return to_out->forward(ctx, attn); + }; + qkv_projection = ggml_reshape_4d(ctx->ggml_ctx, + qkv_projection, + 3 * head_dim, + num_head, + sequence, + batch_size); + if (auto fused = forward_head_major(ctx, qkv_projection, pe, rope_a, rope_b)) { + return project_out(fused); + } + auto qkv = ggml_ext_chunk(ctx->ggml_ctx, qkv_projection, 3, 0); + auto q = ggml_reshape_4d(ctx->ggml_ctx, + qkv[0], + head_dim, + num_head, + sequence, + batch_size); + auto k = ggml_reshape_4d(ctx->ggml_ctx, + qkv[1], + head_dim, + num_head, + sequence, + batch_size); + auto v = ggml_reshape_4d(ctx->ggml_ctx, + qkv[2], + head_dim, + num_head, + sequence, + batch_size); + q = ggml_rms_norm(ctx->ggml_ctx, q, 1e-5f); + k = ggml_rms_norm(ctx->ggml_ctx, k, 1e-5f); + if (rope_a != nullptr && rope_b != nullptr) { + q = apply_partial_rope_tables(ctx->ggml_ctx, q, rope_a, rope_b, pe->ne[2] * 2); + k = apply_partial_rope_tables(ctx->ggml_ctx, k, rope_a, rope_b, pe->ne[2] * 2); + } else { + q = apply_partial_rope(ctx->ggml_ctx, q, pe); + k = apply_partial_rope(ctx->ggml_ctx, k, pe); + } + const int64_t tiles = v->ne[3]; + if (tiles > 1) { + // Attention kernels split their work by the total problem size (flash attention's + // stream-k, batched GEMM heuristics), so a batched call can sum in a different order + // than a single-tile call. One call per tile keeps each tile identical to the unbatched decode. + std::vector outs; + for (int64_t n = 0; n < tiles; ++n) { + auto qn = ggml_view_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], num_head, q->nb[1], q->nb[2], n * num_head * q->nb[2]); + auto kn = ggml_view_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], num_head, k->nb[1], k->nb[2], n * num_head * k->nb[2]); + auto vn = ggml_view_4d(ctx->ggml_ctx, v, v->ne[0], v->ne[1], v->ne[2], 1, v->nb[1], v->nb[2], v->nb[3], n * v->nb[3]); + outs.push_back(ggml_ext_attention_ext(ctx, qn, kn, vn, num_head, nullptr, true, ctx->flash_attn_enabled)); + } + while (outs.size() > 1) { + std::vector merged; + for (size_t i = 0; i + 1 < outs.size(); i += 2) { + merged.push_back(ggml_concat(ctx->ggml_ctx, outs[i], outs[i + 1], 2)); + } + if (outs.size() % 2 == 1) { + merged.push_back(outs.back()); + } + outs = std::move(merged); + } + return project_out(outs[0]); + } + auto out = ggml_ext_attention_ext(ctx, + q, + k, + v, + num_head, + nullptr, + true, + ctx->flash_attn_enabled); + return project_out(out); } }; @@ -312,9 +452,13 @@ namespace MiniMaxH3VAE { blocks["w2"] = std::make_shared(kInnerDim, dim, true); } - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { - auto w1 = std::dynamic_pointer_cast(blocks["w1"]); - auto w2 = std::dynamic_pointer_cast(blocks["w2"]); + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, bool fused_glu = false) { + auto w1 = std::dynamic_pointer_cast(blocks["w1"]); + auto w2 = std::dynamic_pointer_cast(blocks["w2"]); + if (fused_glu) { + // silu(first half) * second half in one op; same per-element math as silu + mul. + return w2->forward(ctx, ggml_swiglu(ctx->ggml_ctx, w1->forward(ctx, x))); + } auto gate = ggml_ext_chunk(ctx->ggml_ctx, w1->forward(ctx, x), 2, 0); return w2->forward(ctx, ggml_mul(ctx->ggml_ctx, @@ -344,7 +488,10 @@ namespace MiniMaxH3VAE { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, - ggml_tensor* pe) { + ggml_tensor* pe, + ggml_tensor* rope_a = nullptr, + ggml_tensor* rope_b = nullptr, + int64_t flat_tiles = 1) { auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); auto attn = std::dynamic_pointer_cast(blocks["attn"]); auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); @@ -352,12 +499,12 @@ namespace MiniMaxH3VAE { x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, - attn->forward(ctx, norm1->forward(ctx, x), pe), + attn->forward(ctx, norm1->forward(ctx, x), pe, rope_a, rope_b, flat_tiles), params["scale1"])); return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, - ff->forward(ctx, norm2->forward(ctx, x)), + ff->forward(ctx, norm2->forward(ctx, x), rope_a != nullptr), params["scale2"])); } }; @@ -396,33 +543,61 @@ namespace MiniMaxH3VAE { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* z, - ggml_tensor* pe) { + ggml_tensor* pe, + ggml_tensor* rope_a = nullptr, + ggml_tensor* rope_b = nullptr) { int64_t width = z->ne[0]; int64_t height = z->ne[1]; int64_t num_frames = z->ne[2]; int64_t batch_size = z->ne[3] / 24; - GGML_ASSERT(batch_size == 1); - - z = ggml_cont(ctx->ggml_ctx, - ggml_ext_torch_permute(ctx->ggml_ctx, z, 3, 0, 1, 2)); - z = ggml_reshape_3d(ctx->ggml_ctx, - z, - 24, - width * height * num_frames, - batch_size); - auto x_embedder = std::dynamic_pointer_cast(blocks["x_embedder"]); - auto h = x_embedder->forward(ctx, z); + GGML_ASSERT(batch_size * 24 == z->ne[3]); + + if (batch_size == 1) { + z = ggml_cont(ctx->ggml_ctx, + ggml_ext_torch_permute(ctx->ggml_ctx, z, 3, 0, 1, 2)); + z = ggml_reshape_3d(ctx->ggml_ctx, + z, + 24, + width * height * num_frames, + batch_size); + } else { + // Tiles are stacked along the channel axis, tile-major: [W, H, T, 24 * N]. + z = ggml_reshape_3d(ctx->ggml_ctx, z, width * height * num_frames, 24, batch_size); + z = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 1, 0, 2, 3)); + } + // Batched tiles run every projection as one 2D matmul over all tiles' tokens. A 3D + // activation would broadcast the weight through cublasGemmBatchedEx, whose results + // did not match the single-tile GEMM. + const bool flat = batch_size > 1; + auto x_embedder = std::dynamic_pointer_cast(blocks["x_embedder"]); + ggml_tensor* h = nullptr; + if (flat) { + h = x_embedder->forward(ctx, ggml_reshape_2d(ctx->ggml_ctx, z, z->ne[0], z->ne[1] * z->ne[2])); + h = ggml_reshape_3d(ctx->ggml_ctx, h, h->ne[0], z->ne[1], z->ne[2]); + } else { + h = x_embedder->forward(ctx, z); + } int64_t num_patches = h->ne[1]; - h = ggml_concat(ctx->ggml_ctx, h, params["register_tokens"], 1); - auto zero = ggml_ext_scale(ctx->ggml_ctx, - ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, 1), - 0.f); - h = ggml_concat(ctx->ggml_ctx, h, zero, 1); - + auto registers = params["register_tokens"]; + if (batch_size > 1) { + registers = ggml_repeat(ctx->ggml_ctx, + registers, + ggml_new_tensor_3d(ctx->ggml_ctx, registers->type, dim, num_register_tokens, batch_size)); + } + h = ggml_concat(ctx->ggml_ctx, h, registers, 1); + auto zero = ggml_ext_scale(ctx->ggml_ctx, + ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, 1), + 0.f); + h = ggml_concat(ctx->ggml_ctx, h, zero, 1); + + const int64_t tokens = h->ne[1]; + if (flat) { + h = ggml_reshape_2d(ctx->ggml_ctx, h, h->ne[0], tokens * batch_size); + } for (int i = 0; i < num_layers; ++i) { auto block = std::dynamic_pointer_cast( blocks["transformer_blocks." + std::to_string(i)]); - h = block->forward(ctx, h, pe); + h = block->forward(ctx, h, pe, rope_a, rope_b, flat ? batch_size : 1); sd::ggml_graph_cut::mark_graph_cut(h, "minimax_h3_vae.decoder.blocks." + std::to_string(i), "hidden_states"); @@ -431,7 +606,10 @@ namespace MiniMaxH3VAE { auto norm_out = std::dynamic_pointer_cast(blocks["norm_out"]); auto proj_out = std::dynamic_pointer_cast(blocks["proj_out"]); h = proj_out->forward(ctx, norm_out->forward(ctx, h)); - h = ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, num_patches); + if (flat) { + h = ggml_reshape_3d(ctx->ggml_ctx, h, h->ne[0], tokens, batch_size); + } + h = ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, num_patches); return DiT::unpatchify_3d(ctx->ggml_ctx, h, num_frames, @@ -448,8 +626,8 @@ namespace MiniMaxH3VAE { MiniMaxH3VideoVAE() { blocks["encoder"] = std::make_shared(); blocks["quant_conv"] = std::make_shared(48, - 48, - std::tuple{1, 1, 1}); + 48, + std::tuple{1, 1, 1}); blocks["post_quant_conv"] = std::make_shared(24, 24, std::tuple{1, 1, 1}); @@ -473,10 +651,12 @@ namespace MiniMaxH3VAE { ggml_tensor* latent, ggml_tensor* pe, ggml_tensor* pixel_mean, - ggml_tensor* pixel_std) { + ggml_tensor* pixel_std, + ggml_tensor* rope_a = nullptr, + ggml_tensor* rope_b = nullptr) { auto post_quant = std::dynamic_pointer_cast(blocks["post_quant_conv"]); auto decoder = std::dynamic_pointer_cast(blocks["decoder"]); - auto pixels = decoder->forward(ctx, post_quant->forward(ctx, latent), pe); + auto pixels = decoder->forward(ctx, post_quant->forward(ctx, latent), pe, rope_a, rope_b); pixels = ggml_add(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, pixels, pixel_std), pixel_mean); @@ -491,6 +671,9 @@ namespace MiniMaxH3VAE { sd::Tensor latents_mean; sd::Tensor latents_std; sd::Tensor rope_cache; + sd::Tensor rope_a_cache; + sd::Tensor rope_b_cache; + std::array rope_key_ = {0, 0, 0, 0}; MiniMaxH3VideoVAERunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map, @@ -562,6 +745,13 @@ namespace MiniMaxH3VAE { params.tile_size_x = 16; params.tile_size_y = 16; params.target_overlap = 0.25f; + // SD_H3_VAE_TILE=N: opt-in N x N latent tiles; the seams move, so frames differ from the default + if (const char* tile = getenv("SD_H3_VAE_TILE")) { + const int n = atoi(tile); + if (n >= 8) { + params.tile_size_x = params.tile_size_y = n; + } + } return params; } @@ -576,27 +766,31 @@ namespace MiniMaxH3VAE { } static sd::Tensor blend_temporal(const sd::Tensor& previous, - const sd::Tensor& current, + sd::Tensor current, int64_t extent) { - auto output = current; extent = std::min({extent, previous.shape()[2], current.shape()[2]}); int64_t previous_start = previous.shape()[2] - extent; + GGML_ASSERT(previous.dim() == 5 && current.dim() == 5 && previous.shape()[0] == current.shape()[0] && + previous.shape()[1] == current.shape()[1] && previous.shape()[3] == current.shape()[3] && + previous.shape()[4] == current.shape()[4]); + const int64_t plane = current.shape()[0] * current.shape()[1]; + const int64_t previous_frames = previous.shape()[2]; + const int64_t current_frames = current.shape()[2]; for (int64_t b = 0; b < current.shape()[4]; ++b) { for (int64_t c = 0; c < current.shape()[3]; ++c) { + const int64_t channel = b * current.shape()[3] + c; for (int64_t t = 0; t < extent; ++t) { - float wb = static_cast(t) / extent; - float wa = 1.f - wb; - for (int64_t h = 0; h < current.shape()[1]; ++h) { - for (int64_t w = 0; w < current.shape()[0]; ++w) { - output.index(w, h, t, c, b) = - previous.index(w, h, previous_start + t, c, b) * wa + - current.index(w, h, t, c, b) * wb; - } + float wb = static_cast(t) / extent; + float wa = 1.f - wb; + const float* pa = previous.data() + (channel * previous_frames + previous_start + t) * plane; + float* pb = current.data() + (channel * current_frames + t) * plane; + for (int64_t i = 0; i < plane; ++i) { + pb[i] = pa[i] * wa + pb[i] * wb; } } } } - return output; + return current; } sd::Tensor encode(int n_threads, @@ -688,25 +882,27 @@ namespace MiniMaxH3VAE { input.shape()[2], {static_cast(tokens_per_chunk + token_overlap), static_cast(token_overlap)}); GGML_ASSERT(plan.tiles.size() == static_cast(num_chunks)); - auto result = process_vae_temporal_tiles(input, plan, [&](const sd::Tensor& chunk, const VAETemporalTile& tile) { - auto decoded = VAE::decode(n_threads, - chunk, - tiling, - true, - circular_x, - circular_y, - silent); - if (decoded.empty()) { - return sd::Tensor(); + const bool keep_resident = env_int("SD_H3_VAE_KEEP_RESIDENT", 1) != 0; + // SD_H3_VAE_REUSE_MEMQUERY=0 queries the device for every capacity check. + set_reuse_device_query(env_int("SD_H3_VAE_REUSE_MEMQUERY", 1) != 0); + struct ReuseQueryGuard { + GGMLRunner& runner; + ~ReuseQueryGuard() { + runner.set_reuse_device_query(false); } - + } reuse_query_guard{*this}; + tile_batch_ = 0; + per_tile_compute_bytes_ = 0; + std::vector> pieces; + // Runs in chunk order, on a worker thread unless SD_H3_VAE_ASYNC_ASSEMBLY=0. + auto assemble = [&](const sd::Tensor& decoded, const VAETemporalTile& tile) { int64_t first_end = std::min(frames_per_chunk, decoded.shape()[2]); auto first = sd::ops::slice(decoded, 2, std::min(frame_pre_padding, first_end), first_end); if (!overlap.empty()) { - first = blend_temporal(overlap, first, frame_overlap); + first = blend_temporal(overlap, std::move(first), frame_overlap); overlap = {}; } @@ -721,19 +917,276 @@ namespace MiniMaxH3VAE { overlap = {}; } return first; - }); - if (result.empty()) { - return {}; - } - + }; + const bool async_assembly = env_int("SD_H3_VAE_ASYNC_ASSEMBLY", 1) != 0; + std::thread assembler; + bool assembly_failed = false; int64_t expected_frames = input.shape()[2] <= 1 ? 1 : ((x.shape()[2] - 2) / 5) * 17 + 5; expected_frames = std::max(1, expected_frames); + // keeps the same frames concat_frames + slice would + sd::Tensor result; + int64_t frame_offset = 0; + auto append_piece = [&](const sd::Tensor& piece) { + GGML_ASSERT(piece.dim() == 5); + if (result.empty()) { + std::vector shape = piece.shape(); + shape[2] = expected_frames; + result = sd::Tensor(std::move(shape)); + } + GGML_ASSERT(piece.shape()[0] == result.shape()[0] && piece.shape()[1] == result.shape()[1] && + piece.shape()[3] == result.shape()[3] && piece.shape()[4] == result.shape()[4]); + const int64_t piece_frames = piece.shape()[2]; + const int64_t keep = std::min(piece_frames, expected_frames - frame_offset); + const int64_t plane = piece.shape()[0] * piece.shape()[1]; + const int64_t planes = piece.shape()[3] * piece.shape()[4]; + for (int64_t p = 0; keep > 0 && p < planes; ++p) { + memcpy(result.data() + (p * expected_frames + frame_offset) * plane, + piece.data() + p * piece_frames * plane, + sizeof(float) * keep * plane); + } + frame_offset += std::max(keep, 0); + }; + auto run_assembly = [&](const sd::Tensor& decoded, const VAETemporalTile& tile) { + try { + if (async_assembly) { + append_piece(assemble(decoded, tile)); + } else { + pieces.push_back(assemble(decoded, tile)); + } + } catch (const std::exception& error) { + LOG_ERROR("MiniMax-H3 video VAE: chunk assembly failed: %s", error.what()); + assembly_failed = true; + } + }; + auto join_assembler = [&]() { + if (assembler.joinable()) { + assembler.join(); + } + }; + struct AssemblerGuard { + std::thread& thread; + ~AssemblerGuard() { + if (thread.joinable()) { + thread.join(); + } + } + } assembler_guard{assembler}; + bool failed = false; + for (const auto& tile : plan.tiles) { + auto decoded = decode_spatial_tiles(n_threads, + sd::ops::slice(input, 2, tile.start, tile.end), + tiling, + circular_x, + circular_y, + silent); + if (!keep_resident) { + runner_end(); + } + join_assembler(); + if (decoded.empty() || assembly_failed) { + failed = true; + break; + } + if (async_assembly && !tile.last) { + auto chunk = std::make_shared>(std::move(decoded)); + try { + assembler = std::thread([&run_assembly, chunk, tile]() { + run_assembly(*chunk, tile); + }); + } catch (const std::system_error&) { + run_assembly(*chunk, tile); // no thread available: assemble inline + } + } else { + run_assembly(decoded, tile); + } + } + join_assembler(); + runner_end(); + if (failed || assembly_failed) { + return {}; + } + if (async_assembly) { + if (result.empty()) { + return {}; + } + if (frame_offset < expected_frames) { + result = sd::ops::slice(result, 2, 0, frame_offset); + } + return result; + } + if (pieces.empty()) { + return {}; + } + result = concat_frames(pieces); if (result.shape()[2] > expected_frames) { result = sd::ops::slice(result, 2, 0, expected_frames); } return result; } + int tile_batch_ = 0; + size_t per_tile_compute_bytes_ = 0; + + static int env_int(const char* name, int fallback) { + const char* value = getenv(name); + if (value == nullptr || value[0] == '\0') { + return fallback; + } + return atoi(value); + } + + static sd::Tensor concat_frames(const std::vector>& pieces) { + std::vector shape = pieces[0].shape(); + int64_t frames = 0; + for (const auto& piece : pieces) { + GGML_ASSERT(piece.dim() == 5 && piece.shape()[0] == shape[0] && piece.shape()[1] == shape[1] && + piece.shape()[3] == shape[3] && piece.shape()[4] == shape[4]); + frames += piece.shape()[2]; + } + shape[2] = frames; + sd::Tensor output(shape); + const int64_t plane = shape[0] * shape[1]; + const int64_t planes = shape[3] * shape[4]; + int64_t frame_offset = 0; + for (const auto& piece : pieces) { + const int64_t piece_frames = piece.shape()[2]; + for (int64_t p = 0; p < planes; ++p) { + memcpy(output.data() + (p * frames + frame_offset) * plane, + piece.data() + p * piece_frames * plane, + sizeof(float) * piece_frames * plane); + } + frame_offset += piece_frames; + } + return output; + } + + // How many spatial tiles go into one decoder graph. Default 1 (one graph per tile): the + // batched graph runs each projection as one matmul over all tiles' tokens, and cuBLAS may + // pick a different kernel for the larger M (seen on RTX PRO 6000 Blackwell: 53 dB PSNR vs + // the per-tile decode), while it measured no faster there. SD_H3_VAE_TILE_BATCH=N forces N, + // SD_H3_VAE_TILE_BATCH=auto runs the first tile alone, measures its compute buffer and sizes + // the batch from the free device memory. + int resolve_tile_batch() { + const char* mode = getenv("SD_H3_VAE_TILE_BATCH"); + if (mode == nullptr || mode[0] == '\0') { + return 1; + } + int forced = strcmp(mode, "auto") == 0 ? 0 : atoi(mode); + if (forced > 0) { + return forced; + } + if (sd_backend_is_cpu(runtime_backend)) { + return 1; + } + if (per_tile_compute_bytes_ == 0) { + return 1; + } + if (tile_batch_ > 0) { + return tile_batch_; + } + size_t free_bytes = 0; + size_t total_bytes = 0; + ggml_backend_dev_t dev = ggml_backend_get_device(runtime_backend); + if (dev == nullptr) { + tile_batch_ = 1; + return tile_batch_; + } + ggml_backend_dev_memory(dev, &free_bytes, &total_bytes); + const size_t reserve = std::max(size_t(1) << 30, total_bytes / 10); + const size_t available = free_bytes > reserve ? free_bytes - reserve : 0; + // Sized as if the single-tile buffer were still held while the batched one is allocated. + int batch = static_cast(available / per_tile_compute_bytes_); + int limit = std::max(1, env_int("SD_H3_VAE_TILE_BATCH_MAX", 4)); + tile_batch_ = std::max(1, std::min(batch, limit)); + LOG_INFO("MiniMax-H3 video VAE: %d tile(s) per decoder graph (%.0f MB per tile, %.0f MB free)", + tile_batch_, + per_tile_compute_bytes_ / (1024.0 * 1024.0), + free_bytes / (1024.0 * 1024.0)); + return tile_batch_; + } + + // Tiles are stacked along the channel axis, so decoder output tile i is the i-th contiguous + // block of the [W, H, T, 3 * N, 1] result. + TileBatchOutput compute_tile_batch(int n_threads, + const std::vector>& tiles) { + TileBatchOutput output; + output.stack_dim = 3; + if (tiles.size() == 1) { + output.data = _compute(n_threads, tiles[0], true); + if (per_tile_compute_bytes_ == 0) { + per_tile_compute_bytes_ = reusable_compute_buffer_bytes(); + } + return output; + } + auto batched = ensure_video_shape(tiles[0]); + for (size_t i = 1; i < tiles.size(); ++i) { + batched = sd::ops::concat(batched, ensure_video_shape(tiles[i]), 3); + } + output.data = _compute(n_threads, batched, true); + return output; + } + + sd::Tensor decode_spatial_tiles(int n_threads, + const sd::Tensor& chunk, + const sd_tiling_params_t& tiling, + bool circular_x, + bool circular_y, + bool silent) { + const int scale_factor = get_scale_factor(); + float tile_overlap; + int tile_size_x, tile_size_y; + get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling, chunk.shape()[0], chunk.shape()[1]); + return process_tiles_2d_batched( + chunk, + static_cast(chunk.shape()[0] * scale_factor), + static_cast(chunk.shape()[1] * scale_factor), + scale_factor, + tile_size_x, + tile_size_y, + tile_overlap, + circular_x, + circular_y, + [&](int) { return resolve_tile_batch(); }, + [&](const std::vector>& tiles) { + auto output = compute_tile_batch(n_threads, tiles); + if (output.data.empty()) { + LOG_ERROR("vae decode compute failed while processing a tile"); + } + return output; + }, + silent); + } + + // rope_cache is [2, 2, half, L] = [[cos, -sin], [sin, cos]] per (pair, position). Rotated + // output d = i + half * j is x_lo[i] * pe[0][j][i] + x_hi[i] * pe[1][j][i]; tables hold the + // coefficient of x[d] (a) and of its partner x[d +- half] (b) for every head channel d. + void build_rope_tables() { + constexpr int64_t head_dim = DecoderAttention::head_dim; + const int64_t half = rope_cache.shape()[2]; + const int64_t positions = rope_cache.shape()[3]; + GGML_ASSERT(rope_cache.shape()[0] == 2 && rope_cache.shape()[1] == 2 && half * 2 <= head_dim); + rope_a_cache = sd::Tensor({head_dim, 1, positions}); + rope_b_cache = sd::Tensor({head_dim, 1, positions}); + const float* pe = rope_cache.data(); + auto pe_at = [&](int64_t a, int64_t j, int64_t i, int64_t l) { + return pe[((l * half + i) * 2 + j) * 2 + a]; + }; + for (int64_t l = 0; l < positions; ++l) { + float* a = rope_a_cache.data() + l * head_dim; + float* b = rope_b_cache.data() + l * head_dim; + for (int64_t i = 0; i < half; ++i) { + a[i] = pe_at(0, 0, i, l); + b[i] = pe_at(1, 0, i, l); + a[half + i] = pe_at(1, 1, i, l); + b[half + i] = pe_at(0, 1, i, l); + } + for (int64_t d = half * 2; d < head_dim; ++d) { + a[d] = 1.f; + b[d] = 0.f; + } + } + } + sd::Tensor build_rope(int64_t width, int64_t height, int64_t num_frames) { @@ -767,11 +1220,19 @@ namespace MiniMaxH3VAE { sd::Tensor _compute(const int n_threads, const sd::Tensor& z, bool decode_graph) override { - auto input = ensure_video_shape(z); + auto input = ensure_video_shape(z); + const bool graph_opt = decode_graph && env_int("SD_H3_VAE_GRAPH_OPT", 1) != 0; if (decode_graph) { - rope_cache = build_rope(input.shape()[0], - input.shape()[1], - input.shape()[2]); + const std::array rope_key = {input.shape()[0], input.shape()[1], input.shape()[2], graph_opt ? 1 : 0}; + if (rope_key != rope_key_ || rope_cache.empty()) { + rope_cache = build_rope(input.shape()[0], + input.shape()[1], + input.shape()[2]); + if (graph_opt) { + build_rope_tables(); + } + rope_key_ = rope_key; + } } auto get_graph = [&]() -> ggml_cgraph* { auto value = make_input(input); @@ -780,8 +1241,10 @@ namespace MiniMaxH3VAE { auto runner_ctx = get_context(); ggml_tensor* out = nullptr; if (decode_graph) { - auto pe = make_input(rope_cache); - out = model.decode(&runner_ctx, value, pe, mean, std); + auto pe = make_input(rope_cache); + ggml_tensor* ropea = graph_opt ? make_input(rope_a_cache) : nullptr; + ggml_tensor* ropeb = graph_opt ? make_input(rope_b_cache) : nullptr; + out = model.decode(&runner_ctx, value, pe, mean, std, ropea, ropeb); } else { out = model.encode(&runner_ctx, value, mean, std); } diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 9a9fe06287..e02270940b 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -1602,6 +1602,9 @@ void ModelManager::set_workspace_reclaimer(uintptr_t owner_id, std::functionfirst.first == owner_id ? device_query_cache_.erase(it) : std::next(it); + } for (auto it = runtime_residencies_.begin(); it != runtime_residencies_.end();) { if (it->first.first == owner_id) { it = runtime_residencies_.erase(it); @@ -1630,7 +1633,21 @@ ModelManager::CapacityCheck ModelManager::check_capacity( return SIZE_MAX; } size_t free_bytes = 0, total_bytes = 0; - ggml_backend_dev_memory(device, &free_bytes, &total_bytes); + // only checks that allocate nothing reuse a reading, so it never predates this owner's growth + const bool reusable = request.reuse_device_query && request.pending_allocation_bytes == 0 && missing == 0; + const auto cache_key = std::make_pair(request.owner_id, device); + auto cached = device_query_cache_.find(cache_key); + if (reusable && cached != device_query_cache_.end()) { + free_bytes = cached->second.first; + total_bytes = cached->second.second; + } else { + ggml_backend_dev_memory(device, &free_bytes, &total_bytes); + if (reusable) { + device_query_cache_[cache_key] = {free_bytes, total_bytes}; + } else if (cached != device_query_cache_.end()) { + device_query_cache_.erase(cached); + } + } if (free_bytes == 0 && total_bytes == 0) { return SIZE_MAX; } @@ -1647,6 +1664,19 @@ ModelManager::CapacityCheck ModelManager::check_capacity( return free_bytes; }; result.available_device_bytes = available_device_bytes(request.compute_backend); + struct DropReadingIfNoFit { + const CapacityCheck& result; + std::map, std::pair>& cache; + uintptr_t owner; + ~DropReadingIfNoFit() { + // A failed check leads to reclaiming / evicting; the retries must see fresh readings. + if (!result.fits()) { + for (auto it = cache.begin(); it != cache.end();) { + it = it->first.first == owner ? cache.erase(it) : std::next(it); + } + } + } + } drop_reading_if_no_fit{result, device_query_cache_, request.owner_id}; if (request.max_backend_bytes > 0) { const size_t resident = add(compute_backend_resident_bytes(request.compute_backend), other_runtime_resident_bytes(request.owner_id, request.compute_backend)); diff --git a/src/model_manager.h b/src/model_manager.h index 7f8df59c9b..e45a329536 100644 --- a/src/model_manager.h +++ b/src/model_manager.h @@ -94,6 +94,7 @@ class ModelManager : public DeviceResidencyManager { std::vector> tensor_states_; std::map tensor_states_by_tensor_; mutable std::list resolved_tensor_states_; + mutable std::map, std::pair> device_query_cache_; std::vector> params_storage_blocks_; std::vector> compute_staging_blocks_; std::map split_buffer_types_; diff --git a/src/runtime/tiling.cpp b/src/runtime/tiling.cpp index 37622253f1..2749b290e8 100644 --- a/src/runtime/tiling.cpp +++ b/src/runtime/tiling.cpp @@ -2,6 +2,11 @@ #include #include +#include +#include +#include +#include +#include #include #include @@ -124,7 +129,32 @@ static sd::Tensor sd_tensor_split_2d(const sd::Tensor& input, int return output; } -static void sd_tensor_merge_2d(const sd::Tensor& input, +template +static void sd_parallel_for_planes(int64_t plane_count, int64_t plane_elements, Fn&& fn) { + // Planes are independent and each element is still produced by one thread with the same + // arithmetic, so splitting by plane does not change the result. + const int64_t min_elements = 1 << 18; + int64_t threads = std::min(plane_count, std::max(1, plane_count * plane_elements / min_elements)); + threads = std::min(threads, std::max(1u, std::min(16u, std::thread::hardware_concurrency()))); + if (threads <= 1) { + fn(0, plane_count); + return; + } + std::vector workers; + const int64_t per_thread = (plane_count + threads - 1) / threads; + for (int64_t begin = per_thread; begin < plane_count; begin += per_thread) { + workers.emplace_back([&fn, begin, end = std::min(plane_count, begin + per_thread)]() { fn(begin, end); }); + } + fn(0, std::min(plane_count, per_thread)); + for (auto& worker : workers) { + worker.join(); + } +} + +static void sd_tensor_merge_2d(const float* input_data, + int64_t width, + int64_t height, + int64_t plane_count, sd::Tensor* output, int x, int y, @@ -135,13 +165,10 @@ static void sd_tensor_merge_2d(const sd::Tensor& input, int x_skip, int y_skip) { GGML_ASSERT(output != nullptr); - int64_t width = input.shape()[0]; - int64_t height = input.shape()[1]; int64_t img_width = output->shape()[0]; int64_t img_height = output->shape()[1]; - int64_t input_plane = sd_tensor_plane_size(input); + int64_t input_plane = width * height; int64_t output_plane = sd_tensor_plane_size(*output); - int64_t plane_count = input.numel() / input_plane; GGML_ASSERT(output->numel() / output_plane == plane_count); // unclamped -> expects x in the range [0-1] @@ -171,25 +198,28 @@ static void sd_tensor_merge_2d(const sd::Tensor& input, wy[iy] = smootherstep_f32(std::min(std::min(y_f_0, y_f_1), 1.f)); } } - for (int64_t plane = 0; plane < plane_count; ++plane) { - const float* src = input.data() + plane * input_plane; - float* dst = output->data() + plane * output_plane; - for (int64_t iy = y_skip; iy < height; iy++) { - const float* src_row = src + width * iy; - float* dst_row = dst + img_width * ((y + iy) % img_height); - if (blend) { - const float sy = wy[iy]; - for (int64_t ix = x_skip; ix < width; ix++) { - float& out = dst_row[dst_x[ix]]; - out = out + src_row[ix] * sy * wx[ix]; - } - } else { - for (int64_t ix = x_skip; ix < width; ix++) { - dst_row[dst_x[ix]] = src_row[ix]; + float* output_data = output->data(); + sd_parallel_for_planes(plane_count, input_plane, [&](int64_t first_plane, int64_t last_plane) { + for (int64_t plane = first_plane; plane < last_plane; ++plane) { + const float* src = input_data + plane * input_plane; + float* dst = output_data + plane * output_plane; + for (int64_t iy = y_skip; iy < height; iy++) { + const float* src_row = src + width * iy; + float* dst_row = dst + img_width * ((y + iy) % img_height); + if (blend) { + const float sy = wy[iy]; + for (int64_t ix = x_skip; ix < width; ix++) { + float& out = dst_row[dst_x[ix]]; + out = out + src_row[ix] * sy * wx[ix]; + } + } else { + for (int64_t ix = x_skip; ix < width; ix++) { + dst_row[dst_x[ix]] = src_row[ix]; + } } } } - } + }); } sd::Tensor process_tiles_2d(const sd::Tensor& input, @@ -203,6 +233,40 @@ sd::Tensor process_tiles_2d(const sd::Tensor& input, bool circular_y, const TileProcessCallback& on_processing, bool silent) { + return process_tiles_2d_batched( + input, + output_width, + output_height, + scale, + p_tile_size_x, + p_tile_size_y, + tile_overlap_factor, + circular_x, + circular_y, + [](int) { return 1; }, + [&](const std::vector>& tiles) { + TileBatchOutput output; + output.data = on_processing(tiles[0]); + if (!output.data.empty()) { + output.stack_dim = static_cast(output.data.dim() - 1); + } + return output; + }, + silent); +} + +sd::Tensor process_tiles_2d_batched(const sd::Tensor& input, + int output_width, + int output_height, + int scale, + int p_tile_size_x, + int p_tile_size_y, + float tile_overlap_factor, + bool circular_x, + bool circular_y, + const TileBatchSizeCallback& batch_size, + const TileBatchProcessCallback& on_processing, + bool silent) { sd::Tensor output; int input_width = static_cast(input.shape()[0]); int input_height = static_cast(input.shape()[1]); @@ -246,17 +310,17 @@ sd::Tensor process_tiles_2d(const sd::Tensor& input, input_tile_size_y *= scale; } - int num_tiles = num_tiles_x * num_tiles_y; - int tile_count = 1; - bool last_y = false; - bool last_x = false; - float last_time = 0.0f; - if (!silent) { - LOG_VERBOSE("num tiles : %d, %d ", num_tiles_x, num_tiles_y); - LOG_VERBOSE("optimal overlap : %f, %f (targeting %f)", tile_overlap_factor_x, tile_overlap_factor_y, tile_overlap_factor); - LOG_VERBOSE("processing %i tiles", num_tiles); - pretty_progress(0, num_tiles, 0.0f); - } + struct TilePlacement { + int x_in; + int y_in; + int x_out; + int y_out; + int dx; + int dy; + }; + std::vector placements; + bool last_y = false; + bool last_x = false; for (int y = 0; y < small_height && !last_y; y += non_tile_overlap_y) { int dy = 0; if (!circular_y && y + tile_size_y >= small_height) { @@ -279,38 +343,136 @@ sd::Tensor process_tiles_2d(const sd::Tensor& input, } last_x = true; } + placements.push_back({decode ? x : scale * x, + decode ? y : scale * y, + decode ? x * scale : x, + decode ? y * scale : y, + dx, + dy}); + } + last_x = false; + } - int x_in = decode ? x : scale * x; - int y_in = decode ? y : scale * y; - int x_out = decode ? x * scale : x; - int y_out = decode ? y * scale : y; - - int overlap_x_out = decode ? tile_overlap_x * scale : tile_overlap_x; - int overlap_y_out = decode ? tile_overlap_y * scale : tile_overlap_y; - - int64_t t1 = ggml_time_ms(); - auto input_tile = sd_tensor_split_2d(input, input_tile_size_x, input_tile_size_y, x_in, y_in); - auto output_tile = on_processing(input_tile); - if (output_tile.empty()) { - return {}; + int overlap_x_out = decode ? tile_overlap_x * scale : tile_overlap_x; + int overlap_y_out = decode ? tile_overlap_y * scale : tile_overlap_y; + int num_tiles = num_tiles_x * num_tiles_y; + int tile_count = 1; + float last_time = 0.0f; + if (!silent) { + LOG_VERBOSE("num tiles : %d, %d ", num_tiles_x, num_tiles_y); + LOG_VERBOSE("optimal overlap : %f, %f (targeting %f)", tile_overlap_factor_x, tile_overlap_factor_y, tile_overlap_factor); + LOG_VERBOSE("processing %i tiles", num_tiles); + pretty_progress(0, num_tiles, 0.0f); + } + // Blend on a worker while the next batch computes, still in tile order; SD_TILE_ASYNC_MERGE=0 blends inline. + static const bool async_merge = []() { + const char* value = getenv("SD_TILE_ASYNC_MERGE"); + return value == nullptr || value[0] == '\0' || atoi(value) != 0; + }(); + std::thread merger; + bool merge_failed = false; + auto join_merger = [&]() { + if (merger.joinable()) { + merger.join(); + } + }; + struct MergerGuard { + std::thread& thread; + ~MergerGuard() { + if (thread.joinable()) { + thread.join(); } - GGML_ASSERT(output_tile.shape()[0] == output_tile_size_x && output_tile.shape()[1] == output_tile_size_y); + } + } merger_guard{merger}; + std::vector> input_tiles; + for (size_t first = 0; first < placements.size();) { + const int remaining = static_cast(placements.size() - first); + const int count = std::max(1, std::min(remaining, batch_size(remaining))); + const size_t last = first + static_cast(count); + int64_t t1 = ggml_time_ms(); + input_tiles.clear(); + for (size_t i = first; i < last; ++i) { + input_tiles.push_back(sd_tensor_split_2d(input, input_tile_size_x, input_tile_size_y, placements[i].x_in, placements[i].y_in)); + } + auto batch = on_processing(input_tiles); + if (batch.data.empty()) { + return {}; + } + const int64_t batch_tiles = static_cast(last - first); + GGML_ASSERT(batch.stack_dim < static_cast(batch.data.dim()) && batch.data.shape()[batch.stack_dim] % batch_tiles == 0); + for (size_t d = batch.stack_dim + 1; d < static_cast(batch.data.dim()); ++d) { + GGML_ASSERT(batch.data.shape()[d] == 1); + } + std::vector tile_shape = batch.data.shape(); + tile_shape[batch.stack_dim] /= batch_tiles; + GGML_ASSERT(tile_shape[0] == output_tile_size_x && tile_shape[1] == output_tile_size_y); + auto merge_batch = [&, first, last, tile_shape](const sd::Tensor& data) { + const int64_t tile_numel = data.numel() / static_cast(last - first); + const int64_t plane_count = tile_numel / (tile_shape[0] * tile_shape[1]); if (output.empty()) { - std::vector output_shape = output_tile.shape(); + std::vector output_shape = tile_shape; output_shape[0] = output_width; output_shape[1] = output_height; output = sd::Tensor::zeros(std::move(output_shape)); } - sd_tensor_merge_2d(output_tile, &output, x_out, y_out, overlap_x_out, overlap_y_out, circular_x, circular_y, dx, dy); + for (size_t i = first; i < last; ++i) { + const auto& placement = placements[i]; + sd_tensor_merge_2d(data.data() + static_cast(i - first) * tile_numel, + tile_shape[0], + tile_shape[1], + plane_count, + &output, + placement.x_out, + placement.y_out, + overlap_x_out, + overlap_y_out, + circular_x, + circular_y, + placement.dx, + placement.dy); + } + }; + join_merger(); + if (merge_failed) { + return {}; + } + bool merged_async = false; + if (async_merge && last < placements.size()) { + auto data = std::make_shared>(std::move(batch.data)); + try { + merger = std::thread([&merge_failed, merge_batch, data]() { + try { + merge_batch(*data); + } catch (const std::exception& error) { + LOG_ERROR("tile merge failed: %s", error.what()); + merge_failed = true; + } + }); + merged_async = true; + } catch (const std::system_error&) { + merge_batch(*data); // no thread available: merge inline + merged_async = true; + } + } + if (!merged_async) { + merge_batch(batch.data); + } - if (!silent) { - int64_t t2 = ggml_time_ms(); - last_time = (t2 - t1) / 1000.0f; + if (!silent) { + int64_t t2 = ggml_time_ms(); + last_time = (t2 - t1) / 1000.0f / static_cast(last - first); + for (size_t i = first; i < last; ++i) { pretty_progress(tile_count, num_tiles, last_time); + tile_count++; } - tile_count++; + } else { + tile_count += static_cast(last - first); } - last_x = false; + first = last; + } + join_merger(); + if (merge_failed) { + return {}; } if (!silent && tile_count < num_tiles) { pretty_progress(num_tiles, num_tiles, last_time); diff --git a/src/runtime/tiling.h b/src/runtime/tiling.h index d97949ce59..433ec08b93 100644 --- a/src/runtime/tiling.h +++ b/src/runtime/tiling.h @@ -2,10 +2,20 @@ #define __SD_RUNTIME_TILING_H__ #include +#include #include "core/tensor.hpp" -using TileProcessCallback = std::function(const sd::Tensor&)>; +using TileProcessCallback = std::function(const sd::Tensor&)>; +using TileBatchSizeCallback = std::function; +// Output tiles of one batch stacked along stack_dim, which must be the highest axis with an extent +// above 1, so tile i is the i-th contiguous block of data. Each tile's shape is data.shape() with +// shape[stack_dim] divided by the number of tiles. +struct TileBatchOutput { + sd::Tensor data; + size_t stack_dim = 0; +}; +using TileBatchProcessCallback = std::function>&)>; sd::Tensor process_tiles_2d(const sd::Tensor& input, int output_width, @@ -19,6 +29,21 @@ sd::Tensor process_tiles_2d(const sd::Tensor& input, const TileProcessCallback& on_processing, bool silent = false); +// Same tile plan and merge order as process_tiles_2d, but hands up to batch_size(remaining) +// consecutive tiles to the callback at once. +sd::Tensor process_tiles_2d_batched(const sd::Tensor& input, + int output_width, + int output_height, + int scale, + int p_tile_size_x, + int p_tile_size_y, + float tile_overlap_factor, + bool circular_x, + bool circular_y, + const TileBatchSizeCallback& batch_size, + const TileBatchProcessCallback& on_processing, + bool silent = false); + // Returns tile_size, adjusted if needed so a non-circular axis of small_dim is never split into tiles that // barely overlap (see the definition for why that shows as a seam). int sd_tiling_seam_safe_tile_size(int small_dim, int tile_size, float tile_overlap_factor);