Merge the MiniMax-H3 speed stack (#18, #19, #20) into master

.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

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.")

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: