DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-11 11:26:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
leejet authored and GitHub committed 2026-11-08 15:50:25 +08:00
1 parent 7bcd4ce8ce
commit de00bafc3d
7 files changed
+103 -43
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/core/cuda/i8_gated_residual.cu
+1 -1
View File
@@ -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
+31 -31
View File
@@ -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>
src/core/cuda/int8.cu
+2 -2
View File
@@ -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*) {
src/core/cuda/rope.cu
+1 -1
View File
@@ -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
+54 -7
View File
@@ -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;
}
src/model/integrations/minimax_h3.cpp
+2 -1
View File
@@ -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;
}