From 7552d8ec917e9cea67cf15dc1e9e2b64d95895c8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 30 Sep 2026 23:56:35 -0700 Subject: [PATCH 01/30] MiniMax-H3: flash attention in the video VAE under --diffusion-fa; carry ggml-cuda speed patch The H3 video VAE decoder is a 36-layer ViT run per 16x16 latent tile. With only --diffusion-fa (what Studio passes) it used mul_mat + scale + softmax attention, materialising a 32 x L x L f32 score matrix per tile and layer. --diffusion-fa now also covers it; SD_H3_VAE_FLASH_ATTN=0 restores the old path. The text encoder is unchanged (that is what --fa would also switch, which moves the conditioning). scripts/unsloth/ggml-patches/ carries two ggml-cuda commits on top of the pinned submodule (BF16 cuBLAS for large-batch quantized matmul behind GGML_CUDA_QUANT_CUBLAS_MIN_BATCH, a cached device integrated flag, a vectorized row copy); the prebuilt workflow applies them after the submodule checkout. --- .github/workflows/unsloth-sd-prebuilt.yml | 6 + docs/minimax_h3.md | 4 + .../0001-ggml-cuda-h3-speed.patch | 255 ++++++++++++++++++ src/pipeline/diffusion_engine.cpp | 13 + 4 files changed, 278 insertions(+) create mode 100644 scripts/unsloth/ggml-patches/0001-ggml-cuda-h3-speed.patch diff --git a/.github/workflows/unsloth-sd-prebuilt.yml b/.github/workflows/unsloth-sd-prebuilt.yml index 291ee0f7ae..81708c3f11 100644 --- a/.github/workflows/unsloth-sd-prebuilt.yml +++ b/.github/workflows/unsloth-sd-prebuilt.yml @@ -222,6 +222,12 @@ jobs: # (we build with SD_SERVER_BUILD_FRONTEND / SD_WEBP / SD_WEBM off), so fetch only # ggml to keep the source tarball small and the build self-contained. git submodule update --init --recursive --depth 1 ggml + # ggml changes we carry on top of the pinned submodule commit (the submodule points at + # leejet/ggml, which we do not push to). A patch that no longer applies fails the build. + for p in scripts/unsloth/ggml-patches/*.patch; do + [ -e "$p" ] || continue + git -C ggml apply --verbose "../$p" + done COMMIT="$(git rev-parse HEAD)" SRC_ARTIFACT="sd-source-${TAG}" if [ "$EXISTS" != "true" ] || [ "${{ github.event_name }}" = "workflow_dispatch" ]; then diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index 82d0e2ffb4..0417d0011a 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -46,6 +46,10 @@ are detected from their weights. Omitting `--audio-vae` still runs the joint diffusion model but produces video without a decoded audio track. +`--diffusion-fa` also enables flash attention in the video VAE decoder (a ViT), which removes +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. + ## 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/0001-ggml-cuda-h3-speed.patch b/scripts/unsloth/ggml-patches/0001-ggml-cuda-h3-speed.patch new file mode 100644 index 0000000000..f8b379f189 --- /dev/null +++ b/scripts/unsloth/ggml-patches/0001-ggml-cuda-h3-speed.patch @@ -0,0 +1,255 @@ +From 22b560daa5e90bdb72c492f0b55667468ad3fd6b Mon Sep 17 00:00:00 2001 +From: Daniel Han +Date: Wed, 30 Sep 2026 22:28:21 -0700 +Subject: [PATCH 1/2] ggml-cuda: BF16 cuBLAS for large-batch quantized matmul, + cache device integrated flag + +GGML_CUDA_QUANT_CUBLAS_MIN_BATCH= dequantizes a quantized weight to BF16 +and runs a cuBLAS GEMM (FP32 accumulate, FP32 output) instead of MMQ when src1 +has at least that many rows; 0 or unset keeps MMQ everywhere. +GGML_CUDA_QUANT_CUBLAS_TYPE=f32 dequantizes to FP32 instead (accuracy oracle). +At diffusion-transformer sizes (MiniMax-H3 960x544x124: 19108 tokens per call) +MMQ reaches 130-240 TFLOP/s on a B200 while the BF16 GEMM reaches ~1300, and +the result is closer to an FP32 reference because activations stay BF16 +instead of being quantized to q8_1. + +ggml_backend_cuda_device_get_memory / get_type called cudaGetDeviceProperties +(~1-2 ms) on every call; stable-diffusion.cpp polls them per graph, 2798 times +in one tiled H3 VAE decode. The flag is now read once at registration. +--- + src/ggml-cuda/ggml-cuda.cu | 55 ++++++++++++++++++++++++++++++++------ + 1 file changed, 47 insertions(+), 8 deletions(-) + +diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu +index e15fdc5e..517d0894 100644 +--- a/src/ggml-cuda/ggml-cuda.cu ++++ b/src/ggml-cuda/ggml-cuda.cu +@@ -2427,6 +2427,25 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { + return use_mul_mat_vec_q; + } + ++// Large-batch matmul with a quantized weight: dequantize the weight to BF16 once and run a cuBLAS ++// GEMM (FP32 accumulate, FP32 output) instead of MMQ. MMQ quantizes the activations to q8_1 and runs ++// int8 MMA, which is the right call for LLM prompt sizes, but at diffusion-transformer sizes (10^4 ++// tokens per call) a tensor-core BF16 GEMM is faster on GPUs whose BF16 rate is not far below their ++// int8 rate, and it is more accurate (activations stay BF16 instead of 8-bit). ++// GGML_CUDA_QUANT_CUBLAS_MIN_BATCH=: use the BF16 path when src1 has at least this many rows; ++// 0 disables it (stock MMQ everywhere). Unset = the per-architecture default below. ++static int64_t ggml_cuda_quant_cublas_min_batch(int cc) { ++ static const int64_t env = [] { ++ const char * e = getenv("GGML_CUDA_QUANT_CUBLAS_MIN_BATCH"); ++ return e != nullptr ? (int64_t) atoll(e) : (int64_t) -1; ++ }(); ++ if (env >= 0) { ++ return env; ++ } ++ GGML_UNUSED(cc); ++ return 0; ++} ++ + static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + GGML_TENSOR_BINARY_OP_LOCALS + +@@ -2492,6 +2511,23 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst); + return; + } ++ if (ggml_is_quantized(src0->type) && ne11 > MMVQ_MAX_BATCH_SIZE && GGML_CUDA_CC_IS_NVIDIA(cc) && ++ bf16_mma_hardware_available(cc) && ggml_get_to_bf16_cuda(src0->type) != nullptr) { ++ const int64_t min_batch = ggml_cuda_quant_cublas_min_batch(cc); ++ if (min_batch > 0 && ne11*ne12*ne13 >= min_batch) { ++ // GGML_CUDA_QUANT_CUBLAS_TYPE=f32 dequantizes to FP32 instead: slow, for accuracy oracles only. ++ static const bool f32_oracle = [] { ++ const char * e = getenv("GGML_CUDA_QUANT_CUBLAS_TYPE"); ++ return e != nullptr && (strcmp(e, "f32") == 0 || strcmp(e, "fp32") == 0); ++ }(); ++ if (f32_oracle) { ++ ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst); ++ } else { ++ ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst); ++ } ++ return; ++ } ++ } + if (ggml_cuda_should_use_mmq(src0->type, cc, ne11, /*n_experts =*/ 0)) { + ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst); + return; +@@ -5200,6 +5236,10 @@ struct ggml_backend_cuda_device_context { + std::string description; + std::string pci_bus_id; + int op_offload_min_batch_size; ++ // cudaDeviceProp::integrated, read once at registration. cudaGetDeviceProperties costs ~1-2 ms per call, and ++ // get_memory / get_type are polled per graph by callers that budget VRAM (stable-diffusion.cpp asks thousands ++ // of times per tiled VAE decode, ~5 s of host time per H3 video on a B200 host). ++ bool integrated = false; + }; + + static const char * ggml_backend_cuda_device_get_name(ggml_backend_dev_t dev) { +@@ -5303,12 +5343,9 @@ static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t * + // ref: https://github.com/ggml-org/llama.cpp/pull/17368 + #if defined(__linux__) + // Check if this is a UMA (Unified Memory Architecture) system +- cudaDeviceProp prop; +- CUDA_CHECK(cudaGetDeviceProperties(&prop, ggml_cuda_get_physical_device(ctx->device))); +- + // Check if UMA is explicitly enabled via environment variable + bool uma_env = getenv("GGML_CUDA_ENABLE_UNIFIED_MEMORY") != nullptr; +- bool is_uma = prop.integrated > 0 || uma_env; ++ bool is_uma = ctx->integrated || uma_env; + + if (is_uma) { + // For UMA systems (like DGX Spark), use system memory info +@@ -5332,10 +5369,7 @@ static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t * + static enum ggml_backend_dev_type ggml_backend_cuda_device_get_type(ggml_backend_dev_t dev) { + ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *) dev->context; + +- cudaDeviceProp prop; +- CUDA_CHECK(cudaGetDeviceProperties(&prop, ggml_cuda_get_physical_device(ctx->device))); +- +- return prop.integrated ++ return ctx->integrated + ? GGML_BACKEND_DEVICE_TYPE_IGPU + : GGML_BACKEND_DEVICE_TYPE_GPU; + } +@@ -6095,6 +6129,11 @@ ggml_backend_reg_t ggml_backend_cuda_reg() { + dev_ctx->device = i; + dev_ctx->name = GGML_CUDA_NAME + std::to_string(i); + dev_ctx->description = ggml_cuda_device_description(i); ++ { ++ cudaDeviceProp prop; ++ CUDA_CHECK(cudaGetDeviceProperties(&prop, physical_id)); ++ dev_ctx->integrated = prop.integrated > 0; ++ } + + char pci_bus_id[32] = {}; + CUDA_CHECK(cudaDeviceGetPCIBusId(pci_bus_id, sizeof(pci_bus_id), physical_id)); +-- +2.43.0 + + +From 0bf02a4e324ef837685c63ebaeddf2b4a4c0f317 Mon Sep 17 00:00:00 2001 +From: Daniel Han +Date: Wed, 30 Sep 2026 22:54:01 -0700 +Subject: [PATCH 2/2] ggml-cuda: vectorized row copy for strided F32->F32/F16 + copies with contiguous rows + +cpy_scalar recomputes the 4-D source and destination index (64-bit div/mod) for +every element, so the permuted ggml_cont copies of a DiT's Q/K/V (rows of +head_dim floats) run at a few percent of memory bandwidth: a 411 MB +[128,42,19108] permute copy takes 1.67 ms on a B200, 0.30 ms as one float4 per +thread with the index math done once per row. Pure copy / same per-element +conversion, so the output is bit-identical; GGML_CUDA_CPY_ROWS=0 restores +cpy_scalar. +--- + src/ggml-cuda/cpy.cu | 83 ++++++++++++++++++++++++++++++++++++++++++++ + 1 file changed, 83 insertions(+) + +diff --git a/src/ggml-cuda/cpy.cu b/src/ggml-cuda/cpy.cu +index 4c6fb4a2..c26a081c 100644 +--- a/src/ggml-cuda/cpy.cu ++++ b/src/ggml-cuda/cpy.cu +@@ -1,4 +1,5 @@ + #include "cpy.cuh" ++#include + #include "dequantize.cuh" + #include "cpy-utils.cuh" + #include "fp8.cuh" +@@ -203,6 +204,84 @@ cudaStream_t stream) { + ggml_cuda_kernel_launch(cpy_scalar_contiguous, launch_params, cx, cdst, ne); + } + ++// Strided copy whose rows (dim 0) are contiguous on both sides and of equal length: one float4 per thread, the ++// 4-D index math done once per row instead of per element. cpy_scalar spends most of its time in the 64-bit ++// div/mod of that per-element math; the permuted ggml_cont copies of a DiT's Q/K/V (rows of head_dim floats) run ++// at a few percent of memory bandwidth through it. A pure copy (or the same per-element conversion), so the ++// result is bit-identical to cpy_scalar. GGML_CUDA_CPY_ROWS=0 disables it. ++template ++static __global__ void cpy_rows_vec4(const char * cx, char * cdst, const int64_t nrows, const int64_t ne00_4, ++ const int64_t ne01, const int64_t ne02, const int64_t nb01, const int64_t nb02, const int64_t nb03, ++ const int64_t ne11, const int64_t ne12, const int64_t nb11, const int64_t nb12, const int64_t nb13) { ++ ggml_cuda_pdl_lc(); ++ const int64_t row = (int64_t) blockIdx.x*blockDim.y + threadIdx.y; ++ if (row >= nrows) { ++ return; ++ } ++ const int64_t i01 = row % ne01; ++ const int64_t i02 = (row / ne01) % ne02; ++ const int64_t i03 = row / (ne01*ne02); ++ const int64_t i11 = row % ne11; ++ const int64_t i12 = (row / ne11) % ne12; ++ const int64_t i13 = row / (ne11*ne12); ++ const float4 * src = (const float4 *) (cx + i01*nb01 + i02*nb02 + i03*nb03); ++ char * dst_row = cdst + i11*nb11 + i12*nb12 + i13*nb13; ++ ggml_cuda_pdl_sync(); ++ for (int64_t c = threadIdx.x; c < ne00_4; c += blockDim.x) { ++ const float4 v = src[c]; ++ if constexpr (std::is_same_v) { ++ ((float4 *) dst_row)[c] = v; ++ } else { ++ dst_t * d = (dst_t *) dst_row + 4*c; ++ cpy_1_scalar((const char *) &v.x, (char *) (d + 0)); ++ cpy_1_scalar((const char *) &v.y, (char *) (d + 1)); ++ cpy_1_scalar((const char *) &v.z, (char *) (d + 2)); ++ cpy_1_scalar((const char *) &v.w, (char *) (d + 3)); ++ } ++ } ++} ++ ++static bool ggml_cuda_cpy_rows_enabled() { ++ static const bool enabled = [] { ++ const char * e = getenv("GGML_CUDA_CPY_ROWS"); ++ return e == nullptr || atoi(e) != 0; ++ }(); ++ return enabled; ++} ++ ++// True (and launched) when src0 -> src1 is a row copy cpy_rows_vec4 can do; false leaves the copy to the caller. ++template ++static bool ggml_cpy_rows_vec4_cuda(const ggml_tensor * src0, ggml_tensor * src1, cudaStream_t stream) { ++ if (!ggml_cuda_cpy_rows_enabled() || src0->type != GGML_TYPE_F32) { ++ return false; ++ } ++ const int64_t ne00 = src0->ne[0]; ++ if (src0->nb[0] != (int64_t) sizeof(float) || src1->nb[0] != (int64_t) sizeof(dst_t) || src1->ne[0] != ne00 || ne00 % 4 != 0) { ++ return false; ++ } ++ const size_t dst_align = std::is_same_v ? 16 : 2*sizeof(dst_t); ++ if (((uintptr_t) src0->data | src0->nb[1] | src0->nb[2] | src0->nb[3]) % 16 != 0 || ++ ((uintptr_t) src1->data | src1->nb[1] | src1->nb[2] | src1->nb[3]) % dst_align != 0) { ++ return false; ++ } ++ const int64_t nrows = ggml_nrows(src0); ++ if (nrows != ggml_nrows(src1) || src0->ne[1] == 0 || src0->ne[2] == 0 || src1->ne[1] == 0 || src1->ne[2] == 0) { ++ return false; ++ } ++ const int64_t ne00_4 = ne00 / 4; ++ const int threads_x = ne00_4 >= 32 ? 32 : (int) ne00_4; ++ const int threads_y = 256 / threads_x; ++ const int64_t blocks = (nrows + threads_y - 1) / threads_y; ++ if (blocks > INT_MAX) { ++ return false; ++ } ++ const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(dim3((unsigned) blocks), dim3(threads_x, threads_y), 0, stream); ++ ggml_cuda_kernel_launch(cpy_rows_vec4, launch_params, (const char *) src0->data, (char *) src1->data, nrows, ne00_4, ++ src0->ne[1], src0->ne[2], (int64_t) src0->nb[1], (int64_t) src0->nb[2], (int64_t) src0->nb[3], ++ src1->ne[1], src1->ne[2], (int64_t) src1->nb[1], (int64_t) src1->nb[2], (int64_t) src1->nb[3]); ++ return true; ++} ++ + template + static void ggml_cpy_scalar_cuda( + const char * cx, char * cdst, const int64_t ne, +@@ -509,6 +588,10 @@ void ggml_cuda_cpy(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, gg + ggml_cpy_scalar_cuda + (src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream); + } ++ } else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && ggml_cpy_rows_vec4_cuda(src0, src1, main_stream)) { ++ // row copy done ++ } else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F16 && !contiguous_srcs && ggml_cpy_rows_vec4_cuda(src0, src1, main_stream)) { ++ // row copy with conversion done + } else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32) { + if (can_be_transposed) { + ggml_cpy_scalar_cuda +-- +2.43.0 + diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index e41c1151a8..eba4cca4da 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -1177,6 +1178,18 @@ bool StableDiffusionGGML::validate_and_load_runners() { high_noise_diffusion_model->set_flash_attention_enabled(true); } } + // The MiniMax-H3 video VAE decoder is a 36-layer ViT (32 x 64 heads, ~1.5k tokens per 16x16 latent tile). + // Without flash attention every tile and layer materialises a 32 x L x L f32 score matrix (mul_mat + scale + + // softmax + mul_mat), which is most of the decode's GPU time. --diffusion-fa therefore also covers it, the + // way --fa would, without switching the text encoder's attention. SD_H3_VAE_FLASH_ATTN=0 restores the old path. + if (!sd_ctx_params->flash_attn && sd_ctx_params->diffusion_flash_attn && first_stage_model && + sd_version_is_minimax_h3(version)) { + const char* env = getenv("SD_H3_VAE_FLASH_ATTN"); + if (env == nullptr || strcmp(env, "0") != 0) { + LOG_INFO("Using flash attention in the MiniMax-H3 video VAE"); + first_stage_model->set_flash_attention_enabled(true); + } + } if (sd_ctx_params->sage_attn && !set_sage_attention_enabled(true)) { return false; } From b01506b950a67ec841d4eb02a0d49c170def2b48 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 00:07:32 -0700 Subject: [PATCH 02/30] tiling: walk tile planes contiguously when splitting and merging process_tiles_2d copied and blended each pixel across all planes in the inner loop, striding a whole output plane (960x544 floats for H3) per element. A video tile has frames x channels planes, so every element was a cache and TLB miss: on MiniMax-H3 960x544x124 (105 tiles) this was about 10 s of single-threaded host time per decode. Loop planes outermost and precompute the per-column and per-row blend weights; each element still computes old + new * wy * wx in the same order, so the output is bitwise identical. --- src/runtime/tiling.cpp | 77 ++++++++++++++++++++++++++++-------------- 1 file changed, 51 insertions(+), 26 deletions(-) diff --git a/src/runtime/tiling.cpp b/src/runtime/tiling.cpp index f0a9dbc73b..767e9c4d19 100644 --- a/src/runtime/tiling.cpp +++ b/src/runtime/tiling.cpp @@ -74,12 +74,20 @@ static sd::Tensor sd_tensor_split_2d(const sd::Tensor& input, int int64_t input_plane = sd_tensor_plane_size(input); int64_t output_plane = sd_tensor_plane_size(output); int64_t plane_count = input.numel() / input_plane; - for (int iy = 0; iy < height; iy++) { - for (int ix = 0; ix < width; ix++) { - int64_t src_xy = (ix + x) % input_width + input_width * ((iy + y) % input_height); - int64_t dst_xy = ix + width * iy; - for (int64_t plane = 0; plane < plane_count; ++plane) { - output[plane * output_plane + dst_xy] = input[plane * input_plane + src_xy]; + // Plane-outer order: a plane is contiguous, so the inner loop streams instead of striding a + // whole plane per element (a video tile has frames x channels planes). + std::vector src_x(static_cast(width)); + for (int ix = 0; ix < width; ix++) { + src_x[ix] = (ix + x) % input_width; + } + 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 (int iy = 0; iy < height; iy++) { + const float* src_row = src + input_width * ((iy + y) % input_height); + float* dst_row = dst + static_cast(width) * iy; + for (int ix = 0; ix < width; ix++) { + dst_row[ix] = src_row[src_x[ix]]; } } } @@ -112,26 +120,43 @@ static void sd_tensor_merge_2d(const sd::Tensor& input, return x * x * x * (x * (6.0f * x - 15.0f) + 10.0f); }; - for (int iy = y_skip; iy < height; iy++) { - for (int ix = x_skip; ix < width; ix++) { - int64_t src_xy = ix + width * iy; - int64_t ox = (x + ix) % img_width; - int64_t oy = (y + iy) % img_height; - int64_t dst_xy = ox + img_width * oy; - for (int64_t plane = 0; plane < plane_count; ++plane) { - float new_value = input[plane * input_plane + src_xy]; - if (overlap_x > 0 || overlap_y > 0) { - float old_value = (*output)[plane * output_plane + dst_xy]; - const float x_f_0 = (circular_x || (overlap_x > 0 && x > 0)) ? (ix - x_skip) / float(overlap_x) : 1.f; - const float x_f_1 = (circular_x || (overlap_x > 0 && x < (img_width - width))) ? (width - ix) / float(overlap_x) : 1.f; - const float y_f_0 = (circular_y || (overlap_y > 0 && y > 0)) ? (iy - y_skip) / float(overlap_y) : 1.f; - const float y_f_1 = (circular_y || (overlap_y > 0 && y < (img_height - height))) ? (height - iy) / float(overlap_y) : 1.f; - const float x_f = std::min(std::min(x_f_0, x_f_1), 1.f); - const float y_f = std::min(std::min(y_f_0, y_f_1), 1.f); - (*output)[plane * output_plane + dst_xy] = - old_value + new_value * smootherstep_f32(y_f) * smootherstep_f32(x_f); - } else { - (*output)[plane * output_plane + dst_xy] = new_value; + // Weights depend on ix or iy only; precompute them and walk each contiguous plane row by row. + // Same arithmetic per element as the per-pixel form, so the result is bit-identical. + const bool blend = overlap_x > 0 || overlap_y > 0; + std::vector wx, wy; + std::vector dst_x(static_cast(width)); + for (int64_t ix = x_skip; ix < width; ix++) { + dst_x[ix] = (x + ix) % img_width; + } + if (blend) { + wx.resize(static_cast(width)); + wy.resize(static_cast(height)); + for (int64_t ix = x_skip; ix < width; ix++) { + const float x_f_0 = (circular_x || (overlap_x > 0 && x > 0)) ? (ix - x_skip) / float(overlap_x) : 1.f; + const float x_f_1 = (circular_x || (overlap_x > 0 && x < (img_width - width))) ? (width - ix) / float(overlap_x) : 1.f; + wx[ix] = smootherstep_f32(std::min(std::min(x_f_0, x_f_1), 1.f)); + } + for (int64_t iy = y_skip; iy < height; iy++) { + const float y_f_0 = (circular_y || (overlap_y > 0 && y > 0)) ? (iy - y_skip) / float(overlap_y) : 1.f; + const float y_f_1 = (circular_y || (overlap_y > 0 && y < (img_height - height))) ? (height - iy) / float(overlap_y) : 1.f; + 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]; } } } From ac399008a8c897de22c585b2110d81ad698bcc4c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 06:14:05 -0700 Subject: [PATCH 03/30] 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 0000000000..5157cba948 --- /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 04/30] 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 5157cba948..564831ee70 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 05/30] 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 0417d0011a..abbd675edd 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 06/30] 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 767e9c4d19..1bfb291a70 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 a832a644e3..352c4abb23 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 07/30] 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 abbd675edd..0e16f8ae8d 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 777566438c..1ae87ec87d 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 08/30] 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 0000000000..1e14c53584 --- /dev/null +++ b/scripts/unsloth/ggml-patches/0003-ggml-cuda-h3-vae-epilogue-fusion.patch @@ -0,0 +1,348 @@ +From 76b49e564c23b99758642d5bc5013d09774a1e24 Mon Sep 17 00:00:00 2001 +From: Daniel Han +Date: Thu, 1 Oct 2026 05:46:18 -0700 +Subject: [PATCH] ggml-cuda: fuse bias epilogues into F16 cuBLAS matmuls, + narrow blocks for short-row RMS norm + +An F16 matmul on the cuBLAS path writes an F16 temporary, converts it to F32, and the +graph then adds a bias (and in transformer blocks multiplies by a per-channel scale and +adds the residual, or applies SWIGLU) in separate passes over the F32 result. When those +nodes directly follow the matmul and have no other users, they now run inside the F16 -> F32 +conversion: MUL_MAT+ADD, MUL_MAT+ADD+MUL+ADD and MUL_MAT+ADD+GLU(SWIGLU). Every element +goes through the same rounded operations in the same order (explicit _rn intrinsics stop +nvcc from contracting them into FMAs), so results are bit-identical. +GGML_CUDA_CUBLAS_EPILOGUE_FUSION=0 disables it. + +RMS norm over rows of at most 128 columns (per-head q/k norms) launched 256-thread blocks +in which most threads had no element. It now uses the next multiple of 32 threads: each +thread still owns the same element and each warp the same partial sum, and the wider block +only added zero partials, so the result is unchanged. GGML_CUDA_NORM_SMALL_ROWS=0 restores +the 256-thread launch. +--- + src/ggml-cuda/ggml-cuda.cu | 200 ++++++++++++++++++++++++++++++++++++- + src/ggml-cuda/norm.cu | 63 ++++++++++++ + 2 files changed, 262 insertions(+), 1 deletion(-) + +diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu +index 517d089..0e61996 100644 +--- a/src/ggml-cuda/ggml-cuda.cu ++++ b/src/ggml-cuda/ggml-cuda.cu +@@ -2017,8 +2017,73 @@ struct batched_mul_mat_traits { + static inline auto convert_nc(ggml_type src_type) { return ggml_get_to_fp16_nc_cuda(src_type); } + }; + ++// Elementwise ops that directly follow an F16 cuBLAS matmul, applied while converting its F16 ++// output to F32 instead of in separate passes over the F32 result. Each element sees the same ++// rounded operations in the same order as the unfused ADD / MUL / ADD or ADD / SWIGLU nodes ++// (explicit _rn intrinsics keep nvcc from contracting them into FMAs), so results are bit-identical. ++enum ggml_cuda_cublas_epilogue_mode { ++ GGML_CUDA_CUBLAS_EPILOGUE_BIAS, // dst = mm + bias ++ GGML_CUDA_CUBLAS_EPILOGUE_BIAS_SCALE_RESIDUAL, // dst = residual + (mm + bias) * scale ++ GGML_CUDA_CUBLAS_EPILOGUE_BIAS_SWIGLU, // dst = silu(mm[:n] + bias[:n]) * (mm[n:] + bias[n:]) ++}; ++ ++struct ggml_cuda_cublas_epilogue { ++ ggml_cuda_cublas_epilogue_mode mode; ++ const float * bias; ++ const float * scale; ++ const float * residual; ++ float * dst; ++ int64_t ne0_out; ++ int64_t n_out; ++}; ++ ++template ++static __global__ void k_cublas_f16_epilogue(const half * __restrict__ mm, const float * __restrict__ bias, ++ const float * __restrict__ scale, const float * residual, float * dst, ++ const int64_t ne0_out, const int64_t n_out) { ++ const int64_t i = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; ++ if (i >= n_out) { ++ return; ++ } ++ const int64_t row = i / ne0_out; ++ const int64_t col = i - row*ne0_out; ++ if constexpr (mode == GGML_CUDA_CUBLAS_EPILOGUE_BIAS_SWIGLU) { ++ const int64_t j = row*2*ne0_out + col; ++ const float x = __fadd_rn(__half2float(mm[j]), bias[col]); ++ const float g = __fadd_rn(__half2float(mm[j + ne0_out]), bias[col + ne0_out]); ++ dst[i] = __fmul_rn(ggml_cuda_op_silu_single(x), g); ++ } else { ++ float v = __fadd_rn(__half2float(mm[i]), bias[col]); ++ if constexpr (mode == GGML_CUDA_CUBLAS_EPILOGUE_BIAS_SCALE_RESIDUAL) { ++ v = __fadd_rn(residual[i], __fmul_rn(v, scale[col])); ++ } ++ dst[i] = v; ++ } ++} ++ ++static void ggml_cuda_cublas_f16_epilogue(const half * mm, const ggml_cuda_cublas_epilogue & epi, cudaStream_t stream) { ++ const int block = 256; ++ const int64_t grid = (epi.n_out + block - 1) / block; ++ switch (epi.mode) { ++ case GGML_CUDA_CUBLAS_EPILOGUE_BIAS: ++ k_cublas_f16_epilogue<<>>( ++ mm, epi.bias, epi.scale, epi.residual, epi.dst, epi.ne0_out, epi.n_out); ++ break; ++ case GGML_CUDA_CUBLAS_EPILOGUE_BIAS_SCALE_RESIDUAL: ++ k_cublas_f16_epilogue<<>>( ++ mm, epi.bias, epi.scale, epi.residual, epi.dst, epi.ne0_out, epi.n_out); ++ break; ++ case GGML_CUDA_CUBLAS_EPILOGUE_BIAS_SWIGLU: ++ k_cublas_f16_epilogue<<>>( ++ mm, epi.bias, epi.scale, epi.residual, epi.dst, epi.ne0_out, epi.n_out); ++ break; ++ } ++ CUDA_CHECK(cudaGetLastError()); ++} ++ + template +-static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { ++static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, ++ const ggml_cuda_cublas_epilogue * epilogue = nullptr) { + using traits = batched_mul_mat_traits; + using cuda_t = typename traits::cuda_type; + +@@ -2226,8 +2291,17 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const + + // Convert output back to F32 if needed + if (cu_data_type != CUDA_R_32F) { ++ if constexpr (compute_type == GGML_TYPE_F16) { ++ if (epilogue != nullptr) { ++ ggml_cuda_cublas_f16_epilogue((const half *) dst_temp.get(), *epilogue, main_stream); ++ return; ++ } ++ } ++ GGML_ASSERT(epilogue == nullptr); + const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(traits::ggml_type_val); + to_fp32_cuda(dst_temp.get(), dst_ddf, ne_dst, main_stream); ++ } else { ++ GGML_ASSERT(epilogue == nullptr); + } + } + +@@ -3838,6 +3912,123 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, + } + + // try and fuse nodes and return the number of nodes to skip ++static bool ggml_cuda_cublas_epilogue_enabled() { ++ static const bool enabled = [] { ++ const char * e = getenv("GGML_CUDA_CUBLAS_EPILOGUE_FUSION"); ++ return e == nullptr || atoi(e) != 0; ++ }(); ++ return enabled; ++} ++ ++// True when ggml_cuda_mul_mat sends this F16 matmul to cuBLAS with F16 compute and an F16 ++// temporary output (the path the epilogue fusion replaces the F16 -> F32 conversion of). ++static bool ggml_cuda_mul_mat_uses_cublas_f16_temp(const ggml_backend_cuda_context & ctx, const ggml_tensor * mm) { ++ const ggml_tensor * src0 = mm->src[0]; ++ const ggml_tensor * src1 = mm->src[1]; ++ if (src0->type != GGML_TYPE_F16 || src1->type != GGML_TYPE_F32 || mm->type != GGML_TYPE_F32 || ++ !ggml_is_contiguous(mm) || ggml_get_op_params_i32(mm, 1) == GGML_HINT_SRC0_IS_HADAMARD || ++ mm->op_params[0] == GGML_PREC_F32 || getenv("GGML_CUDA_CUBLAS_COMPUTE_TYPE") != nullptr) { ++ return false; ++ } ++ const int cc = ggml_cuda_info().devices[ctx.device].cc; ++ const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size; ++ if (!fast_fp16_hardware_available(cc) || cc == GGML_CUDA_CC_VOLTA || GGML_CUDA_CC_IS_RDNA4(cc) || GGML_CUDA_CC_IS_CDNA(cc)) { ++ return false; ++ } ++ const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE ++ && ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && src0->view_src; ++ if (bad_padding_clear) { ++ return true; ++ } ++ const int64_t ne11 = src1->ne[1]; ++ return !ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, ne11) && ++ !ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, ne11, /*mul_mat_id =*/ false); ++} ++ ++static const ggml_tensor * ggml_cuda_other_src(const ggml_tensor * node, const ggml_tensor * operand) { ++ if (node->src[0] == operand) { ++ return node->src[1]; ++ } ++ if (node->src[1] == operand) { ++ return node->src[0]; ++ } ++ return nullptr; ++} ++ ++// Like ggml_node_has_n_uses(cgraph, idx, 1), but also accepts an in-place node whose view ++// source is the (single-use) node before it, e.g. the bias ADD that ggml_add_inplace builds. ++static bool ggml_cuda_single_use_intermediate(const ggml_cgraph * cgraph, int idx, const ggml_tensor * prev) { ++ const ggml_tensor * node = cgraph->nodes[idx]; ++ if (ggml_node_get_use_count(cgraph, idx) != 1 || (node->flags & GGML_TENSOR_FLAG_OUTPUT)) { ++ return false; ++ } ++ return node->view_src == nullptr || node->view_src == prev; ++} ++ ++static bool ggml_cuda_is_row_vector(const ggml_tensor * t, int64_t ne0) { ++ return t != nullptr && t->type == GGML_TYPE_F32 && ggml_is_contiguous(t) && t->ne[0] == ne0 && ggml_nelements(t) == ne0; ++} ++ ++// MUL_MAT (F16 weights, cuBLAS) followed by a row bias ADD, optionally then MUL by a row scale and ++// ADD of a residual, or SWIGLU. Returns the number of extra nodes consumed, 0 if not fused. ++static int ggml_cuda_try_fuse_cublas_epilogue(ggml_backend_cuda_context & ctx, ggml_cgraph * cgraph, int i) { ++ if (!ggml_cuda_cublas_epilogue_enabled() || i + 1 >= cgraph->n_nodes) { ++ return 0; ++ } ++ ggml_tensor * mm = cgraph->nodes[i]; ++ ggml_tensor * add = cgraph->nodes[i + 1]; ++ if (mm->op != GGML_OP_MUL_MAT || add->op != GGML_OP_ADD || (add->flags & GGML_TENSOR_FLAG_COMPUTE) == 0 || ++ !ggml_node_has_n_uses(cgraph, i, 1) || !ggml_are_same_shape(mm, add) || add->type != GGML_TYPE_F32 || ++ !ggml_is_contiguous(add)) { ++ return 0; ++ } ++ const ggml_tensor * bias = ggml_cuda_other_src(add, mm); ++ if (!ggml_cuda_is_row_vector(bias, mm->ne[0]) || !ggml_cuda_mul_mat_uses_cublas_f16_temp(ctx, mm)) { ++ return 0; ++ } ++ ++ ggml_cuda_cublas_epilogue epi = {}; ++ epi.mode = GGML_CUDA_CUBLAS_EPILOGUE_BIAS; ++ epi.bias = (const float *) bias->data; ++ epi.dst = (float *) add->data; ++ epi.ne0_out = add->ne[0]; ++ epi.n_out = ggml_nelements(add); ++ int consumed = 1; ++ ++ if (i + 2 < cgraph->n_nodes && ggml_cuda_single_use_intermediate(cgraph, i + 1, mm)) { ++ ggml_tensor * mul = cgraph->nodes[i + 2]; ++ ggml_tensor * res = i + 3 < cgraph->n_nodes ? cgraph->nodes[i + 3] : nullptr; ++ const ggml_tensor * glu = cgraph->nodes[i + 2]; ++ if (res != nullptr && mul->op == GGML_OP_MUL && res->op == GGML_OP_ADD && (mul->flags & GGML_TENSOR_FLAG_COMPUTE) && ++ (res->flags & GGML_TENSOR_FLAG_COMPUTE) && ggml_node_has_n_uses(cgraph, i + 2, 1) && ++ ggml_are_same_shape(mul, add) && ggml_are_same_shape(res, add) && mul->type == GGML_TYPE_F32 && ++ res->type == GGML_TYPE_F32 && ggml_is_contiguous(res)) { ++ const ggml_tensor * scale = ggml_cuda_other_src(mul, add); ++ const ggml_tensor * residual = ggml_cuda_other_src(res, mul); ++ if (mul->src[0] == add && ggml_cuda_is_row_vector(scale, add->ne[0]) && residual != nullptr && ++ residual->type == GGML_TYPE_F32 && ggml_is_contiguous(residual) && ggml_are_same_shape(residual, res)) { ++ epi.mode = GGML_CUDA_CUBLAS_EPILOGUE_BIAS_SCALE_RESIDUAL; ++ epi.scale = (const float *) scale->data; ++ epi.residual = (const float *) residual->data; ++ epi.dst = (float *) res->data; ++ consumed = 3; ++ } ++ } else if (glu->op == GGML_OP_GLU && (glu->flags & GGML_TENSOR_FLAG_COMPUTE) && glu->src[0] == add && ++ glu->src[1] == nullptr && ggml_get_glu_op(glu) == GGML_GLU_OP_SWIGLU && ++ ggml_get_op_params_i32(glu, 1) == 0 && glu->type == GGML_TYPE_F32 && ggml_is_contiguous(glu) && ++ add->ne[0] % 2 == 0 && glu->ne[0] * 2 == add->ne[0] && ggml_nrows(glu) == ggml_nrows(add)) { ++ epi.mode = GGML_CUDA_CUBLAS_EPILOGUE_BIAS_SWIGLU; ++ epi.dst = (float *) glu->data; ++ epi.ne0_out = glu->ne[0]; ++ epi.n_out = ggml_nelements(glu); ++ consumed = 2; ++ } ++ } ++ ++ ggml_cuda_mul_mat_cublas_impl(ctx, mm->src[0], mm->src[1], mm, &epi); ++ return consumed; ++} ++ + static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { + + static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); +@@ -4523,6 +4714,13 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph + return fused_node_count - 1; + } + ++ if (node->op == GGML_OP_MUL_MAT) { ++ const int consumed = ggml_cuda_try_fuse_cublas_epilogue(*cuda_ctx, cgraph, i); ++ if (consumed > 0) { ++ return consumed; ++ } ++ } ++ + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD }, {})) { + ggml_cuda_op_rms_norm_fused_add(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]); + return 2; +diff --git a/src/ggml-cuda/norm.cu b/src/ggml-cuda/norm.cu +index c3758cd..9c9385f 100644 +--- a/src/ggml-cuda/norm.cu ++++ b/src/ggml-cuda/norm.cu +@@ -301,10 +301,68 @@ static void group_norm_f32_cuda( + } + } + ++// Rows of at most 128 columns (per-head q/k norms) leave most of a 256-thread block idle. A block ++// of the next multiple of 32 gives every thread the same element and every warp the same partial ++// sum; the 256-thread reduction only adds zero partials on top, so the result is bit-identical. ++// GGML_CUDA_NORM_SMALL_ROWS=0 restores the 256-thread launch. ++static bool rms_norm_small_rows_enabled() { ++#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) ++ return false; ++#else ++ static const bool enabled = [] { ++ const char * e = getenv("GGML_CUDA_NORM_SMALL_ROWS"); ++ return e == nullptr || atoi(e) != 0; ++ }(); ++ return enabled && WARP_SIZE == 32; ++#endif ++} ++ ++template ++static void rms_norm_f32_launch_small( ++ const float * x, float * dst, const int ncols, const dim3 blocks_num, ++ const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, ++ const float * mul, const int64_t mul_stride_row, const int64_t mul_stride_channel, const int64_t mul_stride_sample, ++ const uint3 mul_ncols_packed, const uint3 mul_nrows_packed, const uint3 mul_nchannels_packed, const uint3 mul_nsamples_packed, ++ cudaStream_t stream) { ++ const dim3 block_dims(block_size, 1, 1); ++ const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, block_size > WARP_SIZE ? 32 * sizeof(float) : 0, stream}; ++ ggml_cuda_kernel_launch(rms_norm_f32, launch_params, ++ x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, ++ mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, ++ nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); ++} ++ ++template ++static bool rms_norm_f32_try_small( ++ const float * x, float * dst, const int ncols, const dim3 blocks_num, ++ const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, ++ const float * mul, const int64_t mul_stride_row, const int64_t mul_stride_channel, const int64_t mul_stride_sample, ++ const uint3 mul_ncols_packed, const uint3 mul_nrows_packed, const uint3 mul_nchannels_packed, const uint3 mul_nsamples_packed, ++ cudaStream_t stream) { ++ if (ncols > 128 || !rms_norm_small_rows_enabled()) { ++ return false; ++ } ++ if (ncols <= 32) { ++ rms_norm_f32_launch_small<32, do_multiply>(x, dst, ncols, blocks_num, stride_row, stride_channel, stride_sample, eps, mul, ++ mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, stream); ++ } else if (ncols <= 64) { ++ rms_norm_f32_launch_small<64, do_multiply>(x, dst, ncols, blocks_num, stride_row, stride_channel, stride_sample, eps, mul, ++ mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, stream); ++ } else { ++ rms_norm_f32_launch_small<128, do_multiply>(x, dst, ncols, blocks_num, stride_row, stride_channel, stride_sample, eps, mul, ++ mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, stream); ++ } ++ return true; ++} ++ + static void rms_norm_f32_cuda( + const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples, + const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, cudaStream_t stream) { + const dim3 blocks_num(nrows, nchannels, nsamples); ++ if (rms_norm_f32_try_small(x, dst, ncols, blocks_num, stride_row, stride_channel, stride_sample, eps, ++ nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), stream)) { ++ return; ++ } + if (ncols < 1024) { + const dim3 block_dims(256, 1, 1); + const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; +@@ -360,6 +418,11 @@ static void rms_norm_mul_f32_cuda(const float * x, + const uint3 mul_nrows_packed = init_fastdiv_values(mul_nrows); + const uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels); + const uint3 mul_nsamples_packed = init_fastdiv_values(mul_nsamples); ++ if (rms_norm_f32_try_small(x, dst, ncols, blocks_num, stride_row, stride_channel, stride_sample, eps, ++ mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_ncols_packed, mul_nrows_packed, ++ mul_nchannels_packed, mul_nsamples_packed, stream)) { ++ return; ++ } + if (ncols < 1024) { + const dim3 block_dims(256, 1, 1); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; +-- +2.43.0 + From dbb4106b4bdc578497d61c77747f9f6b6e164a08 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 06:48:15 -0700 Subject: [PATCH 09/30] 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 74324169ee..58130582d4 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -207,15 +207,10 @@ ggml_tensor* ggml_ext_gelu_quick(ggml_context* ctx, return x; } -ggml_tensor* ggml_ext_linear(ggml_context* ctx, - ggml_tensor* x, - ggml_tensor* w, - ggml_tensor* b, - bool force_prec_f32, - float scale) { - if (scale != 1.f) { - x = ggml_ext_scale(ctx, x, scale); - } +ggml_tensor* ggml_ext_linear_matmul(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + bool force_prec_f32) { if (x->ne[2] * x->ne[3] > 1024) { // workaround: avoid ggml cuda error int64_t ne2 = x->ne[2]; @@ -232,6 +227,19 @@ ggml_tensor* ggml_ext_linear(ggml_context* ctx, ggml_mul_mat_set_prec(x, GGML_PREC_F32); } } + return x; +} + +ggml_tensor* ggml_ext_linear(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* b, + bool force_prec_f32, + float scale) { + if (scale != 1.f) { + x = ggml_ext_scale(ctx, x, scale); + } + x = ggml_ext_linear_matmul(ctx, x, w, force_prec_f32); if (scale != 1.f) { x = ggml_ext_scale(ctx, x, 1.f / scale); } @@ -796,6 +804,47 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, return kqv; } +ggml_tensor* ggml_ext_attention_prepared(ggml_context* ctx, + ggml_backend_t backend, + ggml_tensor* q, + ggml_tensor* k, + ggml_tensor* v, + int64_t n_head, + int64_t N, + float kv_scale) { + GGML_ASSERT(q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16); + const int64_t L_q = q->ne[1]; + const int64_t d_head = v->ne[0]; + if (backend == nullptr || d_head < 64) { + return nullptr; + } + + // Same expressions as ggml_ext_attention_ext so the softmax scale rounds identically. + float scale = (1.0f / sqrt((float)d_head)); + + auto kqv = ggml_flash_attn_ext(ctx, q, k, v, nullptr, scale / kv_scale, 0, 0); + if (!ggml_backend_supports_op(backend, kqv)) { + return nullptr; + } + ggml_flash_attn_ext_set_prec(kqv, GGML_PREC_F32); + if (kv_scale != 1.0f) { + kqv = ggml_ext_scale(ctx, kqv, 1.0f / kv_scale); + } + kqv = ggml_view_4d(ctx, + kqv, + d_head, + n_head, + L_q, + N, + kqv->nb[1], + kqv->nb[2], + kqv->nb[1] * n_head, + 0); + kqv = ggml_ext_cont(ctx, kqv); + kqv = ggml_reshape_3d(ctx, kqv, d_head * n_head, L_q, N); // [N, L_q, C] + return kqv; +} + ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx, ggml_tensor* x, ggml_tensor* w, diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h index 33d8c3bf21..a5076262ff 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -65,6 +65,12 @@ ggml_tensor* ggml_ext_gelu_quick(ggml_context* ctx, ggml_tensor* x, bool inplace = false); +// The matrix multiply of ggml_ext_linear without its input/output scaling and bias. +ggml_tensor* ggml_ext_linear_matmul(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + bool force_prec_f32); + ggml_tensor* ggml_ext_linear(ggml_context* ctx, ggml_tensor* x, ggml_tensor* w, @@ -223,6 +229,21 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, float kv_scale = 1.0f, bool sage_attn = false); +// Flash attention on inputs that are already in the head-major layout the kernel reads: +// q: F32 [d_head, L_q, n_head * N], k/v: F16 [d_head, L_k, n_kv_head * N] with the kv_scale +// already applied. Builds the same flash_attn_ext + output scaling as ggml_ext_attention_ext +// (skip_reshape, unmasked), so the result is identical; returns nullptr when the backend lacks +// flash attention for these inputs (the caller then uses ggml_ext_attention_ext). +// return: [N, L_q, n_head * d_head] +ggml_tensor* ggml_ext_attention_prepared(ggml_context* ctx, + ggml_backend_t backend, + ggml_tensor* q, + ggml_tensor* k, + ggml_tensor* v, + int64_t n_head, + int64_t N, + float kv_scale); + ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx, ggml_tensor* x, ggml_tensor* w, diff --git a/src/model/common/ggml_block.hpp b/src/model/common/ggml_block.hpp index 991b88f28b..b4c94b2d94 100644 --- a/src/model/common/ggml_block.hpp +++ b/src/model/common/ggml_block.hpp @@ -204,6 +204,30 @@ class Linear : public UnaryBlock { force_prec_f32 = force_prec_f32_; } + // For callers that fold this layer's input/output scaling into neighbouring fused ops: the + // weight when forward() is exactly ggml_ext_linear(x, weight, nullptr, prec_f32(), + // effective_scale()) (no bias, weight scale, adapter, INT8 or FP8 path), nullptr otherwise. + ggml_tensor* foldable_weight(GGMLRunnerContext* ctx) { + ggml_tensor* w = params["weight"]; + if (bias || has_weight_scale || ctx->weight_adapter || w->type == GGML_TYPE_I8) { + return nullptr; + } +#ifndef SD_USE_UPSTREAM_GGML + if (w->type == GGML_TYPE_F8_E4M3 || w->type == GGML_TYPE_F8_E5M2) { + return nullptr; + } +#endif + return w; + } + + float effective_scale(GGMLRunnerContext* ctx) const { + return ctx->linear_scale > 0.f ? ctx->linear_scale : scale; + } + + bool prec_f32() const { + return force_prec_f32; + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { ggml_tensor* w = params["weight"]; const float scale = ctx->linear_scale > 0.f ? ctx->linear_scale : this->scale; diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index 80c29ab528..3b75408aad 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 10/30] 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 0e16f8ae8d..c852492eae 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 0000000000..689dfa5f67 --- /dev/null +++ b/scripts/unsloth/ggml-patches/0004-ggml-h3-dit-fused-ops.patch @@ -0,0 +1,1020 @@ +From 04e9fe6c85b680c2f8e29c22e084dab6012babbe Mon Sep 17 00:00:00 2001 +From: Daniel Han +Date: Thu, 1 Oct 2026 06:48:01 -0700 +Subject: [PATCH] ggml : fused value-preserving ops for adaLN DiT blocks + +Three ops that replace long chains of elementwise ops, copies and concats in +diffusion transformers such as MiniMax-H3 while producing the same bits as the +chains they replace (every product, sum, scale and conversion is rounded +separately, never contracted into an FMA): + +- ggml_rope_pe_permute: rotary embedding with a precomputed [cos, -sin; sin, + cos] table on the first n_rot channels (non-interleaved pairing), pass-through + for the rest, written in the head-major [d, tokens, heads] layout attention + reads, with an optional scale and F16 cast. Reads a strided view, so the + q/k/v chunks of a fused qkv projection need no copies. +- ggml_modulate_rows / ggml_gated_add_rows: (x + x*scale) + shift and + base + u*gate on a row range, written in place into an accumulator so a + sequence split into segments with different modulation vectors needs no + slice copies and no concat. Optional output / input scale. +- ggml_swiglu_scaled: silu(s_in*a) * (s_in*b) scaled by s_out, i.e. a Linear + output scale, swiglu and the next Linear input scale in one pass. + +CPU and CUDA implementations; other backends report the ops unsupported so +callers keep their unfused graphs. test-backend-ops covers strided inputs, +partial and full rotation, F32/F16 outputs, chained segments and the scales. +--- + include/ggml.h | 62 ++++++++++++ + src/ggml-backend-meta.cpp | 3 +- + src/ggml-cpu/ggml-cpu.c | 15 +++ + src/ggml-cpu/ops.cpp | 154 +++++++++++++++++++++++++++++ + src/ggml-cpu/ops.h | 3 + + src/ggml-cuda/ggml-cuda.cu | 17 ++++ + src/ggml-cuda/modulate-rows.cu | 168 ++++++++++++++++++++++++++++++++ + src/ggml-cuda/modulate-rows.cuh | 7 ++ + src/ggml-cuda/rope-pe.cu | 102 +++++++++++++++++++ + src/ggml-cuda/rope-pe.cuh | 5 + + src/ggml.c | 129 +++++++++++++++++++++++- + tests/test-backend-ops.cpp | 139 ++++++++++++++++++++++++++ + 12 files changed, 801 insertions(+), 3 deletions(-) + create mode 100644 src/ggml-cuda/modulate-rows.cu + create mode 100644 src/ggml-cuda/modulate-rows.cuh + create mode 100644 src/ggml-cuda/rope-pe.cu + create mode 100644 src/ggml-cuda/rope-pe.cuh + +diff --git a/include/ggml.h b/include/ggml.h +index 991cc4fc..153eed2a 100644 +--- a/include/ggml.h ++++ b/include/ggml.h +@@ -594,6 +594,9 @@ extern "C" { + + GGML_OP_QUANTIZE_I8_CONVROT, + GGML_OP_SAGE_ATTN, ++ GGML_OP_ROPE_PE_PERMUTE, ++ GGML_OP_MODULATE_ROWS, ++ GGML_OP_SWIGLU_SCALED, + + GGML_OP_COUNT, + }; +@@ -2461,6 +2464,65 @@ extern "C" { + float scale, + enum ggml_sage_attn_mode mode); + ++ // Rotary embedding with a precomputed cos/sin table, fused with the head-major relayout ++ // attention kernels consume, an optional scale and an optional cast to F16. ++ // a: F32 [head_dim, n_head, n_tokens, n_batch], rows contiguous (other strides free) ++ // pe: F32 contiguous [2, 2, n_rot/2, n_tokens], pe[c, r, j, t] = M[r][c] of the 2x2 rotation ++ // for pair j of token t (non-interleaved pairing j <-> j + n_rot/2); NULL when n_rot == 0 ++ // result: `type` (F32 or F16) contiguous [head_dim, n_tokens, n_head, n_batch] ++ // i < n_rot: y = a[j]*pe[0, r, j, t] + a[j + n_rot/2]*pe[1, r, j, t] (r = i / (n_rot/2), j = i % (n_rot/2)) ++ // i >= n_rot: y = a[i] ++ // y = scale*y when scale != 1, then y is converted to `type`. ++ // Each product, the sum, the scale and the conversion are rounded separately, so the result ++ // matches the equivalent chain of ggml_mul/ggml_add/ggml_concat/ggml_scale/ggml_cpy bit for bit. ++ GGML_API struct ggml_tensor * ggml_rope_pe_permute( ++ struct ggml_context * ctx, ++ struct ggml_tensor * a, ++ struct ggml_tensor * pe, ++ int n_rot, ++ float scale, ++ enum ggml_type type); ++ ++ // Row-range modulation for adaLN-style blocks whose sequence is split into segments that use ++ // different modulation vectors. Writes rows [row0, row1) along dim 1 (all of dims 2 and 3) of the ++ // result and leaves every other row as it is in `acc`: chain one call per segment, passing the ++ // previous result as `acc` (written in place); the first call passes acc = NULL and gets a new ++ // tensor shaped like x / base. x, base, u: F32 with contiguous rows; scale, shift, gate: F32 ++ // with ne[0] == x->ne[0] elements, broadcast over rows. ++ // ggml_modulate_rows: y = (x + x*scale) + shift; then y = out_scale*y when out_scale != 1 ++ // ggml_gated_add_rows: w = u_scale*u when u_scale != 1 (else u); y = base + w*gate ++ // Each product, sum and scale is rounded separately, so the result equals the slice + ++ // ggml_mul/ggml_add/ggml_scale + ggml_concat chain bit for bit. ++ GGML_API struct ggml_tensor * ggml_modulate_rows( ++ struct ggml_context * ctx, ++ struct ggml_tensor * acc, ++ struct ggml_tensor * x, ++ struct ggml_tensor * scale, ++ struct ggml_tensor * shift, ++ int64_t row0, ++ int64_t row1, ++ float out_scale); ++ ++ GGML_API struct ggml_tensor * ggml_gated_add_rows( ++ struct ggml_context * ctx, ++ struct ggml_tensor * acc, ++ struct ggml_tensor * base, ++ struct ggml_tensor * u, ++ struct ggml_tensor * gate, ++ int64_t row0, ++ int64_t row1, ++ float u_scale); ++ ++ // swiglu on two separately scaled halves, then scaled: with a' = in_scale*a, b' = in_scale*b, ++ // y = out_scale*(silu(a')*b') (each scale skipped when it is 1). a, b: F32, same shape, rows ++ // contiguous. Matches ggml_scale -> ggml_swiglu_split -> ggml_scale bit for bit. ++ GGML_API struct ggml_tensor * ggml_swiglu_scaled( ++ struct ggml_context * ctx, ++ struct ggml_tensor * a, ++ struct ggml_tensor * b, ++ float in_scale, ++ float out_scale); ++ + GGML_API void ggml_flash_attn_ext_set_prec( + struct ggml_tensor * a, + enum ggml_prec prec); +diff --git a/src/ggml-backend-meta.cpp b/src/ggml-backend-meta.cpp +index 539b4f70..5a34e748 100644 +--- a/src/ggml-backend-meta.cpp ++++ b/src/ggml-backend-meta.cpp +@@ -151,7 +151,8 @@ static ggml_backend_buffer_type_t ggml_backend_meta_device_get_host_buffer_type( + + static bool ggml_backend_meta_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { + GGML_ASSERT(ggml_backend_dev_is_meta(dev)); +- if (op->op == GGML_OP_SAGE_ATTN) { ++ if (op->op == GGML_OP_SAGE_ATTN || op->op == GGML_OP_ROPE_PE_PERMUTE || ++ op->op == GGML_OP_MODULATE_ROWS || op->op == GGML_OP_SWIGLU_SCALED) { + return false; + } + const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; +diff --git a/src/ggml-cpu/ggml-cpu.c b/src/ggml-cpu/ggml-cpu.c +index 350777c1..d83f15f9 100644 +--- a/src/ggml-cpu/ggml-cpu.c ++++ b/src/ggml-cpu/ggml-cpu.c +@@ -2190,6 +2190,18 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm + { + ggml_compute_forward_glu(params, tensor); + } break; ++ case GGML_OP_ROPE_PE_PERMUTE: ++ { ++ ggml_compute_forward_rope_pe_permute(params, tensor); ++ } break; ++ case GGML_OP_MODULATE_ROWS: ++ { ++ ggml_compute_forward_modulate_rows(params, tensor); ++ } break; ++ case GGML_OP_SWIGLU_SCALED: ++ { ++ ggml_compute_forward_swiglu_scaled(params, tensor); ++ } break; + case GGML_OP_GET_REL_POS: + { + ggml_compute_forward_get_rel_pos(params, tensor); +@@ -2522,6 +2534,9 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { + case GGML_OP_ROPE: + case GGML_OP_ROPE_BACK: + case GGML_OP_ADD_REL_POS: ++ case GGML_OP_ROPE_PE_PERMUTE: ++ case GGML_OP_MODULATE_ROWS: ++ case GGML_OP_SWIGLU_SCALED: + { + n_tasks = n_threads; + } break; +diff --git a/src/ggml-cpu/ops.cpp b/src/ggml-cpu/ops.cpp +index 727a4058..40fe1949 100644 +--- a/src/ggml-cpu/ops.cpp ++++ b/src/ggml-cpu/ops.cpp +@@ -12,6 +12,7 @@ + #include + #include + #include ++#include + + struct ggml_fp8_e4m3_cpu { + uint8_t value; +@@ -12136,3 +12137,156 @@ void ggml_compute_forward_quantize_i8_convrot(const ggml_compute_params * params + memset(dst_data + rows * dst->nb[1], 0, (size_t)(rows_padded - rows) * dst->nb[1]); + } + } ++ ++// ggml_compute_forward_rope_pe_permute ++ ++void ggml_compute_forward_rope_pe_permute( ++ const ggml_compute_params * params, ++ ggml_tensor * dst) { ++ const ggml_tensor * src0 = dst->src[0]; ++ const ggml_tensor * pe = dst->src[1]; ++ ++ GGML_ASSERT(src0->type == GGML_TYPE_F32); ++ GGML_ASSERT(dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_F16); ++ ++ const int n_rot = ggml_get_op_params_i32(dst, 0); ++ const float scale = ggml_get_op_params_f32(dst, 1); ++ const int half = n_rot / 2; ++ ++ const int64_t D = src0->ne[0]; ++ const int64_t H = src0->ne[1]; ++ const int64_t L = src0->ne[2]; ++ const int64_t N = src0->ne[3]; ++ ++ const int64_t nrows = H * L * N; ++ const int64_t dr = (nrows + params->nth - 1) / params->nth; ++ const int64_t ir0 = dr * params->ith; ++ const int64_t ir1 = MIN(ir0 + dr, nrows); ++ ++ // Same rounding points as the unfused graph: each product, the sum, the scale and the ++ // conversion are separate passes, so no contraction into fused multiply-adds can happen. ++ std::vector p0(n_rot), p1(n_rot), y(D); ++ ++ for (int64_t ir = ir0; ir < ir1; ++ir) { ++ const int64_t h = ir % H; ++ const int64_t l = (ir / H) % L; ++ const int64_t n = ir / (H * L); ++ ++ const float * x = (const float *) ((const char *) src0->data + h * src0->nb[1] + l * src0->nb[2] + n * src0->nb[3]); ++ ++ if (n_rot > 0) { ++ const float * t = (const float *) pe->data + l * (int64_t) half * 4; ++ for (int i = 0; i < n_rot; ++i) { ++ const int r = i / half; ++ const int j = i - r * half; ++ p0[i] = x[j] * t[4 * j + 2 * r + 0]; ++ } ++ for (int i = 0; i < n_rot; ++i) { ++ const int r = i / half; ++ const int j = i - r * half; ++ p1[i] = x[j + half] * t[4 * j + 2 * r + 1]; ++ } ++ ggml_vec_add_f32(n_rot, y.data(), p0.data(), p1.data()); ++ } ++ if (D > n_rot) { ++ memcpy(y.data() + n_rot, x + n_rot, (D - n_rot) * sizeof(float)); ++ } ++ if (scale != 1.0f) { ++ ggml_vec_scale_f32((int) D, y.data(), scale); ++ } ++ ++ char * out = (char *) dst->data + l * dst->nb[1] + h * dst->nb[2] + n * dst->nb[3]; ++ if (dst->type == GGML_TYPE_F32) { ++ memcpy(out, y.data(), D * sizeof(float)); ++ } else { ++ ggml_cpu_fp32_to_fp16(y.data(), (ggml_fp16_t *) out, D); ++ } ++ } ++} ++ ++// ggml_compute_forward_modulate_rows ++ ++void ggml_compute_forward_modulate_rows( ++ const ggml_compute_params * params, ++ ggml_tensor * dst) { ++ const ggml_tensor * a = dst->src[0]; ++ const ggml_tensor * b = dst->src[1]; ++ const ggml_tensor * c = dst->src[2]; ++ ++ const int mode = ggml_get_op_params_i32(dst, 0); ++ const int64_t row0 = ggml_get_op_params_i32(dst, 1); ++ const int64_t row1 = ggml_get_op_params_i32(dst, 2); ++ const float s = ggml_get_op_params_f32(dst, 3); ++ ++ const int64_t nc = a->ne[0]; ++ const int64_t nr1 = row1 - row0; ++ const int64_t nrows = nr1 * a->ne[2] * a->ne[3]; ++ const int64_t dr = (nrows + params->nth - 1) / params->nth; ++ const int64_t ir0 = dr * params->ith; ++ const int64_t ir1 = MIN(ir0 + dr, nrows); ++ ++ // separate passes keep every product, sum and scale individually rounded, as in the unfused graph ++ std::vector t(nc), w(nc); ++ for (int64_t ir = ir0; ir < ir1; ++ir) { ++ const int64_t i1 = row0 + ir % nr1; ++ const int64_t i2 = (ir / nr1) % a->ne[2]; ++ const int64_t i3 = ir / (nr1 * a->ne[2]); ++ ++ const float * x = (const float *) ((const char *) a->data + i1 * a->nb[1] + i2 * a->nb[2] + i3 * a->nb[3]); ++ float * y = (float *) ((char *) dst->data + i1 * dst->nb[1] + i2 * dst->nb[2] + i3 * dst->nb[3]); ++ ++ if (mode == 0) { ++ ggml_vec_mul_f32((int) nc, t.data(), x, (const float *) b->data); ++ ggml_vec_add_f32((int) nc, w.data(), x, t.data()); ++ ggml_vec_add_f32((int) nc, y, w.data(), (const float *) c->data); ++ if (s != 1.0f) { ++ ggml_vec_scale_f32((int) nc, y, s); ++ } ++ } else { ++ const float * u = (const float *) ((const char *) b->data + i1 * b->nb[1] + i2 * b->nb[2] + i3 * b->nb[3]); ++ memcpy(w.data(), u, nc * sizeof(float)); ++ if (s != 1.0f) { ++ ggml_vec_scale_f32((int) nc, w.data(), s); ++ } ++ ggml_vec_mul_f32((int) nc, t.data(), w.data(), (const float *) c->data); ++ ggml_vec_add_f32((int) nc, y, x, t.data()); ++ } ++ } ++} ++ ++// ggml_compute_forward_swiglu_scaled ++ ++void ggml_compute_forward_swiglu_scaled( ++ const ggml_compute_params * params, ++ ggml_tensor * dst) { ++ const ggml_tensor * a = dst->src[0]; ++ const ggml_tensor * b = dst->src[1]; ++ ++ const float in_scale = ggml_get_op_params_f32(dst, 0); ++ const float out_scale = ggml_get_op_params_f32(dst, 1); ++ ++ const int64_t nc = a->ne[0]; ++ const int64_t nrows = ggml_nrows(a); ++ const int64_t dr = (nrows + params->nth - 1) / params->nth; ++ const int64_t ir0 = dr * params->ith; ++ const int64_t ir1 = MIN(ir0 + dr, nrows); ++ ++ std::vector xa(nc), xb(nc); ++ for (int64_t ir = ir0; ir < ir1; ++ir) { ++ const int64_t i1 = ir % a->ne[1]; ++ const int64_t i2 = (ir / a->ne[1]) % a->ne[2]; ++ const int64_t i3 = ir / (a->ne[1] * a->ne[2]); ++ ++ memcpy(xa.data(), (const char *) a->data + i1 * a->nb[1] + i2 * a->nb[2] + i3 * a->nb[3], nc * sizeof(float)); ++ memcpy(xb.data(), (const char *) b->data + i1 * b->nb[1] + i2 * b->nb[2] + i3 * b->nb[3], nc * sizeof(float)); ++ if (in_scale != 1.0f) { ++ ggml_vec_scale_f32((int) nc, xa.data(), in_scale); ++ ggml_vec_scale_f32((int) nc, xb.data(), in_scale); ++ } ++ float * y = (float *) ((char *) dst->data + i1 * dst->nb[1] + i2 * dst->nb[2] + i3 * dst->nb[3]); ++ ggml_vec_swiglu_f32((int) nc, y, xa.data(), xb.data()); ++ if (out_scale != 1.0f) { ++ ggml_vec_scale_f32((int) nc, y, out_scale); ++ } ++ } ++} +diff --git a/src/ggml-cpu/ops.h b/src/ggml-cpu/ops.h +index 3b579167..e09cc11d 100644 +--- a/src/ggml-cpu/ops.h ++++ b/src/ggml-cpu/ops.h +@@ -98,6 +98,9 @@ void ggml_compute_forward_win_part(const struct ggml_compute_params * params, st + void ggml_compute_forward_win_unpart(const struct ggml_compute_params * params, struct ggml_tensor * dst); + void ggml_compute_forward_unary(const struct ggml_compute_params * params, struct ggml_tensor * dst); + void ggml_compute_forward_glu(const struct ggml_compute_params * params, struct ggml_tensor * dst); ++void ggml_compute_forward_rope_pe_permute(const struct ggml_compute_params * params, struct ggml_tensor * dst); ++void ggml_compute_forward_modulate_rows(const struct ggml_compute_params * params, struct ggml_tensor * dst); ++void ggml_compute_forward_swiglu_scaled(const struct ggml_compute_params * params, struct ggml_tensor * dst); + void ggml_compute_forward_get_rel_pos(const struct ggml_compute_params * params, struct ggml_tensor * dst); + void ggml_compute_forward_add_rel_pos(const struct ggml_compute_params * params, struct ggml_tensor * dst); + void ggml_compute_forward_rwkv_wkv6(const struct ggml_compute_params * params, struct ggml_tensor * dst); +diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu +index 517d0894..2c0984a0 100644 +--- a/src/ggml-cuda/ggml-cuda.cu ++++ b/src/ggml-cuda/ggml-cuda.cu +@@ -44,6 +44,8 @@ + #include "ggml-cuda/rope.cuh" + #include "ggml-cuda/roll.cuh" + #include "ggml-cuda/scale.cuh" ++#include "ggml-cuda/rope-pe.cuh" ++#include "ggml-cuda/modulate-rows.cuh" + #include "ggml-cuda/snake.cuh" + #include "ggml-cuda/softcap.cuh" + #include "ggml-cuda/softmax.cuh" +@@ -3002,6 +3004,15 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg + case GGML_OP_SAGE_ATTN: + ggml_cuda_sage_attn(ctx, dst); + break; ++ case GGML_OP_ROPE_PE_PERMUTE: ++ ggml_cuda_op_rope_pe_permute(ctx, dst); ++ break; ++ case GGML_OP_MODULATE_ROWS: ++ ggml_cuda_op_modulate_rows(ctx, dst); ++ break; ++ case GGML_OP_SWIGLU_SCALED: ++ ggml_cuda_op_swiglu_scaled(ctx, dst); ++ break; + case GGML_OP_CROSS_ENTROPY_LOSS: + ggml_cuda_cross_entropy_loss(ctx, dst); + break; +@@ -5886,6 +5897,12 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g + return ggml_cuda_flash_attn_ext_supported(dev_ctx->device, op); + case GGML_OP_SAGE_ATTN: + return ggml_cuda_sage_attn_supported(dev_ctx->device, op); ++ case GGML_OP_ROPE_PE_PERMUTE: ++ return ggml_cuda_rope_pe_permute_supported(op); ++ case GGML_OP_MODULATE_ROWS: ++ return ggml_cuda_modulate_rows_supported(op); ++ case GGML_OP_SWIGLU_SCALED: ++ return ggml_cuda_swiglu_scaled_supported(op); + case GGML_OP_CROSS_ENTROPY_LOSS: + case GGML_OP_CROSS_ENTROPY_LOSS_BACK: + case GGML_OP_OPT_STEP_ADAMW: +diff --git a/src/ggml-cuda/modulate-rows.cu b/src/ggml-cuda/modulate-rows.cu +new file mode 100644 +index 00000000..1b639d28 +--- /dev/null ++++ b/src/ggml-cuda/modulate-rows.cu +@@ -0,0 +1,168 @@ ++#include "modulate-rows.cuh" ++#include "unary.cuh" ++ ++// Rounding matches the unfused graph: __fmul_rn/__fadd_rn are the separately rounded MUL and ADD ++// kernels (never contracted into an FMA), `s*y + bias` is the expression the SCALE kernel ++// evaluates, and the swiglu product is the GLU kernel's silu(a)*b. ++ ++template ++static __global__ void modulate_rows_kernel( ++ const char * __restrict__ a, const char * __restrict__ b, const char * __restrict__ c, char * __restrict__ dst, ++ const int nc, const int64_t nr1, const int64_t ne2, const int64_t row0, ++ const int64_t nba1, const int64_t nba2, const int64_t nba3, ++ const int64_t nbb1, const int64_t nbb2, const int64_t nbb3, ++ const int64_t nbd1, const int64_t nbd2, const int64_t nbd3, ++ const float s, const float bias) { ++ const int64_t ir = blockIdx.x; ++ const int64_t i1 = row0 + ir % nr1; ++ const int64_t i2 = (ir / nr1) % ne2; ++ const int64_t i3 = ir / (nr1 * ne2); ++ ++ const float * xa = (const float *) (a + i1 * nba1 + i2 * nba2 + i3 * nba3); ++ float * y = (float *) (dst + i1 * nbd1 + i2 * nbd2 + i3 * nbd3); ++ ++ if constexpr (mode == 0) { ++ const float * scale = (const float *) b; ++ const float * shift = (const float *) c; ++ for (int i = threadIdx.x; i < nc; i += blockDim.x) { ++ const float x = xa[i]; ++ float v = __fadd_rn(__fadd_rn(x, __fmul_rn(x, scale[i])), shift[i]); ++ if constexpr (has_scale) { ++ v = s * v + bias; ++ } ++ y[i] = v; ++ } ++ } else { ++ const float * u = (const float *) (b + i1 * nbb1 + i2 * nbb2 + i3 * nbb3); ++ const float * gate = (const float *) c; ++ for (int i = threadIdx.x; i < nc; i += blockDim.x) { ++ float w = u[i]; ++ if constexpr (has_scale) { ++ w = s * w + bias; ++ } ++ y[i] = __fadd_rn(xa[i], __fmul_rn(w, gate[i])); ++ } ++ } ++} ++ ++void ggml_cuda_op_modulate_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ++ const ggml_tensor * a = dst->src[0]; ++ const ggml_tensor * b = dst->src[1]; ++ const ggml_tensor * c = dst->src[2]; ++ ++ const int mode = ggml_get_op_params_i32(dst, 0); ++ const int64_t row0 = ggml_get_op_params_i32(dst, 1); ++ const int64_t row1 = ggml_get_op_params_i32(dst, 2); ++ const float s = ggml_get_op_params_f32(dst, 3); ++ ++ const int64_t nr1 = row1 - row0; ++ const int64_t nrows = nr1 * a->ne[2] * a->ne[3]; ++ const int nc = (int) a->ne[0]; ++ if (nrows == 0) { ++ return; ++ } ++ ++ cudaStream_t stream = ctx.stream(); ++ const dim3 grid((unsigned int) nrows, 1, 1); ++ const dim3 block(256, 1, 1); ++ ++#define MODULATE_ROWS_LAUNCH(MODE, HAS_SCALE) \ ++ modulate_rows_kernel<<>>( \ ++ (const char *) a->data, (const char *) b->data, (const char *) c->data, (char *) dst->data, nc, nr1, \ ++ a->ne[2], row0, a->nb[1], a->nb[2], a->nb[3], b->nb[1], b->nb[2], b->nb[3], dst->nb[1], dst->nb[2], \ ++ dst->nb[3], s, 0.0f) ++ ++ if (mode == 0) { ++ if (s != 1.0f) { ++ MODULATE_ROWS_LAUNCH(0, true); ++ } else { ++ MODULATE_ROWS_LAUNCH(0, false); ++ } ++ } else { ++ if (s != 1.0f) { ++ MODULATE_ROWS_LAUNCH(1, true); ++ } else { ++ MODULATE_ROWS_LAUNCH(1, false); ++ } ++ } ++#undef MODULATE_ROWS_LAUNCH ++} ++ ++bool ggml_cuda_modulate_rows_supported(const ggml_tensor * op) { ++ const int64_t nrows = (int64_t) (ggml_get_op_params_i32(op, 2) - ggml_get_op_params_i32(op, 1)) * op->src[0]->ne[2] * op->src[0]->ne[3]; ++ return op->type == GGML_TYPE_F32 && ggml_is_contiguous_rows(op) && op->src[0]->ne[0] <= INT_MAX && nrows <= INT_MAX && ++ ggml_is_contiguous(op->src[ggml_get_op_params_i32(op, 0) == 0 ? 1 : 2]) && ++ (ggml_get_op_params_i32(op, 0) != 0 || ggml_is_contiguous(op->src[2])); ++} ++ ++template ++static __global__ void swiglu_scaled_kernel( ++ const char * __restrict__ a, const char * __restrict__ b, float * __restrict__ dst, ++ const int nc, const int64_t ne1, const int64_t ne2, ++ const int64_t nba1, const int64_t nba2, const int64_t nba3, ++ const int64_t nbb1, const int64_t nbb2, const int64_t nbb3, ++ const float in_scale, const float out_scale, const float bias) { ++ const int64_t ir = blockIdx.x; ++ const int64_t i1 = ir % ne1; ++ const int64_t i2 = (ir / ne1) % ne2; ++ const int64_t i3 = ir / (ne1 * ne2); ++ ++ const float * xa = (const float *) (a + i1 * nba1 + i2 * nba2 + i3 * nba3); ++ const float * xb = (const float *) (b + i1 * nbb1 + i2 * nbb2 + i3 * nbb3); ++ float * y = dst + ir * nc; ++ ++ for (int i = threadIdx.x; i < nc; i += blockDim.x) { ++ float va = xa[i]; ++ float vb = xb[i]; ++ if constexpr (has_in) { ++ va = in_scale * va + bias; ++ vb = in_scale * vb + bias; ++ } ++ float v = ggml_cuda_op_silu_single(va) * vb; ++ if constexpr (has_out) { ++ v = out_scale * v + bias; ++ } ++ y[i] = v; ++ } ++} ++ ++void ggml_cuda_op_swiglu_scaled(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ++ const ggml_tensor * a = dst->src[0]; ++ const ggml_tensor * b = dst->src[1]; ++ ++ const float in_scale = ggml_get_op_params_f32(dst, 0); ++ const float out_scale = ggml_get_op_params_f32(dst, 1); ++ ++ GGML_ASSERT(ggml_is_contiguous(dst)); ++ const int64_t nrows = ggml_nrows(a); ++ const int nc = (int) a->ne[0]; ++ if (nrows == 0) { ++ return; ++ } ++ ++ cudaStream_t stream = ctx.stream(); ++ const dim3 grid((unsigned int) nrows, 1, 1); ++ const dim3 block(256, 1, 1); ++ ++#define SWIGLU_SCALED_LAUNCH(HAS_IN, HAS_OUT) \ ++ swiglu_scaled_kernel<<>>( \ ++ (const char *) a->data, (const char *) b->data, (float *) dst->data, nc, a->ne[1], a->ne[2], a->nb[1], \ ++ a->nb[2], a->nb[3], b->nb[1], b->nb[2], b->nb[3], in_scale, out_scale, 0.0f) ++ ++ const bool has_in = in_scale != 1.0f; ++ const bool has_out = out_scale != 1.0f; ++ if (has_in && has_out) { ++ SWIGLU_SCALED_LAUNCH(true, true); ++ } else if (has_in) { ++ SWIGLU_SCALED_LAUNCH(true, false); ++ } else if (has_out) { ++ SWIGLU_SCALED_LAUNCH(false, true); ++ } else { ++ SWIGLU_SCALED_LAUNCH(false, false); ++ } ++#undef SWIGLU_SCALED_LAUNCH ++} ++ ++bool ggml_cuda_swiglu_scaled_supported(const ggml_tensor * op) { ++ return op->type == GGML_TYPE_F32 && op->src[0]->ne[0] <= INT_MAX && ggml_nrows(op->src[0]) <= INT_MAX; ++} +diff --git a/src/ggml-cuda/modulate-rows.cuh b/src/ggml-cuda/modulate-rows.cuh +new file mode 100644 +index 00000000..216d0767 +--- /dev/null ++++ b/src/ggml-cuda/modulate-rows.cuh +@@ -0,0 +1,7 @@ ++#include "common.cuh" ++ ++void ggml_cuda_op_modulate_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst); ++bool ggml_cuda_modulate_rows_supported(const ggml_tensor * op); ++ ++void ggml_cuda_op_swiglu_scaled(ggml_backend_cuda_context & ctx, ggml_tensor * dst); ++bool ggml_cuda_swiglu_scaled_supported(const ggml_tensor * op); +diff --git a/src/ggml-cuda/rope-pe.cu b/src/ggml-cuda/rope-pe.cu +new file mode 100644 +index 00000000..3afe4423 +--- /dev/null ++++ b/src/ggml-cuda/rope-pe.cu +@@ -0,0 +1,102 @@ ++#include "rope-pe.cuh" ++#include "convert.cuh" ++ ++// Fused rotary embedding (precomputed cos/sin table) + head-major relayout + scale + cast. ++// One block per (token, batch); the block walks every head so the table entries are read once. ++// Rounding matches the unfused graph exactly: __fmul_rn/__fadd_rn keep the two products and the ++// sum as separately rounded operations (no FMA contraction), the scale is the same `scale*y + bias` ++// expression the SCALE kernel evaluates, and the conversion uses the same ggml_cuda_cast as CPY. ++template ++static __global__ void rope_pe_permute_kernel( ++ const char * __restrict__ x, const float * __restrict__ pe, dst_t * __restrict__ dst, ++ const int D, const int H, const int64_t L, const int n_rot, ++ const int64_t nb01, const int64_t nb02, const int64_t nb03, ++ const float scale, const float bias) { ++ const int64_t l = blockIdx.x; ++ const int64_t n = blockIdx.y; ++ const int half = n_rot / 2; ++ ++ for (int i = threadIdx.x; i < D; i += blockDim.x) { ++ const bool rot = i < n_rot; ++ int j = i; ++ float c0 = 0.0f; ++ float c1 = 0.0f; ++ if (rot) { ++ const int r = i >= half; ++ j = i - r * half; ++ const float * t = pe + (l * half + j) * 4 + 2 * r; ++ c0 = t[0]; ++ c1 = t[1]; ++ } ++ for (int h = 0; h < H; ++h) { ++ const float * xr = (const float *) (x + h * nb01 + l * nb02 + n * nb03); ++ float y; ++ if (rot) { ++ y = __fadd_rn(__fmul_rn(xr[j], c0), __fmul_rn(xr[j + half], c1)); ++ } else { ++ y = xr[i]; ++ } ++ if constexpr (has_scale) { ++ y = scale * y + bias; ++ } ++ dst[((n * H + h) * L + l) * D + i] = ggml_cuda_cast(y); ++ } ++ } ++} ++ ++template ++static void rope_pe_permute_cuda(const ggml_tensor * src0, const float * pe, dst_t * dst, const int n_rot, ++ const float scale, cudaStream_t stream) { ++ const int D = (int) src0->ne[0]; ++ const int H = (int) src0->ne[1]; ++ const int64_t L = src0->ne[2]; ++ const int64_t N = src0->ne[3]; ++ ++ const dim3 grid((unsigned int) L, (unsigned int) N, 1); ++ const dim3 block(D >= 128 ? 128 : ((D + 31) / 32) * 32, 1, 1); ++ const char * x = (const char *) src0->data; ++ if (scale != 1.0f) { ++ rope_pe_permute_kernel<<>>( ++ x, pe, dst, D, H, L, n_rot, src0->nb[1], src0->nb[2], src0->nb[3], scale, 0.0f); ++ } else { ++ rope_pe_permute_kernel<<>>( ++ x, pe, dst, D, H, L, n_rot, src0->nb[1], src0->nb[2], src0->nb[3], scale, 0.0f); ++ } ++} ++ ++void ggml_cuda_op_rope_pe_permute(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ++ const ggml_tensor * src0 = dst->src[0]; ++ const ggml_tensor * pe = dst->src[1]; ++ ++ const int n_rot = ggml_get_op_params_i32(dst, 0); ++ const float scale = ggml_get_op_params_f32(dst, 1); ++ ++ GGML_ASSERT(src0->type == GGML_TYPE_F32 && src0->nb[0] == sizeof(float)); ++ GGML_ASSERT(ggml_is_contiguous(dst)); ++ GGML_ASSERT(n_rot == 0 || (pe != nullptr && ggml_is_contiguous(pe))); ++ ++ cudaStream_t stream = ctx.stream(); ++ const float * pe_d = pe ? (const float *) pe->data : nullptr; ++ ++ if (dst->type == GGML_TYPE_F16) { ++ rope_pe_permute_cuda(src0, pe_d, (half *) dst->data, n_rot, scale, stream); ++ } else { ++ GGML_ASSERT(dst->type == GGML_TYPE_F32); ++ rope_pe_permute_cuda(src0, pe_d, (float *) dst->data, n_rot, scale, stream); ++ } ++} ++ ++bool ggml_cuda_rope_pe_permute_supported(const ggml_tensor * op) { ++ const ggml_tensor * src0 = op->src[0]; ++ const ggml_tensor * pe = op->src[1]; ++ if (src0->type != GGML_TYPE_F32 || src0->nb[0] != sizeof(float)) { ++ return false; ++ } ++ if (op->type != GGML_TYPE_F32 && op->type != GGML_TYPE_F16) { ++ return false; ++ } ++ if (src0->ne[2] > INT_MAX || src0->ne[3] > 65535 || src0->ne[1] > INT_MAX || src0->ne[0] > INT_MAX) { ++ return false; ++ } ++ return pe == nullptr || (pe->type == GGML_TYPE_F32 && ggml_is_contiguous(pe)); ++} +diff --git a/src/ggml-cuda/rope-pe.cuh b/src/ggml-cuda/rope-pe.cuh +new file mode 100644 +index 00000000..b68b6af4 +--- /dev/null ++++ b/src/ggml-cuda/rope-pe.cuh +@@ -0,0 +1,5 @@ ++#include "common.cuh" ++ ++void ggml_cuda_op_rope_pe_permute(ggml_backend_cuda_context & ctx, ggml_tensor * dst); ++ ++bool ggml_cuda_rope_pe_permute_supported(const ggml_tensor * op); +diff --git a/src/ggml.c b/src/ggml.c +index dc4a29f9..fec933bf 100644 +--- a/src/ggml.c ++++ b/src/ggml.c +@@ -1115,9 +1115,12 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { + + "QUANTIZE_I8_CONVROT", + "SAGE_ATTN", ++ "ROPE_PE_PERMUTE", ++ "MODULATE_ROWS", ++ "SWIGLU_SCALED", + }; + +-static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103"); ++static_assert(GGML_OP_COUNT == 106, "GGML_OP_COUNT != 106"); + + static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { + "none", +@@ -1233,9 +1236,12 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { + + "quantize_i8_convrot(x)", + "sage_attn(x,y,z)", ++ "rope_pe_permute(x,pe)", ++ "modulate_rows(x,a,b)", ++ "swiglu_scaled(x,y)", + }; + +-static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103"); ++static_assert(GGML_OP_COUNT == 106, "GGML_OP_COUNT != 106"); + + static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); + +@@ -5557,6 +5563,125 @@ struct ggml_tensor * ggml_sage_attn( + return result; + } + ++struct ggml_tensor * ggml_rope_pe_permute( ++ struct ggml_context * ctx, ++ struct ggml_tensor * a, ++ struct ggml_tensor * pe, ++ int n_rot, ++ float scale, ++ enum ggml_type type) { ++ GGML_ASSERT(a->type == GGML_TYPE_F32); ++ GGML_ASSERT(a->nb[0] == ggml_type_size(a->type)); ++ GGML_ASSERT(type == GGML_TYPE_F32 || type == GGML_TYPE_F16); ++ GGML_ASSERT(n_rot >= 0 && n_rot % 2 == 0 && n_rot <= a->ne[0]); ++ if (n_rot > 0) { ++ GGML_ASSERT(pe != NULL && pe->type == GGML_TYPE_F32 && ggml_is_contiguous(pe)); ++ GGML_ASSERT(pe->ne[0] == 2 && pe->ne[1] == 2 && pe->ne[2] == n_rot / 2 && pe->ne[3] == a->ne[2]); ++ } else { ++ GGML_ASSERT(pe == NULL); ++ } ++ GGML_ASSERT(isfinite(scale)); ++ ++ struct ggml_tensor * result = ggml_new_tensor_4d(ctx, type, a->ne[0], a->ne[2], a->ne[1], a->ne[3]); ++ ++ ggml_set_op_params_i32(result, 0, n_rot); ++ ggml_set_op_params_f32(result, 1, scale); ++ ++ result->op = GGML_OP_ROPE_PE_PERMUTE; ++ result->src[0] = a; ++ result->src[1] = pe; ++ ++ return result; ++} ++ ++static struct ggml_tensor * ggml_modulate_rows_impl( ++ struct ggml_context * ctx, ++ struct ggml_tensor * acc, ++ struct ggml_tensor * a, ++ struct ggml_tensor * b, ++ struct ggml_tensor * c, ++ int64_t row0, ++ int64_t row1, ++ float s, ++ int32_t mode) { ++ GGML_ASSERT(a->type == GGML_TYPE_F32 && a->nb[0] == sizeof(float)); ++ GGML_ASSERT(row0 >= 0 && row0 < row1 && row1 <= a->ne[1] && row1 <= INT32_MAX); ++ GGML_ASSERT(isfinite(s)); ++ // mode 0: a = x, b = scale, c = shift; mode 1: a = base, b = u, c = gate ++ struct ggml_tensor * vec0 = mode == 0 ? b : c; ++ GGML_ASSERT(vec0->type == GGML_TYPE_F32 && vec0->nb[0] == sizeof(float) && ggml_nelements(vec0) == a->ne[0]); ++ if (mode == 0) { ++ GGML_ASSERT(c->type == GGML_TYPE_F32 && c->nb[0] == sizeof(float) && ggml_nelements(c) == a->ne[0]); ++ } else { ++ GGML_ASSERT(b->type == GGML_TYPE_F32 && b->nb[0] == sizeof(float) && ggml_are_same_shape(a, b)); ++ } ++ if (acc != NULL) { ++ GGML_ASSERT(acc->type == GGML_TYPE_F32 && ggml_is_contiguous(acc) && ggml_are_same_shape(acc, a)); ++ } ++ ++ struct ggml_tensor * result = acc != NULL ? ggml_view_tensor(ctx, acc) : ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, a->ne); ++ ++ ggml_set_op_params_i32(result, 0, mode); ++ ggml_set_op_params_i32(result, 1, (int32_t) row0); ++ ggml_set_op_params_i32(result, 2, (int32_t) row1); ++ ggml_set_op_params_f32(result, 3, s); ++ ++ result->op = GGML_OP_MODULATE_ROWS; ++ result->src[0] = a; ++ result->src[1] = b; ++ result->src[2] = c; ++ result->src[3] = acc; ++ ++ return result; ++} ++ ++struct ggml_tensor * ggml_modulate_rows( ++ struct ggml_context * ctx, ++ struct ggml_tensor * acc, ++ struct ggml_tensor * x, ++ struct ggml_tensor * scale, ++ struct ggml_tensor * shift, ++ int64_t row0, ++ int64_t row1, ++ float out_scale) { ++ return ggml_modulate_rows_impl(ctx, acc, x, scale, shift, row0, row1, out_scale, 0); ++} ++ ++struct ggml_tensor * ggml_gated_add_rows( ++ struct ggml_context * ctx, ++ struct ggml_tensor * acc, ++ struct ggml_tensor * base, ++ struct ggml_tensor * u, ++ struct ggml_tensor * gate, ++ int64_t row0, ++ int64_t row1, ++ float u_scale) { ++ return ggml_modulate_rows_impl(ctx, acc, base, u, gate, row0, row1, u_scale, 1); ++} ++ ++struct ggml_tensor * ggml_swiglu_scaled( ++ struct ggml_context * ctx, ++ struct ggml_tensor * a, ++ struct ggml_tensor * b, ++ float in_scale, ++ float out_scale) { ++ GGML_ASSERT(a->type == GGML_TYPE_F32 && b->type == GGML_TYPE_F32); ++ GGML_ASSERT(a->nb[0] == sizeof(float) && b->nb[0] == sizeof(float)); ++ GGML_ASSERT(ggml_are_same_shape(a, b)); ++ GGML_ASSERT(isfinite(in_scale) && isfinite(out_scale)); ++ ++ struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, a->ne); ++ ++ ggml_set_op_params_f32(result, 0, in_scale); ++ ggml_set_op_params_f32(result, 1, out_scale); ++ ++ result->op = GGML_OP_SWIGLU_SCALED; ++ result->src[0] = a; ++ result->src[1] = b; ++ ++ return result; ++} ++ + void ggml_flash_attn_ext_set_prec( + struct ggml_tensor * a, + enum ggml_prec prec) { +diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp +index dd290caf..1e46361c 100644 +--- a/tests/test-backend-ops.cpp ++++ b/tests/test-backend-ops.cpp +@@ -2125,6 +2125,118 @@ struct test_glu : public test_case { + } + }; + ++// GGML_OP_ROPE_PE_PERMUTE ++struct test_rope_pe_permute : public test_case { ++ const ggml_type type; ++ const std::array ne_a; // [head_dim, n_head, n_tokens, n_batch] ++ const int n_rot; ++ const float scale; ++ const bool view; // a is a strided view (one third of a fused qkv projection) ++ ++ std::string vars() override { ++ return VARS_TO_STR5(type, ne_a, n_rot, scale, view); ++ } ++ ++ test_rope_pe_permute(ggml_type type = GGML_TYPE_F32, ++ std::array ne_a = {128, 4, 9, 1}, ++ int n_rot = 96, float scale = 1.0f, bool view = false) ++ : type(type), ne_a(ne_a), n_rot(n_rot), scale(scale), view(view) {} ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * a; ++ if (view) { ++ ggml_tensor * qkv = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 3 * ne_a[0] * ne_a[1], ne_a[2], ne_a[3], 1); ++ ggml_set_name(qkv, "qkv"); ++ a = ggml_view_4d(ctx, qkv, ne_a[0], ne_a[1], ne_a[2], ne_a[3], ++ ne_a[0] * sizeof(float), qkv->nb[1], qkv->nb[2], ne_a[0] * ne_a[1] * sizeof(float)); ++ ggml_set_name(a, "view_of_qkv"); ++ } else { ++ a = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne_a.data()); ++ ggml_set_name(a, "a"); ++ } ++ ggml_tensor * pe = nullptr; ++ if (n_rot > 0) { ++ pe = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 2, 2, n_rot / 2, ne_a[2]); ++ ggml_set_name(pe, "pe"); ++ } ++ ggml_tensor * out = ggml_rope_pe_permute(ctx, a, pe, n_rot, scale, type); ++ ggml_set_name(out, "out"); ++ return out; ++ } ++}; ++ ++// GGML_OP_MODULATE_ROWS (a chain of segments written in place) ++struct test_modulate_rows : public test_case { ++ const int mode; // 0: modulate, 1: gated add ++ const std::array ne; ++ const float s; ++ const bool view; ++ ++ std::string vars() override { ++ return VARS_TO_STR4(mode, ne, s, view); ++ } ++ ++ test_modulate_rows(int mode = 0, std::array ne = {64, 12, 2, 1}, float s = 1.0f, bool view = false) ++ : mode(mode), ne(ne), s(s), view(view) {} ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ auto make = [&](const char * name) { ++ ggml_tensor * t; ++ if (view) { ++ ggml_tensor * big = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 2 * ne[0], ne[1], ne[2], ne[3]); ++ t = ggml_view_4d(ctx, big, ne[0], ne[1], ne[2], ne[3], big->nb[1], big->nb[2], big->nb[3], ne[0] * sizeof(float)); ++ } else { ++ t = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); ++ } ++ ggml_set_name(t, name); ++ return t; ++ }; ++ ggml_tensor * a = make("a"); ++ ggml_tensor * u = mode == 1 ? make("u") : nullptr; ++ // modulation table: 3 rows of 2 vectors each, taken as views like an adaLN projection ++ ggml_tensor * table = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2 * ne[0], 3); ++ ggml_set_name(table, "table"); ++ const int64_t bounds[4] = {0, ne[1] / 4, ne[1] - 2, ne[1]}; ++ ggml_tensor * out = nullptr; ++ for (int seg = 0; seg < 3; ++seg) { ++ const int row = (seg + 1) % 3; ++ ggml_tensor * v0 = ggml_view_1d(ctx, table, ne[0], row * table->nb[1]); ++ ggml_tensor * v1 = ggml_view_1d(ctx, table, ne[0], row * table->nb[1] + ne[0] * sizeof(float)); ++ if (mode == 0) { ++ out = ggml_modulate_rows(ctx, out, a, v0, v1, bounds[seg], bounds[seg + 1], s); ++ } else { ++ out = ggml_gated_add_rows(ctx, out, a, u, v0, bounds[seg], bounds[seg + 1], s); ++ } ++ } ++ ggml_set_name(out, "out"); ++ return out; ++ } ++}; ++ ++// GGML_OP_SWIGLU_SCALED ++struct test_swiglu_scaled : public test_case { ++ const std::array ne; ++ const float in_scale; ++ const float out_scale; ++ ++ std::string vars() override { ++ return VARS_TO_STR3(ne, in_scale, out_scale); ++ } ++ ++ test_swiglu_scaled(std::array ne = {64, 5, 3, 1}, float in_scale = 1.0f, float out_scale = 1.0f) ++ : ne(ne), in_scale(in_scale), out_scale(out_scale) {} ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * h = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 2 * ne[0], ne[1], ne[2], ne[3]); ++ ggml_set_name(h, "h"); ++ ggml_tensor * a = ggml_view_4d(ctx, h, ne[0], ne[1], ne[2], ne[3], h->nb[1], h->nb[2], h->nb[3], 0); ++ ggml_tensor * b = ggml_view_4d(ctx, h, ne[0], ne[1], ne[2], ne[3], h->nb[1], h->nb[2], h->nb[3], ne[0] * sizeof(float)); ++ ggml_tensor * out = ggml_swiglu_scaled(ctx, a, b, in_scale, out_scale); ++ ggml_set_name(out, "out"); ++ return out; ++ } ++}; ++ + struct test_glu_split : public test_case { + const ggml_glu_op op; + const ggml_type type; +@@ -8101,6 +8213,33 @@ static std::vector> make_test_cases_eval() { + } + } + ++ for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_F16}) { ++ for (bool view : {false, true}) { ++ for (float scale : {1.0f, 1.0f / 128.0f}) { ++ test_cases.emplace_back(new test_rope_pe_permute(type, {128, 4, 9, 1}, 96, scale, view)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {128, 3, 5, 2}, 128, scale, view)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {128, 4, 9, 1}, 0, scale, view)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {64, 3, 7, 2}, 32, scale, view)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {200, 2, 3, 1}, 100, scale, view)); ++ } ++ } ++ } ++ ++ for (int mode : {0, 1}) { ++ for (bool view : {false, true}) { ++ for (float sc : {1.0f, mode == 0 ? 1.0f / 128.0f : 128.0f}) { ++ test_cases.emplace_back(new test_modulate_rows(mode, {64, 12, 2, 1}, sc, view)); ++ test_cases.emplace_back(new test_modulate_rows(mode, {300, 9, 1, 2}, sc, view)); ++ } ++ } ++ } ++ for (float in_scale : {1.0f, 128.0f}) { ++ for (float out_scale : {1.0f, 1.0f / 128.0f}) { ++ test_cases.emplace_back(new test_swiglu_scaled({64, 5, 3, 1}, in_scale, out_scale)); ++ test_cases.emplace_back(new test_swiglu_scaled({300, 7, 1, 2}, in_scale, out_scale)); ++ } ++ } ++ + for (int v : {0, 1}) { + for (float alpha : {.5f, 1.702f}) { + for (float limit : {2.0f, 7.0f}) { +-- +2.43.0 + From 36963a4a83f5ae28e813ea51f9483eb7a1b7ad01 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 06:55:12 -0700 Subject: [PATCH 11/30] 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 3b75408aad..07a244a8ad 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 12/30] 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 c852492eae..7d5aa2447f 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 564831ee70..32b2c5378f 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 13/30] 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 7d5aa2447f..4b967e0e36 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 1ae87ec87d..604a05fb97 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 c0eb17539fc0d336928f86498a7ede9f7fe9d569 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 22:42:50 -0700 Subject: [PATCH 14/30] MiniMax-H3: sage attention takes the fused RoPE / head-major q/k/v op With --sage-attn, forward_head_major returned nullptr, so every block fell back to the unfused chunk / slice / rope / concat / scale / cast chain (about 0.75 s per 960x544x124 step on B200). ggml_rope_pe_permute already writes [head_dim, tokens, heads, batch], the layout ggml_sage_attn reads: emit F32 Q, F32 K * kv_scale and F16 V * kv_scale and call sage directly. Same arithmetic and rounding points as the unfused sage branch; the rendered video and audio are bit-identical. SD_H3_FAST_SAGE_QKV=0 restores the old chain. --- docs/minimax_h3.md | 5 ++++- src/model/diffusion/minimax_h3.hpp | 30 ++++++++++++++++++++++++++++-- 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index 4b967e0e36..cc40daa143 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -70,7 +70,10 @@ partial RoPE with the attention relayout and the K/V scale and F16 cast, per-seg 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. +`SD_H3_FAST_VIEWS=0` turn off one part each. With `--sage-attn` the same fused RoPE op writes the +F32 Q/K and F16 V layout the sage kernel reads (K and V carry the kv scale), so the sage path no +longer pays for the chunk / slice / rope / concat / scale / cast chain; the output is bit-identical +to the unfused sage graph. `SD_H3_FAST_SAGE_QKV=0` restores that chain. ## First/last-frame conditioning diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index 07a244a8ad..cc2c24fc49 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -153,6 +153,13 @@ namespace MiniMaxH3 { return enabled; } + // SD_H3_FAST_SAGE_QKV=0: with --sage-attn, q/k/v go through the unfused chunk / slice / rope / + // concat / scale / cast chain instead of one fused RoPE + head-major op per tensor. + static bool fast_sage_qkv() { + static const bool enabled = fast_qkv() && env_flag("SD_H3_FAST_SAGE_QKV", true); + return enabled; + } + static bool fast_segments() { static const bool enabled = graph_fast() && env_flag("SD_H3_FAST_SEGMENTS", true); return enabled; @@ -293,7 +300,12 @@ namespace MiniMaxH3 { #ifdef SD_USE_UPSTREAM_GGML return nullptr; #else - if (!ctx->flash_attn_enabled || ctx->sage_attn_enabled || ctx->backend == nullptr) { + const bool sage = ctx->sage_attn_enabled; + if (ctx->backend == nullptr || (sage ? !fast_sage_qkv() : !ctx->flash_attn_enabled)) { + return nullptr; + } + // the sage kernel takes 64 or 128 channel heads without padding + if (sage && head_dim != 64 && head_dim != 128) { return nullptr; } const int64_t inner = heads * head_dim; @@ -324,13 +336,27 @@ namespace MiniMaxH3 { }; auto q = q_norm->forward(ctx, part(0)); auto k = k_norm->forward(ctx, part(1)); + // Sage reads F32 Q/K and F16 V, all [head_dim, tokens, heads, batch]: the layout this op + // writes. The unfused sage path scales K and V in F32 before V's F16 cast, as here. 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); + k = ggml_rope_pe_permute(ctx->ggml_ctx, k, pe, n_rot, kv_scale, sage ? GGML_TYPE_F32 : GGML_TYPE_F16); auto v = ggml_rope_pe_permute(ctx->ggml_ctx, part(2), nullptr, 0, kv_scale, GGML_TYPE_F16); if (!ggml_backend_supports_op(ctx->backend, q) || !ggml_backend_supports_op(ctx->backend, k) || !ggml_backend_supports_op(ctx->backend, v)) { return nullptr; } + if (sage) { + // same softmax scale expression as ggml_ext_attention_ext's sage branch + const float scale = 1.0f / sqrt((float)head_dim); + auto out = ggml_sage_attn(ctx->ggml_ctx, q, k, v, scale / kv_scale, GGML_SAGE_ATTN_AUTO); + if (!ggml_backend_supports_op(ctx->backend, out)) { + return nullptr; + } + if (kv_scale != 1.0f) { + out = ggml_ext_scale(ctx->ggml_ctx, out, 1.0f / kv_scale); + } + return ggml_reshape_3d(ctx->ggml_ctx, out, head_dim * heads, sequence, batch); // [N, L, C] + } q = ggml_reshape_3d(ctx->ggml_ctx, q, head_dim, sequence, heads * batch); k = ggml_reshape_3d(ctx->ggml_ctx, k, head_dim, sequence, heads * batch); v = ggml_reshape_3d(ctx->ggml_ctx, v, head_dim, sequence, heads * batch); From b93d5b2908f01ae506d0f3a1ae94d8af6d48f653 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 22:58:07 -0700 Subject: [PATCH 15/30] MiniMax-H3 video VAE: keep the device busy between tiles The 960x544x124 decode left the B200 idle for about 2 s between its 105 tile graphs: the host blended each tile, zero-filled and assembled the frames, rebuilt the rotary tables and queried free device memory four times per tile while the GPU waited. - tiling: the blend of a batch into the output runs on a worker thread while the next batch is split and computed; merges stay serial and in tile order with the same arithmetic (SD_TILE_ASYNC_MERGE=0 merges inline). - H3 decode: each temporal chunk is trimmed, cross-faded and copied straight into the trimmed output on a worker thread while the next chunk decodes, instead of concatenating and slicing at the end (SD_H3_VAE_ASYNC_ASSEMBLY=0 restores it). - rotary tables are rebuilt only when the tile shape changes. - DeviceMemoryRequest::reuse_device_query: while a runner repeats identical computes, a capacity check that allocates nothing new (no pending bytes, all parameters resident) reuses the owner's earlier free-memory reading; runner_end() drops it. The H3 decode turns it on for its duration (SD_H3_VAE_REUSE_MEMQUERY=0 disables). Warm decode on B200 6.8-7.0 s -> 5.06 s; frames and audio bit-identical. --- docs/minimax_h3.md | 9 +++ src/core/ggml_runner.cpp | 2 +- src/core/ggml_runner.h | 7 ++ src/device_residency_manager.h | 4 + src/model/vae/minimax_h3_vae.hpp | 135 ++++++++++++++++++++++++++----- src/model_manager.cpp | 18 ++++- src/model_manager.h | 2 + src/runtime/tiling.cpp | 91 ++++++++++++++++----- 8 files changed, 222 insertions(+), 46 deletions(-) diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index cc40daa143..97bcaf8cc9 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -65,6 +65,15 @@ opt-in. The decoder weights stay on the device across temporal chunks table-based rotary embedding and a fused SwiGLU (`SD_H3_VAE_GRAPH_OPT=0` restores the previous graph). Each batched tile still goes through its own attention call. +The host side of the decode overlaps the device: each tile's blend into the frame runs on a +worker thread while the next tile computes (`SD_TILE_ASYNC_MERGE=0` blends inline; this applies +to every tiled VAE decode), each temporal chunk is trimmed, cross-faded and copied into the final +frames on a worker thread while the next chunk decodes (`SD_H3_VAE_ASYNC_ASSEMBLY=0` restores +the concatenate-at-the-end path), the rotary tables are built once per tile shape, and a tile +that allocates nothing new reuses the device free-memory reading taken earlier in the decode +instead of querying the device for every capacity check (`SD_H3_VAE_REUSE_MEMQUERY=0`). The +decoded frames are bit-identical either way. + The DiT blocks use fused ggml ops (CPU and CUDA) for the work around the matmuls and attention: partial RoPE with the attention relayout and the K/V scale and F16 cast, per-segment adaLN modulation and gated residuals written in place, and the MLP Linear scales folded into those ops diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp index abfa8639bd..3ac8e0b2dc 100644 --- a/src/core/ggml_runner.cpp +++ b/src/core/ggml_runner.cpp @@ -725,7 +725,7 @@ std::vector GGMLRunner::memory_requests( } } requests.push_back({size.backend, reinterpret_cast(this), pending, - retained, limit}); + retained, limit, reuse_device_query_}); } return requests; } diff --git a/src/core/ggml_runner.h b/src/core/ggml_runner.h index 29a6808bd8..9b503b407c 100644 --- a/src/core/ggml_runner.h +++ b/src/core/ggml_runner.h @@ -156,6 +156,7 @@ struct GGMLRunner { bool graph_active_ = false; size_t max_graph_vram_bytes = 0; + bool reuse_device_query_ = false; bool graph_cut_layer_split_enabled = false; std::vector graph_cut_layer_split_backend_vram_limits_; @@ -369,6 +370,12 @@ struct GGMLRunner { weight_adapter = adapter; } + // See DeviceMemoryRequest::reuse_device_query. Callers turn it on around a run of identical + // computes (spatial tiles) and off afterwards; runner_end() drops the saved reading. + void set_reuse_device_query(bool enabled) { + reuse_device_query_ = enabled; + } + void set_max_graph_vram_bytes(size_t max_vram_bytes) { max_graph_vram_bytes = max_vram_bytes; } diff --git a/src/device_residency_manager.h b/src/device_residency_manager.h index 0a2a8895f3..fcdb5c389c 100644 --- a/src/device_residency_manager.h +++ b/src/device_residency_manager.h @@ -27,6 +27,10 @@ struct DeviceMemoryRequest { size_t pending_allocation_bytes = 0; size_t runtime_resident_bytes = 0; size_t max_backend_bytes = 0; + // The caller repeats an identical compute (e.g. one tile after another). When nothing new is + // allocated (no pending bytes, every parameter already resident), the manager may answer from + // the device free-memory figure it read for this owner earlier instead of querying again. + bool reuse_device_query = false; // Runtime buffers only; the manager accounts for weights separately. size_t runtime_peak_bytes() const { diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index 604a05fb97..fd09f4a744 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -1,6 +1,7 @@ #ifndef __SD_MODEL_VAE_MINIMAX_H3_VAE_HPP__ #define __SD_MODEL_VAE_MINIMAX_H3_VAE_HPP__ +#include #include #include #include @@ -605,6 +606,7 @@ namespace MiniMaxH3VAE { sd::Tensor rope_cache; sd::Tensor rope_a_cache; sd::Tensor rope_b_cache; + std::array rope_key_ = {0, 0, 0, 0}; MiniMaxH3VideoVAERunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map, @@ -807,18 +809,22 @@ namespace MiniMaxH3VAE { {static_cast(tokens_per_chunk + token_overlap), static_cast(token_overlap)}); GGML_ASSERT(plan.tiles.size() == static_cast(num_chunks)); const bool keep_resident = env_int("SD_H3_VAE_KEEP_RESIDENT", 1) != 0; + // One device free-memory reading serves every tile that allocates nothing new + // (SD_H3_VAE_REUSE_MEMQUERY=0 queries the device for every capacity check). + set_reuse_device_query(env_int("SD_H3_VAE_REUSE_MEMQUERY", 1) != 0); + struct ReuseQueryGuard { + GGMLRunner& runner; + ~ReuseQueryGuard() { + runner.set_reuse_device_query(false); + } + } reuse_query_guard{*this}; tile_batch_ = 0; per_tile_compute_bytes_ = 0; std::vector> pieces; - 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(); - } - + // Trims the pre-padding frames of a decoded chunk and cross-fades it with the previous + // chunk's tail. Runs in chunk order; on a worker thread unless SD_H3_VAE_ASYNC_ASSEMBLY=0, + // so it overlaps the next chunk's decode (same arithmetic, bit-identical result). + auto assemble = [&](const sd::Tensor& decoded, const VAETemporalTile& tile) { int64_t first_end = std::min(frames_per_chunk, decoded.shape()[2]); auto first = sd::ops::slice(decoded, 2, @@ -841,23 +847,102 @@ namespace MiniMaxH3VAE { } return first; }; + const bool async_assembly = env_int("SD_H3_VAE_ASYNC_ASSEMBLY", 1) != 0; + std::thread assembler; + bool assembly_failed = false; + int64_t expected_frames = input.shape()[2] <= 1 ? 1 : ((x.shape()[2] - 2) / 5) * 17 + 5; + expected_frames = std::max(1, expected_frames); + // Pieces are copied straight into the trimmed output as they are assembled (the same + // frames concat_frames + slice would keep), instead of concatenated at the end. + sd::Tensor result; + int64_t frame_offset = 0; + auto append_piece = [&](const sd::Tensor& piece) { + GGML_ASSERT(piece.dim() == 5); + if (result.empty()) { + std::vector shape = piece.shape(); + shape[2] = expected_frames; + result = sd::Tensor(std::move(shape)); + } + GGML_ASSERT(piece.shape()[0] == result.shape()[0] && piece.shape()[1] == result.shape()[1] && + piece.shape()[3] == result.shape()[3] && piece.shape()[4] == result.shape()[4]); + const int64_t piece_frames = piece.shape()[2]; + const int64_t keep = std::min(piece_frames, expected_frames - frame_offset); + const int64_t plane = piece.shape()[0] * piece.shape()[1]; + const int64_t planes = piece.shape()[3] * piece.shape()[4]; + for (int64_t p = 0; keep > 0 && p < planes; ++p) { + memcpy(result.data() + (p * expected_frames + frame_offset) * plane, + piece.data() + p * piece_frames * plane, + sizeof(float) * keep * plane); + } + frame_offset += std::max(keep, 0); + }; + auto run_assembly = [&](const sd::Tensor& decoded, const VAETemporalTile& tile) { + try { + if (async_assembly) { + append_piece(assemble(decoded, tile)); + } else { + pieces.push_back(assemble(decoded, tile)); + } + } catch (const std::exception& error) { + LOG_ERROR("MiniMax-H3 video VAE: chunk assembly failed: %s", error.what()); + assembly_failed = true; + } + }; + auto join_assembler = [&]() { + if (assembler.joinable()) { + assembler.join(); + } + }; + struct AssemblerGuard { + std::thread& thread; + ~AssemblerGuard() { + if (thread.joinable()) { + thread.join(); + } + } + } assembler_guard{assembler}; bool failed = false; for (const auto& tile : plan.tiles) { - auto piece = collect_pieces(sd::ops::slice(input, 2, tile.start, tile.end), tile); - if (piece.empty()) { + auto decoded = decode_spatial_tiles(n_threads, + sd::ops::slice(input, 2, tile.start, tile.end), + tiling, + circular_x, + circular_y, + silent); + if (!keep_resident) { + runner_end(); + } + join_assembler(); + if (decoded.empty() || assembly_failed) { failed = true; break; } - pieces.push_back(std::move(piece)); + if (async_assembly && !tile.last) { + assembler = std::thread([&run_assembly, decoded = std::move(decoded), tile]() { + run_assembly(decoded, tile); + }); + } else { + run_assembly(decoded, tile); + } } + join_assembler(); runner_end(); - if (failed || pieces.empty()) { + if (failed || assembly_failed) { 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); + if (async_assembly) { + if (result.empty()) { + return {}; + } + if (frame_offset < expected_frames) { + result = sd::ops::slice(result, 2, 0, frame_offset); + } + return result; + } + if (pieces.empty()) { + return {}; + } + result = concat_frames(pieces); if (result.shape()[2] > expected_frames) { result = sd::ops::slice(result, 2, 0, expected_frames); } @@ -1063,11 +1148,17 @@ namespace MiniMaxH3VAE { 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(); + // Every spatial tile of a decode has the same extent, so the tables are rebuilt only + // when the tile shape (or the graph variant) changes. + const std::array rope_key = {input.shape()[0], input.shape()[1], input.shape()[2], graph_opt ? 1 : 0}; + if (rope_key != rope_key_ || rope_cache.empty()) { + rope_cache = build_rope(input.shape()[0], + input.shape()[1], + input.shape()[2]); + if (graph_opt) { + build_rope_tables(); + } + rope_key_ = rope_key; } } auto get_graph = [&]() -> ggml_cgraph* { diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 9a9fe06287..d6cd88a480 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -1602,6 +1602,9 @@ void ModelManager::set_workspace_reclaimer(uintptr_t owner_id, std::functionfirst.first == owner_id ? device_query_cache_.erase(it) : std::next(it); + } for (auto it = runtime_residencies_.begin(); it != runtime_residencies_.end();) { if (it->first.first == owner_id) { it = runtime_residencies_.erase(it); @@ -1630,7 +1633,20 @@ ModelManager::CapacityCheck ModelManager::check_capacity( return SIZE_MAX; } size_t free_bytes = 0, total_bytes = 0; - ggml_backend_dev_memory(device, &free_bytes, &total_bytes); + const auto cache_key = std::make_pair(request.owner_id, device); + auto cached = device_query_cache_.find(cache_key); + if (request.reuse_device_query && request.pending_allocation_bytes == 0 && missing == 0 && + cached != device_query_cache_.end()) { + free_bytes = cached->second.first; + total_bytes = cached->second.second; + } else { + ggml_backend_dev_memory(device, &free_bytes, &total_bytes); + if (request.reuse_device_query) { + device_query_cache_[cache_key] = {free_bytes, total_bytes}; + } else if (cached != device_query_cache_.end()) { + device_query_cache_.erase(cached); + } + } if (free_bytes == 0 && total_bytes == 0) { return SIZE_MAX; } diff --git a/src/model_manager.h b/src/model_manager.h index 7f8df59c9b..394391de8f 100644 --- a/src/model_manager.h +++ b/src/model_manager.h @@ -94,6 +94,8 @@ class ModelManager : public DeviceResidencyManager { std::vector> tensor_states_; std::map tensor_states_by_tensor_; mutable std::list resolved_tensor_states_; + // Last device free-memory reading per (owner, device), for DeviceMemoryRequest::reuse_device_query. + mutable std::map, std::pair> device_query_cache_; std::vector> params_storage_blocks_; std::vector> compute_staging_blocks_; std::map split_buffer_types_; diff --git a/src/runtime/tiling.cpp b/src/runtime/tiling.cpp index 1bfb291a70..9511789fc5 100644 --- a/src/runtime/tiling.cpp +++ b/src/runtime/tiling.cpp @@ -2,6 +2,8 @@ #include #include +#include +#include #include #include #include @@ -331,6 +333,29 @@ sd::Tensor process_tiles_2d_batched(const sd::Tensor& input, LOG_VERBOSE("processing %i tiles", num_tiles); pretty_progress(0, num_tiles, 0.0f); } + // The host-side blend of a batch runs on a worker thread while the next batch is split and + // computed, so the device does not idle through it. Merges still run one at a time in tile + // order with the same arithmetic, so the result is bit-identical. SD_TILE_ASYNC_MERGE=0 merges + // inline. + static const bool async_merge = []() { + const char* value = getenv("SD_TILE_ASYNC_MERGE"); + return value == nullptr || value[0] == '\0' || atoi(value) != 0; + }(); + std::thread merger; + bool merge_failed = false; + auto join_merger = [&]() { + if (merger.joinable()) { + merger.join(); + } + }; + struct MergerGuard { + std::thread& thread; + ~MergerGuard() { + if (thread.joinable()) { + thread.join(); + } + } + } merger_guard{merger}; std::vector> input_tiles; for (size_t first = 0; first < placements.size();) { const int remaining = static_cast(placements.size() - first); @@ -353,29 +378,47 @@ sd::Tensor process_tiles_2d_batched(const sd::Tensor& input, 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)); + auto merge_batch = [&, first, last, tile_shape](const sd::Tensor& data) { + const int64_t tile_numel = data.numel() / static_cast(last - first); + const int64_t plane_count = tile_numel / (tile_shape[0] * tile_shape[1]); + if (output.empty()) { + std::vector output_shape = 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(data.data() + static_cast(i - first) * tile_numel, + tile_shape[0], + tile_shape[1], + plane_count, + &output, + placement.x_out, + placement.y_out, + overlap_x_out, + overlap_y_out, + circular_x, + circular_y, + placement.dx, + placement.dy); + } + }; + join_merger(); + if (merge_failed) { + return {}; } - 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 (async_merge && last < placements.size()) { + merger = std::thread([&merge_failed, merge_batch, data = std::move(batch.data)]() { + try { + merge_batch(data); + } catch (const std::exception& error) { + LOG_ERROR("tile merge failed: %s", error.what()); + merge_failed = true; + } + }); + } else { + merge_batch(batch.data); } if (!silent) { @@ -390,6 +433,10 @@ sd::Tensor process_tiles_2d_batched(const sd::Tensor& input, } first = last; } + join_merger(); + if (merge_failed) { + return {}; + } if (!silent && tile_count < num_tiles) { pretty_progress(num_tiles, num_tiles, last_time); } From a6cd7826f31aa894cf231d16144b7ca30b3179b1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 22:45:04 -0700 Subject: [PATCH 16/30] MiniMax-H3: audio VAE anti-aliased activations use direct depthwise convs The up/down-sampling filters of every Activation1D went through ggml_conv_1d / ggml_conv_1d_dw: an F16 im2col plus a single-column matmul per call. Route them through CONV_2D_DW (H=1, F32 per-channel kernel) instead, for the MiniMax-H3 audio VAE only. Falls back to the im2col graph when the backend lacks the op. SD_H3_AUDIO_DIRECT_DW=0 restores the previous graph. --- docs/minimax_h3.md | 6 ++ src/model/vae/ltx_audio_vae.hpp | 84 ++++++++++++++++++++++++-- src/model/vae/minimax_h3_audio_vae.hpp | 16 ++++- 3 files changed, 97 insertions(+), 9 deletions(-) diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index 97bcaf8cc9..634128da29 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -84,6 +84,12 @@ F32 Q/K and F16 V layout the sage kernel reads (K and V carry the kv scale), so longer pays for the chunk / slice / rope / concat / scale / cast chain; the output is bit-identical to the unfused sage graph. `SD_H3_FAST_SAGE_QKV=0` restores that chain. +The audio VAE's anti-aliased activations run their up/down-sampling filters as direct F32 +depthwise convolutions (`CONV_2D_DW`) instead of an F16 im2col plus a matrix-vector product, +which makes the audio decode several times faster. The input is no longer rounded to F16, so +the waveform differs slightly from the previous graph (about 41 dB SNR on a 5 s clip). +`SD_H3_AUDIO_DIRECT_DW=0` restores the previous graph; backends without `CONV_2D_DW` keep it. + ## First/last-frame conditioning Add `--init-img` for I2VA, or both `--init-img` and `--end-img` for FL2VA: diff --git a/src/model/vae/ltx_audio_vae.hpp b/src/model/vae/ltx_audio_vae.hpp index d2314b56a7..7f3098fb3b 100644 --- a/src/model/vae/ltx_audio_vae.hpp +++ b/src/model/vae/ltx_audio_vae.hpp @@ -396,6 +396,74 @@ namespace LTXV { return ggml_reshape_4d(ctx, out, out_time, channels, 1, 1); } + // Direct depthwise conv (one F32 kernel per channel, no im2col, no F16 rounding of the input) for + // the anti-aliased activations. x: [time, channels, 1, 1]. Returns nullptr when the backend lacks + // CONV_2D_DW so callers fall back to the im2col graph. + static ggml_tensor* depthwise_conv1d_direct(GGMLRunnerContext* runner_ctx, + ggml_tensor* x, + ggml_tensor* filter, + int stride, + int padding, + bool reverse) { + auto ctx = runner_ctx->ggml_ctx; + if (runner_ctx->backend == nullptr || x->type != GGML_TYPE_F32 || x->ne[2] != 1 || x->ne[3] != 1) { + return nullptr; + } + const int64_t time = x->ne[0]; + const int64_t channels = x->ne[1]; + ggml_tensor* f = ggml_reshape_4d(ctx, filter, filter->ne[0], 1, 1, 1); + if (reverse) { + f = reverse_1d_filter(ctx, f); + } + if (f->type != GGML_TYPE_F32) { + f = ggml_cast(ctx, f, GGML_TYPE_F32); + } + auto kernel = repeat_with_vulkan_f32_workaround(runner_ctx->backend, ctx, f, f->ne[0], 1, 1, channels); // [K, 1, 1, C] + auto input = ggml_reshape_4d(ctx, ggml_ext_cont(ctx, x), time, 1, channels, 1); // [W, H=1, C, N=1] + auto out = ggml_conv_2d_dw_direct(ctx, kernel, input, stride, 1, padding, 0, 1, 1); + if (!ggml_backend_supports_op(runner_ctx->backend, out)) { + return nullptr; + } + return ggml_reshape_4d(ctx, out, out->ne[0], channels, 1, 1); + } + + // depthwise_conv_transpose1d through the direct depthwise conv: same zero-stuffing, reversed + // filter and full padding, then the same stride gain. + static ggml_tensor* depthwise_conv_transpose1d_direct(GGMLRunnerContext* runner_ctx, + ggml_tensor* x, + ggml_tensor* filter, + int stride) { + auto ctx = runner_ctx->ggml_ctx; + if (x->ne[2] != 1 || x->ne[3] != 1 || filter->ne[1] != 1 || filter->ne[2] != 1 || filter->ne[3] != 1) { + return nullptr; + } + const int64_t time = x->ne[0]; + const int64_t channels = x->ne[1]; + const int64_t kernel_size = filter->ne[0]; + const int64_t out_time = (time - 1) * stride + kernel_size; + + auto x_flat = ggml_reshape_3d(ctx, x, 1, time, channels); + if (stride > 1) { + auto zero_unit = ggml_ext_scale(ctx, x_flat, 0.0f); + auto zero_tail = zero_unit; + for (int i = 1; i < stride - 1; ++i) { + zero_tail = ggml_concat(ctx, zero_tail, zero_unit, 0); + } + x_flat = ggml_concat(ctx, x_flat, zero_tail, 0); + } + x_flat = ggml_reshape_4d(ctx, x_flat, time * stride, channels, 1, 1); + auto out = depthwise_conv1d_direct(runner_ctx, x_flat, filter, 1, static_cast(kernel_size - 1), true); + if (out == nullptr) { + return nullptr; + } + if (out->ne[0] > out_time) { + out = ggml_ext_slice(ctx, out, 0, 0, out_time); + } + GGML_ASSERT(out->ne[0] == out_time); + out = ggml_ext_scale(ctx, out, static_cast(stride)); + return ggml_reshape_4d(ctx, out, out_time, channels, 1, 1); + } + static ggml_tensor* upsample_waveform_hann(GGMLRunnerContext* runner_ctx, ggml_tensor* waveform, ggml_tensor* filter, @@ -667,8 +735,11 @@ namespace LTXV { int up_kernel_size = 12; int down_kernel_size = 12; - explicit Activation1D(int64_t channels) - : channels(channels) { + // direct depthwise convs instead of the im2col graph; opt-in per model (MiniMax-H3 sets it) + bool direct_dw = false; + + explicit Activation1D(int64_t channels, bool direct_dw = false) + : channels(channels), direct_dw(direct_dw) { blocks["act"] = std::make_shared(channels); } @@ -690,16 +761,17 @@ namespace LTXV { int up_pad_right = up_pad * up_ratio + (up_kernel_size - up_ratio + 1) / 2; x = replicate_pad_1d(ctx, x, up_pad, up_pad); - x = depthwise_conv_transpose1d(ctx, x, up_filter, up_ratio); - x = ggml_ext_slice(ctx->ggml_ctx, x, 0, up_pad_left, x->ne[0] - up_pad_right); + ggml_tensor* up = direct_dw ? depthwise_conv_transpose1d_direct(ctx, x, up_filter, up_ratio) : nullptr; + x = up != nullptr ? up : depthwise_conv_transpose1d(ctx, x, up_filter, up_ratio); + x = ggml_ext_slice(ctx->ggml_ctx, x, 0, up_pad_left, x->ne[0] - up_pad_right); x = act->forward(ctx, x); int down_pad_left = down_kernel_size / 2 - (down_kernel_size % 2 == 0 ? 1 : 0); int down_pad_right = down_kernel_size / 2; x = replicate_pad_1d(ctx, x, down_pad_left, down_pad_right); - x = depthwise_conv1d(ctx, x, down_filter, down_ratio, 0); - return x; + ggml_tensor* down = direct_dw ? depthwise_conv1d_direct(ctx, x, down_filter, down_ratio, 0, false) : nullptr; + return down != nullptr ? down : depthwise_conv1d(ctx, x, down_filter, down_ratio, 0); } }; diff --git a/src/model/vae/minimax_h3_audio_vae.hpp b/src/model/vae/minimax_h3_audio_vae.hpp index a4fc144736..300826a030 100644 --- a/src/model/vae/minimax_h3_audio_vae.hpp +++ b/src/model/vae/minimax_h3_audio_vae.hpp @@ -2,6 +2,7 @@ #define __SD_MODEL_VAE_MINIMAX_H3_AUDIO_VAE_HPP__ #include +#include #include #include @@ -10,6 +11,15 @@ namespace MiniMaxH3 { + // SD_H3_AUDIO_DIRECT_DW=0: anti-aliased activations go back to the im2col + mat-vec depthwise convs. + static bool audio_direct_dw() { + static const bool enabled = [] { + const char* v = std::getenv("SD_H3_AUDIO_DIRECT_DW"); + return v == nullptr || v[0] == '\0' || std::atoi(v) != 0; + }(); + return enabled; + } + struct AudioSnake1D : public UnaryBlock { int64_t channels; @@ -229,9 +239,9 @@ namespace MiniMaxH3 { : channels(channels) { for (int i = 0; i < 3; ++i) { blocks["activations." + std::to_string(i * 2)] = - std::make_shared(channels); + std::make_shared(channels, audio_direct_dw()); blocks["activations." + std::to_string(i * 2 + 1)] = - std::make_shared(channels); + std::make_shared(channels, audio_direct_dw()); blocks["convs1." + std::to_string(i)] = std::make_shared(channels, channels, @@ -298,7 +308,7 @@ namespace MiniMaxH3 { } channels = next_channels; } - blocks["activation_post"] = std::make_shared(channels); + blocks["activation_post"] = std::make_shared(channels, audio_direct_dw()); blocks["conv_post"] = std::make_shared(channels, 1, 7, From dcb9dde9c918c0f713f6faf70aca7fbd1c36a46e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 23:02:50 -0700 Subject: [PATCH 17/30] MiniMax-H3 video VAE: fused RoPE / head-major q/k/v for decoder attention The decoder attention chunk-copied q/k/v out of the fused projection, applied the partial RoPE as table mul/add/concat chains and then permuted and cast K/V for flash attention. With one tile per graph, normalise q/k in place on the projection and let ggml_rope_pe_permute (already used by the DiT) write the F32 Q and F16 K/V head-major tensors ggml_ext_attention_prepared reads. Same products, sums and casts as the table path, so the decoded frames are bit-identical. SD_H3_VAE_FUSED_QKV=0 restores the old graph; batched tiles, sage and non-flash decodes keep it. --- docs/minimax_h3.md | 6 +++- src/model/vae/minimax_h3_vae.hpp | 53 ++++++++++++++++++++++++++++++++ 2 files changed, 58 insertions(+), 1 deletion(-) diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index 634128da29..eed2229d2a 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -63,7 +63,11 @@ batched projections can round differently from the per-tile decode on some GPUs, opt-in. The decoder weights stay on the device across temporal chunks (`SD_H3_VAE_KEEP_RESIDENT=0` releases them after every chunk), and the decoder blocks use a table-based rotary embedding and a fused SwiGLU (`SD_H3_VAE_GRAPH_OPT=0` restores the previous -graph). Each batched tile still goes through its own attention call. +graph). Each batched tile still goes through its own attention call. With one tile per graph and +flash attention, each decoder attention normalises q/k in place on the projection and one fused +RoPE op per tensor writes the head-major Q (F32) and K/V (F16) that the attention kernel reads, +replacing the table RoPE, chunk copies, permutes and casts; the frames are bit-identical +(`SD_H3_VAE_FUSED_QKV=0` restores the unfused graph). The host side of the decode overlaps the device: each tile's blend into the frame runs on a worker thread while the next tile computes (`SD_TILE_ASYNC_MERGE=0` blends inline; this applies diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index fd09f4a744..d569090a21 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -289,6 +289,56 @@ namespace MiniMaxH3VAE { blocks["to_out"] = std::make_shared(dim, dim, true); } + // Single-tile flash-attention decode: q/k RMS-normed in place on the projection, then one + // ggml_rope_pe_permute per tensor applies the partial RoPE, writes the head-major layout the + // attention kernel reads and casts K/V to F16. Same products, sums and casts as the table + // RoPE + chunk / permute / cast chain, so the result is bit-identical. Returns nullptr when + // it does not apply (SD_H3_VAE_FUSED_QKV=0 turns it off). + ggml_tensor* forward_head_major(GGMLRunnerContext* ctx, + ggml_tensor* qkv, + ggml_tensor* pe, + ggml_tensor* rope_a, + ggml_tensor* rope_b) { +#ifdef SD_USE_UPSTREAM_GGML + return nullptr; +#else + static const bool enabled = [] { + const char* v = getenv("SD_H3_VAE_FUSED_QKV"); + return v == nullptr || v[0] == '\0' || atoi(v) != 0; + }(); + // rope_a / rope_b are only passed with SD_H3_VAE_GRAPH_OPT on; that path is what this replaces + if (!enabled || rope_a == nullptr || rope_b == nullptr || pe == nullptr || ctx->backend == nullptr || + !ctx->flash_attn_enabled || ctx->sage_attn_enabled || qkv->type != GGML_TYPE_F32 || + qkv->ne[3] != 1 || qkv->nb[0] != sizeof(float)) { + return nullptr; + } + const int64_t sequence = qkv->ne[2]; + const int n_rot = static_cast(pe->ne[2] * 2); + if (pe->type != GGML_TYPE_F32 || !ggml_is_contiguous(pe) || pe->ne[0] != 2 || pe->ne[1] != 2 || + pe->ne[3] != sequence || n_rot > head_dim) { + return nullptr; + } + const float kv_scale = ctx->attn_scale > 0.f ? ctx->attn_scale : 1.f; + auto part = [&](int index) { + return ggml_view_4d(ctx->ggml_ctx, qkv, head_dim, num_head, sequence, 1, + qkv->nb[1], qkv->nb[2], qkv->nb[3], index * head_dim * sizeof(float)); + }; + auto q = ggml_rms_norm(ctx->ggml_ctx, part(0), 1e-5f); + auto k = ggml_rms_norm(ctx->ggml_ctx, part(1), 1e-5f); + 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, num_head); + k = ggml_reshape_3d(ctx->ggml_ctx, k, head_dim, sequence, num_head); + v = ggml_reshape_3d(ctx->ggml_ctx, v, head_dim, sequence, num_head); + return ggml_ext_attention_prepared(ctx->ggml_ctx, ctx->backend, q, k, v, num_head, 1, kv_scale); +#endif + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, @@ -312,6 +362,9 @@ namespace MiniMaxH3VAE { num_head, sequence, batch_size); + if (auto fused = forward_head_major(ctx, qkv_projection, pe, rope_a, rope_b)) { + return project_out(fused); + } auto qkv = ggml_ext_chunk(ctx->ggml_ctx, qkv_projection, 3, 0); auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv[0], From 82ea4be74b4d231b6c19f2d09a3ca4c30675251a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 23:16:40 -0700 Subject: [PATCH 18/30] MiniMax-H3 video VAE: fold the q/k RMS norm into the fused RoPE op Carry ggml patch 0005 (ggml_rope_pe_permute_rms) and use it for the decoder's q/k: the per-head norm was 7560 launches of a 64-thread kernel, about 0.45 s per 960x544x124 decode on B200. The fused op normalises each head row with the same block reduction as the standalone norm, so the frames stay bit-identical. Backends without it (HIP, Vulkan, GGML_CUDA_NORM_SMALL_ROWS=0) keep the separate norm; SD_H3_VAE_FUSED_QK_NORM=0 forces that. --- docs/minimax_h3.md | 4 +- .../0005-ggml-rope-pe-permute-rms.patch | 304 ++++++++++++++++++ src/model/vae/minimax_h3_vae.hpp | 23 +- 3 files changed, 326 insertions(+), 5 deletions(-) create mode 100644 scripts/unsloth/ggml-patches/0005-ggml-rope-pe-permute-rms.patch diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index eed2229d2a..03ce876d9c 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -67,7 +67,9 @@ graph). Each batched tile still goes through its own attention call. With one ti flash attention, each decoder attention normalises q/k in place on the projection and one fused RoPE op per tensor writes the head-major Q (F32) and K/V (F16) that the attention kernel reads, replacing the table RoPE, chunk copies, permutes and casts; the frames are bit-identical -(`SD_H3_VAE_FUSED_QKV=0` restores the unfused graph). +(`SD_H3_VAE_FUSED_QKV=0` restores the unfused graph). On CUDA the q/k RMS norm runs inside that +RoPE op with the same reduction as the standalone norm (`SD_H3_VAE_FUSED_QK_NORM=0` keeps it +separate). The host side of the decode overlaps the device: each tile's blend into the frame runs on a worker thread while the next tile computes (`SD_TILE_ASYNC_MERGE=0` blends inline; this applies diff --git a/scripts/unsloth/ggml-patches/0005-ggml-rope-pe-permute-rms.patch b/scripts/unsloth/ggml-patches/0005-ggml-rope-pe-permute-rms.patch new file mode 100644 index 0000000000..65b336de54 --- /dev/null +++ b/scripts/unsloth/ggml-patches/0005-ggml-rope-pe-permute-rms.patch @@ -0,0 +1,304 @@ +From 092f5d51d0844f7cf282ca115a81fccdde9fafad Mon Sep 17 00:00:00 2001 +From: Daniel Han +Date: Thu, 1 Oct 2026 23:16:25 -0700 +Subject: [PATCH] ggml : rope_pe_permute with the per-head RMS norm folded in + +ggml_rope_pe_permute_rms(a, pe, n_rot, scale, type, eps) is ggml_rope_pe_permute of +ggml_rms_norm(a, eps) without a weight. CPU uses rms_norm's arithmetic (double sum, +1/sqrtf, vec scale); CUDA normalises each head row with the block size and block +reduction rms_norm_f32_cuda's small-row path uses for that row length, so both match +the two-op chain bit for bit. CUDA supports it for head_dim <= 128 when that small-row +path is active (not HIP/MUSA, GGML_CUDA_NORM_SMALL_ROWS not 0); elsewhere the op is +reported unsupported and callers keep the separate norm. test-backend-ops cases added. +--- + include/ggml.h | 12 ++++ + src/ggml-cpu/ops.cpp | 17 +++++- + src/ggml-cuda/rope-pe.cu | 120 +++++++++++++++++++++++++++++++++++++ + src/ggml.c | 15 +++++ + tests/test-backend-ops.cpp | 15 +++-- + 5 files changed, 174 insertions(+), 5 deletions(-) + +diff --git a/include/ggml.h b/include/ggml.h +index 153eed2a..8134626f 100644 +--- a/include/ggml.h ++++ b/include/ggml.h +@@ -2483,6 +2483,18 @@ extern "C" { + float scale, + enum ggml_type type); + ++ // ggml_rope_pe_permute of ggml_rms_norm(a, eps) (no weight): every head row of `a` is ++ // RMS-normalised first, with the same arithmetic as GGML_OP_RMS_NORM on the same backend, so ++ // the result matches the two-op chain bit for bit where the backend supports it. ++ GGML_API struct ggml_tensor * ggml_rope_pe_permute_rms( ++ struct ggml_context * ctx, ++ struct ggml_tensor * a, ++ struct ggml_tensor * pe, ++ int n_rot, ++ float scale, ++ enum ggml_type type, ++ float eps); ++ + // Row-range modulation for adaLN-style blocks whose sequence is split into segments that use + // different modulation vectors. Writes rows [row0, row1) along dim 1 (all of dims 2 and 3) of the + // result and leaves every other row as it is in `acc`: chain one call per segment, passing the +diff --git a/src/ggml-cpu/ops.cpp b/src/ggml-cpu/ops.cpp +index 40fe1949..1d67b482 100644 +--- a/src/ggml-cpu/ops.cpp ++++ b/src/ggml-cpu/ops.cpp +@@ -12151,6 +12151,8 @@ void ggml_compute_forward_rope_pe_permute( + + const int n_rot = ggml_get_op_params_i32(dst, 0); + const float scale = ggml_get_op_params_f32(dst, 1); ++ const bool norm = ggml_get_op_params_i32(dst, 2) != 0; ++ const float eps = ggml_get_op_params_f32(dst, 3); + const int half = n_rot / 2; + + const int64_t D = src0->ne[0]; +@@ -12165,7 +12167,7 @@ void ggml_compute_forward_rope_pe_permute( + + // Same rounding points as the unfused graph: each product, the sum, the scale and the + // conversion are separate passes, so no contraction into fused multiply-adds can happen. +- std::vector p0(n_rot), p1(n_rot), y(D); ++ std::vector p0(n_rot), p1(n_rot), y(D), xn(norm ? D : 0); + + for (int64_t ir = ir0; ir < ir1; ++ir) { + const int64_t h = ir % H; +@@ -12174,6 +12176,19 @@ void ggml_compute_forward_rope_pe_permute( + + const float * x = (const float *) ((const char *) src0->data + h * src0->nb[1] + l * src0->nb[2] + n * src0->nb[3]); + ++ if (norm) { ++ // ggml_compute_forward_rms_norm_f32 without a weight ++ ggml_float sum = 0.0; ++ for (int64_t i = 0; i < D; i++) { ++ sum += (ggml_float)(x[i] * x[i]); ++ } ++ const float mean = sum/D; ++ const float s = 1.0f/sqrtf(mean + eps); ++ memcpy(xn.data(), x, D * sizeof(float)); ++ ggml_vec_scale_f32((int) D, xn.data(), s); ++ x = xn.data(); ++ } ++ + if (n_rot > 0) { + const float * t = (const float *) pe->data + l * (int64_t) half * 4; + for (int i = 0; i < n_rot; ++i) { +diff --git a/src/ggml-cuda/rope-pe.cu b/src/ggml-cuda/rope-pe.cu +index 3afe4423..b6dd7a69 100644 +--- a/src/ggml-cuda/rope-pe.cu ++++ b/src/ggml-cuda/rope-pe.cu +@@ -44,6 +44,112 @@ static __global__ void rope_pe_permute_kernel( + } + } + ++// Same op with every head row RMS-normalised first (no weight). One block of block_size threads ++// per (token, batch), the size and reduction rms_norm_f32_cuda's small-row path uses for this row ++// length, so each normalised value is bit-identical to GGML_OP_RMS_NORM's; the normalised row is ++// staged in shared memory for the RoPE partner reads. ++template ++static __global__ void rope_pe_permute_rms_kernel( ++ const char * __restrict__ x, const float * __restrict__ pe, dst_t * __restrict__ dst, ++ const int D, const int H, const int64_t L, const int n_rot, ++ const int64_t nb01, const int64_t nb02, const int64_t nb03, ++ const float scale, const float bias, const float eps) { ++ const int64_t l = blockIdx.x; ++ const int64_t n = blockIdx.y; ++ const int half = n_rot / 2; ++ const int i = threadIdx.x; ++ ++ __shared__ float s_sum[32]; ++ __shared__ float xs[block_size]; ++ ++ const bool rot = i < n_rot; ++ int j = i; ++ float c0 = 0.0f; ++ float c1 = 0.0f; ++ if (rot) { ++ const int r = i >= half; ++ j = i - r * half; ++ const float * t = pe + (l * half + j) * 4 + 2 * r; ++ c0 = t[0]; ++ c1 = t[1]; ++ } ++ for (int h = 0; h < H; ++h) { ++ const float * xr = (const float *) (x + h * nb01 + l * nb02 + n * nb03); ++ float xi = 0.0f; ++ float tmp = 0.0f; ++ if (i < D) { ++ xi = xr[i]; ++ tmp += xi * xi; ++ } ++ tmp = block_reduce(tmp, s_sum); ++ const float mean = tmp / D; ++ const float sc = rsqrtf(mean + eps); ++ if (i < D) { ++ xs[i] = sc * xi; ++ } ++ __syncthreads(); ++ if (i < D) { ++ float y; ++ if (rot) { ++ y = __fadd_rn(__fmul_rn(xs[j], c0), __fmul_rn(xs[j + half], c1)); ++ } else { ++ y = xs[i]; ++ } ++ if constexpr (has_scale) { ++ y = scale * y + bias; ++ } ++ dst[((n * H + h) * L + l) * D + i] = ggml_cuda_cast(y); ++ } ++ __syncthreads(); ++ } ++} ++ ++// rms_norm_f32_cuda takes its small-row path (and so its reduction order) only here. ++static bool rope_pe_rms_small_rows_enabled() { ++#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) ++ return false; ++#else ++ static const bool enabled = [] { ++ const char * e = getenv("GGML_CUDA_NORM_SMALL_ROWS"); ++ return e == nullptr || atoi(e) != 0; ++ }(); ++ return enabled && WARP_SIZE == 32; ++#endif ++} ++ ++template ++static void rope_pe_permute_rms_launch(const ggml_tensor * src0, const float * pe, dst_t * dst, const int n_rot, ++ const float scale, const float eps, cudaStream_t stream) { ++ const int D = (int) src0->ne[0]; ++ const int H = (int) src0->ne[1]; ++ const int64_t L = src0->ne[2]; ++ const int64_t N = src0->ne[3]; ++ const dim3 grid((unsigned int) L, (unsigned int) N, 1); ++ const dim3 block(block_size, 1, 1); ++ const char * x = (const char *) src0->data; ++ if (scale != 1.0f) { ++ rope_pe_permute_rms_kernel<<>>( ++ x, pe, dst, D, H, L, n_rot, src0->nb[1], src0->nb[2], src0->nb[3], scale, 0.0f, eps); ++ } else { ++ rope_pe_permute_rms_kernel<<>>( ++ x, pe, dst, D, H, L, n_rot, src0->nb[1], src0->nb[2], src0->nb[3], scale, 0.0f, eps); ++ } ++} ++ ++template ++static void rope_pe_permute_rms_cuda(const ggml_tensor * src0, const float * pe, dst_t * dst, const int n_rot, ++ const float scale, const float eps, cudaStream_t stream) { ++ const int D = (int) src0->ne[0]; ++ GGML_ASSERT(D <= 128); ++ if (D <= 32) { ++ rope_pe_permute_rms_launch(src0, pe, dst, n_rot, scale, eps, stream); ++ } else if (D <= 64) { ++ rope_pe_permute_rms_launch(src0, pe, dst, n_rot, scale, eps, stream); ++ } else { ++ rope_pe_permute_rms_launch(src0, pe, dst, n_rot, scale, eps, stream); ++ } ++} ++ + template + static void rope_pe_permute_cuda(const ggml_tensor * src0, const float * pe, dst_t * dst, const int n_rot, + const float scale, cudaStream_t stream) { +@@ -78,6 +184,17 @@ void ggml_cuda_op_rope_pe_permute(ggml_backend_cuda_context & ctx, ggml_tensor * + cudaStream_t stream = ctx.stream(); + const float * pe_d = pe ? (const float *) pe->data : nullptr; + ++ if (ggml_get_op_params_i32(dst, 2) != 0) { ++ const float eps = ggml_get_op_params_f32(dst, 3); ++ if (dst->type == GGML_TYPE_F16) { ++ rope_pe_permute_rms_cuda(src0, pe_d, (half *) dst->data, n_rot, scale, eps, stream); ++ } else { ++ GGML_ASSERT(dst->type == GGML_TYPE_F32); ++ rope_pe_permute_rms_cuda(src0, pe_d, (float *) dst->data, n_rot, scale, eps, stream); ++ } ++ return; ++ } ++ + if (dst->type == GGML_TYPE_F16) { + rope_pe_permute_cuda(src0, pe_d, (half *) dst->data, n_rot, scale, stream); + } else { +@@ -98,5 +215,8 @@ bool ggml_cuda_rope_pe_permute_supported(const ggml_tensor * op) { + if (src0->ne[2] > INT_MAX || src0->ne[3] > 65535 || src0->ne[1] > INT_MAX || src0->ne[0] > INT_MAX) { + return false; + } ++ if (ggml_get_op_params_i32(op, 2) != 0 && (src0->ne[0] > 128 || !rope_pe_rms_small_rows_enabled())) { ++ return false; ++ } + return pe == nullptr || (pe->type == GGML_TYPE_F32 && ggml_is_contiguous(pe)); + } +diff --git a/src/ggml.c b/src/ggml.c +index fec933bf..62a157c6 100644 +--- a/src/ggml.c ++++ b/src/ggml.c +@@ -5594,6 +5594,21 @@ struct ggml_tensor * ggml_rope_pe_permute( + return result; + } + ++struct ggml_tensor * ggml_rope_pe_permute_rms( ++ struct ggml_context * ctx, ++ struct ggml_tensor * a, ++ struct ggml_tensor * pe, ++ int n_rot, ++ float scale, ++ enum ggml_type type, ++ float eps) { ++ GGML_ASSERT(eps >= 0.0f); ++ struct ggml_tensor * result = ggml_rope_pe_permute(ctx, a, pe, n_rot, scale, type); ++ ggml_set_op_params_i32(result, 2, 1); ++ ggml_set_op_params_f32(result, 3, eps); ++ return result; ++} ++ + static struct ggml_tensor * ggml_modulate_rows_impl( + struct ggml_context * ctx, + struct ggml_tensor * acc, +diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp +index 1e46361c..b58c815c 100644 +--- a/tests/test-backend-ops.cpp ++++ b/tests/test-backend-ops.cpp +@@ -2132,15 +2132,16 @@ struct test_rope_pe_permute : public test_case { + const int n_rot; + const float scale; + const bool view; // a is a strided view (one third of a fused qkv projection) ++ const float eps; // >= 0: ggml_rope_pe_permute_rms (rows RMS-normalised first) + + std::string vars() override { +- return VARS_TO_STR5(type, ne_a, n_rot, scale, view); ++ return VARS_TO_STR6(type, ne_a, n_rot, scale, view, eps); + } + + test_rope_pe_permute(ggml_type type = GGML_TYPE_F32, + std::array ne_a = {128, 4, 9, 1}, +- int n_rot = 96, float scale = 1.0f, bool view = false) +- : type(type), ne_a(ne_a), n_rot(n_rot), scale(scale), view(view) {} ++ int n_rot = 96, float scale = 1.0f, bool view = false, float eps = -1.0f) ++ : type(type), ne_a(ne_a), n_rot(n_rot), scale(scale), view(view), eps(eps) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * a; +@@ -2159,7 +2160,8 @@ struct test_rope_pe_permute : public test_case { + pe = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 2, 2, n_rot / 2, ne_a[2]); + ggml_set_name(pe, "pe"); + } +- ggml_tensor * out = ggml_rope_pe_permute(ctx, a, pe, n_rot, scale, type); ++ ggml_tensor * out = eps >= 0.0f ? ggml_rope_pe_permute_rms(ctx, a, pe, n_rot, scale, type, eps) ++ : ggml_rope_pe_permute(ctx, a, pe, n_rot, scale, type); + ggml_set_name(out, "out"); + return out; + } +@@ -8221,6 +8223,11 @@ static std::vector> make_test_cases_eval() { + test_cases.emplace_back(new test_rope_pe_permute(type, {128, 4, 9, 1}, 0, scale, view)); + test_cases.emplace_back(new test_rope_pe_permute(type, {64, 3, 7, 2}, 32, scale, view)); + test_cases.emplace_back(new test_rope_pe_permute(type, {200, 2, 3, 1}, 100, scale, view)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {64, 32, 37, 1}, 48, scale, view, 1e-5f)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {64, 3, 7, 2}, 0, scale, view, 1e-5f)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {128, 4, 9, 1}, 96, scale, view, 1e-6f)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {32, 2, 5, 1}, 16, scale, view, 1e-5f)); ++ test_cases.emplace_back(new test_rope_pe_permute(type, {48, 2, 5, 1}, 32, scale, view, 1e-5f)); + } + } + } +-- +2.43.0 + diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index d569090a21..d7cc87b952 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -323,10 +323,25 @@ namespace MiniMaxH3VAE { return ggml_view_4d(ctx->ggml_ctx, qkv, head_dim, num_head, sequence, 1, qkv->nb[1], qkv->nb[2], qkv->nb[3], index * head_dim * sizeof(float)); }; - auto q = ggml_rms_norm(ctx->ggml_ctx, part(0), 1e-5f); - auto k = ggml_rms_norm(ctx->ggml_ctx, part(1), 1e-5f); - 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); + // q/k RMS norm folded into the RoPE op where the backend reproduces GGML_OP_RMS_NORM's + // arithmetic in it (SD_H3_VAE_FUSED_QK_NORM=0 keeps the separate norm). + static const bool fuse_norm = [] { + const char* v = getenv("SD_H3_VAE_FUSED_QK_NORM"); + return v == nullptr || v[0] == '\0' || atoi(v) != 0; + }(); + ggml_tensor* q = nullptr; + ggml_tensor* k = nullptr; + if (fuse_norm) { + q = ggml_rope_pe_permute_rms(ctx->ggml_ctx, part(0), pe, n_rot, 1.f, GGML_TYPE_F32, 1e-5f); + k = ggml_rope_pe_permute_rms(ctx->ggml_ctx, part(1), pe, n_rot, kv_scale, GGML_TYPE_F16, 1e-5f); + if (!ggml_backend_supports_op(ctx->backend, q) || !ggml_backend_supports_op(ctx->backend, k)) { + q = k = nullptr; + } + } + if (q == nullptr) { + q = ggml_rope_pe_permute(ctx->ggml_ctx, ggml_rms_norm(ctx->ggml_ctx, part(0), 1e-5f), pe, n_rot, 1.f, GGML_TYPE_F32); + k = ggml_rope_pe_permute(ctx->ggml_ctx, ggml_rms_norm(ctx->ggml_ctx, part(1), 1e-5f), pe, n_rot, kv_scale, GGML_TYPE_F16); + } auto v = ggml_rope_pe_permute(ctx->ggml_ctx, part(2), nullptr, 0, kv_scale, GGML_TYPE_F16); if (!ggml_backend_supports_op(ctx->backend, q) || !ggml_backend_supports_op(ctx->backend, k) || !ggml_backend_supports_op(ctx->backend, v)) { From 431871f9e275a7777127b21183b5bf91eeb49cf3 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 23:21:13 -0700 Subject: [PATCH 19/30] MiniMax-H3 video VAE: SD_H3_VAE_TILE opt-in latent tile size Fewer, larger tiles cut the per-tile fixed cost (20x20 decodes 960x544x124 in 3.07 s vs 3.77 s at 16x16 on B200) but move the seams: 36.2 dB PSNR / 0.09 LPIPS vs the 16x16 decode, slightly closer to an untiled decode than 16x16 is. Default stays 16. --- docs/minimax_h3.md | 4 +++- src/model/vae/minimax_h3_vae.hpp | 8 ++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index 03ce876d9c..7a8255db71 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -56,7 +56,9 @@ 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 one 16x16 latent tile per decoder graph by default. +The video VAE decodes one 16x16 latent tile per decoder graph by default. `SD_H3_VAE_TILE=N` +uses N x N latent tiles instead (20 decodes about 0.7 s faster at 960x544x124 on B200); the tile +seams move, so the frames differ from the default (36 dB PSNR at 20) and it is opt-in. `SD_H3_VAE_TILE_BATCH=auto` puts several tiles into one graph, sized from free device memory (at most 4 unless `SD_H3_VAE_TILE_BATCH_MAX` raises it), and `SD_H3_VAE_TILE_BATCH=N` forces N; the batched projections can round differently from the per-tile decode on some GPUs, so it is diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index d7cc87b952..8412ebf1a3 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -746,6 +746,14 @@ namespace MiniMaxH3VAE { params.tile_size_x = 16; params.tile_size_y = 16; params.target_overlap = 0.25f; + // SD_H3_VAE_TILE=N: N x N latent tiles instead of 16 x 16 (opt-in; the tile seams move, so + // the frames are not bit-identical to the default) + if (const char* tile = getenv("SD_H3_VAE_TILE")) { + const int n = atoi(tile); + if (n >= 8) { + params.tile_size_x = params.tile_size_y = n; + } + } return params; } From ee09dfa75e3ebb9088fd951ab01766f5c012a7da Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 23:38:09 -0700 Subject: [PATCH 20/30] Tile merge / H3 assembly / capacity reading: harden the reuse paths - The saved device free-memory reading is now taken only by a check that allocates nothing, so it is never the reading from before this owner's compute buffer was allocated, and any check that does not fit drops the owner's readings so the reclaim / evict retries see fresh ones. - If a worker thread cannot be started, the tile merge and the chunk assembly run inline instead of throwing. --- src/model/vae/minimax_h3_vae.hpp | 12 +++++++++--- src/model_manager.cpp | 21 ++++++++++++++++++--- src/runtime/tiling.cpp | 29 ++++++++++++++++++++--------- 3 files changed, 47 insertions(+), 15 deletions(-) diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index 8412ebf1a3..f4cff90af5 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -994,9 +995,14 @@ namespace MiniMaxH3VAE { break; } if (async_assembly && !tile.last) { - assembler = std::thread([&run_assembly, decoded = std::move(decoded), tile]() { - run_assembly(decoded, tile); - }); + auto chunk = std::make_shared>(std::move(decoded)); + try { + assembler = std::thread([&run_assembly, chunk, tile]() { + run_assembly(*chunk, tile); + }); + } catch (const std::system_error&) { + run_assembly(*chunk, tile); // no thread available: assemble inline + } } else { run_assembly(decoded, tile); } diff --git a/src/model_manager.cpp b/src/model_manager.cpp index d6cd88a480..a3554da546 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -1633,15 +1633,17 @@ ModelManager::CapacityCheck ModelManager::check_capacity( return SIZE_MAX; } size_t free_bytes = 0, total_bytes = 0; + // A reading is saved and reused only for checks that allocate nothing (no pending bytes, + // every parameter resident), so it is never one taken before this owner's buffers grew. + const bool reusable = request.reuse_device_query && request.pending_allocation_bytes == 0 && missing == 0; const auto cache_key = std::make_pair(request.owner_id, device); auto cached = device_query_cache_.find(cache_key); - if (request.reuse_device_query && request.pending_allocation_bytes == 0 && missing == 0 && - cached != device_query_cache_.end()) { + if (reusable && cached != device_query_cache_.end()) { free_bytes = cached->second.first; total_bytes = cached->second.second; } else { ggml_backend_dev_memory(device, &free_bytes, &total_bytes); - if (request.reuse_device_query) { + if (reusable) { device_query_cache_[cache_key] = {free_bytes, total_bytes}; } else if (cached != device_query_cache_.end()) { device_query_cache_.erase(cached); @@ -1663,6 +1665,19 @@ ModelManager::CapacityCheck ModelManager::check_capacity( return free_bytes; }; result.available_device_bytes = available_device_bytes(request.compute_backend); + struct DropReadingIfNoFit { + const CapacityCheck& result; + std::map, std::pair>& cache; + uintptr_t owner; + ~DropReadingIfNoFit() { + // A failed check leads to reclaiming / evicting; the retries must see fresh readings. + if (!result.fits()) { + for (auto it = cache.begin(); it != cache.end();) { + it = it->first.first == owner ? cache.erase(it) : std::next(it); + } + } + } + } drop_reading_if_no_fit{result, device_query_cache_, request.owner_id}; if (request.max_backend_bytes > 0) { const size_t resident = add(compute_backend_resident_bytes(request.compute_backend), other_runtime_resident_bytes(request.owner_id, request.compute_backend)); diff --git a/src/runtime/tiling.cpp b/src/runtime/tiling.cpp index 9511789fc5..1dcf35cd55 100644 --- a/src/runtime/tiling.cpp +++ b/src/runtime/tiling.cpp @@ -4,6 +4,8 @@ #include #include #include +#include +#include #include #include #include @@ -408,16 +410,25 @@ sd::Tensor process_tiles_2d_batched(const sd::Tensor& input, if (merge_failed) { return {}; } + bool merged_async = false; if (async_merge && last < placements.size()) { - merger = std::thread([&merge_failed, merge_batch, data = std::move(batch.data)]() { - try { - merge_batch(data); - } catch (const std::exception& error) { - LOG_ERROR("tile merge failed: %s", error.what()); - merge_failed = true; - } - }); - } else { + auto data = std::make_shared>(std::move(batch.data)); + try { + merger = std::thread([&merge_failed, merge_batch, data]() { + try { + merge_batch(*data); + } catch (const std::exception& error) { + LOG_ERROR("tile merge failed: %s", error.what()); + merge_failed = true; + } + }); + merged_async = true; + } catch (const std::system_error&) { + merge_batch(*data); // no thread available: merge inline + merged_async = true; + } + } + if (!merged_async) { merge_batch(batch.data); } From def3cb3f6b2c508f7bc4475b7655484bd1a02c44 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 3 Oct 2026 19:23:21 -0700 Subject: [PATCH 21/30] MiniMax-H3 / ggml-cuda: cuDNN fused attention on Ampere and newer New ggml patch 0006 adds an optional cuDNN path to ggml-cuda's FLASH_ATTN_EXT and SAGE_ATTN (AUTO mode): GGML_CUDA_CUDNN=ON compiles it against cudnn.h and the header-only cudnn-frontend v1.26.0 (MIT, downloaded by CMake and pinned by checksum); libcudnn.so.9 is opened at run time, so the binaries keep no cuDNN dependency and run the existing kernels wherever cuDNN cannot be loaded, on Turing and older, for small attention calls and for shapes cuDNN declines. GGML_CUDA_CUDNN_ATTN=0 restores the previous kernels everywhere and GGML_CUDA_CUDNN_SAGE=0 keeps the sage kernel for --sage-attn. On a B200 the H3 DiT attention at 960x544x124 (19315 tokens x 56 heads) takes about 8 ms per call against 28 ms for the sage kernel and 43 ms for the ggml flash attention kernel, and is closer to an fp64 reference than either. The Linux CUDA prebuilt fetches the pinned cuDNN 9.27 headers package, builds with GGML_CUDA_CUDNN=ON and fails if either binary ends up linked to libcudnn. docs/minimax_h3.md describes the run-time requirements. --- .github/workflows/unsloth-sd-prebuilt.yml | 33 + docs/minimax_h3.md | 18 + .../0006-ggml-cuda-cudnn-attention.patch | 972 ++++++++++++++++++ 3 files changed, 1023 insertions(+) create mode 100644 scripts/unsloth/ggml-patches/0006-ggml-cuda-cudnn-attention.patch diff --git a/.github/workflows/unsloth-sd-prebuilt.yml b/.github/workflows/unsloth-sd-prebuilt.yml index 81708c3f11..a27a2c004a 100644 --- a/.github/workflows/unsloth-sd-prebuilt.yml +++ b/.github/workflows/unsloth-sd-prebuilt.yml @@ -433,6 +433,29 @@ jobs: # keeps what it compiled. save: false + - name: cuDNN headers (build time only) + run: | + set -euo pipefail + # ggml-cuda's cuDNN attention path (GGML_CUDA_CUDNN) compiles against cudnn.h, but the + # library itself is opened with dlopen when the first attention op runs and is NOT + # bundled: a host without a loadable libcudnn.so.9 keeps the ggml attention kernels. + # Only the headers package is fetched (80 KB), pinned by version and checksum. + deb=libcudnn9-headers-cuda-12_9.27.0.42-1_amd64.deb + curl -fsSL --retry 3 -o "$RUNNER_TEMP/$deb" \ + "https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/$deb" + echo "68243ab04422da565e0bafec00f64ac08cfd82d805c8015c0c503af7ab4dd1d6 $RUNNER_TEMP/$deb" | sha256sum -c - + dpkg -x "$RUNNER_TEMP/$deb" "$RUNNER_TEMP/cudnn-deb" + # The package ships cudnn*_v9.h; the unversioned names normally come from the -dev + # package's alternatives, so link them here. + inc="$RUNNER_TEMP/cudnn-include" + mkdir -p "$inc" + for f in "$RUNNER_TEMP"/cudnn-deb/usr/include/x86_64-linux-gnu/cudnn*_v9.h; do + b="$(basename "$f")" + ln -s "$f" "$inc/${b%_v9.h}.h" + done + test -e "$inc/cudnn.h" + echo "CUDNN_INCLUDE_DIR=$inc" >> "$GITHUB_ENV" + - name: Build sd-cli + sd-server (CUDA) working-directory: src run: | @@ -451,6 +474,8 @@ jobs: -DSD_WEBP=OFF -DSD_WEBM=OFF \ -DGGML_NATIVE=OFF \ -DSD_CUDA=ON \ + -DGGML_CUDA_CUDNN=ON \ + -DGGML_CUDA_CUDNN_INCLUDE_DIR="$CUDNN_INCLUDE_DIR" \ -DCMAKE_CUDA_ARCHITECTURES="$CUDA_ARCHS" \ -DCMAKE_C_COMPILER_LAUNCHER=ccache \ -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \ @@ -479,6 +504,14 @@ jobs: patchelf --set-rpath '$ORIGIN' "$BIN/$exe" done ldd "$BIN/sd-cli" | sed -n '1,40p' + # cuDNN is opened at run time only; a link-time dependency would stop the binaries + # loading on every host without it. + for exe in sd-cli sd-server; do + if readelf -d "$BIN/$exe" | grep -q 'NEEDED.*cudnn'; then + echo "ERROR: $exe links libcudnn; it must stay a run-time dlopen" >&2 + exit 1 + fi + done - name: Package bundle env: diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index 7a8255db71..40b3203c7f 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -56,6 +56,24 @@ 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. +Builds configured with `-DGGML_CUDA_CUDNN=ON` (the Linux CUDA prebuilt is) can run that unmasked +attention, and the `--sage-attn` attention, through cuDNN's fused attention on Ampere and newer. +Only `cudnn.h` is needed at build time (`GGML_CUDA_CUDNN_INCLUDE_DIR`; the header-only +cudnn-frontend v1.26.0 is downloaded by CMake unless `GGML_CUDA_CUDNN_FRONTEND_DIR` points at a +checkout). `libcudnn.so.9` is opened when the first attention op runs and is not shipped: put a cuDNN 9 +built for the same CUDA major version as the binary (cuDNN 9 for CUDA 12 for the prebuilt) on the +library path or beside the binary, or point `GGML_CUDA_CUDNN_LIB` at it (cuDNN 9.27 needs nothing +else; older 9.x releases build these kernels with NVRTC and also need that CUDA major's `libnvrtc`). +Without a loadable cuDNN, on Turing and older, for attention calls under about a million scores +(where the extra conversions cost more than cuDNN saves) and for any shape cuDNN declines, the +existing kernels run. The first call of each attention shape builds a cuDNN plan (about 0.4 to 1.3 s +on a B200, logged once). cuDNN runs F16 with F32 accumulation: for the H3 DiT at 960x544x124 it +takes about 8 ms per attention call on a B200 against 28 ms for the sage kernel and 43 ms for the +ggml flash attention kernel, and it is closer to an exact (fp64) reference than either, so frames +and audio differ slightly from the previous kernels. `GGML_CUDA_CUDNN_ATTN=0` restores them, +`GGML_CUDA_CUDNN_SAGE=0` keeps the sage kernel for `--sage-attn` only, and +`GGML_CUDA_CUDNN_ATTN_BF16=1` runs cuDNN in BF16 (less accurate, same speed). + The video VAE decodes one 16x16 latent tile per decoder graph by default. `SD_H3_VAE_TILE=N` uses N x N latent tiles instead (20 decodes about 0.7 s faster at 960x544x124 on B200); the tile seams move, so the frames differ from the default (36 dB PSNR at 20) and it is opt-in. diff --git a/scripts/unsloth/ggml-patches/0006-ggml-cuda-cudnn-attention.patch b/scripts/unsloth/ggml-patches/0006-ggml-cuda-cudnn-attention.patch new file mode 100644 index 0000000000..7027a8e995 --- /dev/null +++ b/scripts/unsloth/ggml-patches/0006-ggml-cuda-cudnn-attention.patch @@ -0,0 +1,972 @@ +From 80d308615f1dc9fc601f556123861a651b7cffa5 Mon Sep 17 00:00:00 2001 +From: Daniel Han +Date: Sat, 3 Oct 2026 18:27:29 -0700 +Subject: [PATCH] ggml-cuda: optional cuDNN attention for unmasked + FLASH_ATTN_EXT and SAGE_ATTN + +GGML_CUDA_CUDNN=ON (default OFF, Linux) builds a cuDNN fused-attention path +through the header-only cudnn-frontend v1.26.0 (MIT; fetched by CMake or taken +from GGML_CUDA_CUDNN_FRONTEND_DIR) with NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING: +only cudnn.h is needed at build time and libcudnn.so.9 is dlopen'ed on first use +(GGML_CUDA_CUDNN_LIB overrides the name), so the binary has no cuDNN link +dependency and runs unchanged where cuDNN is missing. + +On NVIDIA sm80+ it takes unmasked, sink/ALiBi/softcap-free attention with equal +Q/K/V head sizes (multiple of 8, up to 256) and at least 2^20 attention scores, +GQA and batch included, for FLASH_ATTN_EXT and for SAGE_ATTN in AUTO mode (an +explicitly requested Sage variant keeps its kernel). Q (F32) is narrowed to F16 +into dst's own bytes, F16 K/V are read in place with their strides (F32 K/V are +narrowed), cuDNN writes F16 O to scratch after dst and one kernel widens it into +dst. The scratch is part of the op's graph allocation (max with the existing +kernels' needs). Plans are built once per shape and cached; a shape cuDNN +declines, a failed build or execute, stream capture, Turing and older, or a +missing library keep the existing kernels. + +GGML_CUDA_CUDNN_ATTN=0 disables it, GGML_CUDA_CUDNN_SAGE=0 keeps the Sage kernel, +GGML_CUDA_CUDNN_ATTN_BF16=1 runs BF16, GGML_CUDA_CUDNN_ATTN_LOG=1 logs each call. + +test-backend-ops: unmasked cases at the MiniMax-H3 head layout (odd and aligned +lengths, GQA with batch, F32 K/V, permuted Q) and unmasked perf cases. +--- + CMakeLists.txt | 3 + + src/ggml-cuda/CMakeLists.txt | 49 ++++ + src/ggml-cuda/cudnn-attn-graph.cpp | 258 ++++++++++++++++++++ + src/ggml-cuda/cudnn-attn-graph.h | 31 +++ + src/ggml-cuda/cudnn-attn.cu | 367 +++++++++++++++++++++++++++++ + src/ggml-cuda/cudnn-attn.cuh | 25 ++ + src/ggml-cuda/fattn.cu | 7 +- + src/ggml-cuda/ggml-cuda.cu | 2 +- + src/ggml-cuda/sage-attn.cu | 12 +- + src/ggml-cuda/sage-attn.cuh | 2 +- + tests/test-backend-ops.cpp | 25 ++ + 11 files changed, 775 insertions(+), 6 deletions(-) + create mode 100644 src/ggml-cuda/cudnn-attn-graph.cpp + create mode 100644 src/ggml-cuda/cudnn-attn-graph.h + create mode 100644 src/ggml-cuda/cudnn-attn.cu + create mode 100644 src/ggml-cuda/cudnn-attn.cuh + +diff --git a/CMakeLists.txt b/CMakeLists.txt +index b6843578..1c32cf3b 100644 +--- a/CMakeLists.txt ++++ b/CMakeLists.txt +@@ -208,6 +208,9 @@ option(GGML_CUDA_FA "ggml: compile ggml FlashAttention C + option(GGML_CUDA_FA_ALL_QUANTS "ggml: compile all quants for FlashAttention" OFF) + option(GGML_CUDA_GRAPHS "ggml: use CUDA graphs (llama.cpp only)" ${GGML_CUDA_GRAPHS_DEFAULT}) + option(GGML_CUDA_NCCL "ggml: use NVIDIA Collective Comm. Library" ON) ++option(GGML_CUDA_CUDNN "ggml: cuDNN attention (cudnn-frontend, libcudnn opened at run time)" OFF) ++set (GGML_CUDA_CUDNN_FRONTEND_DIR "" CACHE PATH "ggml: cudnn-frontend checkout for GGML_CUDA_CUDNN (fetched when empty)") ++set (GGML_CUDA_CUDNN_INCLUDE_DIR "" CACHE PATH "ggml: directory holding cudnn.h for GGML_CUDA_CUDNN") + set (GGML_CUDA_COMPRESSION_MODE "size" CACHE STRING + "ggml: cuda link binary compression mode; requires cuda 12.8+") + set_property(CACHE GGML_CUDA_COMPRESSION_MODE PROPERTY STRINGS "none;speed;balance;size") +diff --git a/src/ggml-cuda/CMakeLists.txt b/src/ggml-cuda/CMakeLists.txt +index 05b81487..0ddeb81e 100644 +--- a/src/ggml-cuda/CMakeLists.txt ++++ b/src/ggml-cuda/CMakeLists.txt +@@ -124,11 +124,60 @@ if (CUDAToolkit_FOUND) + template-instances/fattn-vec-instance-bf16-bf16.cu) + endif() + ++ if (GGML_CUDA_CUDNN AND WIN32) ++ message(WARNING "GGML_CUDA_CUDNN: the run-time cuDNN loader is Linux only, building without it") ++ set(GGML_CUDA_CUDNN OFF) ++ endif() ++ if (GGML_CUDA_CUDNN) ++ # Headers only at build time: cudnn.h for the types, cudnn-frontend for the graph API. ++ # libcudnn itself is opened with dlopen when the first attention op runs. ++ if (GGML_CUDA_CUDNN_INCLUDE_DIR) ++ # an explicit directory wins over any earlier cached search result ++ set(GGML_CUDNN_HEADER_DIR ${GGML_CUDA_CUDNN_INCLUDE_DIR}) ++ else() ++ find_path(GGML_CUDNN_HEADER_DIR cudnn.h ++ HINTS $ENV{CUDNN_INCLUDE_DIR} $ENV{CUDNN_PATH} ${CUDAToolkit_INCLUDE_DIRS} ++ PATH_SUFFIXES include) ++ endif() ++ if (NOT GGML_CUDNN_HEADER_DIR OR NOT EXISTS ${GGML_CUDNN_HEADER_DIR}/cudnn.h) ++ message(FATAL_ERROR "GGML_CUDA_CUDNN: cudnn.h not found, set GGML_CUDA_CUDNN_INCLUDE_DIR") ++ endif() ++ if (GGML_CUDA_CUDNN_FRONTEND_DIR) ++ set(GGML_CUDNN_FRONTEND_INCLUDE ${GGML_CUDA_CUDNN_FRONTEND_DIR}/include) ++ else() ++ include(FetchContent) ++ # v1.26.0, the last MIT-licensed release (later ones are Apache-2.0). Only the headers are ++ # used; SOURCE_SUBDIR points at a directory that does not exist so its CMake project stays out. ++ FetchContent_Declare( ++ ggml_cudnn_frontend ++ URL https://github.com/NVIDIA/cudnn-frontend/archive/refs/tags/v1.26.0.tar.gz ++ URL_HASH SHA256=f88361fb714589c970bac2ab56d6a7843722be08ca5ba226484de3fa46fd22cf ++ DOWNLOAD_EXTRACT_TIMESTAMP TRUE ++ SOURCE_SUBDIR ggml-headers-only ++ ) ++ FetchContent_MakeAvailable(ggml_cudnn_frontend) ++ set(GGML_CUDNN_FRONTEND_INCLUDE ${ggml_cudnn_frontend_SOURCE_DIR}/include) ++ endif() ++ if (NOT EXISTS ${GGML_CUDNN_FRONTEND_INCLUDE}/cudnn_frontend.h) ++ message(FATAL_ERROR "GGML_CUDA_CUDNN: cudnn_frontend.h not found under ${GGML_CUDNN_FRONTEND_INCLUDE}") ++ endif() ++ message(STATUS "cuDNN attention: cudnn.h in ${GGML_CUDNN_HEADER_DIR}, cudnn-frontend in ${GGML_CUDNN_FRONTEND_INCLUDE}") ++ list(APPEND GGML_SOURCES_CUDA cudnn-attn-graph.cpp) ++ endif() ++ + ggml_add_backend_library(ggml-cuda + ${GGML_HEADERS_CUDA} + ${GGML_SOURCES_CUDA} + ) + ++ if (GGML_CUDA_CUDNN) ++ target_compile_definitions(ggml-cuda PRIVATE GGML_CUDA_USE_CUDNN NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING) ++ # -isystem: third-party headers, kept out of ggml's warning flags. ++ set_source_files_properties(cudnn-attn-graph.cpp PROPERTIES ++ COMPILE_OPTIONS "-isystem;${GGML_CUDNN_FRONTEND_INCLUDE};-isystem;${GGML_CUDNN_HEADER_DIR}") ++ target_link_libraries(ggml-cuda PRIVATE ${CMAKE_DL_LIBS}) ++ endif() ++ + if (CUDAToolkit_VERSION VERSION_GREATER_EQUAL "11.8") + target_compile_definitions(ggml-cuda PRIVATE GGML_CUDA_USE_CUBLASLT_FP8) + if (NOT GGML_STATIC OR WIN32) +diff --git a/src/ggml-cuda/cudnn-attn-graph.cpp b/src/ggml-cuda/cudnn-attn-graph.cpp +new file mode 100644 +index 00000000..b0f1dc11 +--- /dev/null ++++ b/src/ggml-cuda/cudnn-attn-graph.cpp +@@ -0,0 +1,258 @@ ++// cuDNN scaled dot product attention through cudnn-frontend (header-only), with libcudnn opened at ++// run time (NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING): a build with GGML_CUDA_CUDNN still starts and ++// runs on hosts without cuDNN, it just keeps the ggml kernels. ++ ++#include "cudnn-attn-graph.h" ++ ++#include "ggml-impl.h" ++ ++#include ++ ++#include ++ ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++#include ++ ++namespace cudnn_frontend { ++// cudnn-frontend resolves every backend call through this handle when dynamic loading is on. ++void * cudnn_dlhandle = nullptr; ++} // namespace cudnn_frontend ++ ++namespace fe = cudnn_frontend; ++ ++namespace { ++ ++enum : int64_t { UID_Q = 1, UID_K = 2, UID_V = 3, UID_O = 4 }; ++ ++using plan_key = std::array; ++ ++struct plan_entry { ++ std::shared_ptr graph; // null: cuDNN has no plan for this key ++ int64_t workspace = 0; ++}; ++ ++std::mutex g_mutex; ++int g_load_state = 0; // 0 not tried, 1 loaded, -1 unavailable ++std::map g_handles; ++std::map g_plans; ++ ++plan_key make_key(const ggml_cuda_cudnn_sdpa_params & p) { ++ plan_key k{}; ++ int i = 0; ++ k[i++] = p.device; ++ k[i++] = p.b; ++ k[i++] = p.hq; ++ k[i++] = p.hk; ++ k[i++] = p.sq; ++ k[i++] = p.skv; ++ k[i++] = p.d; ++ for (int j = 0; j < 4; ++j) { ++ k[i++] = p.q_stride[j]; ++ } ++ for (int j = 0; j < 4; ++j) { ++ k[i++] = p.k_stride[j]; ++ } ++ for (int j = 0; j < 4; ++j) { ++ k[i++] = p.v_stride[j]; ++ } ++ for (int j = 0; j < 4; ++j) { ++ k[i++] = p.o_stride[j]; ++ } ++ uint32_t scale_bits = 0; ++ std::memcpy(&scale_bits, &p.scale, sizeof(scale_bits)); ++ k[i++] = scale_bits; ++ k[i++] = p.bf16 ? 1 : 0; ++ GGML_ASSERT(i == (int) k.size()); ++ return k; ++} ++ ++bool load_locked() { ++ if (g_load_state != 0) { ++ return g_load_state == 1; ++ } ++ g_load_state = -1; ++ ++ // cudnn-frontend opens libcudart itself for a few runtime calls; point it at the runtime major ++ // this backend was built against so it never picks a second, mismatched copy. ++ char cudart_name[32]; ++ std::snprintf(cudart_name, sizeof(cudart_name), "libcudart.so.%d", CUDART_VERSION / 1000); ++ setenv("CUDNN_FRONTEND_CUDART_LIB_NAME", cudart_name, 0); ++ ++ const char * user = std::getenv("GGML_CUDA_CUDNN_LIB"); ++ const char * name = user != nullptr && user[0] != '\0' ? user : "libcudnn.so.9"; ++ void * h = dlopen(name, RTLD_NOW | RTLD_LOCAL); ++ if (h == nullptr) { ++ GGML_LOG_DEBUG("%s: cuDNN not loaded (%s), attention keeps the ggml kernels\n", __func__, dlerror()); ++ return false; ++ } ++ using get_version_t = size_t (*)(); ++ auto get_version = reinterpret_cast(dlsym(h, "cudnnGetVersion")); ++ const size_t version = get_version != nullptr ? get_version() : 0; ++ if (version < 90000) { ++ GGML_LOG_WARN("%s: %s reports cuDNN %zu, need 9.0 or newer; attention keeps the ggml kernels\n", __func__, ++ name, version); ++ dlclose(h); ++ return false; ++ } ++ fe::cudnn_dlhandle = h; ++ g_load_state = 1; ++ GGML_LOG_INFO("%s: cuDNN %zu loaded from %s for attention\n", __func__, version, name); ++ return true; ++} ++ ++cudnnHandle_t handle_locked(int device) { ++ auto it = g_handles.find(device); ++ if (it != g_handles.end()) { ++ return it->second; ++ } ++ cudnnHandle_t h = nullptr; ++ if (fe::detail::create_handle(&h) != CUDNN_STATUS_SUCCESS) { ++ h = nullptr; ++ } ++ g_handles[device] = h; ++ return h; ++} ++ ++std::shared_ptr build_graph(const ggml_cuda_cudnn_sdpa_params & p, cudnnHandle_t handle, ++ std::string & why) { ++ auto g = std::make_shared(); ++ g->set_io_data_type(p.bf16 ? fe::DataType_t::BFLOAT16 : fe::DataType_t::HALF) ++ .set_intermediate_data_type(fe::DataType_t::FLOAT) ++ .set_compute_data_type(fe::DataType_t::FLOAT); ++ ++ auto tensor = [&](const char * name, int64_t uid, int64_t h, int64_t s, const int64_t * stride) { ++ return g->tensor(fe::graph::Tensor_attributes() ++ .set_name(name) ++ .set_uid(uid) ++ .set_dim({ p.b, h, s, p.d }) ++ .set_stride({ stride[0], stride[1], stride[2], stride[3] })); ++ }; ++ auto Q = tensor("Q", UID_Q, p.hq, p.sq, p.q_stride); ++ auto K = tensor("K", UID_K, p.hk, p.skv, p.k_stride); ++ auto V = tensor("V", UID_V, p.hk, p.skv, p.v_stride); ++ ++ auto options = fe::graph::SDPA_attributes().set_name("ggml_sdpa").set_generate_stats(false).set_attn_scale(p.scale); ++ auto [O, stats] = g->sdpa(Q, K, V, options); ++ (void) stats; ++ O->set_output(true) ++ .set_uid(UID_O) ++ .set_dim({ p.b, p.hq, p.sq, p.d }) ++ .set_stride({ p.o_stride[0], p.o_stride[1], p.o_stride[2], p.o_stride[3] }); ++ ++ fe::error_t st = g->validate(); ++ if (st.is_good()) { ++ st = g->build_operation_graph(handle); ++ } ++ if (st.is_good()) { ++ st = g->create_execution_plans({ fe::HeurMode_t::A, fe::HeurMode_t::FALLBACK }); ++ } ++ if (st.is_good()) { ++ st = g->check_support(handle); ++ } ++ if (st.is_good()) { ++ st = g->build_plans(handle); ++ } ++ if (!st.is_good()) { ++ why = st.get_message(); ++ return nullptr; ++ } ++ return g; ++} ++ ++} // namespace ++ ++bool ggml_cuda_cudnn_available() { ++ std::lock_guard lock(g_mutex); ++ try { ++ return load_locked(); ++ } catch (const std::exception & e) { ++ GGML_LOG_WARN("%s: %s\n", __func__, e.what()); ++ g_load_state = -1; ++ return false; ++ } ++} ++ ++bool ggml_cuda_cudnn_sdpa_prepare(const ggml_cuda_cudnn_sdpa_params & p, size_t * workspace_size) { ++ std::lock_guard lock(g_mutex); ++ if (g_load_state != 1) { ++ return false; ++ } ++ const plan_key key = make_key(p); ++ auto it = g_plans.find(key); ++ if (it == g_plans.end()) { ++ plan_entry entry; ++ std::string why; ++ const auto t0 = std::chrono::steady_clock::now(); ++ try { ++ cudnnHandle_t handle = handle_locked(p.device); ++ if (handle == nullptr) { ++ why = "cudnnCreate failed"; ++ } else { ++ entry.graph = build_graph(p, handle, why); ++ if (entry.graph && !entry.graph->get_workspace_size(entry.workspace).is_good()) { ++ entry.graph = nullptr; ++ why = "workspace query failed"; ++ } ++ } ++ } catch (const std::exception & e) { ++ entry.graph = nullptr; ++ why = e.what(); ++ } ++ const double ms = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); ++ if (entry.graph) { ++ GGML_LOG_INFO("%s: cuDNN SDPA plan b=%lld hq=%lld hk=%lld sq=%lld skv=%lld d=%lld %s, workspace %lld B, " ++ "built in %.0f ms\n", ++ __func__, (long long) p.b, (long long) p.hq, (long long) p.hk, (long long) p.sq, ++ (long long) p.skv, (long long) p.d, p.bf16 ? "bf16" : "f16", (long long) entry.workspace, ++ ms); ++ } else { ++ GGML_LOG_INFO("%s: no cuDNN SDPA plan for b=%lld hq=%lld hk=%lld sq=%lld skv=%lld d=%lld, keeping the ggml " ++ "kernel (%s)\n", ++ __func__, (long long) p.b, (long long) p.hq, (long long) p.hk, (long long) p.sq, ++ (long long) p.skv, (long long) p.d, why.c_str()); ++ } ++ it = g_plans.emplace(key, std::move(entry)).first; ++ } ++ if (!it->second.graph) { ++ return false; ++ } ++ *workspace_size = (size_t) it->second.workspace; ++ return true; ++} ++ ++bool ggml_cuda_cudnn_sdpa_execute(const ggml_cuda_cudnn_sdpa_params & p, cudaStream_t stream, void * q, void * k, ++ void * v, void * o, void * workspace) { ++ std::lock_guard lock(g_mutex); ++ auto it = g_plans.find(make_key(p)); ++ if (it == g_plans.end() || !it->second.graph) { ++ return false; ++ } ++ try { ++ cudnnHandle_t handle = handle_locked(p.device); ++ if (handle == nullptr || fe::detail::set_stream(handle, stream) != CUDNN_STATUS_SUCCESS) { ++ return false; ++ } ++ std::unordered_map pack = { ++ { UID_Q, q }, { UID_K, k }, { UID_V, v }, { UID_O, o } ++ }; ++ fe::error_t st = it->second.graph->execute(handle, pack, workspace); ++ if (!st.is_good()) { ++ GGML_LOG_WARN("%s: cuDNN SDPA execute failed, keeping the ggml kernel for this shape: %s\n", __func__, ++ st.get_message().c_str()); ++ it->second.graph = nullptr; ++ return false; ++ } ++ } catch (const std::exception & e) { ++ GGML_LOG_WARN("%s: cuDNN SDPA execute failed, keeping the ggml kernel for this shape: %s\n", __func__, e.what()); ++ it->second.graph = nullptr; ++ return false; ++ } ++ return true; ++} +diff --git a/src/ggml-cuda/cudnn-attn-graph.h b/src/ggml-cuda/cudnn-attn-graph.h +new file mode 100644 +index 00000000..9749f2ab +--- /dev/null ++++ b/src/ggml-cuda/cudnn-attn-graph.h +@@ -0,0 +1,31 @@ ++#pragma once ++ ++// Host-side half of the cuDNN attention path (cudnn-attn.cu is the device side). Only compiled ++// with GGML_CUDA_CUDNN. libcudnn is opened at run time, so the binary never links against it. ++ ++#include ++#include ++ ++#include ++ ++struct ggml_cuda_cudnn_sdpa_params { ++ int device; ++ int64_t b, hq, hk, sq, skv, d; // batch, query heads, key/value heads, query tokens, key tokens, head size ++ int64_t q_stride[4]; // element strides in (b, h, s, d) order ++ int64_t k_stride[4]; ++ int64_t v_stride[4]; ++ int64_t o_stride[4]; ++ float scale; ++ bool bf16; // io type: BF16 when set, F16 otherwise ++}; ++ ++// Opens libcudnn once per process. False when it cannot be loaded or is older than cuDNN 9. ++bool ggml_cuda_cudnn_available(); ++ ++// Builds (or finds) the execution plan for these parameters. False when cuDNN has no plan for ++// them; the result is cached either way, so a rejected shape costs one attempt per process. ++bool ggml_cuda_cudnn_sdpa_prepare(const ggml_cuda_cudnn_sdpa_params & p, size_t * workspace_size); ++ ++// Runs a plan built by ggml_cuda_cudnn_sdpa_prepare on the given stream. ++bool ggml_cuda_cudnn_sdpa_execute(const ggml_cuda_cudnn_sdpa_params & p, cudaStream_t stream, ++ void * q, void * k, void * v, void * o, void * workspace); +diff --git a/src/ggml-cuda/cudnn-attn.cu b/src/ggml-cuda/cudnn-attn.cu +new file mode 100644 +index 00000000..b54c29af +--- /dev/null ++++ b/src/ggml-cuda/cudnn-attn.cu +@@ -0,0 +1,367 @@ ++#include "cudnn-attn.cuh" ++ ++#if defined(GGML_CUDA_USE_CUDNN) && !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) ++ ++#include "convert.cuh" ++#include "cudnn-attn-graph.h" ++ ++#include ++#include ++#include ++ ++static bool cudnn_attn_env_set(const char * name, bool default_value) { ++ const char * v = std::getenv(name); ++ if (v == nullptr || v[0] == '\0') { ++ return default_value; ++ } ++ return !(std::strcmp(v, "0") == 0 || std::strcmp(v, "off") == 0 || std::strcmp(v, "false") == 0); ++} ++ ++// Where a staged F16 / BF16 tensor lives: inside dst's bytes or in the scratch after dst. ++struct cudnn_attn_slot { ++ bool in_dst; ++ size_t off; ++}; ++ ++// Everything the cuDNN path needs to know about one op, from shapes and types alone. ++struct cudnn_attn_plan { ++ ggml_cuda_cudnn_sdpa_params p; ++ bool convert_k; ++ bool convert_v; ++ cudnn_attn_slot q, k, v, o; ++ size_t scratch; // scratch bytes after dst ++}; ++ ++static bool cudnn_attn_strides_ok(const ggml_tensor * t) { ++ // cuDNN wants a unit stride on the head dimension and 16-byte aligned rows. ++ return t->nb[0] == ggml_type_size(t->type) && t->nb[1] % 16 == 0 && t->nb[2] % 16 == 0 && t->nb[3] % 16 == 0; ++} ++ ++static bool cudnn_attn_make_plan(int device, const ggml_tensor * dst, cudnn_attn_plan & plan) { ++ static const bool enabled = cudnn_attn_env_set("GGML_CUDA_CUDNN_ATTN", true); ++ static const bool sage = cudnn_attn_env_set("GGML_CUDA_CUDNN_SAGE", true); ++ static const bool want_bf16 = cudnn_attn_env_set("GGML_CUDA_CUDNN_ATTN_BF16", false); ++ if (!enabled || (dst->op == GGML_OP_SAGE_ATTN && !sage)) { ++ return false; ++ } ++ if (dst->op != GGML_OP_FLASH_ATTN_EXT && dst->op != GGML_OP_SAGE_ATTN) { ++ return false; ++ } ++ // An explicitly requested SageAttention variant keeps its kernel; AUTO takes the cuDNN path. ++ if (dst->op == GGML_OP_SAGE_ATTN && ggml_get_op_params_i32(dst, 1) != GGML_SAGE_ATTN_AUTO) { ++ return false; ++ } ++ const int cc = ggml_cuda_info().devices[device].cc; ++ if (!GGML_CUDA_CC_IS_NVIDIA(cc) || cc < GGML_CUDA_CC_AMPERE) { ++ return false; ++ } ++ ++ const ggml_tensor * Q = dst->src[0]; ++ const ggml_tensor * K = dst->src[1]; ++ const ggml_tensor * V = dst->src[2]; ++ if (Q == nullptr || K == nullptr || V == nullptr) { ++ return false; ++ } ++ if (dst->op == GGML_OP_FLASH_ATTN_EXT) { ++ float max_bias = 0.0f, softcap = 0.0f; ++ std::memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); ++ std::memcpy(&softcap, (const float *) dst->op_params + 2, sizeof(float)); ++ if (dst->src[3] != nullptr || dst->src[4] != nullptr || max_bias != 0.0f || softcap != 0.0f) { ++ return false; ++ } ++ } ++ ++ const ggml_type io = want_bf16 ? GGML_TYPE_BF16 : GGML_TYPE_F16; ++ auto kv_type_ok = [](ggml_type t) { ++ return t == GGML_TYPE_F16 || t == GGML_TYPE_F32 || t == GGML_TYPE_BF16; ++ }; ++ if (Q->type != GGML_TYPE_F32 || !kv_type_ok(K->type) || !kv_type_ok(V->type) || dst->type != GGML_TYPE_F32) { ++ return false; ++ } ++ // BF16 -> F16 can overflow, so BF16 K/V only run in BF16 mode. F16 K/V in BF16 mode are narrowed ++ // (3 mantissa bits lost); that mode exists to measure BF16 against the F16 default. ++ if ((K->type == GGML_TYPE_BF16 || V->type == GGML_TYPE_BF16) && io != GGML_TYPE_BF16) { ++ return false; ++ } ++ ++ const int64_t d = Q->ne[0]; ++ if (d % 8 != 0 || d > 256 || K->ne[0] != d || V->ne[0] != d) { ++ return false; ++ } ++ const int64_t sq = Q->ne[1], hq = Q->ne[2], b = Q->ne[3]; ++ const int64_t skv = K->ne[1], hk = K->ne[2]; ++ if (V->ne[1] != skv || V->ne[2] != hk || K->ne[3] != b || V->ne[3] != b || hk == 0 || hq % hk != 0) { ++ return false; ++ } ++ if (!ggml_is_contiguous(dst) || dst->ne[0] != d || dst->ne[1] != hq || dst->ne[2] != sq || dst->ne[3] != b) { ++ return false; ++ } ++ if (Q->nb[0] != ggml_type_size(Q->type) || !cudnn_attn_strides_ok(K) || !cudnn_attn_strides_ok(V)) { ++ return false; ++ } ++ if (hq * b > INT32_MAX || sq > INT32_MAX || skv > INT32_MAX) { ++ return false; ++ } ++ // Below about 2^20 attention scores the two extra conversion launches cost more than cuDNN saves ++ // (B200: 31 tokens x 56 heads 10.6 -> 14-16 us, one query row x 32 heads x 4096 keys 12.9 -> 20 us, ++ // break-even near 256 x 256 x 16, 2.3x faster at 512 x 512 x 16), so small shapes keep the ggml kernels. ++ if ((double) sq * (double) skv * (double) (hq * b) < (double) (1 << 20)) { ++ return false; ++ } ++ if (!ggml_cuda_cudnn_available()) { ++ return false; ++ } ++ ++ plan = {}; ++ plan.p.device = device; ++ plan.p.b = b; ++ plan.p.hq = hq; ++ plan.p.hk = hk; ++ plan.p.sq = sq; ++ plan.p.skv = skv; ++ plan.p.d = d; ++ plan.p.bf16 = io == GGML_TYPE_BF16; ++ std::memcpy(&plan.p.scale, (const float *) dst->op_params + 0, sizeof(float)); ++ ++ // Converted tensors are written contiguous: [d, s, h, b], i.e. (b, h, s, d) strides below. ++ auto set_contiguous = [d](int64_t * stride, int64_t s, int64_t h) { ++ stride[3] = 1; ++ stride[2] = d; ++ stride[1] = d * s; ++ stride[0] = d * s * h; ++ }; ++ auto set_from = [](int64_t * stride, const ggml_tensor * t) { ++ const size_t ts = ggml_type_size(t->type); ++ stride[3] = 1; ++ stride[2] = t->nb[1] / ts; ++ stride[1] = t->nb[2] / ts; ++ stride[0] = t->nb[3] / ts; ++ }; ++ ++ // Q (and K/V when they need converting) are staged in dst's own bytes: dst is F32 and at least ++ // twice the size of F16 Q, and it is only written after cuDNN has read them. O and whatever does ++ // not fit go to scratch after dst. ++ const size_t es = 2; // F16 / BF16 ++ const size_t dst_size = ggml_nbytes(dst); ++ size_t dst_off = 0; ++ size_t off = 0; ++ auto take = [&](size_t bytes, cudnn_attn_slot & slot, bool allow_dst) { ++ if (allow_dst && dst_off + bytes <= dst_size) { ++ slot = { true, dst_off }; ++ dst_off += GGML_PAD(bytes, 256); ++ } else { ++ slot = { false, off }; ++ off += GGML_PAD(bytes, 256); ++ } ++ }; ++ ++ take(ggml_nelements(Q) * es, plan.q, true); ++ set_contiguous(plan.p.q_stride, sq, hq); ++ ++ plan.convert_k = K->type != io; ++ if (plan.convert_k) { ++ take(ggml_nelements(K) * es, plan.k, true); ++ set_contiguous(plan.p.k_stride, skv, hk); ++ } else { ++ set_from(plan.p.k_stride, K); ++ } ++ plan.convert_v = V->type != io; ++ if (plan.convert_v) { ++ take(ggml_nelements(V) * es, plan.v, true); ++ set_contiguous(plan.p.v_stride, skv, hk); ++ } else { ++ set_from(plan.p.v_stride, V); ++ } ++ ++ // O is written in dst's layout, [d, hq, sq, b], then widened to F32 into dst. ++ take(ggml_nelements(dst) * es, plan.o, false); ++ plan.p.o_stride[3] = 1; ++ plan.p.o_stride[1] = d; ++ plan.p.o_stride[2] = d * hq; ++ plan.p.o_stride[0] = d * hq * sq; ++ plan.scratch = off; ++ return true; ++} ++ ++size_t ggml_cuda_cudnn_attn_alloc_size(int device, const ggml_tensor * dst) { ++ cudnn_attn_plan plan; ++ if (!cudnn_attn_make_plan(device, dst, plan)) { ++ return 0; ++ } ++ // dst (which also stages Q), up to 256 bytes to align the scratch start, then the scratch. ++ return ggml_nbytes(dst) + 256 + plan.scratch; ++} ++ ++// F32 rows of 8-element groups (strided rows, 16-byte aligned) -> contiguous F16 / BF16, 8 per thread. ++template ++static __global__ void cudnn_attn_narrow_f32(const char * __restrict__ x, dst_t * __restrict__ y, const int64_t n8, ++ const int64_t g0, const int64_t ne1, const int64_t ne2, const int64_t nb1, ++ const int64_t nb2, const int64_t nb3) { ++ for (int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; i < n8; i += (int64_t) gridDim.x * blockDim.x) { ++ const int64_t g = i % g0; ++ const int64_t row = i / g0; ++ const int64_t i1 = row % ne1; ++ const int64_t i2 = (row / ne1) % ne2; ++ const int64_t i3 = row / (ne1 * ne2); ++ const float4 * src = (const float4 *) (x + i3 * nb3 + i2 * nb2 + i1 * nb1) + 2 * g; ++ const float4 a = src[0]; ++ const float4 b = src[1]; ++ union { ++ dst_t h[8]; ++ uint4 u; ++ } out; ++ out.h[0] = ggml_cuda_cast(a.x); ++ out.h[1] = ggml_cuda_cast(a.y); ++ out.h[2] = ggml_cuda_cast(a.z); ++ out.h[3] = ggml_cuda_cast(a.w); ++ out.h[4] = ggml_cuda_cast(b.x); ++ out.h[5] = ggml_cuda_cast(b.y); ++ out.h[6] = ggml_cuda_cast(b.z); ++ out.h[7] = ggml_cuda_cast(b.w); ++ ((uint4 *) y)[i] = out.u; ++ } ++} ++ ++// Contiguous F16 / BF16 -> contiguous F32, 8 per thread. ++template ++static __global__ void cudnn_attn_widen(const src_t * __restrict__ x, float * __restrict__ y, const int64_t n8) { ++ for (int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; i < n8; i += (int64_t) gridDim.x * blockDim.x) { ++ union { ++ src_t h[8]; ++ uint4 u; ++ } in; ++ in.u = ((const uint4 *) x)[i]; ++ float4 a, b; ++ a.x = ggml_cuda_cast(in.h[0]); ++ a.y = ggml_cuda_cast(in.h[1]); ++ a.z = ggml_cuda_cast(in.h[2]); ++ a.w = ggml_cuda_cast(in.h[3]); ++ b.x = ggml_cuda_cast(in.h[4]); ++ b.y = ggml_cuda_cast(in.h[5]); ++ b.z = ggml_cuda_cast(in.h[6]); ++ b.w = ggml_cuda_cast(in.h[7]); ++ ((float4 *) y)[2 * i] = a; ++ ((float4 *) y)[2 * i + 1] = b; ++ } ++} ++ ++static int cudnn_attn_grid(int64_t n) { ++ return (int) std::min((n + 255) / 256, 65535LL * 8); ++} ++ ++static void cudnn_attn_convert(const ggml_tensor * src, void * dst, bool bf16, cudaStream_t stream) { ++ const size_t ts = ggml_type_size(src->type); ++ const bool vec_ok = src->type == GGML_TYPE_F32 && src->ne[0] % 8 == 0 && (uintptr_t) src->data % 16 == 0 && ++ src->nb[1] % 16 == 0 && src->nb[2] % 16 == 0 && src->nb[3] % 16 == 0; ++ if (vec_ok) { ++ const int64_t n8 = ggml_nelements(src) / 8; ++ if (bf16) { ++ cudnn_attn_narrow_f32<<>>( ++ (const char *) src->data, (nv_bfloat16 *) dst, n8, src->ne[0] / 8, src->ne[1], src->ne[2], src->nb[1], ++ src->nb[2], src->nb[3]); ++ } else { ++ cudnn_attn_narrow_f32<<>>( ++ (const char *) src->data, (half *) dst, n8, src->ne[0] / 8, src->ne[1], src->ne[2], src->nb[1], ++ src->nb[2], src->nb[3]); ++ } ++ return; ++ } ++ if (bf16) { ++ to_bf16_nc_cuda_t f = ggml_get_to_bf16_nc_cuda(src->type); ++ GGML_ASSERT(f != nullptr); ++ f(src->data, (nv_bfloat16 *) dst, src->ne[0], src->ne[1], src->ne[2], src->ne[3], src->nb[1] / ts, ++ src->nb[2] / ts, src->nb[3] / ts, stream); ++ } else { ++ to_fp16_nc_cuda_t f = ggml_get_to_fp16_nc_cuda(src->type); ++ GGML_ASSERT(f != nullptr); ++ f(src->data, (half *) dst, src->ne[0], src->ne[1], src->ne[2], src->ne[3], src->nb[1] / ts, src->nb[2] / ts, ++ src->nb[3] / ts, stream); ++ } ++} ++ ++bool ggml_cuda_cudnn_attn(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ++ static const bool log_calls = cudnn_attn_env_set("GGML_CUDA_CUDNN_ATTN_LOG", false); ++ static std::atomic n_calls{ 0 }; ++ ++ cudnn_attn_plan plan; ++ if (!cudnn_attn_make_plan(ctx.device, dst, plan)) { ++ return false; ++ } ++ cudaStream_t stream = ctx.stream(); ++ ++ // Plans are built outside stream capture only; a captured graph keeps the ggml kernels. ++ cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; ++ CUDA_CHECK(cudaStreamIsCapturing(stream, &capture)); ++ if (capture != cudaStreamCaptureStatusNone) { ++ return false; ++ } ++ ++ const ggml_tensor * Q = dst->src[0]; ++ const ggml_tensor * K = dst->src[1]; ++ const ggml_tensor * V = dst->src[2]; ++ if ((uintptr_t) dst->data % 16 != 0 || (!plan.convert_k && (uintptr_t) K->data % 16 != 0) || ++ (!plan.convert_v && (uintptr_t) V->data % 16 != 0)) { ++ return false; ++ } ++ ++ size_t workspace_size = 0; ++ if (!ggml_cuda_cudnn_sdpa_prepare(plan.p, &workspace_size)) { ++ if (log_calls) { ++ GGML_LOG_INFO("%s: %s %s -> ggml kernel (no cuDNN plan)\n", __func__, ggml_op_desc(dst), dst->name); ++ } ++ return false; ++ } ++ ++ char * base = (char *) dst->data; ++ char * scratch = (char *) GGML_PAD((uintptr_t) dst->data + ggml_nbytes(dst), 256); ++ auto where = [&](const cudnn_attn_slot & slot) -> void * { ++ return (slot.in_dst ? base : scratch) + slot.off; ++ }; ++ void * q = where(plan.q); ++ void * k = plan.convert_k ? where(plan.k) : K->data; ++ void * v = plan.convert_v ? where(plan.v) : V->data; ++ void * o = where(plan.o); ++ ++ cudnn_attn_convert(Q, q, plan.p.bf16, stream); ++ if (plan.convert_k) { ++ cudnn_attn_convert(K, k, plan.p.bf16, stream); ++ } ++ if (plan.convert_v) { ++ cudnn_attn_convert(V, v, plan.p.bf16, stream); ++ } ++ ++ ggml_cuda_pool_alloc workspace(ctx.pool()); ++ if (workspace_size > 0) { ++ workspace.alloc(workspace_size); ++ } ++ if (!ggml_cuda_cudnn_sdpa_execute(plan.p, stream, q, k, v, o, workspace_size > 0 ? workspace.get() : nullptr)) { ++ return false; ++ } ++ ++ // dst holds d * hq * sq * b values with d % 8 == 0, and both buffers are 256-byte aligned. ++ const int64_t n8 = ggml_nelements(dst) / 8; ++ if (plan.p.bf16) { ++ cudnn_attn_widen<<>>((const nv_bfloat16 *) o, (float *) dst->data, n8); ++ } else { ++ cudnn_attn_widen<<>>((const half *) o, (float *) dst->data, n8); ++ } ++ CUDA_CHECK(cudaGetLastError()); ++ ++ const long long n = ++n_calls; ++ if (log_calls) { ++ GGML_LOG_INFO("%s: %s %s -> cuDNN (call %lld)\n", __func__, ggml_op_desc(dst), dst->name, n); ++ } ++ return true; ++} ++ ++#else ++ ++size_t ggml_cuda_cudnn_attn_alloc_size(int, const ggml_tensor *) { ++ return 0; ++} ++ ++bool ggml_cuda_cudnn_attn(ggml_backend_cuda_context &, ggml_tensor *) { ++ return false; ++} ++ ++#endif +diff --git a/src/ggml-cuda/cudnn-attn.cuh b/src/ggml-cuda/cudnn-attn.cuh +new file mode 100644 +index 00000000..f59689d5 +--- /dev/null ++++ b/src/ggml-cuda/cudnn-attn.cuh +@@ -0,0 +1,25 @@ ++#pragma once ++ ++#include "common.cuh" ++ ++// cuDNN scaled dot product attention for GGML_OP_FLASH_ATTN_EXT and GGML_OP_SAGE_ATTN on NVIDIA ++// sm80+ when the backend is built with GGML_CUDA_CUDNN and libcudnn 9 loads at run time. ++// Only unmasked, sink-free, ALiBi-free, softcap-free attention with equal Q/K/V head sizes and at ++// least 2^20 attention scores is taken; everything else, and every shape cuDNN rejects, keeps the ++// existing kernels. ++// ++// Environment: ++// GGML_CUDA_CUDNN_ATTN=0 never use cuDNN for attention ++// GGML_CUDA_CUDNN_SAGE=0 keep the SageAttention kernels for GGML_OP_SAGE_ATTN (AUTO mode; ++// an explicitly requested Sage variant always keeps its kernel) ++// GGML_CUDA_CUDNN_ATTN_BF16=1 run cuDNN in BF16 instead of F16 (F16 K/V are narrowed) ++// GGML_CUDA_CUDNN_LIB=path library to open instead of libcudnn.so.9 ++// GGML_CUDA_CUDNN_ATTN_LOG=1 log every call that takes (or falls back from) the cuDNN path ++ ++// Bytes the op needs from the graph allocation when the cuDNN path can run, 0 otherwise. The ++// region starts at dst->data and includes dst itself. ++size_t ggml_cuda_cudnn_attn_alloc_size(int device, const ggml_tensor * dst); ++ ++// Runs dst through cuDNN. Returns false without touching dst when the cuDNN path does not apply, ++// so the caller runs its own kernel. ++bool ggml_cuda_cudnn_attn(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +diff --git a/src/ggml-cuda/fattn.cu b/src/ggml-cuda/fattn.cu +index b8be30a0..ddc93e30 100644 +--- a/src/ggml-cuda/fattn.cu ++++ b/src/ggml-cuda/fattn.cu +@@ -4,6 +4,7 @@ + #include "fattn-tile.cuh" + #include "fattn-vec.cuh" + #include "fattn.cuh" ++#include "cudnn-attn.cuh" + + template + static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { +@@ -573,11 +574,15 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d + const ggml_cuda_flash_attn_ext_f16_extra_data f16_extra = + ggml_cuda_flash_attn_ext_get_f16_extra_data(dst, need_f16_K, need_f16_V); + +- return f16_extra.end - (uintptr_t) dst->data; ++ // The cuDNN path, when it can run, uses the same region after dst for its own scratch. ++ return std::max(f16_extra.end - (uintptr_t) dst->data, ggml_cuda_cudnn_attn_alloc_size(device, dst)); + } + + void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + ggml_cuda_set_device(ctx.device); ++ if (ggml_cuda_cudnn_attn(ctx, dst)) { ++ return; ++ } + switch (ggml_cuda_get_best_fattn_kernel(ggml_cuda_get_device(), dst)) { + case BEST_FATTN_KERNEL_NONE: + GGML_ABORT("fatal error"); +diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu +index 86a6cc5b..749aead9 100644 +--- a/src/ggml-cuda/ggml-cuda.cu ++++ b/src/ggml-cuda/ggml-cuda.cu +@@ -919,7 +919,7 @@ static size_t ggml_backend_cuda_buffer_type_get_alloc_size(ggml_backend_buffer_t + ? ggml_cuda_flash_attn_ext_get_alloc_size(buft_ctx->device, tensor) + : ggml_nbytes(tensor); + if (tensor->op == GGML_OP_SAGE_ATTN) { +- size = ggml_cuda_sage_attn_get_alloc_size(tensor); ++ size = ggml_cuda_sage_attn_get_alloc_size(buft_ctx->device, tensor); + } + int64_t ne0 = tensor->ne[0]; + +diff --git a/src/ggml-cuda/sage-attn.cu b/src/ggml-cuda/sage-attn.cu +index 2c103eb0..ceb29ed2 100644 +--- a/src/ggml-cuda/sage-attn.cu ++++ b/src/ggml-cuda/sage-attn.cu +@@ -1,4 +1,5 @@ + #include "sage-attn.cuh" ++#include "cudnn-attn.cuh" + + #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDACC_VER_MAJOR__ >= 12 + +@@ -332,13 +333,18 @@ bool ggml_cuda_sage_attn_supported(int device, const ggml_tensor* op) { + op->ne[2] == q->ne[1] && op->ne[3] == q->ne[3]; + } + +-size_t ggml_cuda_sage_attn_get_alloc_size(const ggml_tensor* op) { ++size_t ggml_cuda_sage_attn_get_alloc_size(int device, const ggml_tensor* op) { + // Scratch belongs to the graph allocation, so graph capture and VRAM planning see its full lifetime. +- return SageWorkspace(op).size; ++ // The cuDNN path, when it can run, reuses the same region for its own scratch. ++ return std::max(SageWorkspace(op).size, ggml_cuda_cudnn_attn_alloc_size(device, op)); + } + + void ggml_cuda_sage_attn(ggml_backend_cuda_context& ctx, ggml_tensor* dst) { + GGML_ASSERT(ggml_cuda_sage_attn_supported(ctx.device, dst)); ++ ggml_cuda_set_device(ctx.device); ++ if (ggml_cuda_cudnn_attn(ctx, dst)) { ++ return; ++ } + if (dst->src[0]->ne[0] == 64) { + sage_attn<64>(ctx, dst); + } else { +@@ -352,7 +358,7 @@ bool ggml_cuda_sage_attn_supported(int, const ggml_tensor*) { + return false; + } + +-size_t ggml_cuda_sage_attn_get_alloc_size(const ggml_tensor* op) { ++size_t ggml_cuda_sage_attn_get_alloc_size(int, const ggml_tensor* op) { + return ggml_nbytes(op); + } + +diff --git a/src/ggml-cuda/sage-attn.cuh b/src/ggml-cuda/sage-attn.cuh +index 126c7a74..e2648857 100644 +--- a/src/ggml-cuda/sage-attn.cuh ++++ b/src/ggml-cuda/sage-attn.cuh +@@ -4,7 +4,7 @@ + #include "common.cuh" + + bool ggml_cuda_sage_attn_supported(int device, const ggml_tensor* op); +-size_t ggml_cuda_sage_attn_get_alloc_size(const ggml_tensor* op); ++size_t ggml_cuda_sage_attn_get_alloc_size(int device, const ggml_tensor* op); + void ggml_cuda_sage_attn(ggml_backend_cuda_context& ctx, ggml_tensor* dst); + + #endif +diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp +index b58c815c..55db14a0 100644 +--- a/tests/test-backend-ops.cpp ++++ b/tests/test-backend-ops.cpp +@@ -9773,6 +9773,21 @@ static std::vector> make_test_cases_eval() { + GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + } + ++ // unmasked F16 K/V with the MiniMax-H3 DiT head layout (42 heads x 128), the shape class the CUDA ++ // cuDNN attention path takes: odd and tile-aligned lengths, GQA with batch, F32 K/V, head size 64. ++ for (int64_t n : { 1023, 4096 }) { ++ test_cases.emplace_back(new test_flash_attn_ext(128, 128, 42, {1, 1}, n, n, false, false, 0, 0, ++ GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); ++ } ++ test_cases.emplace_back(new test_flash_attn_ext(128, 128, 42, {1, 1}, 4096, 333, false, false, 0, 0, ++ GGML_PREC_DEFAULT, GGML_TYPE_F16, GGML_TYPE_F16)); ++ test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {4, 2}, 777, 333, false, false, 0, 0, ++ GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); ++ test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {1, 1}, 512, 512, false, false, 0, 0, ++ GGML_PREC_F32, GGML_TYPE_F32, GGML_TYPE_F32)); ++ test_cases.emplace_back(new test_flash_attn_ext(128, 128, 42, {1, 1}, 1023, 1023, false, false, 0, 0, ++ GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3})); ++ + test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3})); + test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1})); + test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, { 10, 5, 4, 3})); +@@ -10124,6 +10139,16 @@ static std::vector> make_test_cases_perf() { + // Qwen3-VL-8B https://github.com/ggml-org/llama.cpp/issues/17012 + test_cases.emplace_back(new test_flash_attn_ext(72, 72, 16, {1, 1}, 5776, 5776, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + ++ // unmasked diffusion attention: MiniMax-H3 DiT (19315 tokens x 56 heads at 960x544x124; 22440 x 42), ++ // its token refiner (31 tokens) and video VAE decoder (1797 x 32 x 64), FLUX-class 1024px (4096 x 24), ++ // and a single-query row ++ for (auto [hs, nh, n] : std::vector>{ { 128, 56, 19315 }, { 128, 42, 22440 }, { 128, 56, 31 }, ++ { 64, 32, 1797 }, { 128, 24, 4096 }, { 128, 16, 256 }, ++ { 128, 16, 512 }, { 128, 16, 1024 } }) { ++ test_cases.emplace_back(new test_flash_attn_ext(hs, hs, nh, {1, 1}, n, n, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); ++ } ++ test_cases.emplace_back(new test_flash_attn_ext(128, 128, 8, {4, 1}, 4096, 1, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); ++ + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0)); +-- +2.43.0 + From ca92d95ddf041bf885a36a7fdafbffd88706c28e Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Thu, 8 Oct 2026 00:15:39 -0300 Subject: [PATCH 22/30] Tighten H3 VAE and tiling comments --- src/pipeline/diffusion_engine.cpp | 6 ++---- src/runtime/tiling.cpp | 6 ++---- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index eba4cca4da..ec50c620cc 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -1178,10 +1178,8 @@ bool StableDiffusionGGML::validate_and_load_runners() { high_noise_diffusion_model->set_flash_attention_enabled(true); } } - // The MiniMax-H3 video VAE decoder is a 36-layer ViT (32 x 64 heads, ~1.5k tokens per 16x16 latent tile). - // Without flash attention every tile and layer materialises a 32 x L x L f32 score matrix (mul_mat + scale + - // softmax + mul_mat), which is most of the decode's GPU time. --diffusion-fa therefore also covers it, the - // way --fa would, without switching the text encoder's attention. SD_H3_VAE_FLASH_ATTN=0 restores the old path. + // The H3 video VAE decoder is a ViT whose f32 L x L attention scores dominate decode time, so --diffusion-fa + // covers it too, without switching the text encoder. SD_H3_VAE_FLASH_ATTN=0 restores the old path. if (!sd_ctx_params->flash_attn && sd_ctx_params->diffusion_flash_attn && first_stage_model && sd_version_is_minimax_h3(version)) { const char* env = getenv("SD_H3_VAE_FLASH_ATTN"); diff --git a/src/runtime/tiling.cpp b/src/runtime/tiling.cpp index 767e9c4d19..782dc4083d 100644 --- a/src/runtime/tiling.cpp +++ b/src/runtime/tiling.cpp @@ -74,8 +74,7 @@ static sd::Tensor sd_tensor_split_2d(const sd::Tensor& input, int int64_t input_plane = sd_tensor_plane_size(input); int64_t output_plane = sd_tensor_plane_size(output); int64_t plane_count = input.numel() / input_plane; - // Plane-outer order: a plane is contiguous, so the inner loop streams instead of striding a - // whole plane per element (a video tile has frames x channels planes). + // Plane-outer: planes are contiguous, so the inner loop streams instead of striding a plane per element. std::vector src_x(static_cast(width)); for (int ix = 0; ix < width; ix++) { src_x[ix] = (ix + x) % input_width; @@ -120,8 +119,7 @@ static void sd_tensor_merge_2d(const sd::Tensor& input, return x * x * x * (x * (6.0f * x - 15.0f) + 10.0f); }; - // Weights depend on ix or iy only; precompute them and walk each contiguous plane row by row. - // Same arithmetic per element as the per-pixel form, so the result is bit-identical. + // Same per-element arithmetic as the per-pixel form, so the result stays bit-identical. const bool blend = overlap_x > 0 || overlap_y > 0; std::vector wx, wy; std::vector dst_x(static_cast(width)); 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 23/30] 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 39d77eae35..2413abe2c7 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 07a244a8ad..f82a6445c2 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 24/30] 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 4b967e0e36..f54d465041 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: From 907615c633d4fcab13c11763a9dc7e95202a52f1 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Thu, 8 Oct 2026 11:09:11 -0300 Subject: [PATCH 25/30] Gate the H3 VAE fused attention on ggml patch 0004 and its fused norm on 0005 --- cmake/ggml.cmake | 5 +++++ src/model/vae/minimax_h3_vae.hpp | 8 +++++--- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/cmake/ggml.cmake b/cmake/ggml.cmake index 2413abe2c7..82d4814933 100644 --- a/cmake/ggml.cmake +++ b/cmake/ggml.cmake @@ -26,12 +26,17 @@ set(sd_ggml_header "${sd_ggml_private_include}/../include/ggml.h") if(EXISTS "${sd_ggml_header}") set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS "${sd_ggml_header}") file(STRINGS "${sd_ggml_header}" sd_ggml_h3_fused_ops REGEX "ggml_rope_pe_permute\\(") + file(STRINGS "${sd_ggml_header}" sd_ggml_h3_fused_qk_norm REGEX "ggml_rope_pe_permute_rms\\(") endif() if(sd_ggml_h3_fused_ops) target_compile_definitions(${SD_LIB} PUBLIC SD_GGML_H3_FUSED_OPS) else() message(STATUS "ggml lacks the fused DiT ops patch: MiniMax-H3 uses the unfused graph") endif() +# The fused q/k RMS norm comes from patch 0005. +if(sd_ggml_h3_fused_qk_norm) + target_compile_definitions(${SD_LIB} PUBLIC SD_GGML_H3_FUSED_QK_NORM) +endif() if(SD_USE_UPSTREAM_GGML) target_compile_definitions(${SD_LIB} PUBLIC SD_USE_UPSTREAM_GGML) diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index f4cff90af5..a4ec74fd0a 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -300,7 +300,7 @@ namespace MiniMaxH3VAE { ggml_tensor* pe, ggml_tensor* rope_a, ggml_tensor* rope_b) { -#ifdef SD_USE_UPSTREAM_GGML +#ifndef SD_GGML_H3_FUSED_OPS return nullptr; #else static const bool enabled = [] { @@ -324,14 +324,15 @@ namespace MiniMaxH3VAE { return ggml_view_4d(ctx->ggml_ctx, qkv, head_dim, num_head, sequence, 1, qkv->nb[1], qkv->nb[2], qkv->nb[3], index * head_dim * sizeof(float)); }; + ggml_tensor* q = nullptr; + ggml_tensor* k = nullptr; +#ifdef SD_GGML_H3_FUSED_QK_NORM // q/k RMS norm folded into the RoPE op where the backend reproduces GGML_OP_RMS_NORM's // arithmetic in it (SD_H3_VAE_FUSED_QK_NORM=0 keeps the separate norm). static const bool fuse_norm = [] { const char* v = getenv("SD_H3_VAE_FUSED_QK_NORM"); return v == nullptr || v[0] == '\0' || atoi(v) != 0; }(); - ggml_tensor* q = nullptr; - ggml_tensor* k = nullptr; if (fuse_norm) { q = ggml_rope_pe_permute_rms(ctx->ggml_ctx, part(0), pe, n_rot, 1.f, GGML_TYPE_F32, 1e-5f); k = ggml_rope_pe_permute_rms(ctx->ggml_ctx, part(1), pe, n_rot, kv_scale, GGML_TYPE_F16, 1e-5f); @@ -339,6 +340,7 @@ namespace MiniMaxH3VAE { q = k = nullptr; } } +#endif if (q == nullptr) { q = ggml_rope_pe_permute(ctx->ggml_ctx, ggml_rms_norm(ctx->ggml_ctx, part(0), 1e-5f), pe, n_rot, 1.f, GGML_TYPE_F32); k = ggml_rope_pe_permute(ctx->ggml_ctx, ggml_rms_norm(ctx->ggml_ctx, part(1), 1e-5f), pe, n_rot, kv_scale, GGML_TYPE_F16); From 338079ece82649c7535da4bbba663f29b4bc1a25 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Thu, 8 Oct 2026 11:10:53 -0300 Subject: [PATCH 26/30] Tighten the H3 VAE, sage and tiling comments --- src/core/ggml_runner.h | 3 +-- src/device_residency_manager.h | 4 +--- src/model/diffusion/minimax_h3.hpp | 6 ++---- src/model/vae/ltx_audio_vae.hpp | 8 ++------ src/model/vae/minimax_h3_vae.hpp | 25 +++++++------------------ src/model_manager.cpp | 3 +-- src/model_manager.h | 1 - src/runtime/tiling.cpp | 5 +---- 8 files changed, 15 insertions(+), 40 deletions(-) diff --git a/src/core/ggml_runner.h b/src/core/ggml_runner.h index 9b503b407c..c5e1a65164 100644 --- a/src/core/ggml_runner.h +++ b/src/core/ggml_runner.h @@ -370,8 +370,7 @@ struct GGMLRunner { weight_adapter = adapter; } - // See DeviceMemoryRequest::reuse_device_query. Callers turn it on around a run of identical - // computes (spatial tiles) and off afterwards; runner_end() drops the saved reading. + // See DeviceMemoryRequest::reuse_device_query; runner_end() drops the saved reading. void set_reuse_device_query(bool enabled) { reuse_device_query_ = enabled; } diff --git a/src/device_residency_manager.h b/src/device_residency_manager.h index fcdb5c389c..0be0d2ee42 100644 --- a/src/device_residency_manager.h +++ b/src/device_residency_manager.h @@ -27,9 +27,7 @@ struct DeviceMemoryRequest { size_t pending_allocation_bytes = 0; size_t runtime_resident_bytes = 0; size_t max_backend_bytes = 0; - // The caller repeats an identical compute (e.g. one tile after another). When nothing new is - // allocated (no pending bytes, every parameter already resident), the manager may answer from - // the device free-memory figure it read for this owner earlier instead of querying again. + // Repeated identical compute: a check that allocates nothing may reuse this owner's last free-memory reading. bool reuse_device_query = false; // Runtime buffers only; the manager accounts for weights separately. diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index 9846ba6e7d..b065f02bd1 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -153,8 +153,7 @@ namespace MiniMaxH3 { return enabled; } - // SD_H3_FAST_SAGE_QKV=0: with --sage-attn, q/k/v go through the unfused chunk / slice / rope / - // concat / scale / cast chain instead of one fused RoPE + head-major op per tensor. + // SD_H3_FAST_SAGE_QKV=0: --sage-attn keeps the unfused q/k/v chain. static bool fast_sage_qkv() { static const bool enabled = fast_qkv() && env_flag("SD_H3_FAST_SAGE_QKV", true); return enabled; @@ -336,8 +335,7 @@ namespace MiniMaxH3 { }; auto q = q_norm->forward(ctx, part(0)); auto k = k_norm->forward(ctx, part(1)); - // Sage reads F32 Q/K and F16 V, all [head_dim, tokens, heads, batch]: the layout this op - // writes. The unfused sage path scales K and V in F32 before V's F16 cast, as here. + // sage reads F32 Q/K and F16 V in this layout; K and V are scaled in F32, as on the unfused path q = ggml_rope_pe_permute(ctx->ggml_ctx, q, pe, n_rot, 1.f, GGML_TYPE_F32); k = ggml_rope_pe_permute(ctx->ggml_ctx, k, pe, n_rot, kv_scale, sage ? GGML_TYPE_F32 : GGML_TYPE_F16); auto v = ggml_rope_pe_permute(ctx->ggml_ctx, part(2), nullptr, 0, kv_scale, GGML_TYPE_F16); diff --git a/src/model/vae/ltx_audio_vae.hpp b/src/model/vae/ltx_audio_vae.hpp index 7f3098fb3b..2ff8528ecc 100644 --- a/src/model/vae/ltx_audio_vae.hpp +++ b/src/model/vae/ltx_audio_vae.hpp @@ -396,9 +396,7 @@ namespace LTXV { return ggml_reshape_4d(ctx, out, out_time, channels, 1, 1); } - // Direct depthwise conv (one F32 kernel per channel, no im2col, no F16 rounding of the input) for - // the anti-aliased activations. x: [time, channels, 1, 1]. Returns nullptr when the backend lacks - // CONV_2D_DW so callers fall back to the im2col graph. + // F32 depthwise conv without im2col or F16 input rounding; nullptr when the backend lacks CONV_2D_DW. static ggml_tensor* depthwise_conv1d_direct(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* filter, @@ -427,8 +425,7 @@ namespace LTXV { return ggml_reshape_4d(ctx, out, out->ne[0], channels, 1, 1); } - // depthwise_conv_transpose1d through the direct depthwise conv: same zero-stuffing, reversed - // filter and full padding, then the same stride gain. + // depthwise_conv_transpose1d via the direct conv: same zero-stuffing, reversed filter, padding and gain. static ggml_tensor* depthwise_conv_transpose1d_direct(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* filter, @@ -735,7 +732,6 @@ namespace LTXV { int up_kernel_size = 12; int down_kernel_size = 12; - // direct depthwise convs instead of the im2col graph; opt-in per model (MiniMax-H3 sets it) bool direct_dw = false; explicit Activation1D(int64_t channels, bool direct_dw = false) diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index a4ec74fd0a..68cb9a9a27 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -290,11 +290,8 @@ namespace MiniMaxH3VAE { blocks["to_out"] = std::make_shared(dim, dim, true); } - // Single-tile flash-attention decode: q/k RMS-normed in place on the projection, then one - // ggml_rope_pe_permute per tensor applies the partial RoPE, writes the head-major layout the - // attention kernel reads and casts K/V to F16. Same products, sums and casts as the table - // RoPE + chunk / permute / cast chain, so the result is bit-identical. Returns nullptr when - // it does not apply (SD_H3_VAE_FUSED_QKV=0 turns it off). + // Single-tile flash-attention decode: one ggml_rope_pe_permute per tensor replaces the table RoPE + + // chunk / permute / cast chain, bit-identically. nullptr when it does not apply (SD_H3_VAE_FUSED_QKV=0). ggml_tensor* forward_head_major(GGMLRunnerContext* ctx, ggml_tensor* qkv, ggml_tensor* pe, @@ -327,8 +324,7 @@ namespace MiniMaxH3VAE { ggml_tensor* q = nullptr; ggml_tensor* k = nullptr; #ifdef SD_GGML_H3_FUSED_QK_NORM - // q/k RMS norm folded into the RoPE op where the backend reproduces GGML_OP_RMS_NORM's - // arithmetic in it (SD_H3_VAE_FUSED_QK_NORM=0 keeps the separate norm). + // same reduction as GGML_OP_RMS_NORM where supported (SD_H3_VAE_FUSED_QK_NORM=0 keeps it separate) static const bool fuse_norm = [] { const char* v = getenv("SD_H3_VAE_FUSED_QK_NORM"); return v == nullptr || v[0] == '\0' || atoi(v) != 0; @@ -749,8 +745,7 @@ namespace MiniMaxH3VAE { params.tile_size_x = 16; params.tile_size_y = 16; params.target_overlap = 0.25f; - // SD_H3_VAE_TILE=N: N x N latent tiles instead of 16 x 16 (opt-in; the tile seams move, so - // the frames are not bit-identical to the default) + // SD_H3_VAE_TILE=N: opt-in N x N latent tiles; the seams move, so frames differ from the default if (const char* tile = getenv("SD_H3_VAE_TILE")) { const int n = atoi(tile); if (n >= 8) { @@ -888,8 +883,7 @@ namespace MiniMaxH3VAE { {static_cast(tokens_per_chunk + token_overlap), static_cast(token_overlap)}); GGML_ASSERT(plan.tiles.size() == static_cast(num_chunks)); const bool keep_resident = env_int("SD_H3_VAE_KEEP_RESIDENT", 1) != 0; - // One device free-memory reading serves every tile that allocates nothing new - // (SD_H3_VAE_REUSE_MEMQUERY=0 queries the device for every capacity check). + // SD_H3_VAE_REUSE_MEMQUERY=0 queries the device for every capacity check. set_reuse_device_query(env_int("SD_H3_VAE_REUSE_MEMQUERY", 1) != 0); struct ReuseQueryGuard { GGMLRunner& runner; @@ -900,9 +894,7 @@ namespace MiniMaxH3VAE { tile_batch_ = 0; per_tile_compute_bytes_ = 0; std::vector> pieces; - // Trims the pre-padding frames of a decoded chunk and cross-fades it with the previous - // chunk's tail. Runs in chunk order; on a worker thread unless SD_H3_VAE_ASYNC_ASSEMBLY=0, - // so it overlaps the next chunk's decode (same arithmetic, bit-identical result). + // Runs in chunk order, on a worker thread unless SD_H3_VAE_ASYNC_ASSEMBLY=0. auto assemble = [&](const sd::Tensor& decoded, const VAETemporalTile& tile) { int64_t first_end = std::min(frames_per_chunk, decoded.shape()[2]); auto first = sd::ops::slice(decoded, @@ -931,8 +923,7 @@ namespace MiniMaxH3VAE { bool assembly_failed = false; int64_t expected_frames = input.shape()[2] <= 1 ? 1 : ((x.shape()[2] - 2) / 5) * 17 + 5; expected_frames = std::max(1, expected_frames); - // Pieces are copied straight into the trimmed output as they are assembled (the same - // frames concat_frames + slice would keep), instead of concatenated at the end. + // keeps the same frames concat_frames + slice would sd::Tensor result; int64_t frame_offset = 0; auto append_piece = [&](const sd::Tensor& piece) { @@ -1232,8 +1223,6 @@ namespace MiniMaxH3VAE { auto input = ensure_video_shape(z); const bool graph_opt = decode_graph && env_int("SD_H3_VAE_GRAPH_OPT", 1) != 0; if (decode_graph) { - // Every spatial tile of a decode has the same extent, so the tables are rebuilt only - // when the tile shape (or the graph variant) changes. const std::array rope_key = {input.shape()[0], input.shape()[1], input.shape()[2], graph_opt ? 1 : 0}; if (rope_key != rope_key_ || rope_cache.empty()) { rope_cache = build_rope(input.shape()[0], diff --git a/src/model_manager.cpp b/src/model_manager.cpp index a3554da546..e02270940b 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -1633,8 +1633,7 @@ ModelManager::CapacityCheck ModelManager::check_capacity( return SIZE_MAX; } size_t free_bytes = 0, total_bytes = 0; - // A reading is saved and reused only for checks that allocate nothing (no pending bytes, - // every parameter resident), so it is never one taken before this owner's buffers grew. + // only checks that allocate nothing reuse a reading, so it never predates this owner's growth const bool reusable = request.reuse_device_query && request.pending_allocation_bytes == 0 && missing == 0; const auto cache_key = std::make_pair(request.owner_id, device); auto cached = device_query_cache_.find(cache_key); diff --git a/src/model_manager.h b/src/model_manager.h index 394391de8f..e45a329536 100644 --- a/src/model_manager.h +++ b/src/model_manager.h @@ -94,7 +94,6 @@ class ModelManager : public DeviceResidencyManager { std::vector> tensor_states_; std::map tensor_states_by_tensor_; mutable std::list resolved_tensor_states_; - // Last device free-memory reading per (owner, device), for DeviceMemoryRequest::reuse_device_query. mutable std::map, std::pair> device_query_cache_; std::vector> params_storage_blocks_; std::vector> compute_staging_blocks_; diff --git a/src/runtime/tiling.cpp b/src/runtime/tiling.cpp index 48fd882756..63f7fe9ae7 100644 --- a/src/runtime/tiling.cpp +++ b/src/runtime/tiling.cpp @@ -333,10 +333,7 @@ sd::Tensor process_tiles_2d_batched(const sd::Tensor& input, LOG_VERBOSE("processing %i tiles", num_tiles); pretty_progress(0, num_tiles, 0.0f); } - // The host-side blend of a batch runs on a worker thread while the next batch is split and - // computed, so the device does not idle through it. Merges still run one at a time in tile - // order with the same arithmetic, so the result is bit-identical. SD_TILE_ASYNC_MERGE=0 merges - // inline. + // Blend on a worker while the next batch computes, still in tile order; SD_TILE_ASYNC_MERGE=0 blends inline. static const bool async_merge = []() { const char* value = getenv("SD_TILE_ASYNC_MERGE"); return value == nullptr || value[0] == '\0' || atoi(value) != 0; From aff34f1542edb66acf2f5bb5daae1c895db0afe4 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Thu, 8 Oct 2026 12:24:21 -0300 Subject: [PATCH 27/30] Ship cudnn-frontend's MIT notice in the CUDA bundle --- .github/workflows/unsloth-sd-prebuilt.yml | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/.github/workflows/unsloth-sd-prebuilt.yml b/.github/workflows/unsloth-sd-prebuilt.yml index 86b8cba0b3..b802f4aefe 100644 --- a/.github/workflows/unsloth-sd-prebuilt.yml +++ b/.github/workflows/unsloth-sd-prebuilt.yml @@ -544,8 +544,16 @@ jobs: LABEL: Linux-Ubuntu-22.04-x86_64-cuda12 COMMIT: ${{ needs.resolve.outputs.commit }} SOURCE_REPO: ${{ github.repository }} - LICENSE_FILE: ${{ github.workspace }}/src/LICENSE - run: python3 tooling/scripts/unsloth/package_bundle.py + LICENSE_FILE: ${{ runner.temp }}/LICENSE + run: | + set -euo pipefail + # The binaries compile in cudnn-frontend's header-only templates, so its MIT notice ships too. + { + cat src/LICENSE + printf '\n\n---- cudnn-frontend (https://github.com/NVIDIA/cudnn-frontend) ----\n\n' + cat src/build/_deps/ggml_cudnn_frontend-src/LICENSE.txt + } > "$LICENSE_FILE" + python3 tooling/scripts/unsloth/package_bundle.py - name: Upload bundle uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 From ff941a8a853c951c4f77310d5ba204ea7bf9ac0e Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Thu, 8 Oct 2026 22:51:20 -0300 Subject: [PATCH 28/30] Keep the H3 VAE on its old attention outside CUDA and ROCm --- docs/minimax_h3.md | 8 +++++--- src/pipeline/diffusion_engine.cpp | 11 +++++++---- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index 0417d0011a..bcceb73c47 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -46,9 +46,11 @@ are detected from their weights. Omitting `--audio-vae` still runs the joint diffusion model but produces video without a decoded audio track. -`--diffusion-fa` also enables flash attention in the video VAE decoder (a ViT), which removes -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 CUDA and ROCm, `--diffusion-fa` also enables flash attention in the video VAE decoder (a ViT), +which removes 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. Other backends +keep that attention by default (on Vulkan the flash attention decode was slower); +`SD_H3_VAE_FLASH_ATTN=1` turns it on there. ## First/last-frame conditioning diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index ec50c620cc..db5800991c 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -1178,12 +1178,15 @@ bool StableDiffusionGGML::validate_and_load_runners() { high_noise_diffusion_model->set_flash_attention_enabled(true); } } - // The H3 video VAE decoder is a ViT whose f32 L x L attention scores dominate decode time, so --diffusion-fa - // covers it too, without switching the text encoder. SD_H3_VAE_FLASH_ATTN=0 restores the old path. + // The H3 video VAE decoder is a ViT whose f32 L x L attention scores dominate decode time, so on CUDA and ROCm + // --diffusion-fa covers it too, without switching the text encoder. It decoded slower on Vulkan, so other + // backends keep the old path unless SD_H3_VAE_FLASH_ATTN=1; SD_H3_VAE_FLASH_ATTN=0 restores it everywhere. if (!sd_ctx_params->flash_attn && sd_ctx_params->diffusion_flash_attn && first_stage_model && sd_version_is_minimax_h3(version)) { - const char* env = getenv("SD_H3_VAE_FLASH_ATTN"); - if (env == nullptr || strcmp(env, "0") != 0) { + const char* env = getenv("SD_H3_VAE_FLASH_ATTN"); + ggml_backend_t vae_backend = backend_for(SDBackendModule::VAE); + const bool fa_faster_by_default = sd_backend_is(vae_backend, "CUDA") || sd_backend_is(vae_backend, "ROCm"); + if (env != nullptr ? strcmp(env, "0") != 0 : fa_faster_by_default) { LOG_INFO("Using flash attention in the MiniMax-H3 video VAE"); first_stage_model->set_flash_attention_enabled(true); } From 195712330b05c66836c3ea23cb238556d681a305 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Fri, 9 Oct 2026 00:28:22 -0300 Subject: [PATCH 29/30] Take cuDNN attention by default only where it measured faster --- docs/minimax_h3.md | 10 +++- .../0006-ggml-cuda-cudnn-attention.patch | 58 +++++++++++++------ 2 files changed, 47 insertions(+), 21 deletions(-) diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index d4ae452644..d50b9f8b18 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -59,7 +59,10 @@ stock kernel. `GGML_CUDA_FA_LONGSEQ_NCOLS=128` opts into a wider tile that is fa but not bit-identical. Builds configured with `-DGGML_CUDA_CUDNN=ON` (the Linux CUDA prebuilt is) can run that unmasked -attention, and the `--sage-attn` attention, through cuDNN's fused attention on Ampere and newer. +attention, and the `--sage-attn` attention, through cuDNN's fused attention. By default it replaces +the flash attention kernel on Ada (sm89) and B200-class (sm100) GPUs and the sage kernel on sm100 only, +where it measured faster; on an RTX 3090 (sm86) both existing kernels were faster, on an RTX 6000 Ada +the sage kernel was, and other architectures were not measured. Only `cudnn.h` is needed at build time (`GGML_CUDA_CUDNN_INCLUDE_DIR`; the header-only cudnn-frontend v1.26.0 is downloaded by CMake unless `GGML_CUDA_CUDNN_FRONTEND_DIR` points at a checkout). `libcudnn.so.9` is opened when the first attention op runs and is not shipped: put a cuDNN 9 @@ -72,8 +75,9 @@ existing kernels run. The first call of each attention shape builds a cuDNN plan on a B200, logged once). cuDNN runs F16 with F32 accumulation: for the H3 DiT at 960x544x124 it takes about 8 ms per attention call on a B200 against 28 ms for the sage kernel and 43 ms for the ggml flash attention kernel, and it is closer to an exact (fp64) reference than either, so frames -and audio differ slightly from the previous kernels. `GGML_CUDA_CUDNN_ATTN=0` restores them, -`GGML_CUDA_CUDNN_SAGE=0` keeps the sage kernel for `--sage-attn` only, and +and audio differ slightly from the previous kernels. `GGML_CUDA_CUDNN_ATTN=0` restores them and +`GGML_CUDA_CUDNN_ATTN=1` uses cuDNN for the flash attention op on any Ampere or newer GPU; +`GGML_CUDA_CUDNN_SAGE=0` / `=1` does the same for the sage op only, and `GGML_CUDA_CUDNN_ATTN_BF16=1` runs cuDNN in BF16 (less accurate, same speed). The video VAE decodes one 16x16 latent tile per decoder graph by default. `SD_H3_VAE_TILE=N` diff --git a/scripts/unsloth/ggml-patches/0006-ggml-cuda-cudnn-attention.patch b/scripts/unsloth/ggml-patches/0006-ggml-cuda-cudnn-attention.patch index 7027a8e995..2b23434e88 100644 --- a/scripts/unsloth/ggml-patches/0006-ggml-cuda-cudnn-attention.patch +++ b/scripts/unsloth/ggml-patches/0006-ggml-cuda-cudnn-attention.patch @@ -1,4 +1,4 @@ -From 80d308615f1dc9fc601f556123861a651b7cffa5 Mon Sep 17 00:00:00 2001 +From 784b441f8b098863343e69e0345ef62d2ff580af Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 3 Oct 2026 18:27:29 -0700 Subject: [PATCH] ggml-cuda: optional cuDNN attention for unmasked @@ -11,7 +11,7 @@ only cudnn.h is needed at build time and libcudnn.so.9 is dlopen'ed on first use (GGML_CUDA_CUDNN_LIB overrides the name), so the binary has no cuDNN link dependency and runs unchanged where cuDNN is missing. -On NVIDIA sm80+ it takes unmasked, sink/ALiBi/softcap-free attention with equal +On NVIDIA sm80+ it can take unmasked, sink/ALiBi/softcap-free attention with equal Q/K/V head sizes (multiple of 8, up to 256) and at least 2^20 attention scores, GQA and batch included, for FLASH_ATTN_EXT and for SAGE_ATTN in AUTO mode (an explicitly requested Sage variant keeps its kernel). Q (F32) is narrowed to F16 @@ -22,24 +22,28 @@ kernels' needs). Plans are built once per shape and cached; a shape cuDNN declines, a failed build or execute, stream capture, Turing and older, or a missing library keep the existing kernels. -GGML_CUDA_CUDNN_ATTN=0 disables it, GGML_CUDA_CUDNN_SAGE=0 keeps the Sage kernel, -GGML_CUDA_CUDNN_ATTN_BF16=1 runs BF16, GGML_CUDA_CUDNN_ATTN_LOG=1 logs each call. +By default FLASH_ATTN_EXT takes it on sm89 and sm100 and SAGE_ATTN on sm100, +where it measured faster; the existing kernels were faster on sm86 (both ops) +and sm89 (Sage), and other architectures were not measured. +GGML_CUDA_CUDNN_ATTN=0|1 and GGML_CUDA_CUDNN_SAGE=0|1 turn it off or on for every +sm80+ GPU, GGML_CUDA_CUDNN_ATTN_BF16=1 runs BF16, GGML_CUDA_CUDNN_ATTN_LOG=1 +logs each call. test-backend-ops: unmasked cases at the MiniMax-H3 head layout (odd and aligned lengths, GQA with batch, F32 K/V, permuted Q) and unmasked perf cases. --- CMakeLists.txt | 3 + src/ggml-cuda/CMakeLists.txt | 49 ++++ - src/ggml-cuda/cudnn-attn-graph.cpp | 258 ++++++++++++++++++++ + src/ggml-cuda/cudnn-attn-graph.cpp | 258 +++++++++++++++++++ src/ggml-cuda/cudnn-attn-graph.h | 31 +++ - src/ggml-cuda/cudnn-attn.cu | 367 +++++++++++++++++++++++++++++ - src/ggml-cuda/cudnn-attn.cuh | 25 ++ + src/ggml-cuda/cudnn-attn.cu | 384 +++++++++++++++++++++++++++++ + src/ggml-cuda/cudnn-attn.cuh | 26 ++ src/ggml-cuda/fattn.cu | 7 +- src/ggml-cuda/ggml-cuda.cu | 2 +- src/ggml-cuda/sage-attn.cu | 12 +- src/ggml-cuda/sage-attn.cuh | 2 +- tests/test-backend-ops.cpp | 25 ++ - 11 files changed, 775 insertions(+), 6 deletions(-) + 11 files changed, 793 insertions(+), 6 deletions(-) create mode 100644 src/ggml-cuda/cudnn-attn-graph.cpp create mode 100644 src/ggml-cuda/cudnn-attn-graph.h create mode 100644 src/ggml-cuda/cudnn-attn.cu @@ -427,10 +431,10 @@ index 00000000..9749f2ab + void * q, void * k, void * v, void * o, void * workspace); diff --git a/src/ggml-cuda/cudnn-attn.cu b/src/ggml-cuda/cudnn-attn.cu new file mode 100644 -index 00000000..b54c29af +index 00000000..a5b9de14 --- /dev/null +++ b/src/ggml-cuda/cudnn-attn.cu -@@ -0,0 +1,367 @@ +@@ -0,0 +1,384 @@ +#include "cudnn-attn.cuh" + +#if defined(GGML_CUDA_USE_CUDNN) && !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) @@ -470,11 +474,20 @@ index 00000000..b54c29af + return t->nb[0] == ggml_type_size(t->type) && t->nb[1] % 16 == 0 && t->nb[2] % 16 == 0 && t->nb[3] % 16 == 0; +} + ++// -1 when unset, otherwise 0 or 1. ++static int cudnn_attn_env_mode(const char * name) { ++ const char * v = std::getenv(name); ++ if (v == nullptr || v[0] == '\0') { ++ return -1; ++ } ++ return cudnn_attn_env_set(name, true) ? 1 : 0; ++} ++ +static bool cudnn_attn_make_plan(int device, const ggml_tensor * dst, cudnn_attn_plan & plan) { -+ static const bool enabled = cudnn_attn_env_set("GGML_CUDA_CUDNN_ATTN", true); -+ static const bool sage = cudnn_attn_env_set("GGML_CUDA_CUDNN_SAGE", true); ++ static const int attn_mode = cudnn_attn_env_mode("GGML_CUDA_CUDNN_ATTN"); ++ static const int sage_mode = cudnn_attn_env_mode("GGML_CUDA_CUDNN_SAGE"); + static const bool want_bf16 = cudnn_attn_env_set("GGML_CUDA_CUDNN_ATTN_BF16", false); -+ if (!enabled || (dst->op == GGML_OP_SAGE_ATTN && !sage)) { ++ if (attn_mode == 0 || (dst->op == GGML_OP_SAGE_ATTN && sage_mode == 0)) { + return false; + } + if (dst->op != GGML_OP_FLASH_ATTN_EXT && dst->op != GGML_OP_SAGE_ATTN) { @@ -488,6 +501,14 @@ index 00000000..b54c29af + if (!GGML_CUDA_CC_IS_NVIDIA(cc) || cc < GGML_CUDA_CC_AMPERE) { + return false; + } ++ // Unset switches keep the existing kernels except where cuDNN measured faster at H3 shapes: against the flash ++ // attention kernel on sm89 and sm100 (slower on sm86), against the Sage kernel on sm100 (slower on sm86 and ++ // sm89). Other architectures were not measured. =1 opts in on any sm80+ GPU. ++ const bool sm100 = cc >= 1000 && cc < GGML_CUDA_CC_BLACKWELL; ++ const bool sm89 = cc >= GGML_CUDA_CC_ADA_LOVELACE && cc < GGML_CUDA_CC_HOPPER; ++ if (dst->op == GGML_OP_SAGE_ATTN ? sage_mode < 0 && !sm100 : attn_mode < 0 && !sm89 && !sm100) { ++ return false; ++ } + + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; @@ -800,10 +821,10 @@ index 00000000..b54c29af +#endif diff --git a/src/ggml-cuda/cudnn-attn.cuh b/src/ggml-cuda/cudnn-attn.cuh new file mode 100644 -index 00000000..f59689d5 +index 00000000..067bea63 --- /dev/null +++ b/src/ggml-cuda/cudnn-attn.cuh -@@ -0,0 +1,25 @@ +@@ -0,0 +1,26 @@ +#pragma once + +#include "common.cuh" @@ -814,10 +835,11 @@ index 00000000..f59689d5 +// least 2^20 attention scores is taken; everything else, and every shape cuDNN rejects, keeps the +// existing kernels. +// ++// By default FLASH_ATTN_EXT takes it on sm89 and sm100 and SAGE_ATTN on sm100, where it measured faster. +// Environment: -+// GGML_CUDA_CUDNN_ATTN=0 never use cuDNN for attention -+// GGML_CUDA_CUDNN_SAGE=0 keep the SageAttention kernels for GGML_OP_SAGE_ATTN (AUTO mode; -+// an explicitly requested Sage variant always keeps its kernel) ++// GGML_CUDA_CUDNN_ATTN=0|1 never use cuDNN for attention / use it on every sm80+ GPU ++// GGML_CUDA_CUDNN_SAGE=0|1 keep the SageAttention kernels for GGML_OP_SAGE_ATTN / use cuDNN on every ++// sm80+ GPU (AUTO mode; an explicitly requested Sage variant keeps its kernel) +// GGML_CUDA_CUDNN_ATTN_BF16=1 run cuDNN in BF16 instead of F16 (F16 K/V are narrowed) +// GGML_CUDA_CUDNN_LIB=path library to open instead of libcudnn.so.9 +// GGML_CUDA_CUDNN_ATTN_LOG=1 log every call that takes (or falls back from) the cuDNN path From ae9429f471a0e657a84139ed297018a4cbe2ca31 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Fri, 9 Oct 2026 02:30:09 -0300 Subject: [PATCH 30/30] Install the NVRTC header the cuDNN frontend includes in the CUDA prebuilt leg --- .github/workflows/unsloth-sd-prebuilt.yml | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/.github/workflows/unsloth-sd-prebuilt.yml b/.github/workflows/unsloth-sd-prebuilt.yml index b802f4aefe..4ffdbd95ae 100644 --- a/.github/workflows/unsloth-sd-prebuilt.yml +++ b/.github/workflows/unsloth-sd-prebuilt.yml @@ -425,7 +425,8 @@ jobs: # prefix (it ships as libcublas / libcublas-dev), so it has to go in the # other list or apt cannot find it. cudart-dev is what carries the headers; # cudart alone is the runtime and does not compile anything. - sub-packages: '["nvcc", "cudart", "cudart-dev", "thrust"]' + # nvrtc-dev: cudnn-frontend's headers include nvrtc.h (it dlopens libnvrtc itself, so nothing links it). + sub-packages: '["nvcc", "cudart", "cudart-dev", "thrust", "nvrtc-dev"]' non-cuda-sub-packages: '["libcublas", "libcublas-dev"]' # This leg rebuilt every object on every run: 3903 s of the 4121 s job on @@ -527,11 +528,11 @@ jobs: patchelf --set-rpath '$ORIGIN' "$BIN/$exe" done ldd "$BIN/sd-cli" | sed -n '1,40p' - # cuDNN is opened at run time only; a link-time dependency would stop the binaries + # cuDNN and NVRTC are opened at run time only; a link-time dependency would stop the binaries # loading on every host without it. for exe in sd-cli sd-server; do - if readelf -d "$BIN/$exe" | grep -q 'NEEDED.*cudnn'; then - echo "ERROR: $exe links libcudnn; it must stay a run-time dlopen" >&2 + if readelf -d "$BIN/$exe" | grep -qE 'NEEDED.*(cudnn|nvrtc)'; then + echo "ERROR: $exe links libcudnn or libnvrtc; they must stay run-time dlopens" >&2 exit 1 fi done