Commit de00baf for stable-diffusion.cpp
commit de00bafc3d473069583752c80a0d56140477ede2
Author: leejet <leejet714@gmail.com>
Date: Sun Nov 8 15:50:25 2026 +0800
feat: add configurable attention chunk budget for MiniMax-H3 (#2292)
diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md
index db554dd..bc1f19c 100644
--- a/docs/minimax_h3.md
+++ b/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.
diff --git a/src/core/cuda/i8_gated_residual.cu b/src/core/cuda/i8_gated_residual.cu
index 9e43d5c..71477d7 100644
--- a/src/core/cuda/i8_gated_residual.cu
+++ b/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);
diff --git a/src/core/cuda/i8_gemm.cu b/src/core/cuda/i8_gemm.cu
index c62b95c..9d06c91 100644
--- a/src/core/cuda/i8_gemm.cu
+++ b/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 Tile = typename Config::Tile;
- using Warp = typename Config::Warp;
- using Swizzle = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
- 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 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 ScaleX = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
- 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 Epilogue = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
- using Kernel = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
+ using Weight = std::conditional_t<PackedW4, cutlass::int4b_t, int8_t>;
+ using Tile = typename Config::Tile;
+ using Warp = typename Config::Warp;
+ using Swizzle = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
+ 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 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 ScaleX = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
+ 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 Epilogue = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
+ 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 Tile = typename Config::Tile;
- using Warp = typename Config::Warp;
- using Swizzle = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
- 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 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 ScaleX = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
- using ScaleW = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
- using Store = SwigluStore<Map>;
- 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 VisitorEpilogue = cutlass::epilogue::threadblock::EpilogueWithVisitorCallbacks<typename GemmBase::Epilogue, Epilogue, 1>;
+ using Weight = int8_t;
+ using Tile = typename Config::Tile;
+ using Warp = typename Config::Warp;
+ using Swizzle = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
+ 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 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 ScaleX = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
+ using ScaleW = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
+ using Store = SwigluStore<Map>;
+ 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 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 Kernel = std::conditional_t<StreamK,
+ using Mma = std::conditional_t<Split, SplitMma, BaseMma>;
+ 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>
diff --git a/src/core/cuda/int8.cu b/src/core/cuda/int8.cu
index 93d8f7f..01a6147 100644
--- a/src/core/cuda/int8.cu
+++ b/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));
- *workspace = total;
+ p->output = fresh_output || (p->gemm && p->kind == SWIGLU_QUANTIZE) ? SIZE_MAX : reserve(total, ggml_nbytes(dst));
+ *workspace = total;
return GGML_STATUS_SUCCESS;
}
static void destroy(void* state, void*) {
diff --git a/src/core/cuda/rope.cu b/src/core/cuda/rope.cu
index f3ab216..2bad5df 100644
--- a/src/core/cuda/rope.cu
+++ b/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);
diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp
index 9a40643..97fd230 100644
--- a/src/model/diffusion/minimax_h3.hpp
+++ b/src/model/diffusion/minimax_h3.hpp
@@ -19,9 +19,10 @@
namespace MiniMaxH3 {
- constexpr int H3_GRAPH_SIZE = 131072;
- constexpr float FRAME_RESCALE = 5.f / 3.f;
- constexpr float VISUAL_COND_TIMESTEP = 0.999f;
+ constexpr int H3_GRAPH_SIZE = 131072;
+ constexpr float DEFAULT_ATTN_CHUNK_BUDGET_GIB = 2.f;
+ 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;
}
diff --git a/src/model/integrations/minimax_h3.cpp b/src/model/integrations/minimax_h3.cpp
index 35a1f67..7624a6c 100644
--- a/src/model/integrations/minimax_h3.cpp
+++ b/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;
}