From c0eb17539fc0d336928f86498a7ede9f7fe9d569 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 1 Oct 2026 22:42:50 -0700 Subject: [PATCH 1/7] 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 4b967e0e3..cc40daa14 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 07a244a8a..cc2c24fc4 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 2/7] 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 cc40daa14..97bcaf8cc 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 abfa8639b..3ac8e0b2d 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 29a6808bd..9b503b407 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 0a2a8895f..fcdb5c389 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 604a05fb9..fd09f4a74 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 9a9fe0628..d6cd88a48 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 7f8df59c9..394391de8 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 1bfb291a7..9511789fc 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 3/7] 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 97bcaf8cc..634128da2 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 d2314b56a..7f3098fb3 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 a4fc14473..300826a03 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 4/7] 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 634128da2..eed2229d2 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 fd09f4a74..d569090a2 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 5/7] 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 eed2229d2..03ce876d9 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 000000000..65b336de5 --- /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 d569090a2..d7cc87b95 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 6/7] 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 03ce876d9..7a8255db7 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 d7cc87b95..8412ebf1a 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 7/7] 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 8412ebf1a..f4cff90af 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 d6cd88a48..a3554da54 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 9511789fc..1dcf35cd5 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); }