From ac399008a8c897de22c585b2110d81ad698bcc4c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 06:14:05 -0700 Subject: [PATCH 01/13] Carry the pipelined long-sequence ggml-cuda FlashAttention patch scripts/unsloth/ggml-patches/0002 pipelines the unmasked, non-GQA mma FlashAttention kernels (all full KV tiles via cp.async, only the partial last tile synchronously). The MiniMax-H3 DiT (56 heads x 128, 19108 tokens) and the H3 video VAE (32 heads x 64) attention take this path under --diffusion-fa. GGML_CUDA_FA_LONGSEQ=0 restores the stock kernels. --- .../0002-ggml-cuda-fa-longseq.patch | 387 ++++++++++++++++++ 1 file changed, 387 insertions(+) create mode 100644 scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch 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 000000000..5157cba94 --- /dev/null +++ b/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch @@ -0,0 +1,387 @@ +From ed3f840655078aef3c2bb4e6b9a1d26f8e868e64 Mon Sep 17 00:00:00 2001 +From: Daniel Han +Date: Thu, 1 Oct 2026 05:41:03 -0700 +Subject: [PATCH] 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 ddd8ac45b252afb82bcc6d70068767c7379760f7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 06:27:44 -0700 Subject: [PATCH 02/13] ggml patch 0002: 128-column long-sequence FlashAttention tile on sm80 / sm89 --- .../0002-ggml-cuda-fa-longseq.patch | 66 ++++++++++++++++++- 1 file changed, 65 insertions(+), 1 deletion(-) diff --git a/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch b/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch index 5157cba94..564831ee7 100644 --- a/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch +++ b/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch @@ -1,7 +1,7 @@ From ed3f840655078aef3c2bb4e6b9a1d26f8e868e64 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 05:41:03 -0700 -Subject: [PATCH] ggml-cuda: pipelined long-sequence FlashAttention for +Subject: [PATCH 1/2] ggml-cuda: pipelined long-sequence FlashAttention for unmasked attention without GQA packing The mma FlashAttention kernels without GQA packing (ncols2 == 1, used whenever @@ -385,3 +385,67 @@ index fdf505cd..dd290caf 100644 -- 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/2] 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 fbc3a400a2c6b6f18f41822b10002a14ea78b116 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 06:27:49 -0700 Subject: [PATCH 03/13] docs: MiniMax-H3 long-sequence flash attention switches --- docs/minimax_h3.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index 0417d0011..abbd675ed 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -50,6 +50,11 @@ decoded audio track. most of its attention cost without touching the text encoder. Set `SD_H3_VAE_FLASH_ATTN=0` to decode with the previous mul_mat + softmax attention. +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 (about 2x 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=64|128` overrides the per-architecture tile width. + ## First/last-frame conditioning Add `--init-img` for I2VA, or both `--init-img` and `--end-img` for FL2VA: From a565a22c9cb276b27185ead0b68accacbe80fc7a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 06:15:20 -0700 Subject: [PATCH 04/13] tiling: batched tile callback, merge tile planes on several threads process_tiles_2d_batched walks the same tile plan and merge order as process_tiles_2d but hands up to batch_size(remaining) consecutive tiles to one callback, which returns them stacked as contiguous blocks; the merge reads each tile in place instead of from a copy. process_tiles_2d is now a one-tile wrapper around it. The merge splits the output planes across threads. Each element is still written by one thread with the same arithmetic, so outputs are bitwise identical. --- src/runtime/tiling.cpp | 225 ++++++++++++++++++++++++++++++----------- src/runtime/tiling.h | 27 ++++- 2 files changed, 192 insertions(+), 60 deletions(-) diff --git a/src/runtime/tiling.cpp b/src/runtime/tiling.cpp index 767e9c4d1..1bfb291a7 100644 --- a/src/runtime/tiling.cpp +++ b/src/runtime/tiling.cpp @@ -2,6 +2,7 @@ #include #include +#include #include #include @@ -94,7 +95,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, @@ -105,13 +131,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] @@ -142,25 +165,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, @@ -174,6 +200,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]); @@ -217,17 +277,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) { @@ -250,38 +310,85 @@ 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 {}; - } - GGML_ASSERT(output_tile.shape()[0] == output_tile_size_x && output_tile.shape()[1] == output_tile_size_y); - if (output.empty()) { - std::vector output_shape = output_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); + 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); + } + 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); + const int64_t tile_numel = batch.data.numel() / batch_tiles; + const int64_t plane_count = tile_numel / (tile_shape[0] * tile_shape[1]); + if (output.empty()) { + std::vector output_shape = tile_shape; + output_shape[0] = output_width; + output_shape[1] = output_height; + output = sd::Tensor::zeros(std::move(output_shape)); + } + for (size_t i = first; i < last; ++i) { + const auto& placement = placements[i]; + sd_tensor_merge_2d(batch.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); + } - 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; } 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 a832a644e..352c4abb2 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,4 +29,19 @@ 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); + #endif // __SD_RUNTIME_TILING_H__ From 32861d81c3adbcedbebd29407e86411e7476f321 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 06:15:20 -0700 Subject: [PATCH 05/13] MiniMax-H3: batch video VAE tiles per decoder graph, keep weights resident, leaner rope and SwiGLU The video VAE decoder (a 36-layer ViT) ran one graph per 16x16 latent tile: 105 graph builds, allocations and output readbacks for 960x544x124, and every temporal chunk ended with runner_end(), which evicted the 5.5 GB of weights and reuploaded them for the next chunk. - Several tiles now go through one decoder graph. The first tile runs alone to measure its compute buffer and the batch is sized from free device memory, capped at 4 (SD_H3_VAE_TILE_BATCH_MAX); SD_H3_VAE_TILE_BATCH=N forces N and 1 restores one graph per tile. Projections run as one 2D matmul over all tiles' tokens and each tile keeps its own attention call: a broadcast batched cuBLAS GEMM and a batched flash-attention launch both changed the summation order relative to the single-tile decode. With this layout the decoded frames match the per-tile decode exactly. - Weights stay resident across temporal chunks and runner_end() runs once after the last chunk (SD_H3_VAE_KEEP_RESIDENT=0 restores the per-chunk release). - Rotary embedding reads the two rotated halves as permuted views of the projection layout and multiplies them by precomputed coefficient tables, with the same separate mul/mul/add per element as before, instead of per-half permutes, repeats, a tail permute and concat. The feed-forward uses ggml_swiglu instead of silu + mul. SD_H3_VAE_GRAPH_OPT=0 restores the previous graph. - Chunks are blended in place with plain loops instead of index() lookups, and the decoded chunks are concatenated once at the end instead of growing the result chunk by chunk. --- docs/minimax_h3.md | 8 + src/model/vae/minimax_h3_vae.hpp | 484 +++++++++++++++++++++++++------ 2 files changed, 396 insertions(+), 96 deletions(-) diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index abbd675ed..0e16f8ae8 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -55,6 +55,14 @@ long-sequence variant of the ggml-cuda flash attention kernel (about 2x to 4x fa shapes; the arithmetic per output element is unchanged). `GGML_CUDA_FA_LONGSEQ=0` restores the stock kernel, `GGML_CUDA_FA_LONGSEQ_NCOLS=64|128` overrides the per-architecture tile width. +The video VAE decodes several 16x16 latent tiles per decoder graph, sized from free device +memory (at most 4 unless `SD_H3_VAE_TILE_BATCH_MAX` raises it; `SD_H3_VAE_TILE_BATCH=N` forces +N, and 1 restores one graph per tile). Its 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 tile still goes through its own attention call and the projections run as one +matmul over all tiles' tokens, so the decoded frames are the same as with one tile per graph. + ## First/last-frame conditioning Add `--init-img` for I2VA, or both `--init-img` and `--end-img` for FL2VA: diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index 777566438..1ae87ec87 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -4,6 +4,8 @@ #include #include #include +#include +#include #include #include #include @@ -13,6 +15,7 @@ #include "model/common/rope.hpp" #include "model/diffusion/dit.hpp" #include "model/vae/vae.hpp" +#include "runtime/tiling.h" namespace MiniMaxH3VAE { @@ -244,6 +247,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; @@ -256,50 +290,88 @@ 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 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); + 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 +384,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 +420,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 +431,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 +475,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 +538,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 +558,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 +583,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 +603,8 @@ namespace MiniMaxH3VAE { sd::Tensor latents_mean; sd::Tensor latents_std; sd::Tensor rope_cache; + sd::Tensor rope_a_cache; + sd::Tensor rope_b_cache; MiniMaxH3VideoVAERunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map, @@ -576,27 +690,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,14 +806,15 @@ 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); + const bool keep_resident = env_int("SD_H3_VAE_KEEP_RESIDENT", 1) != 0; + tile_batch_ = 0; + per_tile_compute_bytes_ = 0; + std::vector> pieces; + auto collect_pieces = [&](const sd::Tensor& chunk, const VAETemporalTile& tile) { + auto decoded = decode_spatial_tiles(n_threads, chunk, tiling, circular_x, circular_y, silent); + if (!keep_resident) { + runner_end(); + } if (decoded.empty()) { return sd::Tensor(); } @@ -706,7 +825,7 @@ namespace MiniMaxH3VAE { 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,10 +840,21 @@ namespace MiniMaxH3VAE { overlap = {}; } return first; - }); - if (result.empty()) { + }; + bool failed = false; + for (const auto& tile : plan.tiles) { + auto piece = collect_pieces(sd::ops::slice(input, 2, tile.start, tile.end), tile); + if (piece.empty()) { + failed = true; + break; + } + pieces.push_back(std::move(piece)); + } + runner_end(); + if (failed || pieces.empty()) { return {}; } + auto result = concat_frames(pieces); int64_t expected_frames = input.shape()[2] <= 1 ? 1 : ((x.shape()[2] - 2) / 5) * 17 + 5; expected_frames = std::max(1, expected_frames); @@ -734,6 +864,162 @@ namespace MiniMaxH3VAE { 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. SD_H3_VAE_TILE_BATCH=N forces N + // (1 restores one graph per tile); otherwise the first tile runs alone, its compute + // buffer is measured, and the batch is sized from the free device memory. + int resolve_tile_batch() { + int forced = env_int("SD_H3_VAE_TILE_BATCH", 0); + 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 +1053,15 @@ 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]); + if (graph_opt) { + build_rope_tables(); + } } auto get_graph = [&]() -> ggml_cgraph* { auto value = make_input(input); @@ -780,8 +1070,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); } From 1f1bef0d4127d55eb4ff0a5fc36cbe42ddf3543b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 06:15:20 -0700 Subject: [PATCH 06/13] ggml-patches: 0003 fused cuBLAS bias epilogues and narrow short-row RMS norm blocks Carried on top of the pinned ggml submodule like 0001 (the prebuilt workflow applies scripts/unsloth/ggml-patches/*.patch in order). In the H3 video VAE decoder every linear layer is an F16 cuBLAS matmul followed by a bias add and then either a scale + residual add or SWIGLU; these now run inside the F16 -> F32 conversion of the GEMM output instead of as separate passes over F32 activations, bitwise identical (GGML_CUDA_CUBLAS_EPILOGUE_FUSION=0 disables). The per-head q/k RMS norm (64 columns) now launches 64-thread blocks instead of 256, also bitwise identical (GGML_CUDA_NORM_SMALL_ROWS=0 restores 256). Applies cleanly with or without 0002. --- ...003-ggml-cuda-h3-vae-epilogue-fusion.patch | 348 ++++++++++++++++++ 1 file changed, 348 insertions(+) create mode 100644 scripts/unsloth/ggml-patches/0003-ggml-cuda-h3-vae-epilogue-fusion.patch 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 000000000..1e14c5358 --- /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 + From dbb4106b4bdc578497d61c77747f9f6b6e164a08 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 06:48:15 -0700 Subject: [PATCH 07/13] MiniMax-H3: value-preserving DiT graph rewrites The DiT spent about 1.9 s of an 8.6 s B200 step (960x544x124, Q3) outside matmul and attention: chunk copies of q/k/v and of the MLP gate, partial RoPE as slice + permute copies + repeat + mul/add + concat, K/V scale then cast, per-segment adaLN modulation and gated residuals as slice copies + mul/add + concat, and separate scale passes around every MLP Linear. Each lever builds the same arithmetic with the same rounding points, so the output is bit-identical: - q/k/v are read in place from the fused projection; one ggml_rope_pe_permute per tensor applies the partial RoPE, writes the head-major layout flash attention reads and, for K/V, the kv scale and F16 cast (SD_H3_FAST_QKV=0 restores the old graph). - The MLP gate is one swiglu on two views of fc1's output instead of two chunk copies + silu + mul (SD_H3_FAST_MLP=0). - adaLN modulation and gated residuals use one in-place ggml_modulate_rows / ggml_gated_add_rows per sequence segment, and the MLP Linear scales are folded into them and into ggml_swiglu_scaled (SD_H3_FAST_SEGMENTS=0). - Modulation vectors and segment slices are views, not copies (SD_H3_FAST_VIEWS=0). SD_H3_GRAPH_FAST=0 turns all of them off. Backends without the fused ops (and upstream ggml builds) keep the original graph. --- src/core/ggml_extend.cpp | 67 ++++++-- src/core/ggml_extend.h | 21 +++ src/model/common/ggml_block.hpp | 24 +++ src/model/diffusion/minimax_h3.hpp | 248 ++++++++++++++++++++++++++++- 4 files changed, 344 insertions(+), 16 deletions(-) diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 74324169e..58130582d 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 33d8c3bf2..a5076262f 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/model/common/ggml_block.hpp b/src/model/common/ggml_block.hpp index 991b88f28..b4c94b2d9 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 80c29ab52..3b75408aa 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,48 @@ 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; + } + + 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 +188,54 @@ 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); +#ifdef SD_USE_UPSTREAM_GGML + 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); +#endif + return ggml_ext_linear_matmul(ctx->ggml_ctx, gated, fc2->foldable_weight(ctx), fc2->prec_f32()); + } + 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 +279,62 @@ 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) { +#ifdef SD_USE_UPSTREAM_GGML + return nullptr; +#else + if (!ctx->flash_attn_enabled || ctx->sage_attn_enabled || ctx->backend == nullptr) { + 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)); + 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, 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, 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 +345,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 +481,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 +502,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 +521,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 +550,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) { +#ifdef SD_USE_UPSTREAM_GGML + 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"]); From 49244dfc0ada20f66ac0af06ea43c6ad0dc8ba6c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 06:48:30 -0700 Subject: [PATCH 08/13] Carry the fused DiT ops ggml patch; document the H3 graph switches scripts/unsloth/ggml-patches/0004 adds ggml_rope_pe_permute, ggml_modulate_rows / ggml_gated_add_rows and ggml_swiglu_scaled (CPU + CUDA, test-backend-ops cases) on top of 0001 and 0002. The MiniMax-H3 DiT uses them when the backend supports them and keeps the unfused graph otherwise. --- docs/minimax_h3.md | 7 + .../0004-ggml-h3-dit-fused-ops.patch | 1020 +++++++++++++++++ 2 files changed, 1027 insertions(+) create mode 100644 scripts/unsloth/ggml-patches/0004-ggml-h3-dit-fused-ops.patch diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index 0e16f8ae8..c852492ea 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -63,6 +63,13 @@ table-based rotary embedding and a fused SwiGLU (`SD_H3_VAE_GRAPH_OPT=0` restore graph). Each tile still goes through its own attention call and the projections run as one matmul over all tiles' tokens, so the decoded frames are the same as with one tile per graph. +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. + ## 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/0004-ggml-h3-dit-fused-ops.patch b/scripts/unsloth/ggml-patches/0004-ggml-h3-dit-fused-ops.patch new file mode 100644 index 000000000..689dfa5f6 --- /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 + From 36963a4a83f5ae28e813ea51f9483eb7a1b7ad01 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 06:55:12 -0700 Subject: [PATCH 09/13] MiniMax-H3: keep the folded MLP path compiling with SD_USE_UPSTREAM_GGML forward_folded is only reachable when the fused ops exist; return from inside the guarded branch so the upstream-ggml build does not reference an undeclared result. --- src/model/diffusion/minimax_h3.hpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index 3b75408aa..07a244a8a 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -215,11 +215,14 @@ namespace MiniMaxH3 { 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); #ifdef SD_USE_UPSTREAM_GGML + 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); -#endif 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 { From 69babafe30eb10c2caccac2a22f53e2117c8ee0a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 06:55:12 -0700 Subject: [PATCH 10/13] ggml patch 0002: long-sequence flash attention keeps the stock grid, wide tile opt-in The pipelined variant now sizes its stream-k grid from the stock kernel so the work split, and with it the result, is identical on every GPU (on L4 the two kernels' occupancy differs). The 128-column tile is opt-in only (GGML_CUDA_FA_LONGSEQ_NCOLS=128). --- docs/minimax_h3.md | 5 +- .../0002-ggml-cuda-fa-longseq.patch | 116 +++++++++++++++++- 2 files changed, 117 insertions(+), 4 deletions(-) diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index c852492ea..7d5aa2447 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -51,9 +51,10 @@ most of its attention cost without touching the text encoder. Set `SD_H3_VAE_FLA to decode with the previous mul_mat + softmax attention. 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 (about 2x to 4x faster at H3 +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=64|128` overrides the per-architecture tile width. +stock kernel. `GGML_CUDA_FA_LONGSEQ_NCOLS=128` opts into a wider tile that is faster on A100 and L4 +but not bit-identical. The video VAE decodes several 16x16 latent tiles per decoder graph, sized from free device memory (at most 4 unless `SD_H3_VAE_TILE_BATCH_MAX` raises it; `SD_H3_VAE_TILE_BATCH=N` forces diff --git a/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch b/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch index 564831ee7..32b2c5378 100644 --- a/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch +++ b/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch @@ -1,7 +1,7 @@ From ed3f840655078aef3c2bb4e6b9a1d26f8e868e64 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 05:41:03 -0700 -Subject: [PATCH 1/2] ggml-cuda: pipelined long-sequence FlashAttention for +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 @@ -389,7 +389,7 @@ index fdf505cd..dd290caf 100644 From 7f7c132c57dbfd8fd44a13dd29031c6724589efc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 06:27:44 -0700 -Subject: [PATCH 2/2] ggml-cuda: per-architecture tile width for the +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 @@ -449,3 +449,115 @@ index 00e9554a..b8be30a0 100644 -- 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 + From b92b2a2cbc0bc6fa8c01d6b23412ba944458cb92 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 07:57:14 -0700 Subject: [PATCH 11/13] MiniMax-H3: video VAE tile batching opt-in (SD_H3_VAE_TILE_BATCH=auto|N) The batched decoder graph runs each projection as one matmul over all tiles' tokens. On an RTX PRO 6000 Blackwell (sm120) cuBLAS picks a different kernel for that larger M and the frames come out at 53 dB PSNR against the per-tile decode (identical on B200), and the batched graph measured no faster there (VAE decode 7.6 s batched vs 7.5 s per tile, the other VAE levers on). Decode one tile per graph by default; SD_H3_VAE_TILE_BATCH=auto or =N opts in. --- docs/minimax_h3.md | 11 ++++++----- src/model/vae/minimax_h3_vae.hpp | 15 +++++++++++---- 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index 7d5aa2447..4b967e0e3 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -56,13 +56,14 @@ shapes; the arithmetic per output element is unchanged). `GGML_CUDA_FA_LONGSEQ=0 stock kernel. `GGML_CUDA_FA_LONGSEQ_NCOLS=128` opts into a wider tile that is faster on A100 and L4 but not bit-identical. -The video VAE decodes several 16x16 latent tiles per decoder graph, sized from free device -memory (at most 4 unless `SD_H3_VAE_TILE_BATCH_MAX` raises it; `SD_H3_VAE_TILE_BATCH=N` forces -N, and 1 restores one graph per tile). Its weights stay on the device across temporal chunks +The video VAE decodes one 16x16 latent tile per decoder graph by default. +`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 tile still goes through its own attention call and the projections run as one -matmul over all tiles' tokens, so the decoded frames are the same as with one tile per graph. +graph). Each batched tile still goes through its own attention call. 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 diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index 1ae87ec87..604a05fb9 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -900,11 +900,18 @@ namespace MiniMaxH3VAE { return output; } - // How many spatial tiles go into one decoder graph. SD_H3_VAE_TILE_BATCH=N forces N - // (1 restores one graph per tile); otherwise the first tile runs alone, its compute - // buffer is measured, and the batch is sized from the free device memory. + // 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() { - int forced = env_int("SD_H3_VAE_TILE_BATCH", 0); + 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; } From 5ab239da00c27904ca484806c53a9ed092505ddc Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Thu, 8 Oct 2026 10:37:48 -0300 Subject: [PATCH 12/13] Build without the fused H3 ops when ggml lacks patch 0004 --- cmake/ggml.cmake | 12 ++++++++++++ src/model/diffusion/minimax_h3.hpp | 6 +++--- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/cmake/ggml.cmake b/cmake/ggml.cmake index 39d77eae3..2413abe2c 100644 --- a/cmake/ggml.cmake +++ b/cmake/ggml.cmake @@ -21,6 +21,18 @@ 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\\(") +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() + 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/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index 07a244a8a..f82a6445c 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -214,7 +214,7 @@ namespace MiniMaxH3 { 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); -#ifdef SD_USE_UPSTREAM_GGML +#ifndef SD_GGML_H3_FUSED_OPS GGML_UNUSED(s1); GGML_UNUSED(s2); GGML_UNUSED(halves); @@ -290,7 +290,7 @@ namespace MiniMaxH3 { int64_t sequence, int64_t batch, ggml_tensor* pe) { -#ifdef SD_USE_UPSTREAM_GGML +#ifndef SD_GGML_H3_FUSED_OPS return nullptr; #else if (!ctx->flash_attn_enabled || ctx->sage_attn_enabled || ctx->backend == nullptr) { @@ -561,7 +561,7 @@ namespace MiniMaxH3 { ggml_tensor* t_emb, const std::vector& segments, ggml_tensor* pe) { -#ifdef SD_USE_UPSTREAM_GGML +#ifndef SD_GGML_H3_FUSED_OPS return nullptr; #else if (ctx->backend == nullptr || segments.empty() || x->type != GGML_TYPE_F32) { From 703f519b0a26c48cef23a8503a1d4d6318ecbd18 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Thu, 8 Oct 2026 10:46:24 -0300 Subject: [PATCH 13/13] Document that the H3 ggml speedups need the carried patches --- docs/minimax_h3.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index 4b967e0e3..f54d46504 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -72,6 +72,12 @@ and the swiglu. The result is bit-identical to the unfused graph. `SD_H3_GRAPH_F 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. +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: