feat: add configurable attention chunk budget for MiniMax-H3 (#2292)

docs/minimax_h3.md

@@ -155,5 +155,17 @@ Notes for this mode:
 - Width and height are aligned upward to a multiple of 32.
 - Frame count is aligned upward to the `17k + 5` grid, with a minimum of 5.
 - MiniMax-H3 runs at 24 fps; another requested value is overridden.
+- Attention head grouping is controlled by
+  `--model-args minimax_h3_attn_chunk_budget=2`, in GiB (default: `2`).
+  The budget estimates the size of a single FP32 Q/K tensor per group:
+  `head_dim * tokens * batch * 4 bytes * heads_per_group`. The group size adapts
+  to the sequence length, with a minimum of one head; inputs that fit use all
+  heads together. Set `0` to disable grouping, or a smaller budget such as
+  `0.25` to reduce peak activation memory for long sequences at the cost of
+  extra launches and concatenation. Each head still attends to the full sequence.
+  This is a grouping target, not a cap on total attention memory or VRAM usage;
+  other intermediates, backend workspace, and weights require additional memory.
+  `--max-vram` and `--offload-to-cpu` control weight residency, but the active
+  block's working memory must also fit within the VRAM budget.
 - The default video flow shift is 12. The audio stream is mapped internally to
   its shift of 3, so the regular samplers can operate on the packed AV latent.

src/core/cuda/i8_gated_residual.cu

@@ -426,7 +426,7 @@ namespace sd_i8_gated {
         return false;
     }
     static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
-        auto* plan   = new Plan;
+        auto* plan    = new Plan;
         *state        = plan;
         plan->device  = device_info(ggml_backend_get_device(backend));
         const auto op = describe(nodes, count);

src/core/cuda/i8_gemm.cu

@@ -179,19 +179,19 @@ namespace sd_i8_cutlass {
     template <class Config, bool StreamK, bool PackedW4 = false>
     struct DenseTypes {
         static constexpr bool kSplit = false;
-        using Weight   = std::conditional_t<PackedW4, cutlass::int4b_t, int8_t>;
+        using Weight                 = std::conditional_t<PackedW4, cutlass::int4b_t, int8_t>;
-        using Tile     = typename Config::Tile;
+        using Tile                   = typename Config::Tile;
-        using Warp     = typename Config::Warp;
+        using Warp                   = typename Config::Warp;
-        using Swizzle  = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
+        using Swizzle                = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
-        using Map      = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
+        using Map                    = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
-        using Col      = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
+        using Col                    = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
-        using Row      = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
+        using Row                    = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
-        using Mul      = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+        using Mul                    = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
-        using ScaleX   = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
+        using ScaleX                 = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
-        using ScaleW   = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
+        using ScaleW                 = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
-        using Store    = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
+        using Store                  = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
-        using Epilogue = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
+        using Epilogue               = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
-        using Kernel   = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
+        using Kernel                 = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
             int8_t,
             cutlass::layout::RowMajor,
             cutlass::ComplexTransform::kNone,
@@ -219,20 +219,20 @@ namespace sd_i8_cutlass {
     template <class Config, bool StreamK, bool Split = false>
     struct SwigluTypes {
         static constexpr bool kSplit = Split;
-        using Weight          = int8_t;
+        using Weight                 = int8_t;
-        using Tile            = typename Config::Tile;
+        using Tile                   = typename Config::Tile;
-        using Warp            = typename Config::Warp;
+        using Warp                   = typename Config::Warp;
-        using Swizzle         = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
+        using Swizzle                = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
-        using Map             = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
+        using Map                    = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
-        using Col             = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
+        using Col                    = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
-        using Row             = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
+        using Row                    = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
-        using Mul             = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+        using Mul                    = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
-        using ScaleX          = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
+        using ScaleX                 = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
-        using ScaleW          = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
+        using ScaleW                 = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
-        using Store           = SwigluStore<Map>;
+        using Store                  = SwigluStore<Map>;
-        using Epilogue        = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
+        using Epilogue               = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
-        using GemmBase        = typename cutlass::gemm::kernel::DefaultGemmUniversal<int8_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 16, int8_t, cutlass::layout::ColumnMajor, cutlass::ComplexTransform::kNone, 16, float, cutlass::layout::RowMajor, int, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm80, Tile, Warp, cutlass::gemm::GemmShape<16, 8, 32>, cutlass::epilogue::thread::LinearCombination<float, 4, int, float>, Swizzle, 4, cutlass::arch::OpMultiplyAddSaturate, cutlass::gemm::SharedMemoryClearOption::kNone, false, false, false, cutlass::layout::NoPermute, cutlass::layout::NoPermute, WeightLayout>::GemmKernel;
+        using GemmBase               = typename cutlass::gemm::kernel::DefaultGemmUniversal<int8_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 16, int8_t, cutlass::layout::ColumnMajor, cutlass::ComplexTransform::kNone, 16, float, cutlass::layout::RowMajor, int, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm80, Tile, Warp, cutlass::gemm::GemmShape<16, 8, 32>, cutlass::epilogue::thread::LinearCombination<float, 4, int, float>, Swizzle, 4, cutlass::arch::OpMultiplyAddSaturate, cutlass::gemm::SharedMemoryClearOption::kNone, false, false, false, cutlass::layout::NoPermute, cutlass::layout::NoPermute, WeightLayout>::GemmKernel;
-        using VisitorEpilogue = cutlass::epilogue::threadblock::EpilogueWithVisitorCallbacks<typename GemmBase::Epilogue, Epilogue, 1>;
+        using VisitorEpilogue        = cutlass::epilogue::threadblock::EpilogueWithVisitorCallbacks<typename GemmBase::Epilogue, Epilogue, 1>;
         using BaseMma                = typename GemmBase::Mma;
         using SplitMma               = cutlass::gemm::threadblock::MmaMultistage<
             typename BaseMma::Shape,
@@ -246,11 +246,11 @@ namespace sd_i8_cutlass {
             typename BaseMma::LayoutC,
             typename BaseMma::Policy,
             BaseMma::Detail::kStages>;
-        using Mma             = std::conditional_t<Split, SplitMma, BaseMma>;
+        using Mma    = std::conditional_t<Split, SplitMma, BaseMma>;
-        using Kernel          = std::conditional_t<StreamK,
+        using Kernel = std::conditional_t<StreamK,
                                           cutlass::gemm::kernel::GemmWithEpilogueVisitorStreamk<Mma, VisitorEpilogue, Swizzle>,
                                           cutlass::gemm::kernel::GemmWithEpilogueVisitor<Mma, VisitorEpilogue, Swizzle>>;
-        using Gemm            = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
+        using Gemm   = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
     };
     struct Plan {
         DeviceInfo device;
@@ -259,7 +259,7 @@ namespace sd_i8_cutlass {
         bool packed_w4;
         bool split;
         size_t kernel_bytes = 0, scratch_offset = 0;
-        int occupancy                                                                                                       = 0;
+        int occupancy                                                                                                                                    = 0;
         ggml_status (*launch)(Plan&, const int8_t*, const int8_t*, const float*, const float*, float*, void*, cudaStream_t, const int8_t*, const float*) = nullptr;
     };
     template <class Types>

src/core/cuda/int8.cu

@@ -552,8 +552,8 @@ namespace SD_CUDA_NAMESPACE {
             const bool fresh_output = !dst->data && !dst->buffer && !dst->view_src &&
                                       (dst->op == GGML_OP_CUSTOM || dst->op == GGML_OP_CONCAT ||
                                        (dst->op == GGML_OP_GLU && !dst->src[1]));
-            p->output  = fresh_output || (p->gemm && p->kind == SWIGLU_QUANTIZE) ? SIZE_MAX : reserve(total, ggml_nbytes(dst));
+            p->output               = fresh_output || (p->gemm && p->kind == SWIGLU_QUANTIZE) ? SIZE_MAX : reserve(total, ggml_nbytes(dst));
-            *workspace = total;
+            *workspace              = total;
             return GGML_STATUS_SUCCESS;
         }
         static void destroy(void* state, void*) {

src/core/cuda/rope.cu

@@ -277,7 +277,7 @@ namespace SD_CUDA_NAMESPACE {
     }
 
     static ggml_status rope_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
-        const auto* dst         = nodes[count - 1];
+        const auto* dst = nodes[count - 1];
         // Owning casts are fresh outputs; copies into existing storage retain view_src.
         const bool fresh_output = !dst->data && !dst->buffer && !dst->view_src;
         *workspace              = fresh_output ? 0 : ggml_nbytes(dst);

src/model/diffusion/minimax_h3.hpp

@@ -19,9 +19,10 @@
 
 namespace MiniMaxH3 {
 
-    constexpr int H3_GRAPH_SIZE          = 131072;
+    constexpr int H3_GRAPH_SIZE                   = 131072;
-    constexpr float FRAME_RESCALE        = 5.f / 3.f;
+    constexpr float DEFAULT_ATTN_CHUNK_BUDGET_GIB = 2.f;
-    constexpr float VISUAL_COND_TIMESTEP = 0.999f;
+    constexpr float FRAME_RESCALE                 = 5.f / 3.f;
+    constexpr float VISUAL_COND_TIMESTEP          = 0.999f;
 
     struct Config {
         int64_t hidden_size              = 5376;
@@ -231,6 +232,7 @@ namespace MiniMaxH3 {
     struct Attention : public GGMLBlock {
         int64_t heads;
         int64_t head_dim;
+        float head_chunk_budget_gib = DEFAULT_ATTN_CHUNK_BUDGET_GIB;
 
         Attention(int64_t hidden_size,
                   int64_t heads,
@@ -315,6 +317,31 @@ namespace MiniMaxH3 {
         ggml_tensor* attend(GGMLRunnerContext* ctx,
                             const std::vector<ggml_tensor*>& qkv,
                             ggml_tensor* pe = nullptr) {
+            // Budget one F32 Q/K tensor per group, not the full attention working set.
+            const double head_bytes  = double(ggml_row_size(GGML_TYPE_F32, head_dim)) * qkv[0]->ne[2] * qkv[0]->ne[3];
+            const int64_t head_chunk = head_chunk_budget_gib > 0
+                                           ? static_cast<int64_t>(std::clamp(double(head_chunk_budget_gib) * (1024ULL * 1024 * 1024) / head_bytes,
+                                                                             1.0, double(heads)))
+                                           : heads;
+            if (head_chunk < heads) {
+                std::vector<ggml_tensor*> outputs;
+                for (int64_t start = 0; start < heads; start += head_chunk) {
+                    std::vector<ggml_tensor*> part;
+                    for (auto tensor : qkv) {
+                        part.push_back(ggml_ext_slice(ctx->ggml_ctx, tensor, 1, start, std::min(start + head_chunk, heads), false));
+                    }
+                    outputs.push_back(attend_heads(ctx, part, pe));
+                    // Keep each head group's temporaries out of the concatenation tree's live set.
+                    ctx->expand_graph(outputs.back());
+                }
+                return ggml_ext_vec_concat(ctx->ggml_ctx, outputs, 0);
+            }
+            return attend_heads(ctx, qkv, pe);
+        }
+
+        ggml_tensor* attend_heads(GGMLRunnerContext* ctx,
+                                  const std::vector<ggml_tensor*>& qkv,
+                                  ggml_tensor* pe) {
             auto q_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
             auto k_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
             auto q      = qkv[0];
@@ -333,7 +360,7 @@ namespace MiniMaxH3 {
                                               q,
                                               k,
                                               v,
-                                              static_cast<int>(heads),
+                                              static_cast<int>(qkv[0]->ne[1]),
                                               nullptr,
                                               true,
                                               ctx->flash_attn_enabled,
@@ -1035,10 +1062,30 @@ namespace MiniMaxH3 {
         MiniMaxH3Runner(ggml_backend_t backend,
                         const String2TensorStorage& tensors,
                         const std::string& prefix                           = "model.diffusion_model",
-                        std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+                        std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
+                        const char* model_args                              = nullptr)
             : DiffusionModelRunner(backend, prefix, weight_manager),
               config(Config::detect_from_weights(tensors, prefix)),
               model(config) {
+            float head_chunk_budget_gib = DEFAULT_ATTN_CHUNK_BUDGET_GIB;
+            for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
+                if (key == "minimax_h3_attn_chunk_budget") {
+                    float parsed = 0.f;
+                    if (parse_strict_float(value, parsed) && std::isfinite(parsed) && parsed >= 0.f) {
+                        head_chunk_budget_gib = parsed;
+                    } else {
+                        LOG_WARN("ignoring invalid MiniMax-H3 model arg '%s=%s': expected a finite non-negative GiB budget", key.c_str(), value.c_str());
+                    }
+                }
+            }
+            std::vector<GGMLBlock*> blocks;
+            model.get_all_blocks(blocks);
+            for (auto block : blocks) {
+                if (auto attention = dynamic_cast<Attention*>(block)) {
+                    attention->head_chunk_budget_gib = head_chunk_budget_gib;
+                }
+            }
+            LOG_INFO("MiniMax-H3 attention chunk budget: %.3g GiB (0 disables grouping)", head_chunk_budget_gib);
             model.init(params_ctx, tensors, prefix);
         }
 
@@ -1178,7 +1225,8 @@ namespace MiniMaxH3 {
                 timestep_features = make_input(timestep_feature_input_cache);
             }
 
-            auto runner_ctx = get_context();
+            auto graph      = new_graph_custom(H3_GRAPH_SIZE);
+            auto runner_ctx = get_context(graph);
             auto output     = model.forward(&runner_ctx,
                                             video,
                                             audio,
@@ -1201,7 +1249,6 @@ namespace MiniMaxH3 {
                                                     output.second,
                                                     1.f + (audio_scale - 1.f) * sigma_a));
             auto merged   = merge_av_latents(compute_ctx, output.first, output.second);
-            auto graph    = new_graph_custom(H3_GRAPH_SIZE);
             ggml_build_forward_expand(graph, merged);
             return graph;
         }

src/model/integrations/minimax_h3.cpp

@@ -18,7 +18,8 @@ static bool build_minimax_h3(const sd::model_builders::Context& ctx, const Token
     result.diffusion               = std::make_shared<MiniMaxH3::MiniMaxH3Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                                   tensor_storage_map,
                                                                                   "model.diffusion_model",
-                                                                                  weight_manager);
+                                                                                  weight_manager,
+                                                                                  sd_ctx_params->model_args);
     return true;
 }